From abcb9ac28b1c0adaaae959e4d86a7d66b29f8a8f Mon Sep 17 00:00:00 2001 From: Sarthak Singh Date: Fri, 10 Jul 2026 02:18:18 +0530 Subject: [PATCH] fix(ai): collect trailing tool approvals --- .../collect-tool-approvals.test.ts | 173 ++++++++++ .../generate-text/collect-tool-approvals.ts | 75 +++-- .../ai/src/generate-text/stream-text.test.ts | 295 ++++++++++++++++++ ...guage-model-prompt.validation.test.ts.snap | 31 ++ .../convert-to-language-model-prompt.ts | 85 ++--- ...o-language-model-prompt.validation.test.ts | 88 ++++++ .../get-latest-assistant-approval-turn.ts | 26 ++ 7 files changed, 706 insertions(+), 67 deletions(-) create mode 100644 packages/ai/src/util/get-latest-assistant-approval-turn.ts diff --git a/packages/ai/src/generate-text/collect-tool-approvals.test.ts b/packages/ai/src/generate-text/collect-tool-approvals.test.ts index 5f0c4c3e61cc..bfc28aef2d8b 100644 --- a/packages/ai/src/generate-text/collect-tool-approvals.test.ts +++ b/packages/ai/src/generate-text/collect-tool-approvals.test.ts @@ -110,6 +110,47 @@ describe('collectToolApprovals', () => { `); }); + it('should return approved approval with approved response followed by user context', () => { + const result = collectToolApprovals({ + messages: [ + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'tool1', + input: { value: 'test-input' }, + }, + { + type: 'tool-approval-request', + approvalId: 'approval-id-1', + toolCallId: 'call-1', + }, + ], + }, + { + role: 'tool', + content: [ + { + type: 'tool-approval-response', + approvalId: 'approval-id-1', + approved: true, + }, + ], + }, + { + role: 'user', + content: 'extra context', + }, + ], + }); + + expect(result.approvedToolApprovals).toHaveLength(1); + expect(result.approvedToolApprovals[0].toolCall.toolCallId).toBe('call-1'); + expect(result.deniedToolApprovals).toHaveLength(0); + }); + it('should return processed approval with approved response and tool result', () => { const result = collectToolApprovals({ messages: [ @@ -156,6 +197,54 @@ describe('collectToolApprovals', () => { `); }); + it('should return processed approval with approved response and tool result followed by user context', () => { + const result = collectToolApprovals({ + messages: [ + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'tool1', + input: { value: 'test-input' }, + }, + { + type: 'tool-approval-request', + approvalId: 'approval-id-1', + toolCallId: 'call-1', + }, + ], + }, + { + role: 'tool', + content: [ + { + type: 'tool-approval-response', + approvalId: 'approval-id-1', + approved: true, + }, + { + type: 'tool-result', + toolCallId: 'call-1', + toolName: 'tool1', + output: { type: 'text', value: 'test-output' }, + }, + ], + }, + { + role: 'user', + content: 'extra context', + }, + ], + }); + + expect(result).toEqual({ + approvedToolApprovals: [], + deniedToolApprovals: [], + }); + }); + it('should return denied approval with denied response', () => { const result = collectToolApprovals({ messages: [ @@ -219,6 +308,48 @@ describe('collectToolApprovals', () => { `); }); + it('should return denied approval with denied response followed by system context', () => { + const result = collectToolApprovals({ + messages: [ + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'tool1', + input: { value: 'test-input' }, + }, + { + type: 'tool-approval-request', + approvalId: 'approval-id-1', + toolCallId: 'call-1', + }, + ], + }, + { + role: 'tool', + content: [ + { + type: 'tool-approval-response', + approvalId: 'approval-id-1', + approved: false, + reason: 'test-reason', + }, + ], + }, + { + role: 'system', + content: 'extra context', + }, + ], + }); + + expect(result.approvedToolApprovals).toHaveLength(0); + expect(result.deniedToolApprovals).toHaveLength(1); + expect(result.deniedToolApprovals[0].toolCall.toolCallId).toBe('call-1'); + }); + it('should return processed approval with denied response and tool result', () => { const result = collectToolApprovals({ messages: [ @@ -266,6 +397,48 @@ describe('collectToolApprovals', () => { `); }); + it('should ignore approval responses before a later assistant response', () => { + const result = collectToolApprovals({ + messages: [ + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'call-1', + toolName: 'tool1', + input: { value: 'test-input' }, + }, + { + type: 'tool-approval-request', + approvalId: 'approval-id-1', + toolCallId: 'call-1', + }, + ], + }, + { + role: 'tool', + content: [ + { + type: 'tool-approval-response', + approvalId: 'approval-id-1', + approved: true, + }, + ], + }, + { + role: 'assistant', + content: 'later response', + }, + ], + }); + + expect(result).toEqual({ + approvedToolApprovals: [], + deniedToolApprovals: [], + }); + }); + it('should throw when approval response has unknown approvalId', () => { expect(() => collectToolApprovals({ diff --git a/packages/ai/src/generate-text/collect-tool-approvals.ts b/packages/ai/src/generate-text/collect-tool-approvals.ts index e6a796a2c8ce..854b137b44e1 100644 --- a/packages/ai/src/generate-text/collect-tool-approvals.ts +++ b/packages/ai/src/generate-text/collect-tool-approvals.ts @@ -6,6 +6,7 @@ import type { } from '@ai-sdk/provider-utils'; import { InvalidToolApprovalError } from '../error/invalid-tool-approval-error'; import { ToolCallNotFoundForApprovalError } from '../error/tool-call-not-found-for-approval-error'; +import { getLatestAssistantApprovalTurn } from '../util/get-latest-assistant-approval-turn'; import type { TypedToolCall } from './tool-call'; import type { TypedToolResult } from './tool-result'; @@ -15,9 +16,20 @@ export type CollectedToolApprovals = { toolCall: TypedToolCall; }; +function createEmptyToolApprovals(): { + approvedToolApprovals: Array>; + deniedToolApprovals: Array>; +} { + return { + approvedToolApprovals: [], + deniedToolApprovals: [], + }; +} + /** - * If the last message is a tool message, this function collects all tool approvals - * from that message. + * Collects unresolved tool approvals for the latest assistant turn. Approval + * responses may be followed by non-assistant context messages, but are not + * replayed once a later assistant response exists. */ export function collectToolApprovals({ messages, @@ -27,15 +39,14 @@ export function collectToolApprovals({ approvedToolApprovals: Array>; deniedToolApprovals: Array>; } { - const lastMessage = messages.at(-1); + const approvalTurn = getLatestAssistantApprovalTurn(messages); - if (lastMessage?.role != 'tool') { - return { - approvedToolApprovals: [], - deniedToolApprovals: [], - }; + if (approvalTurn == null) { + return createEmptyToolApprovals(); } + const { latestAssistantContent, suffixMessages } = approvalTurn; + // gather tool calls and prepare lookup. // // These maps are keyed by client-supplied ids (`toolCallId`, `approvalId`) @@ -49,47 +60,47 @@ export function collectToolApprovals({ string, TypedToolCall > = Object.create(null); - for (const message of messages) { - if (message.role === 'assistant' && typeof message.content !== 'string') { - const content = message.content; - for (const part of content) { - if (part.type === 'tool-call') { - toolCallsByToolCallId[part.toolCallId] = part as TypedToolCall; - } - } + for (const part of latestAssistantContent) { + if (part.type === 'tool-call') { + toolCallsByToolCallId[part.toolCallId] = part as TypedToolCall; } } - // gather approval responses and prepare lookup + // gather approval requests from the latest assistant turn only. If a later + // assistant message exists, older approval responses belong to already + // continued history and must not trigger side effects again. const toolApprovalRequestsByApprovalId: Record = Object.create(null); - for (const message of messages) { - if (message.role === 'assistant' && typeof message.content !== 'string') { - const content = message.content; - for (const part of content) { - if (part.type === 'tool-approval-request') { - toolApprovalRequestsByApprovalId[part.approvalId] = part; - } - } + for (const part of latestAssistantContent) { + if (part.type === 'tool-approval-request') { + toolApprovalRequestsByApprovalId[part.approvalId] = part; } } - // gather tool results from the last tool message + const approvalResponses: ToolApprovalResponse[] = []; + + // gather tool results from the unresolved suffix after the latest assistant + // turn, allowing user/system context to trail approval responses. const toolResults: Record> = Object.create( null, ); - for (const part of lastMessage.content) { - if (part.type === 'tool-result') { - toolResults[part.toolCallId] = part as TypedToolResult; + for (const message of suffixMessages) { + if (message.role !== 'tool') { + continue; + } + + for (const part of message.content) { + if (part.type === 'tool-approval-response') { + approvalResponses.push(part); + } else if (part.type === 'tool-result') { + toolResults[part.toolCallId] = part as TypedToolResult; + } } } const approvedToolApprovals: Array> = []; const deniedToolApprovals: Array> = []; - const approvalResponses = lastMessage.content.filter( - part => part.type === 'tool-approval-response', - ); for (const approvalResponse of approvalResponses) { const approvalRequest = toolApprovalRequestsByApprovalId[approvalResponse.approvalId]; diff --git a/packages/ai/src/generate-text/stream-text.test.ts b/packages/ai/src/generate-text/stream-text.test.ts index 4969b1e23f4e..63694811231f 100644 --- a/packages/ai/src/generate-text/stream-text.test.ts +++ b/packages/ai/src/generate-text/stream-text.test.ts @@ -28261,6 +28261,301 @@ describe('streamText', () => { }); }); + describe('when local tool approval responses are followed by context messages', () => { + function createContinuationModel(prompts: LanguageModelV4Prompt[]) { + return new MockLanguageModelV4({ + doStream: async ({ prompt }) => { + prompts.push(prompt); + return { + stream: convertArrayToReadableStream([ + { type: 'stream-start', warnings: [] }, + { type: 'text-start', id: '1' }, + { + type: 'text-delta', + id: '1', + delta: 'continued', + }, + { type: 'text-end', id: '1' }, + { + type: 'finish', + finishReason: { unified: 'stop', raw: 'stop' }, + usage: testUsage, + }, + ]), + }; + }, + }); + } + + function createApprovedToolMessages({ + approved, + includeToolResult = false, + includeLaterAssistant = false, + }: { + approved: boolean; + includeToolResult?: boolean; + includeLaterAssistant?: boolean; + }): ModelMessage[] { + return [ + { role: 'user', content: 'test-input' }, + { + role: 'assistant', + content: [ + { + input: { value: 'value' }, + providerExecuted: undefined, + providerOptions: undefined, + toolCallId: 'call-1', + toolName: 'tool1', + type: 'tool-call', + }, + { + approvalId: 'id-1', + toolCallId: 'call-1', + type: 'tool-approval-request', + }, + ], + }, + { + role: 'tool', + content: [ + { + approvalId: 'id-1', + type: 'tool-approval-response', + approved, + reason: approved ? undefined : 'not allowed', + }, + ...(includeToolResult + ? [ + { + type: 'tool-result' as const, + toolCallId: 'call-1', + toolName: 'tool1', + output: { + type: 'text' as const, + value: 'existing result', + }, + }, + ] + : []), + ], + }, + ...(includeLaterAssistant + ? [ + { + role: 'assistant' as const, + content: 'already continued', + }, + ] + : []), + { role: 'user', content: 'extra context' }, + ]; + } + + it('should execute an approved local tool exactly once', async () => { + const prompts: LanguageModelV4Prompt[] = []; + const executeFunction = vi.fn().mockReturnValue('result1'); + + const result = streamText({ + model: createContinuationModel(prompts), + tools: { + tool1: tool({ + inputSchema: z.object({ value: z.string() }), + execute: executeFunction, + }), + }, + toolApproval: { + tool1: 'user-approval', + }, + stopWhen: isStepCount(3), + _internal: { + generateId: mockId({ prefix: 'id' }), + generateCallId: () => 'test-telemetry-call-id', + }, + messages: createApprovedToolMessages({ approved: true }), + }); + + await result.consumeStream(); + + expect(executeFunction).toHaveBeenCalledTimes(1); + expect(executeFunction).toHaveBeenCalledWith( + { value: 'value' }, + expect.objectContaining({ + toolCallId: 'call-1', + messages: expect.any(Array), + }), + ); + expect(await result.responseMessages).toEqual([ + { + role: 'tool', + content: [ + { + type: 'tool-result', + toolCallId: 'call-1', + toolName: 'tool1', + output: { + type: 'text', + value: 'result1', + }, + }, + ], + }, + { + role: 'assistant', + content: [ + { + type: 'text', + text: 'continued', + providerOptions: undefined, + }, + ], + }, + ]); + expect(prompts).toHaveLength(1); + }); + + it('should produce a denial result when a denied approval has trailing user context', async () => { + const prompts: LanguageModelV4Prompt[] = []; + const executeFunction = vi.fn().mockReturnValue('result1'); + + const result = streamText({ + model: createContinuationModel(prompts), + tools: { + tool1: tool({ + inputSchema: z.object({ value: z.string() }), + execute: executeFunction, + }), + }, + toolApproval: { + tool1: 'user-approval', + }, + stopWhen: isStepCount(3), + _internal: { + generateId: mockId({ prefix: 'id' }), + generateCallId: () => 'test-telemetry-call-id', + }, + messages: createApprovedToolMessages({ approved: false }), + }); + + await result.consumeStream(); + + expect(executeFunction).not.toHaveBeenCalled(); + expect(await result.responseMessages).toEqual([ + { + role: 'tool', + content: [ + { + type: 'tool-result', + toolCallId: 'call-1', + toolName: 'tool1', + output: { + type: 'execution-denied', + reason: 'not allowed', + }, + }, + ], + }, + { + role: 'assistant', + content: [ + { + type: 'text', + text: 'continued', + providerOptions: undefined, + }, + ], + }, + ]); + expect(prompts).toHaveLength(1); + }); + + it('should not execute an approved local tool again when a tool result already exists', async () => { + const prompts: LanguageModelV4Prompt[] = []; + const executeFunction = vi.fn().mockReturnValue('result1'); + + const result = streamText({ + model: createContinuationModel(prompts), + tools: { + tool1: tool({ + inputSchema: z.object({ value: z.string() }), + execute: executeFunction, + }), + }, + toolApproval: { + tool1: 'user-approval', + }, + stopWhen: isStepCount(3), + _internal: { + generateId: mockId({ prefix: 'id' }), + generateCallId: () => 'test-telemetry-call-id', + }, + messages: createApprovedToolMessages({ + approved: true, + includeToolResult: true, + }), + }); + + await result.consumeStream(); + + expect(executeFunction).not.toHaveBeenCalled(); + expect(await result.responseMessages).toEqual([ + { + role: 'assistant', + content: [ + { + type: 'text', + text: 'continued', + providerOptions: undefined, + }, + ], + }, + ]); + expect(prompts).toHaveLength(1); + }); + + it('should not replay an older approval behind a later assistant response', async () => { + const prompts: LanguageModelV4Prompt[] = []; + const executeFunction = vi.fn().mockReturnValue('result1'); + const onError = vi.fn(); + + const result = streamText({ + model: createContinuationModel(prompts), + tools: { + tool1: tool({ + inputSchema: z.object({ value: z.string() }), + execute: executeFunction, + }), + }, + toolApproval: { + tool1: 'user-approval', + }, + stopWhen: isStepCount(3), + _internal: { + generateId: mockId({ prefix: 'id' }), + generateCallId: () => 'test-telemetry-call-id', + }, + onError, + messages: createApprovedToolMessages({ + approved: true, + includeLaterAssistant: true, + }), + }); + + await result.consumeStream(); + + expect(executeFunction).not.toHaveBeenCalled(); + expect(onError).toHaveBeenCalledWith({ + error: expect.objectContaining({ + message: 'Tool result is missing for tool call call-1.', + }), + }); + await expect(result.text).rejects.toThrow( + 'No output generated. Check the stream for errors.', + ); + expect(prompts).toHaveLength(0); + }); + }); + describe('provider-executed tool (MCP) approval', () => { describe('when a provider-executed tool emits tool-approval-request', () => { let result: StreamTextResult; diff --git a/packages/ai/src/prompt/__snapshots__/convert-to-language-model-prompt.validation.test.ts.snap b/packages/ai/src/prompt/__snapshots__/convert-to-language-model-prompt.validation.test.ts.snap index d9bc580543d6..470f3207e95d 100644 --- a/packages/ai/src/prompt/__snapshots__/convert-to-language-model-prompt.validation.test.ts.snap +++ b/packages/ai/src/prompt/__snapshots__/convert-to-language-model-prompt.validation.test.ts.snap @@ -42,6 +42,37 @@ exports[`tool validation > should pass validation for tool-approval-response 1`] ] `; +exports[`tool validation > should pass validation for tool-approval-response followed by user context 1`] = ` +[ + { + "content": [ + { + "input": { + "action": "delete_db", + }, + "providerExecuted": undefined, + "providerOptions": undefined, + "toolCallId": "call_to_approve", + "toolName": "dangerous_action", + "type": "tool-call", + }, + ], + "providerOptions": undefined, + "role": "assistant", + }, + { + "content": [ + { + "text": "extra context", + "type": "text", + }, + ], + "providerOptions": undefined, + "role": "user", + }, +] +`; + exports[`tool validation > should preserve provider-executed tool-approval-response 1`] = ` [ { diff --git a/packages/ai/src/prompt/convert-to-language-model-prompt.ts b/packages/ai/src/prompt/convert-to-language-model-prompt.ts index 7035c61e7552..eb11f88cfa34 100644 --- a/packages/ai/src/prompt/convert-to-language-model-prompt.ts +++ b/packages/ai/src/prompt/convert-to-language-model-prompt.ts @@ -28,6 +28,7 @@ import { import { convertToLanguageModelV4FilePart } from './file-part-data'; import { logWarnings } from '../logger/log-warnings'; import type { Warning } from '../types/warning'; +import { getLatestAssistantApprovalTurn } from '../util/get-latest-assistant-approval-turn'; import { InvalidMessageRoleError } from './invalid-message-role-error'; import type { StandardizedPrompt } from './standardize-prompt'; import { MissingToolResultsError } from '../error/missing-tool-result-error'; @@ -51,37 +52,8 @@ export async function convertToLanguageModelPrompt({ supportedUrls, ); - const approvalIdToToolCallId = new Map(); - for (const message of prompt.messages) { - if (message.role === 'assistant' && Array.isArray(message.content)) { - for (const part of message.content) { - if ( - part.type === 'tool-approval-request' && - 'approvalId' in part && - 'toolCallId' in part - ) { - approvalIdToToolCallId.set( - part.approvalId as string, - part.toolCallId as string, - ); - } - } - } - } - - const approvedToolCallIds = new Set(); - for (const message of prompt.messages) { - if (message.role === 'tool') { - for (const part of message.content) { - if (part.type === 'tool-approval-response') { - const toolCallId = approvalIdToToolCallId.get(part.approvalId); - if (toolCallId) { - approvedToolCallIds.add(toolCallId); - } - } - } - } - } + const approvalResolvedToolCallIds = + collectApprovalResolvedToolCallIdsFromLatestAssistantTurn(prompt.messages); const messages = [ ...(prompt.instructions != null @@ -136,8 +108,8 @@ export async function convertToLanguageModelPrompt({ } case 'user': case 'system': - // remove approved tool calls from the set before checking: - for (const id of approvedToolCallIds) { + // remove approval-resolved tool calls from the set before checking: + for (const id of approvalResolvedToolCallIds) { toolCallIds.delete(id); } @@ -150,8 +122,8 @@ export async function convertToLanguageModelPrompt({ } } - // remove approved tool calls from the set before checking: - for (const id of approvedToolCallIds) { + // remove approval-resolved tool calls from the set before checking: + for (const id of approvalResolvedToolCallIds) { toolCallIds.delete(id); } @@ -168,6 +140,49 @@ export async function convertToLanguageModelPrompt({ ); } +function collectApprovalResolvedToolCallIdsFromLatestAssistantTurn( + messages: ModelMessage[], +): Set { + const approvalTurn = getLatestAssistantApprovalTurn(messages); + + if (approvalTurn == null) { + return new Set(); + } + + const { latestAssistantContent, suffixMessages } = approvalTurn; + const approvalIdToToolCallId = new Map(); + for (const part of latestAssistantContent) { + if ( + part.type === 'tool-approval-request' && + 'approvalId' in part && + 'toolCallId' in part + ) { + approvalIdToToolCallId.set( + part.approvalId as string, + part.toolCallId as string, + ); + } + } + + const approvalResolvedToolCallIds = new Set(); + for (const message of suffixMessages) { + if (message.role !== 'tool') { + continue; + } + + for (const part of message.content) { + if (part.type === 'tool-approval-response') { + const toolCallId = approvalIdToToolCallId.get(part.approvalId); + if (toolCallId != null) { + approvalResolvedToolCallIds.add(toolCallId); + } + } + } + } + + return approvalResolvedToolCallIds; +} + /** * Convert a ModelMessage to a LanguageModelV4Message. * diff --git a/packages/ai/src/prompt/convert-to-language-model-prompt.validation.test.ts b/packages/ai/src/prompt/convert-to-language-model-prompt.validation.test.ts index a083fe919daf..e2da1ecb1787 100644 --- a/packages/ai/src/prompt/convert-to-language-model-prompt.validation.test.ts +++ b/packages/ai/src/prompt/convert-to-language-model-prompt.validation.test.ts @@ -71,6 +71,94 @@ describe('tool validation', () => { expect(result).toMatchSnapshot(); }); + it('should pass validation for tool-approval-response followed by user context', async () => { + const result = await convertToLanguageModelPrompt({ + prompt: { + instructions: undefined, + messages: [ + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'call_to_approve', + toolName: 'dangerous_action', + input: { action: 'delete_db' }, + }, + { + type: 'tool-approval-request', + toolCallId: 'call_to_approve', + approvalId: 'approval_123', + }, + ], + }, + { + role: 'tool', + content: [ + { + type: 'tool-approval-response', + approvalId: 'approval_123', + approved: true, + }, + ], + }, + { + role: 'user', + content: 'extra context', + }, + ], + }, + supportedUrls: {}, + download: undefined, + }); + + expect(result).toMatchSnapshot(); + }); + + it('should not treat an older approval before a later assistant response as resolving missing results', async () => { + await expect(async () => { + await convertToLanguageModelPrompt({ + prompt: { + instructions: undefined, + messages: [ + { + role: 'assistant', + content: [ + { + type: 'tool-call', + toolCallId: 'call_to_approve', + toolName: 'dangerous_action', + input: { action: 'delete_db' }, + }, + { + type: 'tool-approval-request', + toolCallId: 'call_to_approve', + approvalId: 'approval_123', + }, + ], + }, + { + role: 'tool', + content: [ + { + type: 'tool-approval-response', + approvalId: 'approval_123', + approved: true, + }, + ], + }, + { + role: 'assistant', + content: 'already continued', + }, + ], + }, + supportedUrls: {}, + download: undefined, + }); + }).rejects.toThrow(MissingToolResultsError); + }); + it('should preserve provider-executed tool-approval-response', async () => { const result = await convertToLanguageModelPrompt({ prompt: { diff --git a/packages/ai/src/util/get-latest-assistant-approval-turn.ts b/packages/ai/src/util/get-latest-assistant-approval-turn.ts new file mode 100644 index 000000000000..47e2de911a61 --- /dev/null +++ b/packages/ai/src/util/get-latest-assistant-approval-turn.ts @@ -0,0 +1,26 @@ +import type { ModelMessage } from '@ai-sdk/provider-utils'; + +export function getLatestAssistantApprovalTurn(messages: ModelMessage[]) { + let latestAssistantMessageIndex = -1; + for (let index = messages.length - 1; index >= 0; index--) { + if (messages[index].role === 'assistant') { + latestAssistantMessageIndex = index; + break; + } + } + + if (latestAssistantMessageIndex === -1) { + return undefined; + } + + const latestAssistantMessage = messages[latestAssistantMessageIndex]; + const latestAssistantContent = + typeof latestAssistantMessage.content === 'string' + ? [] + : latestAssistantMessage.content; + + return { + latestAssistantContent, + suffixMessages: messages.slice(latestAssistantMessageIndex + 1), + }; +}