@raoxxxwq/pi-codebuddy-sdk 0.3.0 → 0.4.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@raoxxxwq/pi-codebuddy-sdk",
-  "version": "0.3.0",
+  "version": "0.4.0",
   "description": "Pi extension that uses CodeBuddy (via Agent SDK) as a model provider and adds an AskCodebuddy tool.",
   "author": "raoxxxwq",
   "license": "MIT",

package/src/index.ts

@@ -1,7 +1,7 @@
 import { calculateCost, StringEnum, type AssistantMessage, type AssistantMessageEventStream, type Context, type Model, type SimpleStreamOptions, type Tool } from "@earendil-works/pi-ai";
 import * as piAi from "@earendil-works/pi-ai";
 import { buildSessionContext, compact, keyHint, type CompactionEntry, type ExtensionAPI, type ExtensionUIContext } from "@earendil-works/pi-coding-agent";
-import { createSdkMcpServer, query, type Effort, type Message as CbMessage, type UserMessage as CbUserMessage } from "@tencent-ai/agent-sdk";
+import { createSdkMcpServer, query, unstable_v2_createSession as createSdkSession, type Effort, type Message as CbMessage, type UserMessage as CbUserMessage } from "@tencent-ai/agent-sdk";
 import { Type } from "typebox";
 import { Text } from "@earendil-works/pi-tui";
 import { createSession, deleteSession } from "./cb-session-io.js";
@@ -9,13 +9,13 @@ import { appendFileSync, mkdirSync, realpathSync, statSync } from "fs";
 import { homedir } from "os";
 import { dirname, join } from "path";
 import { PROVIDER_ID, messageContentToText, convertPiMessages } from "./convert.js";
-import { buildModels, codebuddyModelId, FALLBACK_MODELS, rawModelsFromSdk, resolveModel as _resolveModel, type PiModel } from "./models.js";
+import { buildModels, codebuddyModelId, FALLBACK_MODELS, rawModelsFromSdk, rawModelsFromSdkRaw, resolveModel as _resolveModel, type PiModel } from "./models.js";
 import { MCP_SERVER_NAME, MCP_TOOL_PREFIX, buildCodebuddySystemPrompt, enhancePiToolForCodebuddy } from "./skills.js";
 import { verifyWrittenSession as _verifyWrittenSession } from "./session-verify.js";
 import { extractAllToolResults as _extractAllToolResults, type McpResult } from "./extract-tool-results.js";
 import { QueryContext, ctx } from "./query-state.js";
 import { loadConfig, type Config } from "./config.js";
-import { jsonSchemaToZodObject } from "./typebox-to-zod.js";
+import { jsonSchemaToZodObjectForMcp } from "./typebox-to-zod.js";
 import { buildActionSummary, type ToolCallState } from "./askcodebuddy-ui.js";
 import {
 	applyContextWindowCalibrations,
@@ -664,6 +664,8 @@ export const __test = {
 	buildProviderBoundaryOptions,
 	buildProviderQueryOptions,
 	syncSharedSession,
+	isEmptyArgs,
+	hasRequiredParams,
 };
 
 function createDelegationSessionFromContext(
@@ -755,6 +757,30 @@ function mapToolName(name: string, customToolNameToPi?: Map<string, string>): st
 	return name;
 }
 
+// True when a tool-call argument object is empty or absent. Used to detect
+// the parallel-tool-call arg-dropping failure: CodeBuddy's MCP client may
+// dispatch tool_call arguments as {} (especially in parallel batches), and
+// the stream's input_json_delta may also arrive empty. We defer toolcall_end
+// for such blocks until the assistant message (or MCP dispatch) provides
+// real arguments, preventing pi from executing tools with empty args.
+function isEmptyArgs(args: Record<string, unknown> | undefined | null): boolean {
+	if (!args) return true;
+	if (Object.keys(args).length === 0) return true;
+	// Treat an object with only undefined/null values as empty
+	const hasValue = Object.values(args).some((v) => v !== undefined && v !== null);
+	return !hasValue;
+}
+
+// Checks whether a Pi tool has any required parameters. Used by the canUseTool
+// pre-validation to decide whether empty dispatch args should be rejected.
+// If a tool has no required params, empty args are legitimate and should pass.
+function hasRequiredParams(toolName: string, tools: Tool[]): boolean {
+	const tool = tools.find((t) => t.name === toolName);
+	if (!tool?.parameters) return false;
+	const required = (tool.parameters as Record<string, unknown>).required;
+	return Array.isArray(required) && required.length > 0;
+}
+
 // Renames for CodeBuddy SDK param names that differ from pi's native names.
 // Keys not listed here pass through unchanged, so new pi params work automatically.
 const SDK_KEY_RENAMES: Record<string, Record<string, string>> = {
@@ -842,10 +868,79 @@ function buildMcpServers(tools: Tool[], queryCtx: QueryContext): Record<string,
 	const mcpTools = tools.map((tool) => ({
 		name: tool.name,
 		description: tool.description,
-		inputSchema: jsonSchemaToZodObject(tool.parameters),
-		handler: async () => {
-			const toolCallId = queryCtx.turnToolCallIds[queryCtx.nextHandlerIdx++];
-			if (!toolCallId) debug(`WARNING: mcp handler ${tool.name} has no toolCallId (idx=${queryCtx.nextHandlerIdx - 1}, available=${queryCtx.turnToolCallIds.length})`);
+		// MCP schema preserves required constraints so empty {} (from parallel
+		// tool_call arg-dropping) is rejected at MCP validation time — an early
+		// signal that args were lost. The deferred-backfill logic handles the
+		// stream side; this handles the dispatch side. Passthrough allows extra
+		// keys for forward-compat.
+		inputSchema: jsonSchemaToZodObjectForMcp(tool.parameters),
+		handler: async (dispatchArgs?: Record<string, unknown>) => {
+			// Name-based toolCallId matching: find the first toolCallId whose
+			// corresponding turnBlock has a name matching this tool, and hasn't
+			// been claimed by a previous handler invocation. This fixes the
+			// parallel-dispatch race where CodeBuddy calls MCP handlers in a
+			// different order than the stream's content_block_start events —
+			// without this, bash's handler could get read's toolCallId, causing
+			// result misassignment (bash result → read handler, vice versa).
+			//
+			// Fallback to index-based for robustness (e.g. if turnBlocks doesn't
+			// have the block yet, or same-name dedup edge cases).
+			let toolCallId: string | undefined;
+			for (const id of queryCtx.turnToolCallIds) {
+				if (queryCtx.matchedToolCallIds.has(id)) continue;
+				const block = queryCtx.turnBlocks.find((b: any) => b.id === id && b.type === "toolCall");
+				if (block && block.name === tool.name) {
+					toolCallId = id;
+					queryCtx.matchedToolCallIds.add(id);
+					break;
+				}
+			}
+			// Fallback: if no name match, use index (handles edge cases where
+			// block.name hasn't been mapped yet or same-name calls overflow)
+			if (!toolCallId) {
+				toolCallId = queryCtx.turnToolCallIds[queryCtx.nextHandlerIdx];
+				if (toolCallId) {
+					queryCtx.matchedToolCallIds.add(toolCallId);
+					queryCtx.nextHandlerIdx++;
+				}
+			}
+			if (!toolCallId) debug(`WARNING: mcp handler ${tool.name} has no toolCallId (matched=${queryCtx.matchedToolCallIds.size}, available=${queryCtx.turnToolCallIds.length})`);
+
+			debug(`mcp dispatch: ${tool.name} dispatchArgsLen=${JSON.stringify(dispatchArgs ?? {}).length} dispatchArgsEmpty=${isEmptyArgs(dispatchArgs)} matched=${queryCtx.matchedToolCallIds.size}/${queryCtx.turnToolCallIds.length} pendingBlocks=${queryCtx.argsPendingBlocks.length}`);
+
+			// Backfill path: if there are args-pending blocks whose toolcall_end
+			// was deferred (stream args were empty), try to backfill using the
+			// MCP dispatch args. This handles the case where the assistant message
+			// also had empty args but the MCP dispatch carries the real args.
+			if (queryCtx.argsPendingBlocks.length > 0 && toolCallId) {
+				const pendingIdx = queryCtx.argsPendingBlocks.findIndex((p) => p.block.id === toolCallId);
+				if (pendingIdx >= 0) {
+					const pending = queryCtx.argsPendingBlocks[pendingIdx];
+					const backfillArgs = mapToolArgs(tool.name, dispatchArgs);
+					if (!isEmptyArgs(backfillArgs)) {
+						pending.block.arguments = backfillArgs;
+						debug(`mcp handler: backfilled ${tool.name} [${toolCallId}] from MCP dispatch args (argsLen=${JSON.stringify(backfillArgs).length})`);
+					} else {
+						debug(`mcp handler: dispatch args also empty for ${tool.name} [${toolCallId}], emitting with current args`);
+					}
+					queryCtx.turnSawToolCall = true;
+					queryCtx.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: queryCtx.turnOutput });
+					queryCtx.argsPendingBlocks.splice(pendingIdx, 1);
+
+					// If all pending blocks are now resolved and done was deferred, emit it
+					if (queryCtx.argsPendingBlocks.length === 0 && queryCtx.doneDeferredForArgs && queryCtx.currentPiStream && queryCtx.turnOutput) {
+						queryCtx.doneDeferredForArgs = false;
+						queryCtx.turnOutput.stopReason = "toolUse";
+						const stream = queryCtx.currentPiStream;
+						stream.push({ type: "done", reason: "toolUse", message: queryCtx.turnOutput });
+						markStreamComplete(stream);
+						stream.end();
+						queryCtx.currentPiStream = null;
+						debug(`mcp handler: all pending blocks resolved, emitted deferred done event`);
+					}
+				}
+			}
+
 			if (toolCallId && queryCtx.pendingResults.has(toolCallId)) {
 				const result = queryCtx.pendingResults.get(toolCallId)!;
 				queryCtx.pendingResults.delete(toolCallId);
@@ -899,7 +994,7 @@ function observeServedContextWindow(label: string, servedModel: string, observed
 	if (!Number.isFinite(observed) || observed <= 0) return;
 	try {
 		const previousRegistered = MODELS.find((candidate) => candidate.id === model.id)?.contextWindow;
-		const { changed, floorChanged, record } = recordObservedContextWindow(
+		const { changed, record } = recordObservedContextWindow(
 			calibrationCache,
 			model.id,
 			calibrationEnvironment,
@@ -907,11 +1002,11 @@ function observeServedContextWindow(label: string, servedModel: string, observed
 		);
 		if (!changed) return;
 		saveCalibrationCache(calibrationCache);
-		const floor = record.capabilities.contextWindow?.floor;
-		if (floor != null) {
+		const latest = record.capabilities.contextWindow?.latest;
+		if (latest != null) {
 			MODELS = MODELS.map((candidate) => (
 				candidate.id === model.id
-					? { ...candidate, contextWindow: floor }
+					? { ...candidate, contextWindow: latest }
 					: candidate
 			));
 		}
@@ -919,7 +1014,7 @@ function observeServedContextWindow(label: string, servedModel: string, observed
 			`calibration: ${label} observed=${observed} floor=${record.capabilities.contextWindow?.floor ?? "?"} ` +
 			`latest=${record.capabilities.contextWindow?.latest ?? "?"} servedModel=${servedModel} registeredBefore=${previousRegistered ?? "?"}`,
 		);
-		if (floorChanged && floor != null && previousRegistered != null && floor !== previousRegistered) {
+		if (latest != null && previousRegistered != null && latest !== previousRegistered) {
 			scheduleCalibrationRefresh(`contextWindow:${model.id}`);
 		}
 	} catch (err) {
@@ -991,6 +1086,20 @@ function finalizeCurrentStream(c: QueryContext, stopReason?: string): void {
 	if (!c.turnStarted) ensureTurnStarted(c);
 	const reason = stopReason === "length" ? "length" : "stop";
 	const stream = c.currentPiStream;
+
+	// Drain any args-pending blocks: if the SDK ended abnormally (crash, network
+	// error) without yielding an assistant message, these blocks were never
+	// backfilled. Emit them with their current (possibly empty) args so pi can
+	// surface validation errors rather than silently dropping the tool calls.
+	if (c.argsPendingBlocks.length > 0) {
+		debug(`finalizeCurrentStream: draining ${c.argsPendingBlocks.length} pending block(s) with current args`);
+		for (const pending of c.argsPendingBlocks) {
+			c.turnSawToolCall = true;
+			stream!.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
+		}
+		c.argsPendingBlocks = [];
+		c.doneDeferredForArgs = false;
+	}
 	stream!.push({ type: "done", reason, message: c.turnOutput });
 	markStreamComplete(stream);
 	stream!.end();
@@ -1011,7 +1120,8 @@ function processStreamEvent(
 
 	if (event?.type === "message_start") {
 		c.turnToolCallIds = [];
-		c.nextHandlerIdx = 0;
+			c.nextHandlerIdx = 0;
+			c.matchedToolCallIds = new Set();
 		if (event.message?.usage) updateUsage(c.turnOutput, event.message.usage, model);
 		return;
 	}
@@ -1025,7 +1135,7 @@ function processStreamEvent(
 			c.turnBlocks.push({ type: "thinking", thinking: "", thinkingSignature: "", index: event.index });
 			c.currentPiStream!.push({ type: "thinking_start", contentIndex: c.turnBlocks.length - 1, partial: c.turnOutput });
 		} else if (event.content_block?.type === "tool_use") {
-			c.turnSawToolCall = true;
+			c.turnToolCallIds.push(event.content_block.id);
 			c.turnToolCallIds.push(event.content_block.id);
 			c.turnBlocks.push({
 				type: "toolCall", id: event.content_block.id,
@@ -1072,12 +1182,33 @@ function processStreamEvent(
 		} else if (block.type === "thinking") {
 			c.currentPiStream!.push({ type: "thinking_end", contentIndex: index, content: block.thinking, partial: c.turnOutput });
 		} else if (block.type === "toolCall") {
-			c.turnSawToolCall = true;
+			const partialJsonLen = block.partialJson?.length ?? 0;
 			block.arguments = mapToolArgs(
 				block.name, parsePartialJson(block.partialJson, block.arguments),
 			);
 			delete block.partialJson;
-			c.currentPiStream!.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: c.turnOutput });
+
+			debug(`processStreamEvent: content_block_stop ${block.name} [${block.id}] argsSource=${
+				isEmptyArgs(block.arguments) ? "EMPTY" : "stream"
+			} argsLen=${JSON.stringify(block.arguments).length} partialJsonLen=${partialJsonLen}`);
+
+			// Parallel tool-call arg-dropping defense: when the stream delivers
+			// empty args (CodeBuddy's MCP client may dispatch parallel tool_call
+			// arguments as {}), defer toolcall_end until the assistant message
+			// (or MCP dispatch) provides real args. This prevents pi from
+			// executing tools with empty args (e.g. "bash" with no command).
+			//
+			// turnSawToolCall is set only when we actually emit a toolcall_end
+			// (here for non-empty args, or in the backfill path for deferred
+			// blocks). This ensures message_stop's done-deferral check fires
+			// correctly even when ALL tool blocks had empty args.
+			if (isEmptyArgs(block.arguments)) {
+				debug(`processStreamEvent: deferring toolcall_end for ${block.name} [${block.id}] — stream args empty, will backfill from assistant message or MCP dispatch`);
+				c.argsPendingBlocks.push({ block, contentIndex: index });
+			} else {
+				c.turnSawToolCall = true;
+				c.currentPiStream!.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: c.turnOutput });
+			}
 		}
 		return;
 	}
@@ -1088,6 +1219,19 @@ function processStreamEvent(
 		return;
 	}
 
+	// Check args-pending blocks FIRST, before the turnSawToolCall gate.
+	// turnSawToolCall is only set when a toolcall_end is actually emitted (non-empty
+	// args in content_block_stop, or during backfill). When ALL tool blocks had
+	// empty stream args, turnSawToolCall is false here. Without this early check,
+	// doneDeferredForArgs would never be set and the stream would hang forever.
+	// This is defense-in-depth: the backfill path also handles this, but setting
+	// doneDeferredForArgs here ensures the assistant message path knows to emit done.
+	if (event?.type === "message_stop" && c.argsPendingBlocks.length > 0) {
+		debug(`processStreamEvent: message_stop deferring done event — ${c.argsPendingBlocks.length} block(s) awaiting args backfill (turnSawToolCall=${c.turnSawToolCall})`);
+		c.doneDeferredForArgs = true;
+		return;
+	}
+
 	if (event?.type === "message_stop" && c.turnSawToolCall) {
 		// Tool call complete — end this pi stream. The SDK will still yield an
 		// assistant message for this turn, but currentPiStream=null causes
@@ -1117,11 +1261,77 @@ function processStreamEvent(
 // the same stream lifecycle as processStreamEvent — including ending the stream on
 // tool_use to prevent deadlock with the MCP handler.
 function processAssistantMessage(message: CbMessage, model: Model<any>, customToolNameToPi: Map<string, string>, c: QueryContext): void {
-	if (c.turnSawStreamEvent) return;
 	const assistantMsg = (message as any).message;
 	if (!assistantMsg?.content) return;
+
+	// --- Args-backfill path (stream already delivered content, but some tool
+	// blocks had empty args). Use the assistant message's complete tool_use
+	// input to backfill the deferred blocks, then emit the pending toolcall_end
+	// events and the deferred done event.
+	if (c.turnSawStreamEvent && c.argsPendingBlocks.length > 0) {
+		debug(`processAssistantMessage: backfill path — ${c.argsPendingBlocks.length} pending block(s), assistant has ${assistantMsg.content.length} blocks`);
+		const toolUseBlocks = assistantMsg.content.filter((b: any) => b.type === "tool_use");
+
+		// Match pending blocks to assistant message tool_use blocks by position.
+		// Stream order and assistant message order should align (both follow the
+	// model's content_block index). If names mismatch we log a warning but still
+	// backfill by position — the stream is the source of truth for ids.
+		const remaining = [...c.argsPendingBlocks];
+		for (const tuBlock of toolUseBlocks) {
+			if (remaining.length === 0) break;
+			const tuName = mapToolName(tuBlock.name, customToolNameToPi);
+			// Find the first pending block whose name matches (or just the first
+		// pending block if names don't match — positional fallback).
+			let matchIdx = remaining.findIndex((p) => p.block.name === tuName);
+			if (matchIdx < 0) {
+				debug(`processAssistantMessage: backfill name mismatch — assistant '${tuName}' vs pending [${remaining.map((p) => p.block.name).join(",")}], using positional fallback`);
+				matchIdx = 0;
+			}
+			const pending = remaining.splice(matchIdx, 1)[0];
+			const dispatchArgs = mapToolArgs(pending.block.name, tuBlock.input);
+			if (isEmptyArgs(dispatchArgs)) {
+				debug(`processAssistantMessage: backfill for ${pending.block.name} [${pending.block.id}] — assistant args also empty, emitting with empty args (pi will validate)`);
+			} else {
+				debug(`processAssistantMessage: backfilled ${pending.block.name} [${pending.block.id}] from assistant message (argsLen=${JSON.stringify(dispatchArgs).length})`);
+			}
+			pending.block.arguments = dispatchArgs;
+			c.turnSawToolCall = true;
+			c.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
+			// Remove from the original argsPendingBlocks array
+			const origIdx = c.argsPendingBlocks.indexOf(pending);
+			if (origIdx >= 0) c.argsPendingBlocks.splice(origIdx, 1);
+		}
+
+		// Any remaining pending blocks (no matching assistant tool_use) get
+		// emitted with their (empty) args so pi can surface the validation error
+		// rather than hanging forever.
+		for (const pending of remaining) {
+			debug(`processAssistantMessage: backfill — no matching assistant tool_use for ${pending.block.name} [${pending.block.id}], emitting with current args`);
+			c.turnSawToolCall = true;
+			c.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
+			const origIdx = c.argsPendingBlocks.indexOf(pending);
+			if (origIdx >= 0) c.argsPendingBlocks.splice(origIdx, 1);
+		}
+
+		// Emit the deferred done event (message_stop skipped it because of pending blocks)
+		if (c.doneDeferredForArgs && c.turnSawToolCall && c.currentPiStream && c.turnOutput) {
+			c.doneDeferredForArgs = false;
+			c.turnOutput.stopReason = "toolUse";
+			const stream = c.currentPiStream;
+			stream.push({ type: "done", reason: "toolUse", message: c.turnOutput });
+			markStreamComplete(stream);
+			stream.end();
+			c.currentPiStream = null;
+			debug(`processAssistantMessage: backfill complete, emitted deferred done event`);
+		}
+		return;
+	}
+
+	// --- Original path: no stream events, this is the primary content path ---
+	if (c.turnSawStreamEvent) return;
 	c.turnToolCallIds = [];
-	c.nextHandlerIdx = 0;
+		c.nextHandlerIdx = 0;
+		c.matchedToolCallIds = new Set();
 	debug(`processAssistantMessage fallback: ${assistantMsg.content.length} blocks, types=${assistantMsg.content.map((b: any) => b.type).join(",")}`);
 	for (const block of assistantMsg.content) {
 		if (block.type === "text" && block.text) {
@@ -1407,6 +1617,29 @@ function streamCodebuddySdk(model: Model<any>, context: Context, options?: Simpl
 		debugOptions: makeCliDebugOptions("provider"),
 	});
 
+	// P4: canUseTool pre-validation — reject empty dispatch args at MCP layer
+	// before they reach pi. This is an early signal to CodeBuddy that args
+	// were dropped (parallel tool_call dispatch bug), giving the model a chance
+	// to regenerate proper args in the same turn rather than failing at pi-side
+	// validation and requiring a full retry.
+	//
+	// Only rejects when the tool has required params AND dispatch args are empty.
+	// Tools without required params (e.g. some MCP tools) pass through normally.
+	const piTools = context.tools ?? [];
+	(queryOptions as SdkQueryOptions).canUseTool = async (toolName: string, input: Record<string, unknown>, opts) => {
+		const piName = mapToolName(toolName, customToolNameToPi);
+		const mappedArgs = mapToolArgs(piName, input);
+		if (isEmptyArgs(mappedArgs) && hasRequiredParams(piName, piTools)) {
+			debug(`canUseTool: rejecting ${toolName}→${piName} — empty args for tool with required params (toolUseId=${opts.toolUseID})`);
+			return {
+				behavior: "deny" as const,
+				message: `Tool "${piName}" requires arguments but received empty input. This can happen with parallel tool calls. Please provide complete arguments for all required fields.`,
+				toolUseID: opts.toolUseID,
+			};
+		}
+		return { behavior: "allow" as const, toolUseID: opts.toolUseID };
+	};
+
 	debug("provider: fresh query",
 		`model=${cliModel} msgs=${context.messages.length} tools=${mcpTools.length}`,
 		`resume=${resumeSessionId?.slice(0, 8) ?? "none"} effort=${effort ?? "default"}`,
@@ -1716,18 +1949,51 @@ let discoverInFlight: Promise<void> | null = null;
 
 async function discoverModels(pi: ExtensionAPI): Promise<void> {
 	await withSdkGate(async () => {
+		const codebuddyExecutable = providerSettings.pathToCodebuddyCode;
+		const commonOpts = {
+			maxTurns: 0,
+			permissionMode: "bypassPermissions" as const,
+			tools: [] as string[],
+			cwd: process.cwd(),
+			env: { ...process.env, DISABLE_AUTO_COMPACT: "1" },
+			...(codebuddyExecutable ? { pathToCodebuddyCode: codebuddyExecutable } : {}),
+			...makeCliDebugOptions("discover-models"),
+		};
+
+		// Preferred path: Session API getAvailableModelsRaw() returns RawLanguageModel[]
+		// with the real per-model maxInputTokens (context window) and maxOutputTokens,
+		// plus capability flags (supportsImages / supportsReasoning). This eliminates
+		// first-run Window Drift — Pi registers the true context window up front
+		// instead of the 200K conservative default. Runtime calibration remains as a
+		// downward correction layer for entitlement-limited served windows.
+		try {
+			const session = createSdkSession(commonOpts);
+			try {
+				await session.connect();
+				const rawModels = await session.getAvailableModelsRaw();
+				if (rawModels.length) {
+					MODELS = applyModelCalibrations(rawModelsFromSdkRaw(rawModels));
+					registerCurrentProvider(pi);
+					modelsDiscovered = true;
+					debug(`discoverModels: registered ${MODELS.length} models via getAvailableModelsRaw()`);
+					return;
+				}
+				debug("discoverModels: getAvailableModelsRaw() returned empty, falling back to supportedModels()");
+			} finally {
+				session.close();
+			}
+		} catch (err) {
+			debug("discoverModels: getAvailableModelsRaw() failed, falling back to supportedModels()", err);
+		}
+
+		// Fallback path: one-shot query().supportedModels() returns the simplified
+		// ModelInfo (value/displayName/description) without token limits, so we must
+		// use the conservativeContextWindow() heuristic. Used when the CLI doesn't
+		// support the get_available_models control request or the Session API errors.
 		try {
-			const codebuddyExecutable = providerSettings.pathToCodebuddyCode;
 			const q = query({
 				prompt: " ",
-				options: {
-					maxTurns: 0,
-					permissionMode: "bypassPermissions",
-					tools: [],
-					env: { ...process.env, DISABLE_AUTO_COMPACT: "1" },
-					...(codebuddyExecutable ? { pathToCodebuddyCode: codebuddyExecutable } : {}),
-					...makeCliDebugOptions("supported-models"),
-				},
+				options: commonOpts,
 			});
 			const supported = await q.supportedModels();
 			await q.return().catch(() => {});
@@ -1735,9 +2001,9 @@ async function discoverModels(pi: ExtensionAPI): Promise<void> {
 			MODELS = applyModelCalibrations(rawModelsFromSdk(supported as any));
 			registerCurrentProvider(pi);
 			modelsDiscovered = true;
-			debug(`discoverModels: registered ${MODELS.length} models from SDK`);
+			debug(`discoverModels: registered ${MODELS.length} models from supportedModels()`);
 		} catch (err) {
-			debug("discoverModels: failed, using fallback models", err);
+			debug("discoverModels: supportedModels() failed, using fallback models", err);
 		}
 	});
 }

package/src/model-calibration.ts

@@ -90,8 +90,8 @@ export function applyContextWindowCalibrations<T extends { id: string; contextWi
 ): T[] {
 	return models.map((model) => {
 		const metric = getContextWindowCalibration(cache, model.id, environment);
-		if (!metric?.floor) return model;
-		return { ...model, contextWindow: metric.floor };
+		if (!metric?.latest) return model;
+			return { ...model, contextWindow: metric.latest };
 	});
 }
 

package/src/models.ts

@@ -1,5 +1,7 @@
-// Dynamic model list from CodeBuddy SDK supportedModels().
-import type { ModelInfo } from "@tencent-ai/agent-sdk";
+// Dynamic model list from CodeBuddy SDK.
+// - supportedModels(): simplified ModelInfo (no token limits) — legacy fallback.
+// - getAvailableModelsRaw(): RawLanguageModel with real maxInputTokens / maxOutputTokens.
+import type { ModelInfo, RawLanguageModel } from "@tencent-ai/agent-sdk";
 
 export type PiModel = {
 	id: string;
@@ -12,8 +14,8 @@ export type PiModel = {
 	cost: { input: number; output: number; cacheRead: number; cacheWrite: number };
 };
 
-const CONSERVATIVE_DEFAULT_CONTEXT = 65_536;
-const CONSERVATIVE_GEMINI_CONTEXT = 131_072;
+const CONSERVATIVE_DEFAULT_CONTEXT = 200_000;
+const CONSERVATIVE_GEMINI_CONTEXT = 1_048_576;
 const DEFAULT_MAX_TOKENS = 8192;
 
 function detectThinking(id: string): boolean {
@@ -50,6 +52,29 @@ export function rawModelsFromSdk(supported: Array<ModelInfo & { id?: string; nam
 		}));
 }
 
+// Build PiModel[] from the SDK Session API's getAvailableModelsRaw() result.
+// Unlike rawModelsFromSdk() (which only gets the simplified ModelInfo and must
+// fall back to conservativeContextWindow), this reads the real per-model token
+// limits (maxInputTokens = context window, maxOutputTokens = max output) plus
+// capability flags (supportsImages / supportsReasoning) directly from the CLI.
+// Models marked disabled are filtered out so Pi never offers an unusable model.
+export function rawModelsFromSdkRaw(raw: RawLanguageModel[]): PiModel[] {
+	return raw
+		.filter((m) => m.id && !m.disabled)
+		.map((m) => {
+			const id = m.id!;
+			return {
+				id,
+				name: m.name || id,
+				reasoning: m.supportsReasoning ?? detectThinking(id),
+				input: (m.supportsImages ?? detectImages(id)) ? ["text", "image"] as const : ["text"] as const,
+				contextWindow: m.maxInputTokens ?? conservativeContextWindow(id),
+				maxTokens: m.maxOutputTokens ?? conservativeMaxTokens(id),
+				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+			};
+		});
+}
+
 export const FALLBACK_MODELS: PiModel[] = [
 	{
 		id: "hy3-preview-agent-ioa",

package/src/query-state.ts

@@ -23,8 +23,25 @@ export class QueryContext {
 	pendingResults = new Map<string, McpResult>();
 	turnToolCallIds: string[] = [];
 	nextHandlerIdx = 0;
+	// toolCallIds already assigned to an MCP handler. Used by name-based
+	// matching to avoid the index-order race when CodeBuddy dispatches
+	// parallel tool handlers in a different order than the stream's
+	// content_block_start events. Without this, bash's handler could get
+	// read's toolCallId, causing result misassignment.
+	matchedToolCallIds = new Set<string>();
 	deferredUserMessages: string[] = [];
 
+	// Tool-call blocks whose stream args came through empty (parallel tool_call
+	// dispatch dropping args to {}). Their toolcall_end is deferred until either
+	// the assistant message arrives with complete args, or the MCP handler
+	// receives non-empty dispatch args. The done event is also deferred while
+	// this list is non-empty, so pi does not execute tools with empty args.
+	argsPendingBlocks: Array<{ block: any; contentIndex: number }> = [];
+	// Set when message_stop arrived but done was deferred due to pending blocks.
+	// Prevents double-deferral and lets the assistant-message path know it
+	// must emit the done event after backfilling.
+	doneDeferredForArgs = false;
+
 	// Per-turn (reset together)
 	turnOutput: AssistantMessage | null = null;
 	turnStarted = false;
@@ -47,6 +64,9 @@ export class QueryContext {
 		this.turnStarted = false;
 		this.turnSawStreamEvent = false;
 		this.turnSawToolCall = false;
+		this.argsPendingBlocks = [];
+		this.doneDeferredForArgs = false;
+		this.matchedToolCallIds = new Set();
 		// turnToolCallIds and nextHandlerIdx are NOT reset — they persist across
 		// tool-result delivery callbacks within the same assistant message.
 	}

package/src/skills.ts

@@ -71,6 +71,9 @@ export function buildPiToolBridgeInstruction(options: ToolBridgeInstructionOptio
 	if (has("bash")) {
 		lines.push(`- Use \`${mcpName("bash")}\` only when file tools are insufficient, for search/test/build/git information, or when command execution is requested.`);
 	}
+	if (toolNames.size > 1) {
+		lines.push("When calling multiple tools in a single response, ensure each tool_call has complete arguments — do not leave required fields empty or rely on defaults. Incomplete parallel tool calls will be rejected.");
+	}
 	lines.push("Tool arguments must match the Pi tool schema exactly. After a tool result, base the next step on that result; if it is an error, correct the call instead of assuming success.");
 
 	return lines.join("\n");

package/src/typebox-to-zod.ts

@@ -44,24 +44,26 @@ function intersectionSchema(schemas: z.ZodTypeAny[]): z.ZodTypeAny {
 	return schemas.slice(1).reduce((left, right) => z.intersection(left, right), schemas[0]);
 }
 
-function objectSchema(prop: JsonSchema): z.ZodTypeAny {
-	const shape = jsonSchemaToZodShape(prop);
+function objectSchema(prop: JsonSchema, relax = false): z.ZodTypeAny {
+	const shape = jsonSchemaToZodShape(prop, relax);
 	const additionalPropertiesSchema =
 		prop.additionalProperties && typeof prop.additionalProperties === "object"
-			? jsonSchemaPropertyToZod(prop.additionalProperties as JsonSchema)
+			? jsonSchemaPropertyToZod(prop.additionalProperties as JsonSchema, relax)
 			: undefined;
 	if (Object.keys(shape).length > 0) {
 		const base = z.object(shape);
-		if (prop.additionalProperties === false) return base.strict();
+		// relax: never strict — accept {} and extra keys so a CodeBuddy MCP client
+		// that drops parallel tool_call arguments to {} still passes validation.
+		if (!relax && prop.additionalProperties === false) return base.strict();
 		if (additionalPropertiesSchema) return base.catchall(additionalPropertiesSchema);
 		return base.passthrough();
 	}
 	if (additionalPropertiesSchema) return z.record(z.string(), additionalPropertiesSchema);
-	if (prop.additionalProperties === false) return z.object({}).strict();
+	if (!relax && prop.additionalProperties === false) return z.object({}).strict();
 	return z.record(z.string(), z.unknown());
 }
 
-export function jsonSchemaPropertyToZod(prop: JsonSchema): z.ZodTypeAny {
+export function jsonSchemaPropertyToZod(prop: JsonSchema, relax = false): z.ZodTypeAny {
 	let base: z.ZodTypeAny;
 
 	if (prop.const !== undefined) {
@@ -122,7 +124,7 @@ export function jsonSchemaPropertyToZod(prop: JsonSchema): z.ZodTypeAny {
 				: z.array(z.unknown());
 			break;
 		case "object":
-			base = objectSchema(prop);
+			base = objectSchema(prop, relax);
 			break;
 		case "null":
 			base = z.null();
@@ -136,7 +138,7 @@ export function jsonSchemaPropertyToZod(prop: JsonSchema): z.ZodTypeAny {
 	return withDescription(base, prop);
 }
 
-export function jsonSchemaToZodShape(schema: unknown): Record<string, z.ZodTypeAny> {
+export function jsonSchemaToZodShape(schema: unknown, relax = false): Record<string, z.ZodTypeAny> {
 	const s = schema as JsonSchema;
 	const objectSchemas = [s, ...(Array.isArray(s?.allOf) ? s.allOf as JsonSchema[] : [])]
 		.filter((item) => item && typeof item === "object" && ((item as JsonSchema).type === "object" || (item as JsonSchema).properties));
@@ -148,13 +150,51 @@ export function jsonSchemaToZodShape(schema: unknown): Record<string, z.ZodTypeA
 		if (!props) continue;
 		const required = new Set(Array.isArray(objectSchemaPart.required) ? objectSchemaPart.required as string[] : []);
 		for (const [key, prop] of Object.entries(props)) {
-			const zodProp = jsonSchemaPropertyToZod(prop);
-			shape[key] = required.has(key) ? zodProp : zodProp.optional();
+			const zodProp = jsonSchemaPropertyToZod(prop, relax);
+			// relax: make every property optional so an empty {} (args dropped by a
+			// buggy parallel tool_call dispatch) still validates. Descriptions are
+			// preserved so the model keeps parameter guidance.
+			shape[key] = relax || !required.has(key) ? zodProp.optional() : zodProp;
 		}
 	}
 	return shape;
 }
 
-export function jsonSchemaToZodObject(schema: unknown): z.ZodTypeAny {
-	return jsonSchemaPropertyToZod(schema as JsonSchema);
+export function jsonSchemaToZodObject(schema: unknown, relax = false): z.ZodTypeAny {
+	return jsonSchemaPropertyToZod(schema as JsonSchema, relax);
+}
+
+/**
+ * MCP-specific schema builder: preserves required field constraints so that
+ * an empty {} (from a buggy parallel tool_call dispatch) is rejected by MCP
+ * validation, while still allowing extra keys (passthrough) for forward-compat.
+ *
+ * Unlike jsonSchemaToZodObject(relax=true) which makes ALL params optional,
+ * this keeps required params required. This forces CodeBuddy to regenerate
+ * proper args instead of silently passing {} through to the handler.
+ *
+ * The deferred-backfill logic in processStreamEvent/processAssistantMessage
+ * handles the case where stream args arrive empty — it waits for the
+ * assistant message or MCP dispatch to provide real args. So MCP validation
+ * rejecting {} is safe and desirable: it's an early signal that args were
+ * dropped, rather than letting an empty-args call reach pi.
+ */
+export function jsonSchemaToZodObjectForMcp(schema: unknown): z.ZodTypeAny {
+	const s = schema as JsonSchema;
+	// Build shape with relax=false to preserve required constraints
+	const shape = jsonSchemaToZodShape(s, false);
+	const additionalPropertiesSchema =
+		s?.additionalProperties && typeof s.additionalProperties === "object"
+			? jsonSchemaPropertyToZod(s.additionalProperties as JsonSchema, false)
+			: undefined;
+
+	if (Object.keys(shape).length > 0) {
+		const base = z.object(shape);
+		// Always passthrough: accept extra keys (forward-compat with new params)
+		// but enforce required fields.
+		if (additionalPropertiesSchema) return base.catchall(additionalPropertiesSchema);
+		return base.passthrough();
+	}
+	if (s?.additionalProperties === false) return z.object({}).strict();
+	return z.record(z.string(), z.unknown());
 }