From 8a4f19cd42ce9d2f0cc64158b53f3efc5b0ffc21 Mon Sep 17 00:00:00 2001 From: King Star Date: Tue, 21 Jul 2026 11:57:36 +0800 Subject: [PATCH] fix(ai): reject messages after tool approval responses --- .changeset/fix-tool-approval-sequence.md | 5 + .../ai/src/generate-text/stream-text.test.ts | 62 ++++++++ .../convert-to-language-model-prompt.ts | 28 ++-- ...o-language-model-prompt.validation.test.ts | 133 ++++++++++++++++++ 4 files changed, 212 insertions(+), 16 deletions(-) create mode 100644 .changeset/fix-tool-approval-sequence.md diff --git a/.changeset/fix-tool-approval-sequence.md b/.changeset/fix-tool-approval-sequence.md new file mode 100644 index 000000000000..d4083128f56c --- /dev/null +++ b/.changeset/fix-tool-approval-sequence.md @@ -0,0 +1,5 @@ +--- +'ai': patch +--- + +fix: reject trailing messages after tool approval responses diff --git a/packages/ai/src/generate-text/stream-text.test.ts b/packages/ai/src/generate-text/stream-text.test.ts index 8dc5c1bc09a7..949cc81add41 100644 --- a/packages/ai/src/generate-text/stream-text.test.ts +++ b/packages/ai/src/generate-text/stream-text.test.ts @@ -43,6 +43,7 @@ import { import { mockSandboxSessionFileStubs } from '../test/mock-sandbox'; import { z } from 'zod/v4'; import { Output, type LanguageModelCallEndEvent, type Telemetry } from '..'; +import { MissingToolResultsError } from '../error/missing-tool-result-error'; import * as logWarningsModule from '../logger/log-warnings'; import type { Instructions } from '../prompt'; import { MockLanguageModelV4 } from '../test/mock-language-model-v4'; @@ -27061,6 +27062,67 @@ describe('streamText', () => { }); }); + it('should reject a user message after a tool approval response', async () => { + const executeFunction = vi.fn().mockReturnValue('result1'); + const doStream = vi.fn(async () => { + throw new Error('model should not be called'); + }); + const onError = vi.fn(); + + const result = streamText({ + model: new MockLanguageModelV4({ doStream }), + tools: { + tool1: tool({ + inputSchema: z.object({ value: z.string() }), + execute: executeFunction, + }), + }, + toolApproval: { + tool1: 'user-approval', + }, + onError, + messages: [ + { role: 'user', content: 'test-input' }, + { + role: 'assistant', + content: [ + { + input: { value: 'value' }, + 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: true, + }, + ], + }, + { role: 'user', content: 'additional context' }, + ], + }); + + await result.consumeStream(); + + expect(onError).toHaveBeenCalledOnce(); + const error = onError.mock.calls[0][0].error; + expect(MissingToolResultsError.isInstance(error)).toBe(true); + expect(error).toMatchObject({ toolCallIds: ['call-1'] }); + expect(executeFunction).not.toHaveBeenCalled(); + expect(doStream).not.toHaveBeenCalled(); + }); + describe('when a call from a single tool that needs approval is approved', () => { let result: StreamTextResult; let prompts: LanguageModelV4Prompt[]; 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..eaf050c12d6c 100644 --- a/packages/ai/src/prompt/convert-to-language-model-prompt.ts +++ b/packages/ai/src/prompt/convert-to-language-model-prompt.ts @@ -69,15 +69,16 @@ export async function convertToLanguageModelPrompt({ } } - 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); - } + // collectToolApprovals only consumes responses from the final tool message. + // Older responses must not suppress missing-result validation. + const approvalResponseToolCallIds = new Set(); + const lastMessage = prompt.messages.at(-1); + if (lastMessage?.role === 'tool') { + for (const part of lastMessage.content) { + if (part.type === 'tool-approval-response') { + const toolCallId = approvalIdToToolCallId.get(part.approvalId); + if (toolCallId) { + approvalResponseToolCallIds.add(toolCallId); } } } @@ -136,11 +137,6 @@ export async function convertToLanguageModelPrompt({ } case 'user': case 'system': - // remove approved tool calls from the set before checking: - for (const id of approvedToolCallIds) { - toolCallIds.delete(id); - } - if (toolCallIds.size > 0) { throw new MissingToolResultsError({ toolCallIds: Array.from(toolCallIds), @@ -150,8 +146,8 @@ export async function convertToLanguageModelPrompt({ } } - // remove approved tool calls from the set before checking: - for (const id of approvedToolCallIds) { + // The final approval response is consumed before the next model call. + for (const id of approvalResponseToolCallIds) { toolCallIds.delete(id); } 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..776033988d99 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 @@ -1,4 +1,5 @@ import { describe, expect, it } from 'vitest'; +import type { ModelMessage } from '@ai-sdk/provider-utils'; import { convertToLanguageModelPrompt } from './convert-to-language-model-prompt'; import { MissingToolResultsError } from '../error/missing-tool-result-error'; @@ -71,6 +72,138 @@ describe('tool validation', () => { expect(result).toMatchSnapshot(); }); + it.each<{ + decision: string; + approved: boolean; + name: string; + trailingMessage: ModelMessage; + }>([ + { + decision: 'approved', + approved: true, + name: 'user', + trailingMessage: { role: 'user', content: 'additional context' }, + }, + { + decision: 'approved', + approved: true, + name: 'system', + trailingMessage: { role: 'system', content: 'additional instructions' }, + }, + { + decision: 'approved', + approved: true, + name: 'assistant', + trailingMessage: { role: 'assistant', content: 'continued response' }, + }, + { + decision: 'denied', + approved: false, + name: 'user', + trailingMessage: { role: 'user', content: 'additional context' }, + }, + ])( + 'should reject a $name message after an $decision tool approval response', + async ({ approved, trailingMessage }) => { + await expect( + 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, + }, + ], + }, + trailingMessage, + ], + }, + supportedUrls: {}, + download: undefined, + }), + ).rejects.toEqual( + expect.objectContaining({ + name: 'AI_MissingToolResultsError', + toolCallIds: ['call_to_approve'], + }), + ); + }, + ); + + it('should allow later messages after an approved tool call has a result', 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, + }, + { + type: 'tool-result', + toolCallId: 'call_to_approve', + toolName: 'dangerous_action', + output: { type: 'text', value: 'completed' }, + }, + ], + }, + { role: 'assistant', content: 'The action completed.' }, + { role: 'user', content: 'What happened?' }, + ], + }, + supportedUrls: {}, + download: undefined, + }); + + expect(result.map(message => message.role)).toEqual([ + 'assistant', + 'tool', + 'assistant', + 'user', + ]); + }); + it('should preserve provider-executed tool-approval-response', async () => { const result = await convertToLanguageModelPrompt({ prompt: {