diff --git a/docs/design/2026-07-11-tool-call-preparing-events.md b/docs/design/2026-07-11-tool-call-preparing-events.md new file mode 100644 index 00000000000..8a29cfdb6e2 --- /dev/null +++ b/docs/design/2026-07-11-tool-call-preparing-events.md @@ -0,0 +1,227 @@ +# Tool-call preparation events + +## Context + +Qwen Code currently emits a tool call only after the provider has finished +streaming its arguments. For tools with large or complex inputs, generating +those arguments can take much longer than executing the tool itself. ACP +clients therefore show no activity during the expensive part and users can +mistake the turn for a stalled request. + +The provider streams already expose stable tool identity before the arguments +are complete: + +- Anthropic sends `id` and `name` in `content_block_start` for a `tool_use` + block, then sends argument fragments as `input_json_delta`. +- OpenAI-compatible providers normally send `id` and `function.name` in the + first `choice.delta.tool_calls` item, then append argument fragments. + +Qwen Code deliberately waits for `content_block_stop` or `finish_reason` +before constructing a Gemini-compatible `functionCall`. That execution safety +property must remain unchanged. + +## Goal + +Let ACP clients render a tool card while the model is still preparing tool +arguments, with this lifecycle: + +```text +preparing -> in_progress -> completed | failed +``` + +The early event contains only the stable tool-call ID and tool name. It never +contains partial arguments and never starts tool execution. + +## Scope + +This change supports the two provider paths used by the integrating client: + +- Anthropic and Anthropic-compatible streaming responses. +- OpenAI and OpenAI-compatible streaming responses. + +Other providers keep their current behavior. Because preparation metadata is +optional, they naturally degrade to the existing +`in_progress -> completed | failed` lifecycle. + +The change does not alter: + +- tool permission checks; +- hook ordering; +- tool scheduling or execution; +- model conversation history; +- `functionCall` or `functionResponse` construction; +- non-ACP output formats. + +## Design + +### 1. Internal response metadata + +Associate transient tool preparation metadata with each +`GenerateContentResponse` through a module-local `WeakMap`: + +```ts +interface ToolCallPreparation { + callId: string; + toolName: string; +} +``` + +Provider adapters store this metadata against the top-level response chunk. +It is neither an enumerable response property nor a Gemini `Part`, so it is +not serialized and Gemini history assembly continues to see only text, +thought, and complete `functionCall` parts. Shared helpers provide typed store +and read operations, avoiding provider-specific casts in ACP. + +### 2. Anthropic producer + +In `AnthropicContentGenerator.processStream()`, when +`content_block_start(tool_use)` contains a non-empty `id` and `name`, yield an +otherwise empty Gemini response chunk carrying one preparation entry. + +Continue accumulating `input_json_delta` unchanged. At `content_block_stop`, +emit the existing complete `functionCall` with parsed arguments. No argument +data is exposed before that point. + +### 3. OpenAI-compatible producer + +In `convertOpenAIChunkToGemini()`, observe each +`choice.delta.tool_calls` item after passing it to the existing stream-local +tool-call parser. When a stable non-empty ID and name are available for the +first time, attach one preparation entry to the current response chunk. + +Deduplicate by tool-call ID within the request context. Continue emitting the +complete `functionCall` only when `finish_reason` is present. Providers that do +not expose both identity fields early simply keep the existing behavior. + +### 4. ACP consumer and state transitions + +ACP `Session` reads preparation metadata before collecting complete +`functionCalls`. For each new preparation it emits the standard ACP +`tool_call` frame with: + +```ts +{ + status: 'pending', + rawInput: {}, + _meta: { + phase: 'preparing', + toolName, + // existing provenance metadata remains present + }, +} +``` + +The existing execution path later emits the same `toolCallId` with +`status: 'in_progress'` and the complete arguments. Existing result emission +then finishes the card as `completed` or `failed`. + +`TodoWrite` keeps its current special handling and does not emit a tool card. +Preparation emission uses the same filtering rule, so it cannot create a card +that the execution path intentionally suppresses. + +### 5. Retry, fallback, cancellation, and stream failure + +Each active ACP model stream tracks preparations until the stream completes and +hands its parsed calls to tool execution. When an attempt is abandoned by +retry, model fallback, user cancellation, or stream error, ACP emits a terminal +`tool_call_update` for each remaining entry: + +```ts +{ + status: 'failed', + content: [], + _meta: { + phase: 'preparing', + preparationDiscarded: true, + toolName, + }, +} +``` + +`preparationDiscarded` means the model attempt was abandoned before a parsed +tool request reached execution. It is not a tool execution failure. The integrating +client should remove this transient card rather than render a failed tool. +Using a protocol-valid terminal status ensures older clients do not retain an +indefinitely pending card. + +`RETRY` now clears complete `functionCalls` collected from the abandoned +attempt, matching the existing `MODEL_FALLBACK` behavior across all four ACP +stream paths. This prevents a parsed call from the failed attempt from being +executed together with calls from the replacement attempt. + +When a complete `functionCall` with the same ID arrives and the stream finishes +normally, ACP hands it to the existing execution path without a discarded +update. If the stream fails after parsing the call but before execution, the +preparation is still discarded. Normal tool errors therefore continue through +the existing result path and are never marked as discarded. + +## Downstream impact + +- `GeminiChat` and history builders ignore the optional top-level metadata and + continue persisting only candidate content. +- A response containing only preparation metadata is not counted as + user-visible output, so transport retry and model fallback keep their + existing pre-output behavior. +- Preparation IDs use the same cross-turn normalization as complete + `functionCall` IDs, preserving ACP update correlation when a provider reuses + an ID from history. +- Core `Turn`, TUI, and non-interactive JSON consumers keep their current + behavior because no new Gemini `Part` or server event is introduced. +- ACP is the only consumer that opts into the metadata and emits the early UI + state. +- The same metadata contract is shared by Anthropic and OpenAI-compatible + adapters, so ACP has no provider-specific branches. + +## Test plan + +### Core provider tests + +- Anthropic: a `content_block_start(tool_use)` yields preparation metadata + before any `input_json_delta` and before the final `functionCall`. +- Anthropic: missing ID or name does not emit preparation metadata. +- OpenAI-compatible: the first delta with stable ID and name emits one + preparation entry; later argument deltas do not duplicate it. +- OpenAI-compatible: complete calls still appear only at `finish_reason`, with + unchanged parsed arguments. +- OpenAI-compatible: missing early identity fields fall back to current + behavior without an invalid preparation event. +- GeminiChat: preparation-only chunks do not suppress transport retry, primary + model fallback, or continuation through a multi-model fallback chain. +- GeminiChat: cross-turn duplicate provider IDs are normalized consistently in + preparation metadata and complete calls. + +### ACP tests + +- Preparation metadata emits `pending` with `_meta.phase = 'preparing'` and + no partial input. +- The complete call reuses the same ID and transitions to `in_progress` with + complete arguments. +- Retry, fallback, cancellation, and stream error discard preparations that + have not reached tool execution with `_meta.preparationDiscarded = true`. +- Retry and model fallback clear complete calls collected from the abandoned + attempt before accepting replacement chunks. +- A preparation that became a complete call is not discarded after a normally + completed stream, but is discarded if that stream fails before execution. +- `TodoWrite` remains suppressed. + +### Regression verification + +Run the focused provider and ACP suites from their package directories, then +run repository build, typecheck, and lint before completion. The implementation +rebased on v0.19.9 has been verified with: + +- Core provider and stream suites: 649 passed. +- ACP lifecycle suites: 316 passed. +- Repository build, workspace typecheck, and full lint: passed. +- Changed-file Prettier and diff checks: passed. + +## Acceptance criteria + +1. Anthropic and OpenAI-compatible ACP turns emit a pending tool card as soon + as stable tool identity is available. +2. No tool starts before complete arguments and the existing permission and + execution paths run. +3. Complete calls and results retain their current IDs, arguments, ordering, + and history representation. +4. Abandoned attempts leave no indefinitely pending preparation card. +5. Providers without preparation metadata behave exactly as before. diff --git a/packages/cli/src/acp-integration/session/Session.test.ts b/packages/cli/src/acp-integration/session/Session.test.ts index 2bad043012b..d08cadb7833 100644 --- a/packages/cli/src/acp-integration/session/Session.test.ts +++ b/packages/cli/src/acp-integration/session/Session.test.ts @@ -15,7 +15,12 @@ import { resolveHomeLoopResolverRoots, Session, } from './Session.js'; -import type { Content, FunctionCall, Part } from '@google/genai'; +import type { + Content, + FunctionCall, + GenerateContentResponse, + Part, +} from '@google/genai'; import type { ChatRecord, Config, @@ -232,6 +237,16 @@ function createStreamWithChunks( })(); } +/** Builds provider preparation metadata that arrives before complete arguments. */ +function createPreparationResponse( + callId: string, + toolName: string, +): GenerateContentResponse { + const response = {} as GenerateContentResponse; + core.setToolCallPreparations(response, [{ callId, toolName }]); + return response; +} + function expectCompressBeforeSend( compressMock: ReturnType, sendMock: ReturnType, @@ -8324,6 +8339,570 @@ describe('Session', () => { }); }); + describe('tool preparation stream lifecycle', () => { + function registerAllowedTool( + name: string, + execute: ReturnType, + ) { + mockToolRegistry.getTool.mockImplementation((toolName: string) => + toolName === name + ? { + name, + kind: core.Kind.Read, + displayName: name, + build: vi + .fn() + .mockImplementation((args: Record) => ({ + params: args, + execute, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue(name), + toolLocations: vi.fn().mockReturnValue([]), + })), + } + : undefined, + ); + mockConfig.getApprovalMode = vi.fn().mockReturnValue(ApprovalMode.YOLO); + } + + it('emits preparing before execution and keeps resolved preparations out of finally discard', async () => { + const execute = vi.fn().mockResolvedValue({ + llmContent: 'file contents', + returnDisplay: 'file contents', + }); + registerAllowedTool('read_file', execute); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse('call-1', 'read_file'), + }, + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.sql' }, + }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read the file' }], + }); + + const updates = vi + .mocked(mockClient.sessionUpdate) + .mock.calls.map(([params]) => params.update); + const preparingIndex = updates.findIndex( + (update) => + update.sessionUpdate === 'tool_call' && + update.toolCallId === 'call-1' && + update.status === 'pending' && + update._meta?.['phase'] === 'preparing', + ); + const inProgressIndex = updates.findIndex( + (update) => + update.sessionUpdate === 'tool_call_update' && + update.toolCallId === 'call-1' && + update.status === 'in_progress' && + 'rawInput' in update && + update.rawInput?.['file_path'] === 'a.sql', + ); + + expect(preparingIndex).toBeGreaterThanOrEqual(0); + expect(inProgressIndex).toBeGreaterThan(preparingIndex); + expect(execute).toHaveBeenCalledOnce(); + expect(updates).not.toContainEqual( + expect.objectContaining({ + toolCallId: 'call-1', + _meta: expect.objectContaining({ preparationDiscarded: true }), + }), + ); + }); + + it('suppresses TodoWrite preparation updates at the Session boundary', async () => { + mockChat.sendMessageStream = vi.fn().mockResolvedValue( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse( + 'call-todo', + core.ToolNames.TODO_WRITE, + ), + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'update the plan' }], + }); + + const updates = vi + .mocked(mockClient.sessionUpdate) + .mock.calls.map(([params]) => params.update); + expect( + updates.some( + (update) => + (update.sessionUpdate === 'tool_call' || + update.sessionUpdate === 'tool_call_update') && + update.toolCallId === 'call-todo', + ), + ).toBe(false); + }); + + it('discards a resolved preparation when the stream fails before tool execution', async () => { + const execute = vi.fn(); + registerAllowedTool('read_file', execute); + mockChat.sendMessageStream = vi.fn().mockResolvedValue( + (async function* () { + yield { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse( + 'call-stream-error', + 'read_file', + ), + }; + yield { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { + id: 'call-stream-error', + name: 'read_file', + args: { file_path: 'a.sql' }, + }, + ], + }, + }; + throw new Error('stream failed after function call'); + })(), + ); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'read the file' }], + }), + ).rejects.toThrow('stream failed after function call'); + + expect(execute).not.toHaveBeenCalled(); + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + sessionUpdate: 'tool_call_update', + toolCallId: 'call-stream-error', + status: 'failed', + _meta: expect.objectContaining({ + phase: 'preparing', + preparationDiscarded: true, + }), + }), + }); + }); + + it.each([ + core.StreamEventType.RETRY, + core.StreamEventType.MODEL_FALLBACK, + ])( + 'discards preparation and stale function calls on %s before accepting the next attempt', + async (resetEvent) => { + const staleExecute = vi.fn(); + const currentExecute = vi.fn().mockResolvedValue({ + llmContent: 'current result', + returnDisplay: 'current result', + }); + mockToolRegistry.getTool.mockImplementation((toolName: string) => { + const execute = + toolName === 'stale_tool' + ? staleExecute + : toolName === 'current_tool' + ? currentExecute + : undefined; + if (!execute) return undefined; + return { + name: toolName, + kind: core.Kind.Read, + displayName: toolName, + build: vi.fn().mockImplementation((args) => ({ + params: args, + execute, + getDefaultPermission: vi.fn().mockResolvedValue('allow'), + getDescription: vi.fn().mockReturnValue(toolName), + toolLocations: vi.fn().mockReturnValue([]), + })), + }; + }); + mockConfig.getApprovalMode = vi + .fn() + .mockReturnValue(ApprovalMode.YOLO); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse( + 'preparing-stale', + 'read_file', + ), + }, + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { id: 'stale-call', name: 'stale_tool', args: {} }, + ], + }, + }, + { type: resetEvent, value: {} }, + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { id: 'current-call', name: 'current_tool', args: {} }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'retry the tool call' }], + }); + + expect(staleExecute).not.toHaveBeenCalled(); + expect(currentExecute).toHaveBeenCalledOnce(); + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + sessionUpdate: 'tool_call_update', + toolCallId: 'preparing-stale', + status: 'failed', + _meta: expect.objectContaining({ + phase: 'preparing', + preparationDiscarded: true, + }), + }), + }); + }, + ); + + it.each([ + core.StreamEventType.RETRY, + core.StreamEventType.MODEL_FALLBACK, + ])( + 'continues after preparation cleanup fails during %s', + async (resetEvent) => { + debugLoggerWarnSpy.mockClear(); + const execute = vi.fn().mockResolvedValue({ + llmContent: 'current result', + returnDisplay: 'current result', + }); + registerAllowedTool('current_tool', execute); + vi.mocked(mockClient.sessionUpdate).mockImplementation( + async ({ update }) => { + if (update._meta?.['preparationDiscarded'] === true) { + throw new Error('cleanup failed'); + } + }, + ); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse( + 'preparing-stale', + 'read_file', + ), + }, + { type: resetEvent, value: {} }, + { + type: core.StreamEventType.CHUNK, + value: { + functionCalls: [ + { id: 'current-call', name: 'current_tool', args: {} }, + ], + }, + }, + ]), + ) + .mockResolvedValueOnce(createEmptyStream()); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'retry the tool call' }], + }); + + expect(execute).toHaveBeenCalledOnce(); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('cleanup failed'), + ); + }, + ); + + it('preserves the stream error when discarding its unresolved preparation also fails', async () => { + debugLoggerWarnSpy.mockClear(); + vi.mocked(mockClient.sessionUpdate).mockImplementation( + async ({ update }) => { + if (update._meta?.['preparationDiscarded'] === true) { + throw new Error('cleanup failed'); + } + }, + ); + mockChat.sendMessageStream = vi.fn().mockResolvedValue( + (async function* () { + yield { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse('call-error', 'read_file'), + }; + throw new Error('stream failed'); + })(), + ); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'trigger a stream error' }], + }), + ).rejects.toThrow('stream failed'); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('cleanup failed'), + ); + }); + + it('preserves a cancelled return when discarding its unresolved preparation fails', async () => { + debugLoggerWarnSpy.mockClear(); + vi.mocked(mockClient.sessionUpdate).mockImplementation( + async ({ update }) => { + if (update._meta?.['preparationDiscarded'] === true) { + throw new Error('cleanup failed'); + } + }, + ); + mockChat.sendMessageStream = vi + .fn() + .mockImplementation( + async (_model, request: { config: { abortSignal: AbortSignal } }) => + (async function* (signal: AbortSignal) { + yield { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse('call-cancel', 'read_file'), + }; + await new Promise((resolve) => { + signal.addEventListener('abort', () => resolve(), { + once: true, + }); + }); + yield { type: core.StreamEventType.CHUNK, value: {} }; + })(request.config.abortSignal), + ); + + const promptPromise = session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'cancel this stream' }], + }); + await vi.waitFor(() => { + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + sessionUpdate: 'tool_call', + toolCallId: 'call-cancel', + }), + }); + }); + + await session.cancelPendingPrompt(); + + await expect(promptPromise).resolves.toEqual({ + stopReason: 'cancelled', + }); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('cleanup failed'), + ); + }); + + it('preserves a normally completed stream when preparation cleanup fails', async () => { + debugLoggerWarnSpy.mockClear(); + vi.mocked(mockClient.sessionUpdate).mockImplementation( + async ({ update }) => { + if (update._meta?.['preparationDiscarded'] === true) { + throw new Error('cleanup failed'); + } + }, + ); + mockChat.sendMessageStream = vi.fn().mockResolvedValue( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse('call-normal', 'read_file'), + }, + ]), + ); + + await expect( + session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'finish normally' }], + }), + ).resolves.toEqual({ stopReason: 'end_turn' }); + expect(debugLoggerWarnSpy).toHaveBeenCalledWith( + expect.stringContaining('cleanup failed'), + ); + }); + + it('tracks unresolved preparation in a Stop Hook continuation stream', async () => { + const messageBus = { + request: vi.fn().mockImplementation(async (request) => ({ + success: true, + output: + request.eventName === 'Stop' + ? { + decision: 'block', + reason: 'Continue after Stop hook', + } + : {}, + })), + }; + mockConfig.getMessageBus = vi.fn().mockReturnValue(messageBus); + mockConfig.getDisableAllHooks = vi.fn().mockReturnValue(false); + mockConfig.hasHooksForEvent = vi + .fn() + .mockImplementation((eventName: string) => eventName === 'Stop'); + mockConfig.getStopHookBlockingCap = vi.fn().mockReturnValue(2); + mockChat.getLastModelMessageText = vi + .fn() + .mockReturnValue('response text'); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse('call-stop-hook', 'read_file'), + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'run stop hook' }], + }); + + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + sessionUpdate: 'tool_call_update', + toolCallId: 'call-stop-hook', + _meta: expect.objectContaining({ preparationDiscarded: true }), + }), + }); + }); + + it('tracks unresolved preparation in a cron stream', async () => { + const scheduler = { + size: 1, + hasPendingWork: true, + start: vi.fn( + ( + callback: (job: { prompt: string; cronExpr?: string }) => void, + ) => { + callback({ prompt: 'scheduled work', cronExpr: '* * * * *' }); + }, + ), + stop: vi.fn(), + getExitSummary: vi.fn().mockReturnValue(undefined), + }; + mockConfig.isCronEnabled = vi.fn().mockReturnValue(true); + mockConfig.getCronScheduler = vi.fn().mockReturnValue(scheduler); + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse('call-cron', 'read_file'), + }, + ]), + ); + + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start cron' }], + }); + + await vi.waitFor(() => { + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + sessionUpdate: 'tool_call_update', + toolCallId: 'call-cron', + _meta: expect.objectContaining({ preparationDiscarded: true }), + }), + }); + }); + }); + + it('tracks unresolved preparation in a background notification stream', async () => { + mockChat.sendMessageStream = vi + .fn() + .mockResolvedValueOnce(createEmptyStream()) + .mockResolvedValueOnce( + createStreamWithChunks([ + { + type: core.StreamEventType.CHUNK, + value: createPreparationResponse( + 'call-notification', + 'read_file', + ), + }, + ]), + ); + await session.prompt({ + sessionId: 'test-session-id', + prompt: [{ type: 'text', text: 'start background work' }], + }); + const callback = mockBackgroundTaskRegistry.setNotificationCallback.mock + .calls[0][0] as ( + displayText: string, + modelText: string, + meta: { agentId: string; status: string }, + ) => void; + + callback('done', '', { + agentId: 'agent-1', + status: 'completed', + }); + + await vi.waitFor(() => { + expect(mockClient.sessionUpdate).toHaveBeenCalledWith({ + sessionId: 'test-session-id', + update: expect.objectContaining({ + sessionUpdate: 'tool_call_update', + toolCallId: 'call-notification', + _meta: expect.objectContaining({ preparationDiscarded: true }), + }), + }); + }); + }); + }); + it('passes resolved paths to read_many_files tool', async () => { const tempDir = await fs.mkdtemp( path.join(os.tmpdir(), 'qwen-acp-session-'), diff --git a/packages/cli/src/acp-integration/session/Session.ts b/packages/cli/src/acp-integration/session/Session.ts index 8db40a97e11..b59c0c487b9 100644 --- a/packages/cli/src/acp-integration/session/Session.ts +++ b/packages/cli/src/acp-integration/session/Session.ts @@ -209,6 +209,7 @@ import type { } from './types.js'; import { HistoryReplayer } from './history-replayer.js'; import { ToolCallEmitter } from './emitters/tool-call-emitter.js'; +import { ToolCallPreparationTracker } from './tool-call-preparation-tracker.js'; import { PlanEmitter } from './emitters/PlanEmitter.js'; import { MessageEmitter } from './emitters/MessageEmitter.js'; import { SubAgentTracker } from './SubAgentTracker.js'; @@ -227,6 +228,22 @@ const USER_CANCEL_ABORT_REASON = 'qwen:user-cancel'; const DAEMON_RETRY_META_KEY = 'qwen.daemon.retry'; const DAEMON_CONTINUE_META_KEY = 'qwen.daemon.continueLastTurn'; +/** Finalizes preparations without allowing ACP cleanup to change the stream outcome. */ +async function finalizeToolCallPreparations( + tracker: ToolCallPreparationTracker, + includeResolved: boolean, + streamName: string, +): Promise { + try { + await tracker.discard(includeResolved); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + debugLogger.warn( + `Failed to discard tool preparations for ${streamName}; continuing stream: ${message}`, + ); + } +} + function maskApiKeyForDisplay(apiKey: string | undefined): string { const trimmed = apiKey?.trim() ?? ''; if (trimmed.length === 0) return '(not set)'; @@ -1962,6 +1979,9 @@ export class Session implements SessionContext { } const functionCalls: FunctionCall[] = []; + const preparationTracker = new ToolCallPreparationTracker( + this.toolCallEmitter, + ); let usageMetadata: GenerateContentResponseUsageMetadata | null = null; const streamStartTime = Date.now(); @@ -1991,49 +2011,70 @@ export class Session implements SessionContext { const responseStream = sendResult.responseStream; nextMessage = null; - for await (const resp of responseStream) { - if (pendingSend.signal.aborted) { - return { stopReason: 'cancelled' }; - } - - if ( - resp.type === StreamEventType.CHUNK && - resp.value.candidates && - resp.value.candidates.length > 0 - ) { - const candidate = resp.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) { - continue; - } + let streamFailed = false; + try { + for await (const resp of responseStream) { + if (pendingSend.signal.aborted) { + return { stopReason: 'cancelled' }; + } - this.messageEmitter.emitMessage( - part.text, - 'assistant', - part.thought, - ); - if (!part.thought) { - messageDisplay?.addChunk(part.text); + if ( + resp.type === StreamEventType.CHUNK && + resp.value.candidates && + resp.value.candidates.length > 0 + ) { + const candidate = resp.value.candidates[0]; + for (const part of candidate.content?.parts ?? []) { + if (!part.text) { + continue; + } + + this.messageEmitter.emitMessage( + part.text, + 'assistant', + part.thought, + ); + if (!part.thought) { + messageDisplay?.addChunk(part.text); + } } } - } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.usageMetadata - ) { - usageMetadata = resp.value.usageMetadata; - } + if ( + resp.type === StreamEventType.CHUNK && + resp.value.usageMetadata + ) { + usageMetadata = resp.value.usageMetadata; + } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.functionCalls - ) { - functionCalls.push(...resp.value.functionCalls); - } - if (resp.type === StreamEventType.MODEL_FALLBACK) { - functionCalls.length = 0; + if (resp.type === StreamEventType.CHUNK) { + await preparationTracker.observe(resp.value); + if (resp.value.functionCalls) { + preparationTracker.resolve(resp.value.functionCalls); + functionCalls.push(...resp.value.functionCalls); + } + } + if ( + resp.type === StreamEventType.RETRY || + resp.type === StreamEventType.MODEL_FALLBACK + ) { + await finalizeToolCallPreparations( + preparationTracker, + true, + `main prompt ${resp.type}`, + ); + functionCalls.length = 0; + } } + } catch (error) { + streamFailed = true; + throw error; + } finally { + await finalizeToolCallPreparations( + preparationTracker, + streamFailed || pendingSend.signal.aborted, + 'main prompt', + ); } } catch (error) { // Restore the stripped orphan if the send threw before @@ -2303,6 +2344,9 @@ export class Session implements SessionContext { } const functionCalls: FunctionCall[] = []; + const preparationTracker = new ToolCallPreparationTracker( + this.toolCallEmitter, + ); let usageMetadata: GenerateContentResponseUsageMetadata | null = null; const streamStartTime = Date.now(); const messageDisplay = this.#createMessageDisplayDispatcher( @@ -2327,46 +2371,67 @@ export class Session implements SessionContext { const continueResponseStream = continueSendResult.responseStream; nextMessage = null; - for await (const resp of continueResponseStream) { - if (pendingSend.signal.aborted) { - return { stopReason: 'cancelled' }; - } + let streamFailed = false; + try { + for await (const resp of continueResponseStream) { + if (pendingSend.signal.aborted) { + return { stopReason: 'cancelled' }; + } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.candidates && - resp.value.candidates.length > 0 - ) { - const candidate = resp.value.candidates[0]; - for (const part of candidate.content?.parts ?? []) { - if (!part.text) continue; - this.messageEmitter.emitMessage( - part.text, - 'assistant', - part.thought, - ); - if (!part.thought) { - messageDisplay?.addChunk(part.text); + if ( + resp.type === StreamEventType.CHUNK && + resp.value.candidates && + resp.value.candidates.length > 0 + ) { + const candidate = resp.value.candidates[0]; + for (const part of candidate.content?.parts ?? []) { + if (!part.text) continue; + this.messageEmitter.emitMessage( + part.text, + 'assistant', + part.thought, + ); + if (!part.thought) { + messageDisplay?.addChunk(part.text); + } } } - } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.usageMetadata - ) { - usageMetadata = resp.value.usageMetadata; - } + if ( + resp.type === StreamEventType.CHUNK && + resp.value.usageMetadata + ) { + usageMetadata = resp.value.usageMetadata; + } - if ( - resp.type === StreamEventType.CHUNK && - resp.value.functionCalls - ) { - functionCalls.push(...resp.value.functionCalls); - } - if (resp.type === StreamEventType.MODEL_FALLBACK) { - functionCalls.length = 0; + if (resp.type === StreamEventType.CHUNK) { + await preparationTracker.observe(resp.value); + if (resp.value.functionCalls) { + preparationTracker.resolve(resp.value.functionCalls); + functionCalls.push(...resp.value.functionCalls); + } + } + if ( + resp.type === StreamEventType.RETRY || + resp.type === StreamEventType.MODEL_FALLBACK + ) { + await finalizeToolCallPreparations( + preparationTracker, + true, + `Stop Hook continuation ${resp.type}`, + ); + functionCalls.length = 0; + } } + } catch (error) { + streamFailed = true; + throw error; + } finally { + await finalizeToolCallPreparations( + preparationTracker, + streamFailed || pendingSend.signal.aborted, + 'Stop Hook continuation', + ); } } catch (error) { // Fire StopFailure hook (fire-and-forget) @@ -3253,6 +3318,9 @@ export class Session implements SessionContext { if (ac.signal.aborted) return; const functionCalls: FunctionCall[] = []; + const preparationTracker = new ToolCallPreparationTracker( + this.toolCallEmitter, + ); let usageMetadata: GenerateContentResponseUsageMetadata | null = null; const streamStartTime = Date.now(); @@ -3286,6 +3354,7 @@ export class Session implements SessionContext { ac.signal, ); + let streamFailed = false; try { for await (const resp of responseStream) { if (ac.signal.aborted) return; @@ -3316,20 +3385,40 @@ export class Session implements SessionContext { usageMetadata = resp.value.usageMetadata; } + if (resp.type === StreamEventType.CHUNK) { + await preparationTracker.observe(resp.value); + if (resp.value.functionCalls) { + preparationTracker.resolve(resp.value.functionCalls); + functionCalls.push(...resp.value.functionCalls); + } + } if ( - resp.type === StreamEventType.CHUNK && - resp.value.functionCalls + resp.type === StreamEventType.RETRY || + resp.type === StreamEventType.MODEL_FALLBACK ) { - functionCalls.push(...resp.value.functionCalls); - } - if (resp.type === StreamEventType.MODEL_FALLBACK) { + await finalizeToolCallPreparations( + preparationTracker, + true, + `cron/loop tick ${resp.type}`, + ); functionCalls.length = 0; } } + } catch (error) { + streamFailed = true; + throw error; } finally { - // is_final (skipped on abort) delivered and drained on - // every exit path, same as the interactive prompt loops. - await messageDisplay?.finish(); + try { + await finalizeToolCallPreparations( + preparationTracker, + streamFailed || ac.signal.aborted, + 'cron/loop tick', + ); + } finally { + // is_final (skipped on abort) delivered and drained on + // every exit path, same as the interactive prompt loops. + await messageDisplay?.finish(); + } } if (usageMetadata) { @@ -3590,6 +3679,9 @@ export class Session implements SessionContext { } const functionCalls: FunctionCall[] = []; + const preparationTracker = new ToolCallPreparationTracker( + this.toolCallEmitter, + ); let usageMetadata: GenerateContentResponseUsageMetadata | null = null; let responseText = ''; @@ -3617,6 +3709,7 @@ export class Session implements SessionContext { ac.signal, ); + let streamFailed = false; try { for await (const resp of responseStream) { if (ac.signal.aborted) { @@ -3652,20 +3745,40 @@ export class Session implements SessionContext { usageMetadata = resp.value.usageMetadata; } + if (resp.type === StreamEventType.CHUNK) { + await preparationTracker.observe(resp.value); + if (resp.value.functionCalls) { + preparationTracker.resolve(resp.value.functionCalls); + functionCalls.push(...resp.value.functionCalls); + } + } if ( - resp.type === StreamEventType.CHUNK && - resp.value.functionCalls + resp.type === StreamEventType.RETRY || + resp.type === StreamEventType.MODEL_FALLBACK ) { - functionCalls.push(...resp.value.functionCalls); - } - if (resp.type === StreamEventType.MODEL_FALLBACK) { + await finalizeToolCallPreparations( + preparationTracker, + true, + `background notification ${resp.type}`, + ); functionCalls.length = 0; } } + } catch (error) { + streamFailed = true; + throw error; } finally { - // is_final (skipped on abort) delivered and drained on every - // exit path, same as the interactive prompt loops. - await messageDisplay?.finish(); + try { + await finalizeToolCallPreparations( + preparationTracker, + streamFailed || ac.signal.aborted, + 'background notification', + ); + } finally { + // is_final (skipped on abort) delivered and drained on every + // exit path, same as the interactive prompt loops. + await messageDisplay?.finish(); + } } if (responseText.length > 0) { diff --git a/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.test.ts b/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.test.ts index 829f1e3d08a..30d8019f931 100644 --- a/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.test.ts +++ b/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.test.ts @@ -162,6 +162,151 @@ describe('ToolCallEmitter', () => { }); }); + describe('tool preparation lifecycle', () => { + it('emits a preparing tool call without partial input', async () => { + await emitter.emitStart({ + callId: 'call-1', + toolName: 'read_file', + args: {}, + status: 'pending', + phase: 'preparing', + }); + + expect(sendUpdateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + sessionUpdate: 'tool_call', + toolCallId: 'call-1', + status: 'pending', + rawInput: {}, + _meta: expect.objectContaining({ + toolName: 'read_file', + phase: 'preparing', + }), + }), + ); + }); + + it('suppresses duplicate preparing frames for the same call ID', async () => { + const params = { + callId: 'call-1', + toolName: 'read_file', + args: {}, + status: 'pending' as const, + phase: 'preparing' as const, + }; + + const first = await emitter.emitStart(params); + const duplicate = await emitter.emitStart(params); + + expect(first).toBe(true); + expect(duplicate).toBe(false); + expect(sendUpdateSpy).toHaveBeenCalledTimes(1); + }); + + it('updates the prepared tool call when execution starts', async () => { + await emitter.emitStart({ + callId: 'call-1', + toolName: 'read_file', + args: {}, + status: 'pending', + phase: 'preparing', + }); + await emitter.emitStart({ + callId: 'call-1', + toolName: 'read_file', + args: { file_path: 'README.md' }, + status: 'in_progress', + }); + + expect(sendUpdateSpy.mock.calls.map(([update]) => update)).toEqual([ + expect.objectContaining({ + sessionUpdate: 'tool_call', + toolCallId: 'call-1', + _meta: expect.objectContaining({ phase: 'preparing' }), + }), + expect.objectContaining({ + sessionUpdate: 'tool_call_update', + toolCallId: 'call-1', + status: 'in_progress', + rawInput: { file_path: 'README.md' }, + }), + ]); + }); + + it('emits a protocol-valid discarded preparation terminal update', async () => { + await emitter.emitPreparationDiscarded( + 'call-1', + 'mcp__filesystem__read_file', + ); + + expect(sendUpdateSpy).toHaveBeenCalledWith({ + sessionUpdate: 'tool_call_update', + toolCallId: 'call-1', + status: 'failed', + content: [], + _meta: { + toolName: 'mcp__filesystem__read_file', + phase: 'preparing', + preparationDiscarded: true, + provenance: 'mcp', + serverId: 'filesystem', + }, + }); + }); + + it.each(['result', 'error'] as const)( + 'clears prepared state after terminal %s', + async (terminal) => { + const preparation = { + callId: 'call-1', + toolName: 'read_file', + args: {}, + status: 'pending' as const, + phase: 'preparing' as const, + }; + await emitter.emitStart(preparation); + + if (terminal === 'result') { + await emitter.emitResult({ + callId: 'call-1', + toolName: 'read_file', + success: true, + message: [], + }); + } else { + await emitter.emitError('call-1', 'read_file', new Error('failed')); + } + sendUpdateSpy.mockClear(); + + const emitted = await emitter.emitStart(preparation); + + expect(emitted).toBe(true); + expect(sendUpdateSpy).toHaveBeenCalledOnce(); + expect(sendUpdateSpy).toHaveBeenCalledWith( + expect.objectContaining({ + sessionUpdate: 'tool_call', + toolCallId: 'call-1', + _meta: expect.objectContaining({ phase: 'preparing' }), + }), + ); + }, + ); + + it('suppresses preparation lifecycle frames for TodoWrite', async () => { + const emitted = await emitter.emitStart({ + callId: 'call-todo', + toolName: ToolNames.TODO_WRITE, + args: {}, + status: 'pending', + phase: 'preparing', + }); + await emitter.emitPreparationDiscarded('call-todo', ToolNames.TODO_WRITE); + + expect(emitted).toBe(false); + expect(sendUpdateSpy).not.toHaveBeenCalled(); + }); + }); + describe('emitResult', () => { it('should emit tool_call_update with completed status on success', async () => { await emitter.emitResult({ diff --git a/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.ts b/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.ts index 8a0588fd630..d78cc8b5f0e 100644 --- a/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.ts +++ b/packages/cli/src/acp-integration/session/emitters/tool-call-emitter.ts @@ -56,6 +56,7 @@ const KIND_MAP: Record = { */ export class ToolCallEmitter extends BaseEmitter { private readonly planEmitter: PlanEmitter; + private readonly preparedCallIds = new Set(); constructor(ctx: SessionEmitterContext) { super(ctx); @@ -73,6 +74,12 @@ export class ToolCallEmitter extends BaseEmitter { if (this.isTodoWriteTool(params.toolName)) { return false; } + if ( + params.phase === 'preparing' && + this.preparedCallIds.has(params.callId) + ) { + return false; + } const { title, locations, kind } = this.resolveToolMetadata( params.toolName, @@ -82,9 +89,11 @@ export class ToolCallEmitter extends BaseEmitter { params.toolName, params.subagentMeta, ); + const updatesPreparedCall = + params.phase !== 'preparing' && + this.preparedCallIds.delete(params.callId); - await this.sendUpdate({ - sessionUpdate: 'tool_call', + const update = { toolCallId: params.callId, status: params.status || 'pending', title, @@ -94,6 +103,7 @@ export class ToolCallEmitter extends BaseEmitter { rawInput: params.args ?? {}, _meta: { toolName: params.toolName, + ...(params.phase ? { phase: params.phase } : {}), ...params.subagentMeta, provenance: provenance.provenance, ...(provenance.serverId ? { serverId: provenance.serverId } : {}), @@ -101,11 +111,49 @@ export class ToolCallEmitter extends BaseEmitter { timestamp: BaseEmitter.toEpochMs(params.timestamp), }), }, - }); + }; + await this.sendUpdate( + updatesPreparedCall + ? { sessionUpdate: 'tool_call_update', ...update } + : { sessionUpdate: 'tool_call', ...update }, + ); + if (params.phase === 'preparing') { + this.preparedCallIds.add(params.callId); + } return true; } + /** + * Emits a terminal frame when a prepared tool call is discarded before + * execution. TodoWrite remains represented exclusively by plan updates. + * + * @param callId - ID of the prepared tool call + * @param toolName - Name of the prepared tool + */ + async emitPreparationDiscarded( + callId: string, + toolName: string, + ): Promise { + if (this.isTodoWriteTool(toolName)) return; + + this.preparedCallIds.delete(callId); + const provenance = ToolCallEmitter.resolveToolProvenance(toolName); + await this.sendUpdate({ + sessionUpdate: 'tool_call_update', + toolCallId: callId, + status: 'failed', + content: [], + _meta: { + toolName, + phase: 'preparing', + preparationDiscarded: true, + provenance: provenance.provenance, + ...(provenance.serverId ? { serverId: provenance.serverId } : {}), + }, + }); + } + /** * Emits a tool call result event. * Handles TodoWriteTool specially by routing to plan updates. @@ -130,6 +178,8 @@ export class ToolCallEmitter extends BaseEmitter { return; // Skip tool_call_update for TodoWriteTool } + this.preparedCallIds.delete(params.callId); + // Determine content for the update let contentArray: ToolCallContent[] = []; @@ -200,6 +250,7 @@ export class ToolCallEmitter extends BaseEmitter { error: Error, subagentMeta?: SubagentMeta, ): Promise { + this.preparedCallIds.delete(callId); const provenance = ToolCallEmitter.resolveToolProvenance( toolName, subagentMeta, diff --git a/packages/cli/src/acp-integration/session/tool-call-preparation-tracker.test.ts b/packages/cli/src/acp-integration/session/tool-call-preparation-tracker.test.ts new file mode 100644 index 00000000000..a4755d67031 --- /dev/null +++ b/packages/cli/src/acp-integration/session/tool-call-preparation-tracker.test.ts @@ -0,0 +1,151 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import { GenerateContentResponse } from '@google/genai'; +import { setToolCallPreparations } from '@qwen-code/qwen-code-core'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type { ToolCallEmitter } from './emitters/tool-call-emitter.js'; +import { ToolCallPreparationTracker } from './tool-call-preparation-tracker.js'; + +describe('ToolCallPreparationTracker', () => { + let emitStart: ReturnType; + let emitPreparationDiscarded: ReturnType; + let emitter: ToolCallEmitter; + + beforeEach(() => { + emitStart = vi.fn().mockResolvedValue(true); + emitPreparationDiscarded = vi.fn().mockResolvedValue(undefined); + emitter = { + emitStart, + emitPreparationDiscarded, + } as unknown as ToolCallEmitter; + }); + + it('emits a preparation once and does not discard it after resolution', async () => { + const tracker = new ToolCallPreparationTracker(emitter); + const response = new GenerateContentResponse(); + setToolCallPreparations(response, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + + await tracker.observe(response); + await tracker.observe(response); + + expect(emitStart).toHaveBeenCalledTimes(1); + expect(emitStart).toHaveBeenCalledWith({ + callId: 'call-1', + toolName: 'read_file', + args: {}, + status: 'pending', + phase: 'preparing', + }); + + tracker.resolve([ + { id: 'call-1', name: 'read_file', args: { file_path: 'a.sql' } }, + ]); + await tracker.discard(); + + expect(emitPreparationDiscarded).not.toHaveBeenCalled(); + }); + + it('discards every unresolved preparation exactly once', async () => { + const tracker = new ToolCallPreparationTracker(emitter); + const response = new GenerateContentResponse(); + setToolCallPreparations(response, [ + { callId: 'call-1', toolName: 'read_file' }, + { callId: 'call-2', toolName: 'shell' }, + ]); + + await tracker.observe(response); + await tracker.discard(); + await tracker.discard(); + + expect(emitPreparationDiscarded.mock.calls).toEqual([ + ['call-1', 'read_file'], + ['call-2', 'shell'], + ]); + }); + + it('discards a resolved preparation when the stream attempt is abandoned', async () => { + const tracker = new ToolCallPreparationTracker(emitter); + const response = new GenerateContentResponse(); + setToolCallPreparations(response, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + + await tracker.observe(response); + tracker.resolve([ + { id: 'call-1', name: 'read_file', args: { file_path: 'a.sql' } }, + ]); + await tracker.discard(true); + + expect(emitPreparationDiscarded).toHaveBeenCalledOnce(); + expect(emitPreparationDiscarded).toHaveBeenCalledWith( + 'call-1', + 'read_file', + ); + }); + + it('keeps preparations unresolved for missing or empty function call IDs', async () => { + const tracker = new ToolCallPreparationTracker(emitter); + const response = new GenerateContentResponse(); + setToolCallPreparations(response, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + + await tracker.observe(response); + tracker.resolve([ + { id: undefined, name: 'read_file', args: {} }, + { id: '', name: 'read_file', args: {} }, + ]); + await tracker.discard(); + + expect(emitPreparationDiscarded).toHaveBeenCalledOnce(); + expect(emitPreparationDiscarded).toHaveBeenCalledWith( + 'call-1', + 'read_file', + ); + }); + + it('attempts every unresolved discard before surfacing the first cleanup error', async () => { + const cleanupError = new Error('first discard failed'); + emitPreparationDiscarded + .mockRejectedValueOnce(cleanupError) + .mockResolvedValueOnce(undefined); + const tracker = new ToolCallPreparationTracker(emitter); + const response = new GenerateContentResponse(); + setToolCallPreparations(response, [ + { callId: 'call-1', toolName: 'read_file' }, + { callId: 'call-2', toolName: 'shell' }, + ]); + + await tracker.observe(response); + + await expect(tracker.discard()).rejects.toBe(cleanupError); + expect(emitPreparationDiscarded.mock.calls).toEqual([ + ['call-1', 'read_file'], + ['call-2', 'shell'], + ]); + await tracker.discard(); + expect(emitPreparationDiscarded).toHaveBeenCalledTimes(2); + }); + + it('does not retry a preparation when the emitter suppresses its start', async () => { + emitStart.mockResolvedValue(false); + const tracker = new ToolCallPreparationTracker(emitter); + const response = new GenerateContentResponse(); + setToolCallPreparations(response, [ + { callId: 'todo-1', toolName: 'TodoWrite' }, + ]); + + await tracker.observe(response); + await tracker.observe(response); + await tracker.discard(); + + expect(emitStart).toHaveBeenCalledTimes(1); + expect(emitPreparationDiscarded).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/cli/src/acp-integration/session/tool-call-preparation-tracker.ts b/packages/cli/src/acp-integration/session/tool-call-preparation-tracker.ts new file mode 100644 index 00000000000..a600d12de0e --- /dev/null +++ b/packages/cli/src/acp-integration/session/tool-call-preparation-tracker.ts @@ -0,0 +1,93 @@ +/** + * @license + * Copyright 2025 Qwen + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { FunctionCall, GenerateContentResponse } from '@google/genai'; +import { getToolCallPreparations } from '@qwen-code/qwen-code-core'; +import type { ToolCallEmitter } from './emitters/tool-call-emitter.js'; + +/** + * Tracks preparations exposed to ACP before their complete function calls are + * parsed. Each model stream gets its own instance so retries, fallbacks, and + * cancellation cannot leak pending calls into a later attempt. + */ +export class ToolCallPreparationTracker { + /** Contains only calls whose start frame was emitted successfully. */ + private readonly pending = new Map(); + /** Contains calls whose start frame was intentionally suppressed. */ + private readonly suppressed = new Set(); + /** Calls parsed completely but not yet handed to tool execution. */ + private readonly resolved = new Set(); + + constructor(private readonly emitter: ToolCallEmitter) {} + + /** + * Emits at most one preparing frame per call ID before the full call arrives. + */ + async observe(response: GenerateContentResponse): Promise { + for (const preparation of getToolCallPreparations(response)) { + if ( + this.pending.has(preparation.callId) || + this.suppressed.has(preparation.callId) + ) { + continue; + } + + const emitted = await this.emitter.emitStart({ + callId: preparation.callId, + toolName: preparation.toolName, + args: {}, + status: 'pending', + phase: 'preparing', + }); + if (emitted) { + this.pending.set(preparation.callId, preparation.toolName); + } else { + this.suppressed.add(preparation.callId); + } + } + } + + /** Resolves preparations once their complete function calls arrive. */ + resolve(functionCalls: readonly FunctionCall[]): void { + for (const functionCall of functionCalls) { + if (functionCall.id && this.pending.has(functionCall.id)) { + this.resolved.add(functionCall.id); + } + } + } + + /** + * Terminates unresolved preparations. The map is cleared first so repeated + * cleanup, including re-entry after an emission failure, cannot emit twice. + */ + async discard(includeResolved = false): Promise { + const pending = [...this.pending.entries()]; + this.pending.clear(); + const resolved = new Set(this.resolved); + this.resolved.clear(); + let firstError: unknown; + let hasError = false; + + for (const [callId, toolName] of pending) { + if (!includeResolved && resolved.has(callId)) continue; + + try { + await this.emitter.emitPreparationDiscarded(callId, toolName); + } catch (error) { + // One failed ACP update must not prevent the remaining calls from being + // finalized. Preserve the first failure and throw it after all attempts. + if (!hasError) { + firstError = error; + hasError = true; + } + } + } + + if (hasError) { + throw firstError; + } + } +} diff --git a/packages/cli/src/acp-integration/session/types.ts b/packages/cli/src/acp-integration/session/types.ts index 8fc17de4755..868df0a9100 100644 --- a/packages/cli/src/acp-integration/session/types.ts +++ b/packages/cli/src/acp-integration/session/types.ts @@ -95,6 +95,8 @@ export interface ToolCallStartParams { args?: Record; /** Status of the tool call */ status?: 'pending' | 'in_progress' | 'completed' | 'failed'; + /** Transient phase recognized by clients that support tool preparation. */ + phase?: 'preparing'; /** Optional subagent metadata */ subagentMeta?: SubagentMeta; /** Server-side timestamp (ISO string or ms) for message ordering */ diff --git a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts index 450c8b65bac..25e22df8b12 100644 --- a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts +++ b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.test.ts @@ -2954,6 +2954,272 @@ describe('AnthropicContentGenerator', () => { }); describe('generateContentStream', () => { + it('emits tool preparation metadata before the complete function call', async () => { + const { AnthropicContentGenerator } = await importGenerator(); + const { getToolCallPreparations } = await import( + '../tool-call-preparation.js' + ); + let stopEventReached = false; + anthropicState.createImpl.mockResolvedValue( + (async function* toolUseStream() { + yield { + type: 'message_start', + message: { + id: 'msg-1', + model: 'claude-test', + usage: { input_tokens: 1 }, + }, + }; + yield { + type: 'content_block_start', + index: 0, + content_block: { + type: 'tool_use', + id: 'call-1', + name: 'read_file', + input: {}, + }, + }; + yield { + type: 'content_block_delta', + index: 0, + delta: { + type: 'input_json_delta', + partial_json: '{"file_path":', + }, + }; + yield { + type: 'content_block_delta', + index: 0, + delta: { + type: 'input_json_delta', + partial_json: '"a.sql"}', + }, + }; + yield { + get type() { + stopEventReached = true; + return 'content_block_stop' as const; + }, + index: 0, + }; + yield { + type: 'message_delta', + delta: { stop_reason: 'tool_use' }, + usage: { output_tokens: 5 }, + }; + yield { type: 'message_stop' }; + })(), + ); + + const generator = new AnthropicContentGenerator( + { + model: 'claude-test', + apiKey: 'test-key', + timeout: 10_000, + maxRetries: 2, + samplingParams: { max_tokens: 100 }, + schemaCompliance: 'auto', + }, + mockConfig, + ); + + const stream = await generator.generateContentStream({ + model: 'models/ignored', + contents: 'Hello', + } as unknown as GenerateContentParameters); + const chunks: GenerateContentResponse[] = []; + let stopReachedWhenFunctionCallEmitted: boolean | undefined; + for await (const chunk of stream) { + chunks.push(chunk); + if (chunk.functionCalls) { + stopReachedWhenFunctionCallEmitted = stopEventReached; + } + } + + expect(getToolCallPreparations(chunks[0]!)).toEqual([ + { callId: 'call-1', toolName: 'read_file' }, + ]); + const functionCallChunks = chunks.filter((chunk) => chunk.functionCalls); + expect(functionCallChunks).toHaveLength(1); + expect(stopReachedWhenFunctionCallEmitted).toBe(true); + expect(functionCallChunks[0]!.functionCalls).toEqual([ + { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.sql' }, + }, + ]); + }); + + it('emits preparations before both function calls in a multi-tool stream', async () => { + const { AnthropicContentGenerator } = await importGenerator(); + const { getToolCallPreparations } = await import( + '../tool-call-preparation.js' + ); + anthropicState.createImpl.mockResolvedValue( + (async function* multiToolStream() { + yield { + type: 'content_block_start', + index: 0, + content_block: { + type: 'tool_use', + id: 'call-1', + name: 'read_file', + input: {}, + }, + }; + yield { + type: 'content_block_start', + index: 1, + content_block: { + type: 'tool_use', + id: 'call-2', + name: 'run_shell_command', + input: {}, + }, + }; + yield { + type: 'content_block_delta', + index: 0, + delta: { + type: 'input_json_delta', + partial_json: '{"file_path":"a.sql"}', + }, + }; + yield { type: 'content_block_stop', index: 0 }; + yield { + type: 'content_block_delta', + index: 1, + delta: { + type: 'input_json_delta', + partial_json: '{"command":"pwd"}', + }, + }; + yield { type: 'content_block_stop', index: 1 }; + })(), + ); + const generator = new AnthropicContentGenerator( + { + model: 'claude-test', + apiKey: 'test-key', + timeout: 10_000, + maxRetries: 2, + samplingParams: { max_tokens: 100 }, + schemaCompliance: 'auto', + }, + mockConfig, + ); + + const stream = await generator.generateContentStream({ + model: 'models/ignored', + contents: 'Hello', + } as unknown as GenerateContentParameters); + const chunks: GenerateContentResponse[] = []; + for await (const chunk of stream) { + chunks.push(chunk); + } + + const preparations = chunks.flatMap((chunk, index) => + getToolCallPreparations(chunk).map((preparation) => ({ + ...preparation, + index, + })), + ); + const functionCalls = chunks.flatMap((chunk, index) => + (chunk.functionCalls ?? []).map((functionCall) => ({ + ...functionCall, + index, + })), + ); + expect(preparations).toEqual([ + { callId: 'call-1', toolName: 'read_file', index: 0 }, + { callId: 'call-2', toolName: 'run_shell_command', index: 1 }, + ]); + expect(functionCalls).toEqual([ + { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.sql' }, + index: 2, + }, + { + id: 'call-2', + name: 'run_shell_command', + args: { command: 'pwd' }, + index: 3, + }, + ]); + }); + + it.each([ + { label: 'id is missing', contentBlock: { name: 'read_file' } }, + { label: 'name is missing', contentBlock: { id: 'call-1' } }, + { + label: 'id is empty', + contentBlock: { id: '', name: 'read_file' }, + }, + { + label: 'name is empty', + contentBlock: { id: 'call-1', name: '' }, + }, + { + label: 'id is not a string', + contentBlock: { id: 42, name: 'read_file' }, + }, + { + label: 'name is not a string', + contentBlock: { id: 'call-1', name: 42 }, + }, + ])( + 'does not emit tool preparation metadata when $label', + async ({ contentBlock }) => { + const { AnthropicContentGenerator } = await importGenerator(); + const { getToolCallPreparations } = await import( + '../tool-call-preparation.js' + ); + anthropicState.createImpl.mockResolvedValue( + (async function* toolUseStream() { + yield { + type: 'content_block_start', + index: 0, + content_block: { + type: 'tool_use', + ...contentBlock, + input: {}, + }, + }; + yield { type: 'content_block_stop', index: 0 }; + })(), + ); + + const generator = new AnthropicContentGenerator( + { + model: 'claude-test', + apiKey: 'test-key', + timeout: 10_000, + maxRetries: 2, + samplingParams: { max_tokens: 100 }, + schemaCompliance: 'auto', + }, + mockConfig, + ); + + const stream = await generator.generateContentStream({ + model: 'models/ignored', + contents: 'Hello', + } as unknown as GenerateContentParameters); + const chunks: GenerateContentResponse[] = []; + for await (const chunk of stream) { + chunks.push(chunk); + } + + expect( + chunks.every((chunk) => getToolCallPreparations(chunk).length === 0), + ).toBe(true); + }, + ); + it('redacts proxy credentials from stream creation errors', async () => { const { AnthropicContentGenerator } = await importGenerator(); anthropicState.createImpl.mockRejectedValue( @@ -3203,8 +3469,9 @@ describe('AnthropicContentGenerator', () => { thoughtSignature: 'abc', }); - // Tool call chunk. - expect(chunks[3]?.candidates?.[0]?.content?.parts?.[0]).toEqual({ + // The preparation-only chunk precedes the complete tool call chunk. + expect(chunks[3]?.functionCalls).toBeUndefined(); + expect(chunks[4]?.candidates?.[0]?.content?.parts?.[0]).toEqual({ functionCall: { id: 't1', name: 'tool', args: { x: 1 } }, }); diff --git a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts index dea0b258099..51dac3b4f6e 100644 --- a/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts +++ b/packages/core/src/core/anthropicContentGenerator/anthropicContentGenerator.ts @@ -48,6 +48,7 @@ import { reconcileMaxTokens, parsePositiveIntegerEnvValue, } from '../tokenLimits.js'; +import { setToolCallPreparations } from '../tool-call-preparation.js'; const debugLogger = createDebugLogger('ANTHROPIC'); @@ -1055,18 +1056,20 @@ export class AnthropicContentGenerator implements ContentGenerator { case 'content_block_start': { const index = event.index ?? 0; const type = String(event.content_block.type || 'text'); + const id = + 'id' in event.content_block ? event.content_block.id : undefined; + const name = + 'name' in event.content_block + ? event.content_block.name + : undefined; const initialInput = type === 'tool_use' && 'input' in event.content_block ? JSON.stringify(event.content_block.input) : ''; blocks.set(index, { type, - id: - 'id' in event.content_block ? event.content_block.id : undefined, - name: - 'name' in event.content_block - ? event.content_block.name - : undefined, + id, + name, inputJson: initialInput !== '{}' ? initialInput : '', signature: type === 'thinking' && @@ -1075,6 +1078,18 @@ export class AnthropicContentGenerator implements ContentGenerator { ? event.content_block.signature : '', }); + if ( + type === 'tool_use' && + typeof id === 'string' && + id.length > 0 && + typeof name === 'string' && + name.length > 0 + ) { + const chunk = this.buildGeminiChunk(undefined, messageId, model); + setToolCallPreparations(chunk, [{ callId: id, toolName: name }]); + collectedResponses.push(chunk); + yield chunk; + } break; } case 'content_block_delta': { diff --git a/packages/core/src/core/geminiChat.test.ts b/packages/core/src/core/geminiChat.test.ts index 10e5d41aa57..0f63d4969cf 100644 --- a/packages/core/src/core/geminiChat.test.ts +++ b/packages/core/src/core/geminiChat.test.ts @@ -40,6 +40,10 @@ import { } from '../services/tokenEstimation.js'; import { SYSTEM_REMINDER_OPEN } from '../utils/environmentContext.js'; import { SessionStartSource } from '../hooks/types.js'; +import { + getToolCallPreparations, + setToolCallPreparations, +} from './tool-call-preparation.js'; // Mock fs module to prevent actual file system operations during tests const mockFileSystem = new Map(); @@ -1001,6 +1005,88 @@ describe('GeminiChat', async () => { ).resolves.not.toThrow(); }); + it('uses the normalized function call ID for preparation metadata', async () => { + chat.setHistory([ + { role: 'user', parts: [{ text: 'first request' }] }, + { + role: 'model', + parts: [ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'a.txt' }, + }, + }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + id: 'call-1', + name: 'read_file', + response: { output: 'first result' }, + }, + }, + ], + }, + ]); + const preparationResponse = { + candidates: [{ content: { role: 'model', parts: [] } }], + } as unknown as GenerateContentResponse; + setToolCallPreparations(preparationResponse, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + (async function* () { + yield preparationResponse; + yield { + candidates: [ + { + content: { + role: 'model', + parts: [ + { + functionCall: { + id: 'call-1', + name: 'read_file', + args: { file_path: 'b.txt' }, + }, + }, + ], + }, + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'second request' }, + 'prompt-normalized-preparation-id', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + const preparation = events + .filter((event) => event.type === StreamEventType.CHUNK) + .flatMap((event) => getToolCallPreparations(event.value))[0]; + const functionCall = events.find( + (event) => + event.type === StreamEventType.CHUNK && + event.value.functionCalls?.length, + ); + expect(preparation?.callId).toBe('call-1__qwen_dup_2'); + expect( + functionCall?.type === StreamEventType.CHUNK + ? functionCall.value.functionCalls?.[0]?.id + : undefined, + ).toBe(preparation?.callId); + }); + it('persists partial assistant turn when stream throws after a tool_use chunk', async () => { // Weak-network scenario: Anthropic-compatible providers emit the // `functionCall` part on `content_block_stop`; the SSE may then drop @@ -5191,7 +5277,7 @@ describe('GeminiChat', async () => { } }); - it('tries the next fallback when a fallback fails before emitting output', async () => { + it('tries the next fallback when a fallback emits only preparation metadata', async () => { vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ authType: AuthType.USE_GEMINI, model: 'test-model', @@ -5244,7 +5330,18 @@ describe('GeminiChat', async () => { vi.mocked( mockContentGenerator.generateContentStream, ).mockRejectedValueOnce(capacityError); - fallbackAGenerateContentStream.mockRejectedValueOnce(capacityError); + const preparationResponse = { + candidates: [{ content: { parts: [] } }], + } as unknown as GenerateContentResponse; + setToolCallPreparations(preparationResponse, [ + { callId: 'call-fallback-a', toolName: 'read_file' }, + ]); + fallbackAGenerateContentStream.mockResolvedValueOnce( + (async function* () { + yield preparationResponse; + throw capacityError; + })(), + ); fallbackBGenerateContentStream.mockResolvedValueOnce( (async function* () { yield { @@ -6003,6 +6100,136 @@ describe('GeminiChat', async () => { ).toBe(true); }); + it('retries a transport stream error after yielding only tool preparation metadata', async () => { + vi.useFakeTimers(); + try { + const transportError = Object.assign(new TypeError('terminated'), { + cause: Object.assign(new Error('other side closed'), { + code: 'UND_ERR_SOCKET', + }), + }); + const preparationResponse = { + candidates: [{ content: { parts: [] } }], + } as unknown as GenerateContentResponse; + setToolCallPreparations(preparationResponse, [ + { callId: 'call-preparing', toolName: 'read_file' }, + ]); + + vi.mocked(mockContentGenerator.generateContentStream) + .mockResolvedValueOnce( + (async function* () { + yield preparationResponse; + throw transportError; + })(), + ) + .mockResolvedValueOnce( + (async function* () { + yield { + candidates: [ + { + content: { + parts: [{ text: 'Recovered after preparation' }], + }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-transport-after-preparation', + ); + const events = await collectStreamWithFakeTimers(stream, 5_000); + + expect( + mockContentGenerator.generateContentStream, + ).toHaveBeenCalledTimes(2); + expect( + events.filter((event) => event.type === StreamEventType.RETRY), + ).toHaveLength(1); + } finally { + vi.useRealTimers(); + } + }); + + it('falls back after yielding only tool preparation metadata', async () => { + vi.mocked(mockConfig.getContentGeneratorConfig).mockReturnValue({ + authType: AuthType.USE_GEMINI, + model: 'test-model', + maxRetries: 0, + }); + vi.mocked(mockConfig.getModelFallbacks).mockReturnValue([ + 'fallback-model', + ]); + + const fallbackGenerateContentStream = vi.fn().mockResolvedValue( + (async function* () { + yield { + candidates: [ + { + content: { parts: [{ text: 'Recovered with fallback' }] }, + finishReason: 'STOP', + }, + ], + } as unknown as GenerateContentResponse; + })(), + ); + const resolveForModel = vi.fn().mockResolvedValue({ + contentGenerator: { + generateContent: vi.fn(), + generateContentStream: fallbackGenerateContentStream, + countTokens: vi.fn(), + embedContent: vi.fn(), + batchEmbedContents: vi.fn(), + useSummarizedThinking: vi.fn().mockReturnValue(false), + } as unknown as ContentGenerator, + retryAuthType: AuthType.USE_GEMINI, + retryErrorCodes: undefined, + model: 'fallback-model', + }); + vi.mocked(mockConfig.getBaseLlmClient).mockReturnValue({ + resolveForModel, + } as unknown as ReturnType); + + const capacityError = Object.assign( + new StreamContentError( + '{"error":{"code":"429","message":"Throttling"}}', + ), + { status: 429 }, + ); + const preparationResponse = { + candidates: [{ content: { parts: [] } }], + } as unknown as GenerateContentResponse; + setToolCallPreparations(preparationResponse, [ + { callId: 'call-fallback', toolName: 'read_file' }, + ]); + vi.mocked(mockContentGenerator.generateContentStream).mockResolvedValue( + (async function* () { + yield preparationResponse; + throw capacityError; + })(), + ); + + const stream = await chat.sendMessageStream( + 'test-model', + { message: 'test' }, + 'prompt-fallback-after-preparation', + ); + const events: StreamEvent[] = []; + for await (const event of stream) events.push(event); + + expect(resolveForModel).toHaveBeenCalledWith('fallback-model', { + failClosed: true, + }); + expect(fallbackGenerateContentStream).toHaveBeenCalledOnce(); + expect( + events.filter((event) => event.type === StreamEventType.MODEL_FALLBACK), + ).toHaveLength(1); + }); + it('classifies every allow-listed stream transport code as retryable transport', () => { // Drift guard: the stream allow-list is a hand-curated subset of the // classifier's transport codes. If a code is renamed/removed there, or diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index e8c8c8c6613..951dc26a7f9 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -101,13 +101,29 @@ import { RETRYABLE_STREAM_TRANSPORT_CODES } from './stream-transport-retry.js'; import { collectToolCallIdsFromHistory, normalizeModelToolCallIds, + reserveModelToolCallId, } from './toolCallIdUtils.js'; +import { + getToolCallPreparations, + setToolCallPreparations, +} from './tool-call-preparation.js'; import { InvalidStreamError } from './invalid-stream-error.js'; export { InvalidStreamError }; const debugLogger = createDebugLogger('QWEN_CODE_CHAT'); +function isToolCallPreparationOnly(response: GenerateContentResponse): boolean { + if (getToolCallPreparations(response).length === 0) return false; + + const hasCandidateOutput = response.candidates?.some( + (candidate) => + Boolean(candidate.finishReason) || + (candidate.content?.parts?.length ?? 0) > 0, + ); + return !hasCandidateOutput && !response.usageMetadata; +} + function syncFunctionCallsField( response: GenerateContentResponse, parts: readonly Part[], @@ -2248,8 +2264,10 @@ export class GeminiChat { lastFinishReason = undefined; for await (const chunk of stream) { - streamYieldedChunk = true; - streamYieldedAnyChunk = true; + if (!isToolCallPreparationOnly(chunk)) { + streamYieldedChunk = true; + streamYieldedAnyChunk = true; + } const fr = chunk.candidates?.[0]?.finishReason; if (fr) lastFinishReason = fr; yield { type: StreamEventType.CHUNK, value: chunk }; @@ -3004,8 +3022,13 @@ export class GeminiChat { fallbackRetryAuthType, fallbackRetryErrorCodes, )) { - currentFallbackYieldedAnyChunk = true; - fallbackStreamYieldedAnyChunk = true; + const emittedUserVisibleOutput = + event.type !== StreamEventType.CHUNK || + !isToolCallPreparationOnly(event.value); + if (emittedUserVisibleOutput) { + currentFallbackYieldedAnyChunk = true; + fallbackStreamYieldedAnyChunk = true; + } yield event; } @@ -3562,6 +3585,7 @@ export class GeminiChat { const allModelParts: Part[] = []; const usedToolCallIds = collectToolCallIdsFromHistory(this.history); const rawToolCallIdsInCurrentTurn = new Set(); + const reservedToolCallIds = new Map(); let usageMetadata: GenerateContentResponseUsageMetadata | undefined; let coercedUsage: | { @@ -3587,6 +3611,21 @@ export class GeminiChat { try { for await (const chunk of streamResponse) { + const preparations = getToolCallPreparations(chunk); + if (preparations.length > 0) { + setToolCallPreparations( + chunk, + preparations.map((preparation) => ({ + ...preparation, + callId: reserveModelToolCallId( + preparation.callId, + usedToolCallIds, + reservedToolCallIds, + ), + })), + ); + } + // Use ||= to avoid later usage-only chunks (no candidates) overwriting // a finishReason that was already seen in an earlier chunk. hasFinishReason ||= @@ -3615,6 +3654,7 @@ export class GeminiChat { content.parts, usedToolCallIds, rawToolCallIdsInCurrentTurn, + reservedToolCallIds, ); syncFunctionCallsField(chunk, content.parts); diff --git a/packages/core/src/core/openaiContentGenerator/converter.test.ts b/packages/core/src/core/openaiContentGenerator/converter.test.ts index 79e9d427106..c59b1c78a39 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.test.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.test.ts @@ -21,6 +21,7 @@ import { } from '@google/genai'; import type OpenAI from 'openai'; import { convertToFunctionResponse } from '../coreToolScheduler.js'; +import { getToolCallPreparations } from '../tool-call-preparation.js'; import { isOpenAIReasoningThoughtPart } from '../../utils/thoughtUtils.js'; describe('OpenAIContentConverter', () => { @@ -5602,6 +5603,357 @@ describe('Truncated tool call detection in streaming', () => { ); } + it('emits tool preparation metadata before the complete function call', () => { + const context = createStreamingRequestContext(); + const opener = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-open', + created: 100, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: 'call-1', + type: 'function', + function: { name: 'read_file', arguments: '' }, + }, + ], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const args = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-args', + created: 101, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + type: 'function', + function: { arguments: '{"file_path":"a.sql"}' }, + }, + ], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const finish = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-finish', + created: 102, + model: 'test-model', + choices: [ + { + index: 0, + delta: {}, + finish_reason: 'tool_calls', + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + + expect(getToolCallPreparations(opener)).toEqual([ + { callId: 'call-1', toolName: 'read_file' }, + ]); + expect(getToolCallPreparations(args)).toEqual([]); + expect(getToolCallPreparations(finish)).toEqual([]); + expect(opener.functionCalls).toBeUndefined(); + expect(finish.functionCalls).toEqual([ + { id: 'call-1', name: 'read_file', args: { file_path: 'a.sql' } }, + ]); + }); + + it('does not duplicate tool preparation metadata for a replayed opener', () => { + const context = createStreamingRequestContext(); + const opener = { + object: 'chat.completion.chunk', + id: 'chunk-open', + created: 100, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: 'call-1', + type: 'function', + function: { name: 'read_file', arguments: '' }, + }, + ], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk; + + const first = converter.convertOpenAIChunkToGemini(opener, context); + const replay = converter.convertOpenAIChunkToGemini(opener, context); + + expect(getToolCallPreparations(first)).toEqual([ + { callId: 'call-1', toolName: 'read_file' }, + ]); + expect(getToolCallPreparations(replay)).toEqual([]); + }); + + it('emits preparation after split identity deltas using the remapped parser index', () => { + const context = createStreamingRequestContext(); + + const firstCall = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-first-call', + created: 100, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: 'call-1', + type: 'function', + function: { + name: 'read_file', + arguments: '{"file_path":"a.sql"}', + }, + }, + ], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const secondCallId = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-second-id', + created: 101, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [{ index: 0, id: 'call-2', type: 'function' }], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const secondCallName = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-second-name', + created: 102, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + type: 'function', + function: { + name: 'write_file', + arguments: '{"file_path":"b.sql"}', + }, + }, + ], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const thirdCallName = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-third-name', + created: 103, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + type: 'function', + function: { + name: 'delete_file', + arguments: '', + }, + }, + ], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const thirdCallArguments = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-third-arguments', + created: 104, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + type: 'function', + function: { arguments: '{"file_path":"c.sql"}' }, + }, + ], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const thirdCallId = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-third-id', + created: 105, + model: 'test-model', + choices: [ + { + index: 0, + delta: { + tool_calls: [{ index: 0, id: 'call-3', type: 'function' }], + }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + const finish = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-finish', + created: 106, + model: 'test-model', + choices: [ + { + index: 0, + delta: {}, + finish_reason: 'tool_calls', + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + context, + ); + + expect(getToolCallPreparations(firstCall)).toEqual([ + { callId: 'call-1', toolName: 'read_file' }, + ]); + expect(getToolCallPreparations(secondCallId)).toEqual([]); + expect(getToolCallPreparations(secondCallName)).toEqual([ + { callId: 'call-2', toolName: 'write_file' }, + ]); + expect(getToolCallPreparations(thirdCallName)).toEqual([]); + expect(getToolCallPreparations(thirdCallArguments)).toEqual([]); + expect(getToolCallPreparations(thirdCallId)).toEqual([ + { callId: 'call-3', toolName: 'delete_file' }, + ]); + expect(firstCall.functionCalls).toBeUndefined(); + expect(secondCallId.functionCalls).toBeUndefined(); + expect(secondCallName.functionCalls).toBeUndefined(); + expect(thirdCallName.functionCalls).toBeUndefined(); + expect(thirdCallArguments.functionCalls).toBeUndefined(); + expect(thirdCallId.functionCalls).toBeUndefined(); + expect(finish.functionCalls).toEqual([ + { id: 'call-1', name: 'read_file', args: { file_path: 'a.sql' } }, + { id: 'call-2', name: 'write_file', args: { file_path: 'b.sql' } }, + { id: 'call-3', name: 'delete_file', args: { file_path: 'c.sql' } }, + ]); + }); + + it.each([ + { + label: 'call ID is missing', + toolCall: { + index: 0, + type: 'function' as const, + function: { name: 'read_file', arguments: '' }, + }, + }, + { + label: 'tool name is missing', + toolCall: { + index: 0, + id: 'call-1', + type: 'function' as const, + function: { arguments: '' }, + }, + }, + ])('does not emit tool preparation metadata when $label', ({ toolCall }) => { + const response = converter.convertOpenAIChunkToGemini( + { + object: 'chat.completion.chunk', + id: 'chunk-open', + created: 100, + model: 'test-model', + choices: [ + { + index: 0, + delta: { tool_calls: [toolCall] }, + finish_reason: null, + logprobs: null, + }, + ], + } as unknown as OpenAI.Chat.ChatCompletionChunk, + createStreamingRequestContext(), + ); + + expect(getToolCallPreparations(response)).toEqual([]); + }); + it('should override finishReason to MAX_TOKENS when tool call JSON is truncated and provider reports "stop"', () => { // Simulate: write_file call truncated mid-JSON, provider says "stop" const result = feedToolCallChunks( diff --git a/packages/core/src/core/openaiContentGenerator/converter.ts b/packages/core/src/core/openaiContentGenerator/converter.ts index b1c7a81f1c6..e13453f36bf 100644 --- a/packages/core/src/core/openaiContentGenerator/converter.ts +++ b/packages/core/src/core/openaiContentGenerator/converter.ts @@ -28,6 +28,10 @@ import { convertSchema, type SchemaComplianceMode, } from '../../utils/schemaConverter.js'; +import { + setToolCallPreparations, + type ToolCallPreparation, +} from '../tool-call-preparation.js'; import { InvalidStreamError } from '../invalid-stream-error.js'; const debugLogger = createDebugLogger('CONVERTER'); @@ -1224,6 +1228,7 @@ export function convertOpenAIChunkToGemini( ): GenerateContentResponse { const choice = chunk.choices?.[0]; const response = new GenerateContentResponse(); + const preparations: ToolCallPreparation[] = []; const toolCallParser = requestContext.toolCallParser; if (!toolCallParser) { throw new Error( @@ -1360,21 +1365,29 @@ export function convertOpenAIChunkToGemini( const index = toolCall.index ?? 0; // Process the tool call chunk through the streaming parser - if (toolCall.function?.arguments) { - toolCallParser.addChunk( - index, - toolCall.function.arguments, - toolCall.id, - toolCall.function.name, - ); - } else { - // Handle metadata-only chunks (id and/or name without arguments) - toolCallParser.addChunk( - index, - '', // Empty chunk for metadata-only updates - toolCall.id, - toolCall.function?.name, - ); + const parseResult = toolCall.function?.arguments + ? toolCallParser.addChunk( + index, + toolCall.function.arguments, + toolCall.id, + toolCall.function.name, + ) + : toolCallParser.addChunk( + index, + '', // Empty chunk for metadata-only updates + toolCall.id, + toolCall.function?.name, + ); + + const { id: callId, name: toolName } = toolCallParser.getToolCallMeta( + parseResult.actualIndex ?? index, + ); + if (callId && toolName) { + const emitted = (requestContext.preparedToolCallIds ??= new Set()); + if (!emitted.has(callId)) { + emitted.add(callId); + preparations.push({ callId, toolName }); + } } } } @@ -1535,6 +1548,10 @@ export function convertOpenAIChunkToGemini( }; } + if (preparations.length > 0) { + setToolCallPreparations(response, preparations); + } + return response; } diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts index 7e5230d6e1c..0252835d39d 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.test.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.test.ts @@ -27,6 +27,7 @@ import { MAX_STREAM_IDLE_TIMEOUT_MS, QWEN_STREAM_IDLE_TIMEOUT_MS_ENV, } from './constants.js'; +import { setToolCallPreparations } from '../tool-call-preparation.js'; // Mock dependencies vi.mock('./converter.js', () => ({ @@ -1652,6 +1653,48 @@ describe('ContentGenerationPipeline', () => { expect(results[0]).toBe(mockValidResponse); }); + it('should preserve an otherwise empty response with tool preparation metadata', async () => { + const request: GenerateContentParameters = { + model: 'test-model', + contents: [{ parts: [{ text: 'Hello' }], role: 'user' }], + }; + const mockChunk = { + id: 'chunk-tool-opener', + choices: [{ delta: { tool_calls: [] }, finish_reason: null }], + } as unknown as OpenAI.Chat.ChatCompletionChunk; + const mockStream = { + async *[Symbol.asyncIterator]() { + yield mockChunk; + }, + }; + const preparationResponse = new GenerateContentResponse(); + preparationResponse.candidates = [ + { content: { parts: [], role: 'model' } }, + ]; + setToolCallPreparations(preparationResponse, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + + (mockConverter.convertGeminiRequestToOpenAI as Mock).mockReturnValue([]); + (mockConverter.convertOpenAIChunkToGemini as Mock).mockReturnValue( + preparationResponse, + ); + (mockClient.chat.completions.create as Mock).mockResolvedValue( + mockStream, + ); + + const resultGenerator = await pipeline.executeStream( + request, + 'test-prompt-id', + ); + const results = []; + for await (const result of resultGenerator) { + results.push(result); + } + + expect(results).toEqual([preparationResponse]); + }); + it('should handle streaming errors and reset tool calls', async () => { // Arrange const request: GenerateContentParameters = { diff --git a/packages/core/src/core/openaiContentGenerator/pipeline.ts b/packages/core/src/core/openaiContentGenerator/pipeline.ts index f0048e5b8f9..60e0d3d0dad 100644 --- a/packages/core/src/core/openaiContentGenerator/pipeline.ts +++ b/packages/core/src/core/openaiContentGenerator/pipeline.ts @@ -27,6 +27,7 @@ import { QWEN_STREAM_IDLE_TIMEOUT_MS_ENV, } from './constants.js'; import { createDebugLogger } from '../../utils/debugLogger.js'; +import { getToolCallPreparations } from '../tool-call-preparation.js'; import { InvalidStreamError } from '../invalid-stream-error.js'; const debugLogger = createDebugLogger('OPENAI_PIPELINE'); @@ -503,7 +504,9 @@ export class ContentGenerationPipeline { if ( response.candidates?.[0]?.content?.parts?.length === 0 && !response.candidates?.[0]?.finishReason && - !response.usageMetadata + !response.usageMetadata && + // Preparation-only responses must reach ACP before arguments complete. + getToolCallPreparations(response).length === 0 ) { continue; } diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts index f9485a2d9d5..994ee7554be 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.test.ts @@ -944,6 +944,44 @@ describe('StreamingToolCallParser', () => { }); describe('Complex collision scenarios', () => { + it('does not append continuation fragments to a completed remapped slot', () => { + parser.addChunk(0, '{"first":true}', 'call_1', 'function1'); + const remapped = parser.addChunk( + 0, + '{"second":true}', + undefined, + 'function2', + ); + + expect(remapped.actualIndex).toBe(1); + expect(remapped.complete).toBe(true); + + const continuation = parser.addChunk(0, '{"third":true}'); + + expect(continuation.actualIndex).not.toBe(remapped.actualIndex); + expect(parser.getBuffer(remapped.actualIndex!)).toBe('{"second":true}'); + }); + + it('associates a late stable ID with its completed remapped slot', () => { + parser.addChunk(0, '{"first":true}', 'call_1', 'function1'); + const remapped = parser.addChunk( + 0, + '{"second":true}', + undefined, + 'function2', + ); + + const identified = parser.addChunk(0, '', 'call_2'); + + expect(identified.actualIndex).toBe(remapped.actualIndex); + expect(parser.getCompletedToolCalls()).toContainEqual({ + id: 'call_2', + name: 'function2', + args: { second: true }, + index: remapped.actualIndex, + }); + }); + it('should handle rapid tool call switching at same index', () => { // Rapid switching between different tool calls at index 0 parser.addChunk(0, '{"step1":', 'call_1', 'function1'); diff --git a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts index 2c962397b89..abad9005326 100644 --- a/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts +++ b/packages/core/src/core/openaiContentGenerator/streamingToolCallParser.ts @@ -13,6 +13,8 @@ const debugLogger = createDebugLogger('STREAMING_TOOL_CALL_PARSER'); * Type definition for the result of parsing a JSON chunk in tool calls */ export interface ToolCallParseResult { + /** Parser index that received this chunk after collision remapping. */ + actualIndex?: number; /** Whether the JSON parsing is complete */ complete: boolean; /** The parsed JSON value (only present when complete is true) */ @@ -47,6 +49,8 @@ export class StreamingToolCallParser { private namelessToolCallIndices = new Set(); /** Map from tool call ID to actual index used for storage */ private idToIndexMap: Map = new Map(); + /** Remapped slots awaiting a stable ID from a later chunk. */ + private pendingIndexRemaps: Map = new Map(); /** Counter for generating new indices when collisions occur */ private nextAvailableIndex: number = 0; @@ -82,7 +86,10 @@ export class StreamingToolCallParser { let actualIndex = index; const isKnownId = Boolean(id && this.idToIndexMap.has(id)); - const isNameOnlyDelta = Boolean(name && chunk.length === 0); + const existingName = this.toolCallMeta.get(index)?.name; + const isNameOnlyDelta = Boolean( + name && chunk.length === 0 && (!existingName || existingName === name), + ); // Handle tool call ID mapping for collision detection if (id) { @@ -90,6 +97,11 @@ export class StreamingToolCallParser { if (this.idToIndexMap.has(id)) { // We've seen this ID before, use the existing mapped index actualIndex = this.idToIndexMap.get(id)!; + } else if (this.pendingIndexRemaps.has(index)) { + // Some providers stream name or arguments before the stable ID. + actualIndex = this.pendingIndexRemaps.get(index)!; + this.pendingIndexRemaps.delete(index); + this.idToIndexMap.set(id, actualIndex); } else { // New tool call ID // Check if the requested index is already occupied by a different complete tool call @@ -131,7 +143,20 @@ export class StreamingToolCallParser { // No ID provided - this is a continuation chunk // Try to find which tool call this belongs to based on the index // Look for an existing tool call at this index that's not complete - if (this.buffers.has(index)) { + if (this.pendingIndexRemaps.has(index)) { + // Keep later argument chunks on the remapped slot until the ID arrives. + actualIndex = this.pendingIndexRemaps.get(index)!; + const existingBuffer = this.buffers.get(actualIndex)!; + const existingDepth = this.depths.get(actualIndex)!; + if (existingDepth === 0 && existingBuffer.trim()) { + try { + JSON.parse(existingBuffer); + actualIndex = this.findMostRecentIncompleteIndex(); + } catch { + // The remapped buffer is still incomplete; append below. + } + } + } else if (this.buffers.has(index)) { const existingBuffer = this.buffers.get(index)!; const existingDepth = this.depths.get(index)!; @@ -173,7 +198,10 @@ export class StreamingToolCallParser { } else { this.namelessToolCallIndices.delete(actualIndex); } - return { complete: false }; + if (!meta.id && actualIndex !== index) { + this.pendingIndexRemaps.set(index, actualIndex); + } + return { actualIndex, complete: false }; } if (isKnownId && currentDepth === 0) { @@ -183,7 +211,7 @@ export class StreamingToolCallParser { debugLogger.debug( `Ignoring replay chunk for completed toolCall id=${id}`, ); - return { complete: false }; + return { actualIndex, complete: false }; } catch { // Not complete yet; append the incoming chunk below. } @@ -193,6 +221,9 @@ export class StreamingToolCallParser { // Update metadata if (id) meta.id = id; if (name) meta.name = name; + if (!meta.id && actualIndex !== index) { + this.pendingIndexRemaps.set(index, actualIndex); + } // Get current state for the actual index const currentInString = this.inStrings.get(actualIndex)!; @@ -236,13 +267,14 @@ export class StreamingToolCallParser { try { // Standard JSON parsing attempt const parsed = JSON.parse(newBuffer); - return { complete: true, value: parsed }; + return { actualIndex, complete: true, value: parsed }; } catch (e) { // Intelligent repair: try auto-closing unclosed strings if (inString) { try { const repaired = JSON.parse(newBuffer + '"'); return { + actualIndex, complete: true, value: repaired, repaired: true, @@ -252,6 +284,7 @@ export class StreamingToolCallParser { } } return { + actualIndex, complete: false, error: e instanceof Error ? e : new Error(String(e)), }; @@ -259,7 +292,7 @@ export class StreamingToolCallParser { } // JSON structure is incomplete, continue accumulating chunks - return { complete: false }; + return { actualIndex, complete: false }; } /** @@ -460,6 +493,11 @@ export class StreamingToolCallParser { this.escapes.set(index, false); this.toolCallMeta.set(index, {}); this.namelessToolCallIndices.delete(index); + for (const [providerIndex, actualIndex] of this.pendingIndexRemaps) { + if (providerIndex === index || actualIndex === index) { + this.pendingIndexRemaps.delete(providerIndex); + } + } } /** @@ -477,6 +515,7 @@ export class StreamingToolCallParser { this.toolCallMeta.clear(); this.namelessToolCallIndices.clear(); this.idToIndexMap.clear(); + this.pendingIndexRemaps.clear(); this.nextAvailableIndex = 0; } diff --git a/packages/core/src/core/openaiContentGenerator/types.ts b/packages/core/src/core/openaiContentGenerator/types.ts index 3e638cd4da3..31c37bae33b 100644 --- a/packages/core/src/core/openaiContentGenerator/types.ts +++ b/packages/core/src/core/openaiContentGenerator/types.ts @@ -79,6 +79,8 @@ export interface RequestContext { * emitted after the reasoning thought if no tagged thought appears. */ pendingContentParts?: Part[]; + /** Tool IDs whose preparing metadata has already been emitted in this stream. */ + preparedToolCallIds?: Set; pendingUntrustedResponseParts?: Part[]; hasStructuredReasoningContent?: boolean; hasThinkingTagInReasoning?: boolean; diff --git a/packages/core/src/core/tool-call-preparation.test.ts b/packages/core/src/core/tool-call-preparation.test.ts new file mode 100644 index 00000000000..6a068597ca7 --- /dev/null +++ b/packages/core/src/core/tool-call-preparation.test.ts @@ -0,0 +1,39 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { GenerateContentResponse } from '@google/genai'; +import { describe, expect, it } from 'vitest'; +import { + getToolCallPreparations, + setToolCallPreparations, +} from './tool-call-preparation.js'; + +describe('tool-call preparation metadata', () => { + it('attaches metadata without adding enumerable response fields', () => { + const response = new GenerateContentResponse(); + const preparations = [{ callId: 'call-1', toolName: 'read_file' }]; + + setToolCallPreparations(response, preparations); + + expect(getToolCallPreparations(response)).toEqual(preparations); + expect(JSON.stringify(response)).not.toContain('toolCallPreparations'); + }); + + it('returns an empty list when no metadata is attached', () => { + expect(getToolCallPreparations(new GenerateContentResponse())).toEqual([]); + }); + + it('clears attached metadata when set to an empty list', () => { + const response = new GenerateContentResponse(); + setToolCallPreparations(response, [ + { callId: 'call-1', toolName: 'read_file' }, + ]); + + setToolCallPreparations(response, []); + + expect(getToolCallPreparations(response)).toEqual([]); + }); +}); diff --git a/packages/core/src/core/tool-call-preparation.ts b/packages/core/src/core/tool-call-preparation.ts new file mode 100644 index 00000000000..041b9dedada --- /dev/null +++ b/packages/core/src/core/tool-call-preparation.ts @@ -0,0 +1,30 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import type { GenerateContentResponse } from '@google/genai'; + +export interface ToolCallPreparation { + callId: string; + toolName: string; +} + +const preparationsByResponse = new WeakMap< + GenerateContentResponse, + readonly ToolCallPreparation[] +>(); + +export function setToolCallPreparations( + response: GenerateContentResponse, + preparations: readonly ToolCallPreparation[], +): void { + preparationsByResponse.set(response, preparations); +} + +export function getToolCallPreparations( + response: GenerateContentResponse, +): readonly ToolCallPreparation[] { + return preparationsByResponse.get(response) ?? []; +} diff --git a/packages/core/src/core/toolCallIdUtils.test.ts b/packages/core/src/core/toolCallIdUtils.test.ts index 1f51aa347bc..09aa0657e2c 100644 --- a/packages/core/src/core/toolCallIdUtils.test.ts +++ b/packages/core/src/core/toolCallIdUtils.test.ts @@ -11,6 +11,7 @@ import { dedupeToolCallsById, getProviderToolCallId, normalizeModelToolCallIds, + reserveModelToolCallId, } from './toolCallIdUtils.js'; describe('toolCallIdUtils', () => { @@ -100,6 +101,50 @@ describe('toolCallIdUtils', () => { ).toEqual([undefined, undefined]); }); + it('reserves a fresh model tool call id', () => { + const usedIds = new Set(); + const reservedIds = new Map(); + + expect(reserveModelToolCallId('call-1', usedIds, reservedIds)).toBe( + 'call-1', + ); + expect(reservedIds.get('call-1')).toBe('call-1'); + expect(usedIds.has('call-1')).toBe(true); + }); + + it('returns the same id when reserving a raw id repeatedly', () => { + const usedIds = new Set(['call-1']); + const reservedIds = new Map(); + + const first = reserveModelToolCallId('call-1', usedIds, reservedIds); + const second = reserveModelToolCallId('call-1', usedIds, reservedIds); + + expect(first).toBe('call-1__qwen_dup_2'); + expect(second).toBe(first); + expect([...usedIds]).toEqual(['call-1', 'call-1__qwen_dup_2']); + }); + + it('normalizes a colliding raw id to its reserved suffixed id', () => { + const usedIds = new Set(['call-1']); + const reservedIds = new Map(); + const reservedId = reserveModelToolCallId('call-1', usedIds, reservedIds); + + const normalized = normalizeModelToolCallIds( + [ + { + functionCall: { id: 'call-1', name: 'read_file', args: {} }, + }, + ], + usedIds, + new Set(), + reservedIds, + ); + + expect(reservedId).toBe('call-1__qwen_dup_2'); + expect(normalized[0]?.functionCall?.id).toBe(reservedId); + expect(getProviderToolCallId(normalized[0]!.functionCall!)).toBe('call-1'); + }); + it('deduplicates direct function call batches by id', () => { const calls = [ { id: 'call_1', name: 'read_file', args: { file_path: 'a.ts' } }, diff --git a/packages/core/src/core/toolCallIdUtils.ts b/packages/core/src/core/toolCallIdUtils.ts index 886db5345c8..5dab1a13c65 100644 --- a/packages/core/src/core/toolCallIdUtils.ts +++ b/packages/core/src/core/toolCallIdUtils.ts @@ -61,6 +61,7 @@ export function normalizeModelToolCallIds( parts: readonly Part[], usedIds: Set, rawIdsInCurrentTurn: Set, + reservedIds?: ReadonlyMap, ): Part[] { const normalized: Part[] = []; @@ -83,7 +84,7 @@ export function normalizeModelToolCallIds( } const id = rawId - ? nextAvailableDuplicateId(rawId, usedIds) + ? (reservedIds?.get(rawId) ?? nextAvailableDuplicateId(rawId, usedIds)) : nextGeneratedId(usedIds); if (rawId && id !== rawId) { debugLogger.debug( @@ -112,6 +113,20 @@ export function normalizeModelToolCallIds( return normalized; } +export function reserveModelToolCallId( + rawId: string, + usedIds: Set, + reservedIds: Map, +): string { + const existing = reservedIds.get(rawId); + if (existing) return existing; + + const id = nextAvailableDuplicateId(rawId, usedIds); + reservedIds.set(rawId, id); + usedIds.add(id); + return id; +} + export function getProviderToolCallId( functionCall: FunctionCall, ): string | undefined { diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index ae7a247524d..8cd45bf157e 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -72,6 +72,7 @@ export * from './core/nonInteractiveToolExecutor.js'; export * from './core/prompts.js'; export * from './core/session-recovery.js'; export * from './core/tokenLimits.js'; +export * from './core/tool-call-preparation.js'; export * from './core/toolCallIdUtils.js'; export * from './core/turn.js'; export * from './core/turn-interruption.js';