@raoxxxwq/pi-codebuddy-sdk 0.3.1 → 0.4.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.4.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/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(
@@ -755,6 +757,30 @@ function mapToolName(name: string, customToolNameToPi?: Map<string, string>): st
755
757
  return name;
756
758
  }
757
759
 
760
+ // True when a tool-call argument object is empty or absent. Used to detect
761
+ // the parallel-tool-call arg-dropping failure: CodeBuddy's MCP client may
762
+ // dispatch tool_call arguments as {} (especially in parallel batches), and
763
+ // the stream's input_json_delta may also arrive empty. We defer toolcall_end
764
+ // for such blocks until the assistant message (or MCP dispatch) provides
765
+ // real arguments, preventing pi from executing tools with empty args.
766
+ function isEmptyArgs(args: Record<string, unknown> | undefined | null): boolean {
767
+ if (!args) return true;
768
+ if (Object.keys(args).length === 0) return true;
769
+ // Treat an object with only undefined/null values as empty
770
+ const hasValue = Object.values(args).some((v) => v !== undefined && v !== null);
771
+ return !hasValue;
772
+ }
773
+
774
+ // Checks whether a Pi tool has any required parameters. Used by the canUseTool
775
+ // pre-validation to decide whether empty dispatch args should be rejected.
776
+ // If a tool has no required params, empty args are legitimate and should pass.
777
+ function hasRequiredParams(toolName: string, tools: Tool[]): boolean {
778
+ const tool = tools.find((t) => t.name === toolName);
779
+ if (!tool?.parameters) return false;
780
+ const required = (tool.parameters as Record<string, unknown>).required;
781
+ return Array.isArray(required) && required.length > 0;
782
+ }
783
+
758
784
  // Renames for CodeBuddy SDK param names that differ from pi's native names.
759
785
  // Keys not listed here pass through unchanged, so new pi params work automatically.
760
786
  const SDK_KEY_RENAMES: Record<string, Record<string, string>> = {
@@ -842,10 +868,79 @@ function buildMcpServers(tools: Tool[], queryCtx: QueryContext): Record<string,
842
868
  const mcpTools = tools.map((tool) => ({
843
869
  name: tool.name,
844
870
  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})`);
871
+ // MCP schema preserves required constraints so empty {} (from parallel
872
+ // tool_call arg-dropping) is rejected at MCP validation time — an early
873
+ // signal that args were lost. The deferred-backfill logic handles the
874
+ // stream side; this handles the dispatch side. Passthrough allows extra
875
+ // keys for forward-compat.
876
+ inputSchema: jsonSchemaToZodObjectForMcp(tool.parameters),
877
+ handler: async (dispatchArgs?: Record<string, unknown>) => {
878
+ // Name-based toolCallId matching: find the first toolCallId whose
879
+ // corresponding turnBlock has a name matching this tool, and hasn't
880
+ // been claimed by a previous handler invocation. This fixes the
881
+ // parallel-dispatch race where CodeBuddy calls MCP handlers in a
882
+ // different order than the stream's content_block_start events —
883
+ // without this, bash's handler could get read's toolCallId, causing
884
+ // result misassignment (bash result → read handler, vice versa).
885
+ //
886
+ // Fallback to index-based for robustness (e.g. if turnBlocks doesn't
887
+ // have the block yet, or same-name dedup edge cases).
888
+ let toolCallId: string | undefined;
889
+ for (const id of queryCtx.turnToolCallIds) {
890
+ if (queryCtx.matchedToolCallIds.has(id)) continue;
891
+ const block = queryCtx.turnBlocks.find((b: any) => b.id === id && b.type === "toolCall");
892
+ if (block && block.name === tool.name) {
893
+ toolCallId = id;
894
+ queryCtx.matchedToolCallIds.add(id);
895
+ break;
896
+ }
897
+ }
898
+ // Fallback: if no name match, use index (handles edge cases where
899
+ // block.name hasn't been mapped yet or same-name calls overflow)
900
+ if (!toolCallId) {
901
+ toolCallId = queryCtx.turnToolCallIds[queryCtx.nextHandlerIdx];
902
+ if (toolCallId) {
903
+ queryCtx.matchedToolCallIds.add(toolCallId);
904
+ queryCtx.nextHandlerIdx++;
905
+ }
906
+ }
907
+ if (!toolCallId) debug(`WARNING: mcp handler ${tool.name} has no toolCallId (matched=${queryCtx.matchedToolCallIds.size}, available=${queryCtx.turnToolCallIds.length})`);
908
+
909
+ debug(`mcp dispatch: ${tool.name} dispatchArgsLen=${JSON.stringify(dispatchArgs ?? {}).length} dispatchArgsEmpty=${isEmptyArgs(dispatchArgs)} matched=${queryCtx.matchedToolCallIds.size}/${queryCtx.turnToolCallIds.length} pendingBlocks=${queryCtx.argsPendingBlocks.length}`);
910
+
911
+ // Backfill path: if there are args-pending blocks whose toolcall_end
912
+ // was deferred (stream args were empty), try to backfill using the
913
+ // MCP dispatch args. This handles the case where the assistant message
914
+ // also had empty args but the MCP dispatch carries the real args.
915
+ if (queryCtx.argsPendingBlocks.length > 0 && toolCallId) {
916
+ const pendingIdx = queryCtx.argsPendingBlocks.findIndex((p) => p.block.id === toolCallId);
917
+ if (pendingIdx >= 0) {
918
+ const pending = queryCtx.argsPendingBlocks[pendingIdx];
919
+ const backfillArgs = mapToolArgs(tool.name, dispatchArgs);
920
+ if (!isEmptyArgs(backfillArgs)) {
921
+ pending.block.arguments = backfillArgs;
922
+ debug(`mcp handler: backfilled ${tool.name} [${toolCallId}] from MCP dispatch args (argsLen=${JSON.stringify(backfillArgs).length})`);
923
+ } else {
924
+ debug(`mcp handler: dispatch args also empty for ${tool.name} [${toolCallId}], emitting with current args`);
925
+ }
926
+ queryCtx.turnSawToolCall = true;
927
+ queryCtx.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: queryCtx.turnOutput });
928
+ queryCtx.argsPendingBlocks.splice(pendingIdx, 1);
929
+
930
+ // If all pending blocks are now resolved and done was deferred, emit it
931
+ if (queryCtx.argsPendingBlocks.length === 0 && queryCtx.doneDeferredForArgs && queryCtx.currentPiStream && queryCtx.turnOutput) {
932
+ queryCtx.doneDeferredForArgs = false;
933
+ queryCtx.turnOutput.stopReason = "toolUse";
934
+ const stream = queryCtx.currentPiStream;
935
+ stream.push({ type: "done", reason: "toolUse", message: queryCtx.turnOutput });
936
+ markStreamComplete(stream);
937
+ stream.end();
938
+ queryCtx.currentPiStream = null;
939
+ debug(`mcp handler: all pending blocks resolved, emitted deferred done event`);
940
+ }
941
+ }
942
+ }
943
+
849
944
  if (toolCallId && queryCtx.pendingResults.has(toolCallId)) {
850
945
  const result = queryCtx.pendingResults.get(toolCallId)!;
851
946
  queryCtx.pendingResults.delete(toolCallId);
@@ -991,6 +1086,20 @@ function finalizeCurrentStream(c: QueryContext, stopReason?: string): void {
991
1086
  if (!c.turnStarted) ensureTurnStarted(c);
992
1087
  const reason = stopReason === "length" ? "length" : "stop";
993
1088
  const stream = c.currentPiStream;
1089
+
1090
+ // Drain any args-pending blocks: if the SDK ended abnormally (crash, network
1091
+ // error) without yielding an assistant message, these blocks were never
1092
+ // backfilled. Emit them with their current (possibly empty) args so pi can
1093
+ // surface validation errors rather than silently dropping the tool calls.
1094
+ if (c.argsPendingBlocks.length > 0) {
1095
+ debug(`finalizeCurrentStream: draining ${c.argsPendingBlocks.length} pending block(s) with current args`);
1096
+ for (const pending of c.argsPendingBlocks) {
1097
+ c.turnSawToolCall = true;
1098
+ stream!.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
1099
+ }
1100
+ c.argsPendingBlocks = [];
1101
+ c.doneDeferredForArgs = false;
1102
+ }
994
1103
  stream!.push({ type: "done", reason, message: c.turnOutput });
995
1104
  markStreamComplete(stream);
996
1105
  stream!.end();
@@ -1011,7 +1120,8 @@ function processStreamEvent(
1011
1120
 
1012
1121
  if (event?.type === "message_start") {
1013
1122
  c.turnToolCallIds = [];
1014
- c.nextHandlerIdx = 0;
1123
+ c.nextHandlerIdx = 0;
1124
+ c.matchedToolCallIds = new Set();
1015
1125
  if (event.message?.usage) updateUsage(c.turnOutput, event.message.usage, model);
1016
1126
  return;
1017
1127
  }
@@ -1025,7 +1135,7 @@ function processStreamEvent(
1025
1135
  c.turnBlocks.push({ type: "thinking", thinking: "", thinkingSignature: "", index: event.index });
1026
1136
  c.currentPiStream!.push({ type: "thinking_start", contentIndex: c.turnBlocks.length - 1, partial: c.turnOutput });
1027
1137
  } else if (event.content_block?.type === "tool_use") {
1028
- c.turnSawToolCall = true;
1138
+ c.turnToolCallIds.push(event.content_block.id);
1029
1139
  c.turnToolCallIds.push(event.content_block.id);
1030
1140
  c.turnBlocks.push({
1031
1141
  type: "toolCall", id: event.content_block.id,
@@ -1072,12 +1182,33 @@ function processStreamEvent(
1072
1182
  } else if (block.type === "thinking") {
1073
1183
  c.currentPiStream!.push({ type: "thinking_end", contentIndex: index, content: block.thinking, partial: c.turnOutput });
1074
1184
  } else if (block.type === "toolCall") {
1075
- c.turnSawToolCall = true;
1185
+ const partialJsonLen = block.partialJson?.length ?? 0;
1076
1186
  block.arguments = mapToolArgs(
1077
1187
  block.name, parsePartialJson(block.partialJson, block.arguments),
1078
1188
  );
1079
1189
  delete block.partialJson;
1080
- c.currentPiStream!.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: c.turnOutput });
1190
+
1191
+ debug(`processStreamEvent: content_block_stop ${block.name} [${block.id}] argsSource=${
1192
+ isEmptyArgs(block.arguments) ? "EMPTY" : "stream"
1193
+ } argsLen=${JSON.stringify(block.arguments).length} partialJsonLen=${partialJsonLen}`);
1194
+
1195
+ // Parallel tool-call arg-dropping defense: when the stream delivers
1196
+ // empty args (CodeBuddy's MCP client may dispatch parallel tool_call
1197
+ // arguments as {}), defer toolcall_end until the assistant message
1198
+ // (or MCP dispatch) provides real args. This prevents pi from
1199
+ // executing tools with empty args (e.g. "bash" with no command).
1200
+ //
1201
+ // turnSawToolCall is set only when we actually emit a toolcall_end
1202
+ // (here for non-empty args, or in the backfill path for deferred
1203
+ // blocks). This ensures message_stop's done-deferral check fires
1204
+ // correctly even when ALL tool blocks had empty args.
1205
+ if (isEmptyArgs(block.arguments)) {
1206
+ debug(`processStreamEvent: deferring toolcall_end for ${block.name} [${block.id}] — stream args empty, will backfill from assistant message or MCP dispatch`);
1207
+ c.argsPendingBlocks.push({ block, contentIndex: index });
1208
+ } else {
1209
+ c.turnSawToolCall = true;
1210
+ c.currentPiStream!.push({ type: "toolcall_end", contentIndex: index, toolCall: block, partial: c.turnOutput });
1211
+ }
1081
1212
  }
1082
1213
  return;
1083
1214
  }
@@ -1088,6 +1219,19 @@ function processStreamEvent(
1088
1219
  return;
1089
1220
  }
1090
1221
 
1222
+ // Check args-pending blocks FIRST, before the turnSawToolCall gate.
1223
+ // turnSawToolCall is only set when a toolcall_end is actually emitted (non-empty
1224
+ // args in content_block_stop, or during backfill). When ALL tool blocks had
1225
+ // empty stream args, turnSawToolCall is false here. Without this early check,
1226
+ // doneDeferredForArgs would never be set and the stream would hang forever.
1227
+ // This is defense-in-depth: the backfill path also handles this, but setting
1228
+ // doneDeferredForArgs here ensures the assistant message path knows to emit done.
1229
+ if (event?.type === "message_stop" && c.argsPendingBlocks.length > 0) {
1230
+ debug(`processStreamEvent: message_stop deferring done event — ${c.argsPendingBlocks.length} block(s) awaiting args backfill (turnSawToolCall=${c.turnSawToolCall})`);
1231
+ c.doneDeferredForArgs = true;
1232
+ return;
1233
+ }
1234
+
1091
1235
  if (event?.type === "message_stop" && c.turnSawToolCall) {
1092
1236
  // Tool call complete — end this pi stream. The SDK will still yield an
1093
1237
  // assistant message for this turn, but currentPiStream=null causes
@@ -1117,11 +1261,77 @@ function processStreamEvent(
1117
1261
  // the same stream lifecycle as processStreamEvent — including ending the stream on
1118
1262
  // tool_use to prevent deadlock with the MCP handler.
1119
1263
  function processAssistantMessage(message: CbMessage, model: Model<any>, customToolNameToPi: Map<string, string>, c: QueryContext): void {
1120
- if (c.turnSawStreamEvent) return;
1121
1264
  const assistantMsg = (message as any).message;
1122
1265
  if (!assistantMsg?.content) return;
1266
+
1267
+ // --- Args-backfill path (stream already delivered content, but some tool
1268
+ // blocks had empty args). Use the assistant message's complete tool_use
1269
+ // input to backfill the deferred blocks, then emit the pending toolcall_end
1270
+ // events and the deferred done event.
1271
+ if (c.turnSawStreamEvent && c.argsPendingBlocks.length > 0) {
1272
+ debug(`processAssistantMessage: backfill path — ${c.argsPendingBlocks.length} pending block(s), assistant has ${assistantMsg.content.length} blocks`);
1273
+ const toolUseBlocks = assistantMsg.content.filter((b: any) => b.type === "tool_use");
1274
+
1275
+ // Match pending blocks to assistant message tool_use blocks by position.
1276
+ // Stream order and assistant message order should align (both follow the
1277
+ // model's content_block index). If names mismatch we log a warning but still
1278
+ // backfill by position — the stream is the source of truth for ids.
1279
+ const remaining = [...c.argsPendingBlocks];
1280
+ for (const tuBlock of toolUseBlocks) {
1281
+ if (remaining.length === 0) break;
1282
+ const tuName = mapToolName(tuBlock.name, customToolNameToPi);
1283
+ // Find the first pending block whose name matches (or just the first
1284
+ // pending block if names don't match — positional fallback).
1285
+ let matchIdx = remaining.findIndex((p) => p.block.name === tuName);
1286
+ if (matchIdx < 0) {
1287
+ debug(`processAssistantMessage: backfill name mismatch — assistant '${tuName}' vs pending [${remaining.map((p) => p.block.name).join(",")}], using positional fallback`);
1288
+ matchIdx = 0;
1289
+ }
1290
+ const pending = remaining.splice(matchIdx, 1)[0];
1291
+ const dispatchArgs = mapToolArgs(pending.block.name, tuBlock.input);
1292
+ if (isEmptyArgs(dispatchArgs)) {
1293
+ debug(`processAssistantMessage: backfill for ${pending.block.name} [${pending.block.id}] — assistant args also empty, emitting with empty args (pi will validate)`);
1294
+ } else {
1295
+ debug(`processAssistantMessage: backfilled ${pending.block.name} [${pending.block.id}] from assistant message (argsLen=${JSON.stringify(dispatchArgs).length})`);
1296
+ }
1297
+ pending.block.arguments = dispatchArgs;
1298
+ c.turnSawToolCall = true;
1299
+ c.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
1300
+ // Remove from the original argsPendingBlocks array
1301
+ const origIdx = c.argsPendingBlocks.indexOf(pending);
1302
+ if (origIdx >= 0) c.argsPendingBlocks.splice(origIdx, 1);
1303
+ }
1304
+
1305
+ // Any remaining pending blocks (no matching assistant tool_use) get
1306
+ // emitted with their (empty) args so pi can surface the validation error
1307
+ // rather than hanging forever.
1308
+ for (const pending of remaining) {
1309
+ debug(`processAssistantMessage: backfill — no matching assistant tool_use for ${pending.block.name} [${pending.block.id}], emitting with current args`);
1310
+ c.turnSawToolCall = true;
1311
+ c.currentPiStream?.push({ type: "toolcall_end", contentIndex: pending.contentIndex, toolCall: pending.block, partial: c.turnOutput });
1312
+ const origIdx = c.argsPendingBlocks.indexOf(pending);
1313
+ if (origIdx >= 0) c.argsPendingBlocks.splice(origIdx, 1);
1314
+ }
1315
+
1316
+ // Emit the deferred done event (message_stop skipped it because of pending blocks)
1317
+ if (c.doneDeferredForArgs && c.turnSawToolCall && c.currentPiStream && c.turnOutput) {
1318
+ c.doneDeferredForArgs = false;
1319
+ c.turnOutput.stopReason = "toolUse";
1320
+ const stream = c.currentPiStream;
1321
+ stream.push({ type: "done", reason: "toolUse", message: c.turnOutput });
1322
+ markStreamComplete(stream);
1323
+ stream.end();
1324
+ c.currentPiStream = null;
1325
+ debug(`processAssistantMessage: backfill complete, emitted deferred done event`);
1326
+ }
1327
+ return;
1328
+ }
1329
+
1330
+ // --- Original path: no stream events, this is the primary content path ---
1331
+ if (c.turnSawStreamEvent) return;
1123
1332
  c.turnToolCallIds = [];
1124
- c.nextHandlerIdx = 0;
1333
+ c.nextHandlerIdx = 0;
1334
+ c.matchedToolCallIds = new Set();
1125
1335
  debug(`processAssistantMessage fallback: ${assistantMsg.content.length} blocks, types=${assistantMsg.content.map((b: any) => b.type).join(",")}`);
1126
1336
  for (const block of assistantMsg.content) {
1127
1337
  if (block.type === "text" && block.text) {
@@ -1407,6 +1617,29 @@ function streamCodebuddySdk(model: Model<any>, context: Context, options?: Simpl
1407
1617
  debugOptions: makeCliDebugOptions("provider"),
1408
1618
  });
1409
1619
 
1620
+ // P4: canUseTool pre-validation — reject empty dispatch args at MCP layer
1621
+ // before they reach pi. This is an early signal to CodeBuddy that args
1622
+ // were dropped (parallel tool_call dispatch bug), giving the model a chance
1623
+ // to regenerate proper args in the same turn rather than failing at pi-side
1624
+ // validation and requiring a full retry.
1625
+ //
1626
+ // Only rejects when the tool has required params AND dispatch args are empty.
1627
+ // Tools without required params (e.g. some MCP tools) pass through normally.
1628
+ const piTools = context.tools ?? [];
1629
+ (queryOptions as SdkQueryOptions).canUseTool = async (toolName: string, input: Record<string, unknown>, opts) => {
1630
+ const piName = mapToolName(toolName, customToolNameToPi);
1631
+ const mappedArgs = mapToolArgs(piName, input);
1632
+ if (isEmptyArgs(mappedArgs) && hasRequiredParams(piName, piTools)) {
1633
+ debug(`canUseTool: rejecting ${toolName}→${piName} — empty args for tool with required params (toolUseId=${opts.toolUseID})`);
1634
+ return {
1635
+ behavior: "deny" as const,
1636
+ message: `Tool "${piName}" requires arguments but received empty input. This can happen with parallel tool calls. Please provide complete arguments for all required fields.`,
1637
+ toolUseID: opts.toolUseID,
1638
+ };
1639
+ }
1640
+ return { behavior: "allow" as const, toolUseID: opts.toolUseID };
1641
+ };
1642
+
1410
1643
  debug("provider: fresh query",
1411
1644
  `model=${cliModel} msgs=${context.messages.length} tools=${mcpTools.length}`,
1412
1645
  `resume=${resumeSessionId?.slice(0, 8) ?? "none"} effort=${effort ?? "default"}`,
@@ -1716,18 +1949,51 @@ let discoverInFlight: Promise<void> | null = null;
1716
1949
 
1717
1950
  async function discoverModels(pi: ExtensionAPI): Promise<void> {
1718
1951
  await withSdkGate(async () => {
1952
+ const codebuddyExecutable = providerSettings.pathToCodebuddyCode;
1953
+ const commonOpts = {
1954
+ maxTurns: 0,
1955
+ permissionMode: "bypassPermissions" as const,
1956
+ tools: [] as string[],
1957
+ cwd: process.cwd(),
1958
+ env: { ...process.env, DISABLE_AUTO_COMPACT: "1" },
1959
+ ...(codebuddyExecutable ? { pathToCodebuddyCode: codebuddyExecutable } : {}),
1960
+ ...makeCliDebugOptions("discover-models"),
1961
+ };
1962
+
1963
+ // Preferred path: Session API getAvailableModelsRaw() returns RawLanguageModel[]
1964
+ // with the real per-model maxInputTokens (context window) and maxOutputTokens,
1965
+ // plus capability flags (supportsImages / supportsReasoning). This eliminates
1966
+ // first-run Window Drift — Pi registers the true context window up front
1967
+ // instead of the 200K conservative default. Runtime calibration remains as a
1968
+ // downward correction layer for entitlement-limited served windows.
1969
+ try {
1970
+ const session = createSdkSession(commonOpts);
1971
+ try {
1972
+ await session.connect();
1973
+ const rawModels = await session.getAvailableModelsRaw();
1974
+ if (rawModels.length) {
1975
+ MODELS = applyModelCalibrations(rawModelsFromSdkRaw(rawModels));
1976
+ registerCurrentProvider(pi);
1977
+ modelsDiscovered = true;
1978
+ debug(`discoverModels: registered ${MODELS.length} models via getAvailableModelsRaw()`);
1979
+ return;
1980
+ }
1981
+ debug("discoverModels: getAvailableModelsRaw() returned empty, falling back to supportedModels()");
1982
+ } finally {
1983
+ session.close();
1984
+ }
1985
+ } catch (err) {
1986
+ debug("discoverModels: getAvailableModelsRaw() failed, falling back to supportedModels()", err);
1987
+ }
1988
+
1989
+ // Fallback path: one-shot query().supportedModels() returns the simplified
1990
+ // ModelInfo (value/displayName/description) without token limits, so we must
1991
+ // use the conservativeContextWindow() heuristic. Used when the CLI doesn't
1992
+ // support the get_available_models control request or the Session API errors.
1719
1993
  try {
1720
- const codebuddyExecutable = providerSettings.pathToCodebuddyCode;
1721
1994
  const q = query({
1722
1995
  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
- },
1996
+ options: commonOpts,
1731
1997
  });
1732
1998
  const supported = await q.supportedModels();
1733
1999
  await q.return().catch(() => {});
@@ -1735,9 +2001,9 @@ async function discoverModels(pi: ExtensionAPI): Promise<void> {
1735
2001
  MODELS = applyModelCalibrations(rawModelsFromSdk(supported as any));
1736
2002
  registerCurrentProvider(pi);
1737
2003
  modelsDiscovered = true;
1738
- debug(`discoverModels: registered ${MODELS.length} models from SDK`);
2004
+ debug(`discoverModels: registered ${MODELS.length} models from supportedModels()`);
1739
2005
  } catch (err) {
1740
- debug("discoverModels: failed, using fallback models", err);
2006
+ debug("discoverModels: supportedModels() failed, using fallback models", err);
1741
2007
  }
1742
2008
  });
1743
2009
  }
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
@@ -71,6 +71,9 @@ export function buildPiToolBridgeInstruction(options: ToolBridgeInstructionOptio
71
71
  if (has("bash")) {
72
72
  lines.push(`- Use \`${mcpName("bash")}\` only when file tools are insufficient, for search/test/build/git information, or when command execution is requested.`);
73
73
  }
74
+ if (toolNames.size > 1) {
75
+ 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.");
76
+ }
74
77
  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
78
 
76
79
  return lines.join("\n");
@@ -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
  }