@raoxxxwq/pi-codebuddy-sdk 0.3.1 → 0.5.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@raoxxxwq/pi-codebuddy-sdk",
3
- "version": "0.3.1",
3
+ "version": "0.5.0",
4
4
  "description": "Pi extension that uses CodeBuddy (via Agent SDK) as a model provider and adds an AskCodebuddy tool.",
5
5
  "author": "raoxxxwq",
6
6
  "license": "MIT",
package/src/config.ts CHANGED
@@ -26,6 +26,7 @@ export interface Config {
26
26
  appendSystemPrompt?: boolean;
27
27
  settingSources?: SettingSource[];
28
28
  strictMcpConfig?: boolean;
29
+ serialToolCalls?: boolean;
29
30
  pathToCodebuddyCode?: string;
30
31
  };
31
32
  }
package/src/index.ts CHANGED
@@ -1,7 +1,7 @@
1
1
  import { calculateCost, StringEnum, type AssistantMessage, type AssistantMessageEventStream, type Context, type Model, type SimpleStreamOptions, type Tool } from "@earendil-works/pi-ai";
2
2
  import * as piAi from "@earendil-works/pi-ai";
3
3
  import { buildSessionContext, compact, keyHint, type CompactionEntry, type ExtensionAPI, type ExtensionUIContext } from "@earendil-works/pi-coding-agent";
4
- import { createSdkMcpServer, query, type Effort, type Message as CbMessage, type UserMessage as CbUserMessage } from "@tencent-ai/agent-sdk";
4
+ import { createSdkMcpServer, query, unstable_v2_createSession as createSdkSession, type Effort, type Message as CbMessage, type UserMessage as CbUserMessage } from "@tencent-ai/agent-sdk";
5
5
  import { Type } from "typebox";
6
6
  import { Text } from "@earendil-works/pi-tui";
7
7
  import { createSession, deleteSession } from "./cb-session-io.js";
@@ -9,13 +9,13 @@ import { appendFileSync, mkdirSync, realpathSync, statSync } from "fs";
9
9
  import { homedir } from "os";
10
10
  import { dirname, join } from "path";
11
11
  import { PROVIDER_ID, messageContentToText, convertPiMessages } from "./convert.js";
12
- import { buildModels, codebuddyModelId, FALLBACK_MODELS, rawModelsFromSdk, resolveModel as _resolveModel, type PiModel } from "./models.js";
12
+ import { buildModels, codebuddyModelId, FALLBACK_MODELS, rawModelsFromSdk, rawModelsFromSdkRaw, resolveModel as _resolveModel, type PiModel } from "./models.js";
13
13
  import { MCP_SERVER_NAME, MCP_TOOL_PREFIX, buildCodebuddySystemPrompt, enhancePiToolForCodebuddy } from "./skills.js";
14
14
  import { verifyWrittenSession as _verifyWrittenSession } from "./session-verify.js";
15
15
  import { extractAllToolResults as _extractAllToolResults, type McpResult } from "./extract-tool-results.js";
16
16
  import { QueryContext, ctx } from "./query-state.js";
17
17
  import { loadConfig, type Config } from "./config.js";
18
- import { jsonSchemaToZodObject } from "./typebox-to-zod.js";
18
+ import { jsonSchemaToZodObjectForMcp } from "./typebox-to-zod.js";
19
19
  import { buildActionSummary, type ToolCallState } from "./askcodebuddy-ui.js";
20
20
  import {
21
21
  applyContextWindowCalibrations,
@@ -664,6 +664,8 @@ export const __test = {
664
664
  buildProviderBoundaryOptions,
665
665
  buildProviderQueryOptions,
666
666
  syncSharedSession,
667
+ isEmptyArgs,
668
+ hasRequiredParams,
667
669
  };
668
670
 
669
671
  function createDelegationSessionFromContext(
@@ -698,11 +700,13 @@ function stripToolHistoryForDelegation(messages: Context["messages"]): Context["
698
700
  function buildProviderBoundaryOptions(settings: NonNullable<Config["provider"]>) {
699
701
  const appendSystemPrompt = settings.appendSystemPrompt !== false;
700
702
  const strictMcpConfigEnabled = settings.strictMcpConfig !== false;
703
+ const serialToolCalls = settings.serialToolCalls !== false;
701
704
  const extraArgs: Record<string, string | null> = {};
702
705
  if (strictMcpConfigEnabled) extraArgs["strict-mcp-config"] = null;
703
706
  return {
704
707
  appendSystemPrompt,
705
708
  strictMcpConfigEnabled,
709
+ serialToolCalls,
706
710
  tools: [] as string[],
707
711
  extraArgs,
708
712
  settingSources: appendSystemPrompt
@@ -755,6 +759,30 @@ function mapToolName(name: string, customToolNameToPi?: Map<string, string>): st
755
759
  return name;
756
760
  }
757
761
 
762
+ // True when a tool-call argument object is empty or absent. Used to detect
763
+ // the parallel-tool-call arg-dropping failure: CodeBuddy's MCP client may
764
+ // dispatch tool_call arguments as {} (especially in parallel batches), and
765
+ // the stream's input_json_delta may also arrive empty. We defer toolcall_end
766
+ // for such blocks until the assistant message (or MCP dispatch) provides
767
+ // real arguments, preventing pi from executing tools with empty args.
768
+ function isEmptyArgs(args: Record<string, unknown> | undefined | null): boolean {
769
+ if (!args) return true;
770
+ if (Object.keys(args).length === 0) return true;
771
+ // Treat an object with only undefined/null values as empty
772
+ const hasValue = Object.values(args).some((v) => v !== undefined && v !== null);
773
+ return !hasValue;
774
+ }
775
+
776
+ // Checks whether a Pi tool has any required parameters. Used by the canUseTool
777
+ // pre-validation to decide whether empty dispatch args should be rejected.
778
+ // If a tool has no required params, empty args are legitimate and should pass.
779
+ function hasRequiredParams(toolName: string, tools: Tool[]): boolean {
780
+ const tool = tools.find((t) => t.name === toolName);
781
+ if (!tool?.parameters) return false;
782
+ const required = (tool.parameters as Record<string, unknown>).required;
783
+ return Array.isArray(required) && required.length > 0;
784
+ }
785
+
758
786
  // Renames for CodeBuddy SDK param names that differ from pi's native names.
759
787
  // Keys not listed here pass through unchanged, so new pi params work automatically.
760
788
  const SDK_KEY_RENAMES: Record<string, Record<string, string>> = {
@@ -842,10 +870,79 @@ function buildMcpServers(tools: Tool[], queryCtx: QueryContext): Record<string,
842
870
  const mcpTools = tools.map((tool) => ({
843
871
  name: tool.name,
844
872
  description: tool.description,
845
- inputSchema: jsonSchemaToZodObject(tool.parameters),
846
- handler: async () => {
847
- const toolCallId = queryCtx.turnToolCallIds[queryCtx.nextHandlerIdx++];
848
- if (!toolCallId) debug(`WARNING: mcp handler ${tool.name} has no toolCallId (idx=${queryCtx.nextHandlerIdx - 1}, available=${queryCtx.turnToolCallIds.length})`);
873
+ // MCP schema preserves required constraints so empty {} (from parallel
874
+ // tool_call arg-dropping) is rejected at MCP validation time — an early
875
+ // signal that args were lost. The deferred-backfill logic handles the
876
+ // stream side; this handles the dispatch side. Passthrough allows extra
877
+ // keys for forward-compat.
878
+ inputSchema: jsonSchemaToZodObjectForMcp(tool.parameters),
879
+ handler: async (dispatchArgs?: Record<string, unknown>) => {
880
+ // Name-based toolCallId matching: find the first toolCallId whose
881
+ // corresponding turnBlock has a name matching this tool, and hasn't
882
+ // been claimed by a previous handler invocation. This fixes the
883
+ // parallel-dispatch race where CodeBuddy calls MCP handlers in a
884
+ // different order than the stream's content_block_start events —
885
+ // without this, bash's handler could get read's toolCallId, causing
886
+ // result misassignment (bash result → read handler, vice versa).
887
+ //
888
+ // Fallback to index-based for robustness (e.g. if turnBlocks doesn't
889
+ // have the block yet, or same-name dedup edge cases).
890
+ let toolCallId: string | undefined;
891
+ for (const id of queryCtx.turnToolCallIds) {
892
+ if (queryCtx.matchedToolCallIds.has(id)) continue;
893
+ const block = queryCtx.turnBlocks.find((b: any) => b.id === id && b.type === "toolCall");
894
+ if (block && block.name === tool.name) {
895
+ toolCallId = id;
896
+ queryCtx.matchedToolCallIds.add(id);
897
+ break;
898
+ }
899
+ }
900
+ // Fallback: if no name match, use index (handles edge cases where
901
+ // block.name hasn't been mapped yet or same-name calls overflow)
902
+ if (!toolCallId) {
903
+ toolCallId = queryCtx.turnToolCallIds[queryCtx.nextHandlerIdx];
904
+ if (toolCallId) {
905
+ queryCtx.matchedToolCallIds.add(toolCallId);
906
+ queryCtx.nextHandlerIdx++;
907
+ }
908
+ }
909
+ if (!toolCallId) debug(`WARNING: mcp handler ${tool.name} has no toolCallId (matched=${queryCtx.matchedToolCallIds.size}, available=${queryCtx.turnToolCallIds.length})`);
910
+
911
+ debug(`mcp dispatch: ${tool.name} dispatchArgsLen=${JSON.stringify(dispatchArgs ?? {}).length} dispatchArgsEmpty=${isEmptyArgs(dispatchArgs)} matched=${queryCtx.matchedToolCallIds.size}/${queryCtx.turnToolCallIds.length} pendingBlocks=${queryCtx.argsPendingBlocks.length}`);
912
+
913
+ // Backfill path: if there are args-pending blocks whose toolcall_end
914
+ // was deferred (stream args were empty), try to backfill using the
915
+ // MCP dispatch args. This handles the case where the assistant message
916
+ // also had empty args but the MCP dispatch carries the real args.
917
+ if (queryCtx.argsPendingBlocks.length > 0 && toolCallId) {
918
+ const pendingIdx = queryCtx.argsPendingBlocks.findIndex((p) => p.block.id === toolCallId);
919
+ if (pendingIdx >= 0) {
920
+ const pending = queryCtx.argsPendingBlocks[pendingIdx];
921
+ const backfillArgs = mapToolArgs(tool.name, dispatchArgs);
922
+ if (!isEmptyArgs(backfillArgs)) {
923
+ pending.block.arguments = backfillArgs;
924
+ debug(`mcp handler: backfilled ${tool.name} [${toolCallId}] from MCP dispatch args (argsLen=${JSON.stringify(backfillArgs).length})`);
925
+ } else {
926
+ debug(`mcp handler: dispatch args also empty for ${tool.name} [${toolCallId}], emitting with current args`);
927
+ }
928
+ queryCtx.turnSawToolCall = true;
929
+ queryCtx.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: queryCtx.turnOutput });
930
+ queryCtx.argsPendingBlocks.splice(pendingIdx, 1);
931
+
932
+ // If all pending blocks are now resolved and done was deferred, emit it
933
+ if (queryCtx.argsPendingBlocks.length === 0 && queryCtx.doneDeferredForArgs && queryCtx.currentPiStream && queryCtx.turnOutput) {
934
+ queryCtx.doneDeferredForArgs = false;
935
+ queryCtx.turnOutput.stopReason = "toolUse";
936
+ const stream = queryCtx.currentPiStream;
937
+ stream.push({ type: "done", reason: "toolUse", message: queryCtx.turnOutput });
938
+ markStreamComplete(stream);
939
+ stream.end();
940
+ queryCtx.currentPiStream = null;
941
+ debug(`mcp handler: all pending blocks resolved, emitted deferred done event`);
942
+ }
943
+ }
944
+ }
945
+
849
946
  if (toolCallId && queryCtx.pendingResults.has(toolCallId)) {
850
947
  const result = queryCtx.pendingResults.get(toolCallId)!;
851
948
  queryCtx.pendingResults.delete(toolCallId);
@@ -991,6 +1088,20 @@ function finalizeCurrentStream(c: QueryContext, stopReason?: string): void {
991
1088
  if (!c.turnStarted) ensureTurnStarted(c);
992
1089
  const reason = stopReason === "length" ? "length" : "stop";
993
1090
  const stream = c.currentPiStream;
1091
+
1092
+ // Drain any args-pending blocks: if the SDK ended abnormally (crash, network
1093
+ // error) without yielding an assistant message, these blocks were never
1094
+ // backfilled. Emit them with their current (possibly empty) args so pi can
1095
+ // surface validation errors rather than silently dropping the tool calls.
1096
+ if (c.argsPendingBlocks.length > 0) {
1097
+ debug(`finalizeCurrentStream: draining ${c.argsPendingBlocks.length} pending block(s) with current args`);
1098
+ for (const pending of c.argsPendingBlocks) {
1099
+ c.turnSawToolCall = true;
1100
+ stream!.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
1101
+ }
1102
+ c.argsPendingBlocks = [];
1103
+ c.doneDeferredForArgs = false;
1104
+ }
994
1105
  stream!.push({ type: "done", reason, message: c.turnOutput });
995
1106
  markStreamComplete(stream);
996
1107
  stream!.end();
@@ -1011,7 +1122,8 @@ function processStreamEvent(
1011
1122
 
1012
1123
  if (event?.type === "message_start") {
1013
1124
  c.turnToolCallIds = [];
1014
- c.nextHandlerIdx = 0;
1125
+ c.nextHandlerIdx = 0;
1126
+ c.matchedToolCallIds = new Set();
1015
1127
  if (event.message?.usage) updateUsage(c.turnOutput, event.message.usage, model);
1016
1128
  return;
1017
1129
  }
@@ -1025,7 +1137,7 @@ function processStreamEvent(
1025
1137
  c.turnBlocks.push({ type: "thinking", thinking: "", thinkingSignature: "", index: event.index });
1026
1138
  c.currentPiStream!.push({ type: "thinking_start", contentIndex: c.turnBlocks.length - 1, partial: c.turnOutput });
1027
1139
  } else if (event.content_block?.type === "tool_use") {
1028
- c.turnSawToolCall = true;
1140
+ c.turnToolCallIds.push(event.content_block.id);
1029
1141
  c.turnToolCallIds.push(event.content_block.id);
1030
1142
  c.turnBlocks.push({
1031
1143
  type: "toolCall", id: event.content_block.id,
@@ -1072,12 +1184,33 @@ function processStreamEvent(
1072
1184
  } else if (block.type === "thinking") {
1073
1185
  c.currentPiStream!.push({ type: "thinking_end", contentIndex: index, content: block.thinking, partial: c.turnOutput });
1074
1186
  } else if (block.type === "toolCall") {
1075
- c.turnSawToolCall = true;
1187
+ const partialJsonLen = block.partialJson?.length ?? 0;
1076
1188
  block.arguments = mapToolArgs(
1077
1189
  block.name, parsePartialJson(block.partialJson, block.arguments),
1078
1190
  );
1079
1191
  delete block.partialJson;
1080
- c.currentPiStream!.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: c.turnOutput });
1192
+
1193
+ debug(`processStreamEvent: content_block_stop ${block.name} [${block.id}] argsSource=${
1194
+ isEmptyArgs(block.arguments) ? "EMPTY" : "stream"
1195
+ } argsLen=${JSON.stringify(block.arguments).length} partialJsonLen=${partialJsonLen}`);
1196
+
1197
+ // Parallel tool-call arg-dropping defense: when the stream delivers
1198
+ // empty args (CodeBuddy's MCP client may dispatch parallel tool_call
1199
+ // arguments as {}), defer toolcall_end until the assistant message
1200
+ // (or MCP dispatch) provides real args. This prevents pi from
1201
+ // executing tools with empty args (e.g. "bash" with no command).
1202
+ //
1203
+ // turnSawToolCall is set only when we actually emit a toolcall_end
1204
+ // (here for non-empty args, or in the backfill path for deferred
1205
+ // blocks). This ensures message_stop's done-deferral check fires
1206
+ // correctly even when ALL tool blocks had empty args.
1207
+ if (isEmptyArgs(block.arguments)) {
1208
+ debug(`processStreamEvent: deferring toolcall_end for ${block.name} [${block.id}] — stream args empty, will backfill from assistant message or MCP dispatch`);
1209
+ c.argsPendingBlocks.push({ block, contentIndex: index });
1210
+ } else {
1211
+ c.turnSawToolCall = true;
1212
+ c.currentPiStream!.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: c.turnOutput });
1213
+ }
1081
1214
  }
1082
1215
  return;
1083
1216
  }
@@ -1088,6 +1221,19 @@ function processStreamEvent(
1088
1221
  return;
1089
1222
  }
1090
1223
 
1224
+ // Check args-pending blocks FIRST, before the turnSawToolCall gate.
1225
+ // turnSawToolCall is only set when a toolcall_end is actually emitted (non-empty
1226
+ // args in content_block_stop, or during backfill). When ALL tool blocks had
1227
+ // empty stream args, turnSawToolCall is false here. Without this early check,
1228
+ // doneDeferredForArgs would never be set and the stream would hang forever.
1229
+ // This is defense-in-depth: the backfill path also handles this, but setting
1230
+ // doneDeferredForArgs here ensures the assistant message path knows to emit done.
1231
+ if (event?.type === "message_stop" && c.argsPendingBlocks.length > 0) {
1232
+ debug(`processStreamEvent: message_stop deferring done event — ${c.argsPendingBlocks.length} block(s) awaiting args backfill (turnSawToolCall=${c.turnSawToolCall})`);
1233
+ c.doneDeferredForArgs = true;
1234
+ return;
1235
+ }
1236
+
1091
1237
  if (event?.type === "message_stop" && c.turnSawToolCall) {
1092
1238
  // Tool call complete — end this pi stream. The SDK will still yield an
1093
1239
  // assistant message for this turn, but currentPiStream=null causes
@@ -1117,11 +1263,77 @@ function processStreamEvent(
1117
1263
  // the same stream lifecycle as processStreamEvent — including ending the stream on
1118
1264
  // tool_use to prevent deadlock with the MCP handler.
1119
1265
  function processAssistantMessage(message: CbMessage, model: Model<any>, customToolNameToPi: Map<string, string>, c: QueryContext): void {
1120
- if (c.turnSawStreamEvent) return;
1121
1266
  const assistantMsg = (message as any).message;
1122
1267
  if (!assistantMsg?.content) return;
1268
+
1269
+ // --- Args-backfill path (stream already delivered content, but some tool
1270
+ // blocks had empty args). Use the assistant message's complete tool_use
1271
+ // input to backfill the deferred blocks, then emit the pending toolcall_end
1272
+ // events and the deferred done event.
1273
+ if (c.turnSawStreamEvent && c.argsPendingBlocks.length > 0) {
1274
+ debug(`processAssistantMessage: backfill path — ${c.argsPendingBlocks.length} pending block(s), assistant has ${assistantMsg.content.length} blocks`);
1275
+ const toolUseBlocks = assistantMsg.content.filter((b: any) => b.type === "tool_use");
1276
+
1277
+ // Match pending blocks to assistant message tool_use blocks by position.
1278
+ // Stream order and assistant message order should align (both follow the
1279
+ // model's content_block index). If names mismatch we log a warning but still
1280
+ // backfill by position — the stream is the source of truth for ids.
1281
+ const remaining = [...c.argsPendingBlocks];
1282
+ for (const tuBlock of toolUseBlocks) {
1283
+ if (remaining.length === 0) break;
1284
+ const tuName = mapToolName(tuBlock.name, customToolNameToPi);
1285
+ // Find the first pending block whose name matches (or just the first
1286
+ // pending block if names don't match — positional fallback).
1287
+ let matchIdx = remaining.findIndex((p) => p.block.name === tuName);
1288
+ if (matchIdx < 0) {
1289
+ debug(`processAssistantMessage: backfill name mismatch — assistant '${tuName}' vs pending [${remaining.map((p) => p.block.name).join(",")}], using positional fallback`);
1290
+ matchIdx = 0;
1291
+ }
1292
+ const pending = remaining.splice(matchIdx, 1)[0];
1293
+ const dispatchArgs = mapToolArgs(pending.block.name, tuBlock.input);
1294
+ if (isEmptyArgs(dispatchArgs)) {
1295
+ debug(`processAssistantMessage: backfill for ${pending.block.name} [${pending.block.id}] — assistant args also empty, emitting with empty args (pi will validate)`);
1296
+ } else {
1297
+ debug(`processAssistantMessage: backfilled ${pending.block.name} [${pending.block.id}] from assistant message (argsLen=${JSON.stringify(dispatchArgs).length})`);
1298
+ }
1299
+ pending.block.arguments = dispatchArgs;
1300
+ c.turnSawToolCall = true;
1301
+ c.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
1302
+ // Remove from the original argsPendingBlocks array
1303
+ const origIdx = c.argsPendingBlocks.indexOf(pending);
1304
+ if (origIdx >= 0) c.argsPendingBlocks.splice(origIdx, 1);
1305
+ }
1306
+
1307
+ // Any remaining pending blocks (no matching assistant tool_use) get
1308
+ // emitted with their (empty) args so pi can surface the validation error
1309
+ // rather than hanging forever.
1310
+ for (const pending of remaining) {
1311
+ debug(`processAssistantMessage: backfill — no matching assistant tool_use for ${pending.block.name} [${pending.block.id}], emitting with current args`);
1312
+ c.turnSawToolCall = true;
1313
+ c.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
1314
+ const origIdx = c.argsPendingBlocks.indexOf(pending);
1315
+ if (origIdx >= 0) c.argsPendingBlocks.splice(origIdx, 1);
1316
+ }
1317
+
1318
+ // Emit the deferred done event (message_stop skipped it because of pending blocks)
1319
+ if (c.doneDeferredForArgs && c.turnSawToolCall && c.currentPiStream && c.turnOutput) {
1320
+ c.doneDeferredForArgs = false;
1321
+ c.turnOutput.stopReason = "toolUse";
1322
+ const stream = c.currentPiStream;
1323
+ stream.push({ type: "done", reason: "toolUse", message: c.turnOutput });
1324
+ markStreamComplete(stream);
1325
+ stream.end();
1326
+ c.currentPiStream = null;
1327
+ debug(`processAssistantMessage: backfill complete, emitted deferred done event`);
1328
+ }
1329
+ return;
1330
+ }
1331
+
1332
+ // --- Original path: no stream events, this is the primary content path ---
1333
+ if (c.turnSawStreamEvent) return;
1123
1334
  c.turnToolCallIds = [];
1124
- c.nextHandlerIdx = 0;
1335
+ c.nextHandlerIdx = 0;
1336
+ c.matchedToolCallIds = new Set();
1125
1337
  debug(`processAssistantMessage fallback: ${assistantMsg.content.length} blocks, types=${assistantMsg.content.map((b: any) => b.type).join(",")}`);
1126
1338
  for (const block of assistantMsg.content) {
1127
1339
  if (block.type === "text" && block.text) {
@@ -1371,7 +1583,7 @@ function streamCodebuddySdk(model: Model<any>, context: Context, options?: Simpl
1371
1583
  const boundaryOptions = buildProviderBoundaryOptions(providerSettings);
1372
1584
  const appendSystemPrompt = boundaryOptions.appendSystemPrompt;
1373
1585
  const systemPrompt = appendSystemPrompt
1374
- ? buildCodebuddySystemPrompt(context.systemPrompt, { availableToolNames: mcpTools.map((tool) => tool.name) })
1586
+ ? buildCodebuddySystemPrompt(context.systemPrompt, { availableToolNames: mcpTools.map((tool) => tool.name), serialToolCalls: boundaryOptions.serialToolCalls })
1375
1587
  : undefined;
1376
1588
 
1377
1589
  // Provider Path keeps CodeBuddy inside Pi's tool boundary: no built-in SDK
@@ -1407,6 +1619,29 @@ function streamCodebuddySdk(model: Model<any>, context: Context, options?: Simpl
1407
1619
  debugOptions: makeCliDebugOptions("provider"),
1408
1620
  });
1409
1621
 
1622
+ // P4: canUseTool pre-validation — reject empty dispatch args at MCP layer
1623
+ // before they reach pi. This is an early signal to CodeBuddy that args
1624
+ // were dropped (parallel tool_call dispatch bug), giving the model a chance
1625
+ // to regenerate proper args in the same turn rather than failing at pi-side
1626
+ // validation and requiring a full retry.
1627
+ //
1628
+ // Only rejects when the tool has required params AND dispatch args are empty.
1629
+ // Tools without required params (e.g. some MCP tools) pass through normally.
1630
+ const piTools = context.tools ?? [];
1631
+ (queryOptions as SdkQueryOptions).canUseTool = async (toolName: string, input: Record<string, unknown>, opts) => {
1632
+ const piName = mapToolName(toolName, customToolNameToPi);
1633
+ const mappedArgs = mapToolArgs(piName, input);
1634
+ if (isEmptyArgs(mappedArgs) && hasRequiredParams(piName, piTools)) {
1635
+ debug(`canUseTool: rejecting ${toolName}→${piName} — empty args for tool with required params (toolUseId=${opts.toolUseID})`);
1636
+ return {
1637
+ behavior: "deny" as const,
1638
+ message: `Tool "${piName}" requires arguments but received empty input. This can happen with parallel tool calls. Please provide complete arguments for all required fields.`,
1639
+ toolUseID: opts.toolUseID,
1640
+ };
1641
+ }
1642
+ return { behavior: "allow" as const, toolUseID: opts.toolUseID };
1643
+ };
1644
+
1410
1645
  debug("provider: fresh query",
1411
1646
  `model=${cliModel} msgs=${context.messages.length} tools=${mcpTools.length}`,
1412
1647
  `resume=${resumeSessionId?.slice(0, 8) ?? "none"} effort=${effort ?? "default"}`,
@@ -1716,18 +1951,51 @@ let discoverInFlight: Promise<void> | null = null;
1716
1951
 
1717
1952
  async function discoverModels(pi: ExtensionAPI): Promise<void> {
1718
1953
  await withSdkGate(async () => {
1954
+ const codebuddyExecutable = providerSettings.pathToCodebuddyCode;
1955
+ const commonOpts = {
1956
+ maxTurns: 0,
1957
+ permissionMode: "bypassPermissions" as const,
1958
+ tools: [] as string[],
1959
+ cwd: process.cwd(),
1960
+ env: { ...process.env, DISABLE_AUTO_COMPACT: "1" },
1961
+ ...(codebuddyExecutable ? { pathToCodebuddyCode: codebuddyExecutable } : {}),
1962
+ ...makeCliDebugOptions("discover-models"),
1963
+ };
1964
+
1965
+ // Preferred path: Session API getAvailableModelsRaw() returns RawLanguageModel[]
1966
+ // with the real per-model maxInputTokens (context window) and maxOutputTokens,
1967
+ // plus capability flags (supportsImages / supportsReasoning). This eliminates
1968
+ // first-run Window Drift — Pi registers the true context window up front
1969
+ // instead of the 200K conservative default. Runtime calibration remains as a
1970
+ // downward correction layer for entitlement-limited served windows.
1971
+ try {
1972
+ const session = createSdkSession(commonOpts);
1973
+ try {
1974
+ await session.connect();
1975
+ const rawModels = await session.getAvailableModelsRaw();
1976
+ if (rawModels.length) {
1977
+ MODELS = applyModelCalibrations(rawModelsFromSdkRaw(rawModels));
1978
+ registerCurrentProvider(pi);
1979
+ modelsDiscovered = true;
1980
+ debug(`discoverModels: registered ${MODELS.length} models via getAvailableModelsRaw()`);
1981
+ return;
1982
+ }
1983
+ debug("discoverModels: getAvailableModelsRaw() returned empty, falling back to supportedModels()");
1984
+ } finally {
1985
+ session.close();
1986
+ }
1987
+ } catch (err) {
1988
+ debug("discoverModels: getAvailableModelsRaw() failed, falling back to supportedModels()", err);
1989
+ }
1990
+
1991
+ // Fallback path: one-shot query().supportedModels() returns the simplified
1992
+ // ModelInfo (value/displayName/description) without token limits, so we must
1993
+ // use the conservativeContextWindow() heuristic. Used when the CLI doesn't
1994
+ // support the get_available_models control request or the Session API errors.
1719
1995
  try {
1720
- const codebuddyExecutable = providerSettings.pathToCodebuddyCode;
1721
1996
  const q = query({
1722
1997
  prompt: " ",
1723
- options: {
1724
- maxTurns: 0,
1725
- permissionMode: "bypassPermissions",
1726
- tools: [],
1727
- env: { ...process.env, DISABLE_AUTO_COMPACT: "1" },
1728
- ...(codebuddyExecutable ? { pathToCodebuddyCode: codebuddyExecutable } : {}),
1729
- ...makeCliDebugOptions("supported-models"),
1730
- },
1998
+ options: commonOpts,
1731
1999
  });
1732
2000
  const supported = await q.supportedModels();
1733
2001
  await q.return().catch(() => {});
@@ -1735,9 +2003,9 @@ async function discoverModels(pi: ExtensionAPI): Promise<void> {
1735
2003
  MODELS = applyModelCalibrations(rawModelsFromSdk(supported as any));
1736
2004
  registerCurrentProvider(pi);
1737
2005
  modelsDiscovered = true;
1738
- debug(`discoverModels: registered ${MODELS.length} models from SDK`);
2006
+ debug(`discoverModels: registered ${MODELS.length} models from supportedModels()`);
1739
2007
  } catch (err) {
1740
- debug("discoverModels: failed, using fallback models", err);
2008
+ debug("discoverModels: supportedModels() failed, using fallback models", err);
1741
2009
  }
1742
2010
  });
1743
2011
  }
package/src/models.ts CHANGED
@@ -1,5 +1,7 @@
1
- // Dynamic model list from CodeBuddy SDK supportedModels().
2
- import type { ModelInfo } from "@tencent-ai/agent-sdk";
1
+ // Dynamic model list from CodeBuddy SDK.
2
+ // - supportedModels(): simplified ModelInfo (no token limits) — legacy fallback.
3
+ // - getAvailableModelsRaw(): RawLanguageModel with real maxInputTokens / maxOutputTokens.
4
+ import type { ModelInfo, RawLanguageModel } from "@tencent-ai/agent-sdk";
3
5
 
4
6
  export type PiModel = {
5
7
  id: string;
@@ -50,6 +52,29 @@ export function rawModelsFromSdk(supported: Array<ModelInfo & { id?: string; nam
50
52
  }));
51
53
  }
52
54
 
55
+ // Build PiModel[] from the SDK Session API's getAvailableModelsRaw() result.
56
+ // Unlike rawModelsFromSdk() (which only gets the simplified ModelInfo and must
57
+ // fall back to conservativeContextWindow), this reads the real per-model token
58
+ // limits (maxInputTokens = context window, maxOutputTokens = max output) plus
59
+ // capability flags (supportsImages / supportsReasoning) directly from the CLI.
60
+ // Models marked disabled are filtered out so Pi never offers an unusable model.
61
+ export function rawModelsFromSdkRaw(raw: RawLanguageModel[]): PiModel[] {
62
+ return raw
63
+ .filter((m) => m.id && !m.disabled)
64
+ .map((m) => {
65
+ const id = m.id!;
66
+ return {
67
+ id,
68
+ name: m.name || id,
69
+ reasoning: m.supportsReasoning ?? detectThinking(id),
70
+ input: (m.supportsImages ?? detectImages(id)) ? ["text", "image"] as const : ["text"] as const,
71
+ contextWindow: m.maxInputTokens ?? conservativeContextWindow(id),
72
+ maxTokens: m.maxOutputTokens ?? conservativeMaxTokens(id),
73
+ cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
74
+ };
75
+ });
76
+ }
77
+
53
78
  export const FALLBACK_MODELS: PiModel[] = [
54
79
  {
55
80
  id: "hy3-preview-agent-ioa",
@@ -23,8 +23,25 @@ export class QueryContext {
23
23
  pendingResults = new Map<string, McpResult>();
24
24
  turnToolCallIds: string[] = [];
25
25
  nextHandlerIdx = 0;
26
+ // toolCallIds already assigned to an MCP handler. Used by name-based
27
+ // matching to avoid the index-order race when CodeBuddy dispatches
28
+ // parallel tool handlers in a different order than the stream's
29
+ // content_block_start events. Without this, bash's handler could get
30
+ // read's toolCallId, causing result misassignment.
31
+ matchedToolCallIds = new Set<string>();
26
32
  deferredUserMessages: string[] = [];
27
33
 
34
+ // Tool-call blocks whose stream args came through empty (parallel tool_call
35
+ // dispatch dropping args to {}). Their toolcall_end is deferred until either
36
+ // the assistant message arrives with complete args, or the MCP handler
37
+ // receives non-empty dispatch args. The done event is also deferred while
38
+ // this list is non-empty, so pi does not execute tools with empty args.
39
+ argsPendingBlocks: Array<{ block: any; contentIndex: number }> = [];
40
+ // Set when message_stop arrived but done was deferred due to pending blocks.
41
+ // Prevents double-deferral and lets the assistant-message path know it
42
+ // must emit the done event after backfilling.
43
+ doneDeferredForArgs = false;
44
+
28
45
  // Per-turn (reset together)
29
46
  turnOutput: AssistantMessage | null = null;
30
47
  turnStarted = false;
@@ -47,6 +64,9 @@ export class QueryContext {
47
64
  this.turnStarted = false;
48
65
  this.turnSawStreamEvent = false;
49
66
  this.turnSawToolCall = false;
67
+ this.argsPendingBlocks = [];
68
+ this.doneDeferredForArgs = false;
69
+ this.matchedToolCallIds = new Set();
50
70
  // turnToolCallIds and nextHandlerIdx are NOT reset — they persist across
51
71
  // tool-result delivery callbacks within the same assistant message.
52
72
  }
package/src/skills.ts CHANGED
@@ -20,6 +20,7 @@ const BUILTIN_TOOL_GUIDANCE: Record<string, string> = {
20
20
 
21
21
  export interface ToolBridgeInstructionOptions {
22
22
  availableToolNames?: Iterable<string>;
23
+ serialToolCalls?: boolean;
23
24
  }
24
25
 
25
26
  function normalizeToolNames(names: Iterable<string> | undefined): Set<string> {
@@ -71,6 +72,18 @@ export function buildPiToolBridgeInstruction(options: ToolBridgeInstructionOptio
71
72
  if (has("bash")) {
72
73
  lines.push(`- Use \`${mcpName("bash")}\` only when file tools are insufficient, for search/test/build/git information, or when command execution is requested.`);
73
74
  }
75
+ if (options.serialToolCalls !== false) {
76
+ // Force serial tool calls. CodeBuddy's MCP client drops arguments for
77
+ // 2nd+ parallel tool_call blocks in a single response: the
78
+ // content_block_stop for those blocks never arrives, so pi finalizes a
79
+ // dangling toolcall_start with empty {} args and the call fails
80
+ // validation. Instructing one-tool-per-turn avoids the failure at the
81
+ // source. The stream-side backfill defense cannot catch this because it
82
+ // keys off content_block_stop, which never fires for the dropped call.
83
+ lines.push("Call AT MOST ONE tool per turn. Never emit multiple tool_call blocks in a single response. Issue one tool call, wait for its result, then decide the next call in the following turn. Parallel tool calls are unsupported and will fail.");
84
+ } else if (toolNames.size > 1) {
85
+ 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.");
86
+ }
74
87
  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.");
75
88
 
76
89
  return lines.join("\n");
@@ -127,13 +140,13 @@ function applySkillsRewrite(systemPrompt: string): string {
127
140
  /** Pi system prompt as CodeBuddy override (replaces default "CodeBuddy Code" identity). */
128
141
  export function buildCodebuddySystemPrompt(
129
142
  piSystemPrompt: string | undefined,
130
- options?: { includeAgents?: boolean; includeSkills?: boolean; includeToolBridge?: boolean; availableToolNames?: Iterable<string> },
143
+ options?: { includeAgents?: boolean; includeSkills?: boolean; includeToolBridge?: boolean; availableToolNames?: Iterable<string>; serialToolCalls?: boolean },
131
144
  ): string | undefined {
132
145
  const parts: string[] = [];
133
146
  const includeToolBridge = options?.includeToolBridge !== false;
134
147
 
135
148
  if (includeToolBridge) {
136
- parts.push(buildPiToolBridgeInstruction({ availableToolNames: options?.availableToolNames }));
149
+ parts.push(buildPiToolBridgeInstruction({ availableToolNames: options?.availableToolNames, serialToolCalls: options?.serialToolCalls }));
137
150
  }
138
151
 
139
152
  if (piSystemPrompt) {
@@ -44,24 +44,26 @@ function intersectionSchema(schemas: z.ZodTypeAny[]): z.ZodTypeAny {
44
44
  return schemas.slice(1).reduce((left, right) => z.intersection(left, right), schemas[0]);
45
45
  }
46
46
 
47
- function objectSchema(prop: JsonSchema): z.ZodTypeAny {
48
- const shape = jsonSchemaToZodShape(prop);
47
+ function objectSchema(prop: JsonSchema, relax = false): z.ZodTypeAny {
48
+ const shape = jsonSchemaToZodShape(prop, relax);
49
49
  const additionalPropertiesSchema =
50
50
  prop.additionalProperties && typeof prop.additionalProperties === "object"
51
- ? jsonSchemaPropertyToZod(prop.additionalProperties as JsonSchema)
51
+ ? jsonSchemaPropertyToZod(prop.additionalProperties as JsonSchema, relax)
52
52
  : undefined;
53
53
  if (Object.keys(shape).length > 0) {
54
54
  const base = z.object(shape);
55
- if (prop.additionalProperties === false) return base.strict();
55
+ // relax: never strict accept {} and extra keys so a CodeBuddy MCP client
56
+ // that drops parallel tool_call arguments to {} still passes validation.
57
+ if (!relax && prop.additionalProperties === false) return base.strict();
56
58
  if (additionalPropertiesSchema) return base.catchall(additionalPropertiesSchema);
57
59
  return base.passthrough();
58
60
  }
59
61
  if (additionalPropertiesSchema) return z.record(z.string(), additionalPropertiesSchema);
60
- if (prop.additionalProperties === false) return z.object({}).strict();
62
+ if (!relax && prop.additionalProperties === false) return z.object({}).strict();
61
63
  return z.record(z.string(), z.unknown());
62
64
  }
63
65
 
64
- export function jsonSchemaPropertyToZod(prop: JsonSchema): z.ZodTypeAny {
66
+ export function jsonSchemaPropertyToZod(prop: JsonSchema, relax = false): z.ZodTypeAny {
65
67
  let base: z.ZodTypeAny;
66
68
 
67
69
  if (prop.const !== undefined) {
@@ -122,7 +124,7 @@ export function jsonSchemaPropertyToZod(prop: JsonSchema): z.ZodTypeAny {
122
124
  : z.array(z.unknown());
123
125
  break;
124
126
  case "object":
125
- base = objectSchema(prop);
127
+ base = objectSchema(prop, relax);
126
128
  break;
127
129
  case "null":
128
130
  base = z.null();
@@ -136,7 +138,7 @@ export function jsonSchemaPropertyToZod(prop: JsonSchema): z.ZodTypeAny {
136
138
  return withDescription(base, prop);
137
139
  }
138
140
 
139
- export function jsonSchemaToZodShape(schema: unknown): Record<string, z.ZodTypeAny> {
141
+ export function jsonSchemaToZodShape(schema: unknown, relax = false): Record<string, z.ZodTypeAny> {
140
142
  const s = schema as JsonSchema;
141
143
  const objectSchemas = [s, ...(Array.isArray(s?.allOf) ? s.allOf as JsonSchema[] : [])]
142
144
  .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
148
150
  if (!props) continue;
149
151
  const required = new Set(Array.isArray(objectSchemaPart.required) ? objectSchemaPart.required as string[] : []);
150
152
  for (const [key, prop] of Object.entries(props)) {
151
- const zodProp = jsonSchemaPropertyToZod(prop);
152
- shape[key] = required.has(key) ? zodProp : zodProp.optional();
153
+ const zodProp = jsonSchemaPropertyToZod(prop, relax);
154
+ // relax: make every property optional so an empty {} (args dropped by a
155
+ // buggy parallel tool_call dispatch) still validates. Descriptions are
156
+ // preserved so the model keeps parameter guidance.
157
+ shape[key] = relax || !required.has(key) ? zodProp.optional() : zodProp;
153
158
  }
154
159
  }
155
160
  return shape;
156
161
  }
157
162
 
158
- export function jsonSchemaToZodObject(schema: unknown): z.ZodTypeAny {
159
- return jsonSchemaPropertyToZod(schema as JsonSchema);
163
+ export function jsonSchemaToZodObject(schema: unknown, relax = false): z.ZodTypeAny {
164
+ return jsonSchemaPropertyToZod(schema as JsonSchema, relax);
165
+ }
166
+
167
+ /**
168
+ * MCP-specific schema builder: preserves required field constraints so that
169
+ * an empty {} (from a buggy parallel tool_call dispatch) is rejected by MCP
170
+ * validation, while still allowing extra keys (passthrough) for forward-compat.
171
+ *
172
+ * Unlike jsonSchemaToZodObject(relax=true) which makes ALL params optional,
173
+ * this keeps required params required. This forces CodeBuddy to regenerate
174
+ * proper args instead of silently passing {} through to the handler.
175
+ *
176
+ * The deferred-backfill logic in processStreamEvent/processAssistantMessage
177
+ * handles the case where stream args arrive empty — it waits for the
178
+ * assistant message or MCP dispatch to provide real args. So MCP validation
179
+ * rejecting {} is safe and desirable: it's an early signal that args were
180
+ * dropped, rather than letting an empty-args call reach pi.
181
+ */
182
+ export function jsonSchemaToZodObjectForMcp(schema: unknown): z.ZodTypeAny {
183
+ const s = schema as JsonSchema;
184
+ // Build shape with relax=false to preserve required constraints
185
+ const shape = jsonSchemaToZodShape(s, false);
186
+ const additionalPropertiesSchema =
187
+ s?.additionalProperties && typeof s.additionalProperties === "object"
188
+ ? jsonSchemaPropertyToZod(s.additionalProperties as JsonSchema, false)
189
+ : undefined;
190
+
191
+ if (Object.keys(shape).length > 0) {
192
+ const base = z.object(shape);
193
+ // Always passthrough: accept extra keys (forward-compat with new params)
194
+ // but enforce required fields.
195
+ if (additionalPropertiesSchema) return base.catchall(additionalPropertiesSchema);
196
+ return base.passthrough();
197
+ }
198
+ if (s?.additionalProperties === false) return z.object({}).strict();
199
+ return z.record(z.string(), z.unknown());
160
200
  }