diff --git a/CHANGELOG.md b/CHANGELOG.md index 542e07b2..75924988 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ # Changelog +## Unreleased + +- Fix continuing a tool approval when a context handler returns history beyond `recentMessages`; context handlers no longer see approval requests or decisions, and pending approvals are left out of context instead of auto-denied (#416) + ## 0.7.8 - Fix answering a tool approval after an earlier approval in the same thread diff --git a/docs/tool-approval.mdx b/docs/tool-approval.mdx index 220d3af8..00b8cf70 100644 --- a/docs/tool-approval.mdx +++ b/docs/tool-approval.mdx @@ -169,16 +169,15 @@ single-decision helpers so the request history is read only once. If a partial continuation requests another approval, that newer request also remains pending while you resume an earlier sibling. -If an unrelated generation starts while approvals are still unresolved, they -are **automatically denied** with the reason -`"auto-denied: new generation started"`. This prevents broken message -histories where tool calls lack results. These denials are synthetic and only -added to the model's context for that request; they are not saved, so the -approvals remain pending in storage. - -This does not apply when resuming approvals: on that continuation, earlier -approvals that were not part of the resume are omitted from that model request. -No response is written for them, so they remain pending in storage. +If an unrelated generation starts while approvals are still unresolved, their +tool calls are left out of its context. The model does not see them, and they +stay pending, so you can still approve or deny them later. + +Approval requests and decisions are never passed to a +[context handler](./context.mdx), nor are tool calls whose approval has not +run. When generation continues from a decision, the stored calls and decisions +are appended after the handler's messages, so only decisions saved with these +methods can run a tool. ## Client-side flow diff --git a/src/vercel/client/prepareApprovalContext.ts b/src/vercel/client/prepareApprovalContext.ts index e203d08c..71ac9a2e 100644 --- a/src/vercel/client/prepareApprovalContext.ts +++ b/src/vercel/client/prepareApprovalContext.ts @@ -1,340 +1,218 @@ import type { AssistantModelMessage, ModelMessage, ToolModelMessage } from "ai"; -import { convexToJson, type Value } from "convex/values"; +import type { MessageDoc } from "../../validators.js"; +import { docsToModelMessages } from "../mapping.js"; -type AssistantPart = Exclude[number]; -type Request = Extract; -type Call = Extract; type ToolPart = ToolModelMessage["content"][number]; -type Response = Extract; -type Result = Extract; - -type ApprovalPart = Request | Call | Response | Result; - -const validRole = { - "tool-approval-request": "assistant", - "tool-call": "assistant", - "tool-approval-response": "tool", - "tool-result": "tool", -} as const; -const addedNoun = { - "tool-approval-request": "approval requests", - "tool-call": "approved tool calls", - "tool-approval-response": "approval decisions", - "tool-result": "results for approved tool calls", -} as const; -const removedNoun = { - "tool-approval-request": "request", - "tool-call": "call", - "tool-approval-response": "decision", - "tool-result": "result", -} as const; - -// The AI SDK runs the approvals in the context's last message whose call has -// no result in that same message (collectToolApprovals), and never sends a -// local approval's request or decision to the provider. So a local approval -// appears as request and decision only in the final synthesized pair, and -// elsewhere as just its call and result. -export function prepareApprovalContext( - messages: ModelMessage[], - approvalIds: ReadonlySet, - storedContext: ModelMessage[], -): ModelMessage[] { - const storedRequests = new Map(); - for (const part of partsOf(storedContext)) { - if (part.type === "tool-approval-request") { - storedRequests.set(part.approvalId, part); - } - } - const approvalCallIds = new Set( - [...storedRequests.values()].map((request) => request.toolCallId), +type Part = + | Exclude[number] + | ToolPart; +type Request = Extract; +type Response = Extract; + +export type ApprovalPlan = { + hidden: Set; + replay: Map; + unit: ModelMessage[]; +}; + +// The AI SDK runs only the decisions in the prompt's last message, so approval +// requests and decisions are protocol, not context. A context handler never +// sees them or a call whose approval has not run, and the continued approvals +// are appended after it from storage. Providers read their own decisions from +// anywhere in the prompt, so those of provider-executed approvals that ran are +// replayed after their call. +export function planApprovals( + stored: MessageDoc[], + search: MessageDoc[], + continued: string[], +): ApprovalPlan { + const plan: ApprovalPlan = { hidden: new Set(), replay: new Map(), unit: [] }; + const [state, searchState] = [stored, search].map((docs) => + scan(docsToModelMessages(docs)), ); - - const stored = new Map(); - const waitingProviderApprovalIds = new Set(); - const consumedProviderApprovalIds = new Set(); - for (const message of storedContext) { - if (message.role === "assistant") { - // Provider-owned approvals do not always produce a tool-result. A later - // assistant response is the durable evidence that they were consumed. - for (const approvalId of waitingProviderApprovalIds) { - consumedProviderApprovalIds.add(approvalId); - } - waitingProviderApprovalIds.clear(); - } - for (const [key, part] of approvalParts(message, approvalCallIds)) { - const first = stored.get(key); - if (first && !sameApprovalPart(first, part)) { - throw new Error( - `Stored context has more than one ${part.type} for ${partId(part)}`, - ); - } - stored.set(key, first ?? part); - if ( - part.type === "tool-approval-response" && - isProviderCall(stored, storedRequests.get(part.approvalId)!.toolCallId) - ) { - waitingProviderApprovalIds.add(part.approvalId); + for (const { requests, decisions, providerCalls, settled } of [ + state, + searchState, + ]) { + for (const request of [...requests.values()].flat()) { + if (!settled(request)) plan.hidden.add(request.toolCallId); + else if (providerCalls.has(request.toolCallId)) { + const replayed = decisions.get(request.approvalId) ?? []; + plan.replay.set(request.toolCallId, structuredClone(replayed)); } } } - - // Every approval part a handler returns is a copy of a stored one. A handler - // may repeat a part (cloning the context, say), so the first copy is - // canonical and carries any part-level provider options it added. - const supplied = new Map(); - for (const message of messages) { - for (const [key, part] of approvalParts(message, approvalCallIds)) { - if (message.role !== validRole[part.type]) { - throw new Error( - `Continuation context has a ${part.type} in a ${message.role} message`, - ); - } - const original = stored.get(key); - if (!original) { - throw new Error( - `Continuation context cannot add ${addedNoun[part.type]}`, - ); - } - if (!sameApprovalPart(original, part)) { - throw new Error( - `Continuation context must preserve the stored ${part.type} for ${partId(part)}`, - ); - } - if (!supplied.has(key)) supplied.set(key, part); + const calls: Part[] = []; + const decisions: Response[] = []; + for (const id of continued) { + const [request] = only(state.requests, "tool-approval-request", id); + const [decision] = only(state.decisions, "tool-approval-response", id); + if (!decision || state.settled(request)) continue; + const [call, ...others] = state.calls.get(request.toolCallId) ?? []; + if (!call || others.length > 0 || calls.includes(call)) { + throw new Error(`Approval ${id} must name exactly one tool call`); } + calls.push(call, request); + decisions.push(decision); } - // An approval is kept or dropped whole, and the ones being continued are - // kept, so a handler can trim history but not separate a decision or result - // from the call it belongs to. - for (const request of storedRequests.values()) { - const keys = unitKeys(request).filter((key) => stored.has(key)); - if ( - !approvalIds.has(request.approvalId) && - !keys.some((key) => supplied.has(key)) - ) { - continue; - } - const missing = keys.find((key) => !supplied.has(key)); - if (missing) { - const part = stored.get(missing)!; - throw new Error( - `Continuation context removed ${removedNoun[part.type]} for ${partId(part)}`, - ); - } - } - - const activeCallIds = new Set(); - const executable = new Set(); - const undecided = new Set(); - const deferred = new Set(); - for (const [approvalId, { toolCallId }] of storedRequests) { - if (approvalIds.has(approvalId)) activeCallIds.add(toolCallId); - if ( - stored.has(`tool-result:${toolCallId}`) || - consumedProviderApprovalIds.has(approvalId) - ) { - continue; - } - if (!approvalIds.has(approvalId)) deferred.add(toolCallId); - else if (stored.has(`tool-approval-response:${approvalId}`)) { - executable.add(toolCallId); - } else undecided.add(toolCallId); + if (decisions.length > 0) { + plan.unit = structuredClone([ + { role: "assistant", content: calls }, + { role: "tool", content: decisions }, + ] as ModelMessage[]); } + return plan; +} - // A handler may clone or reorder the context. The final call/response pair - // is built once from the canonical parts, so every other occurrence is - // dropped, and a completed call's result is re-emitted directly after each - // message that carries the call, so nothing a handler inserts can separate - // them for providers that require that adjacency. - const relocated = new Map(); - for (const message of messages) { - if (message.role !== "tool") continue; - for (const part of message.content) { - if ( - part.type === "tool-result" && - activeCallIds.has(part.toolCallId) && - !relocated.has(part.toolCallId) - ) { - relocated.set(part.toolCallId, { ...message, content: [part] }); - } - } - } - // Message-level provider options (cache control, say) are the handler's, - // and a handler that targets one message usually means the last one, so the - // synthesized pair takes the options of the last message it absorbs. - let assistantOptions: ModelMessage["providerOptions"]; - let toolOptions: ModelMessage["providerOptions"]; - const projected: ModelMessage[] = []; - for (const message of messages) { - if (!Array.isArray(message.content)) { - projected.push(message); - continue; - } - const content = message.content.filter((part) => { - if (part.type === "tool-call" || part.type === "tool-approval-request") { - if (deferred.has(part.toolCallId) || undecided.has(part.toolCallId)) { - return false; - } - if (executable.has(part.toolCallId)) { - assistantOptions = message.providerOptions; - return false; +// What the context handler is given. Siblings decided in separate batches +// store each result after its own decision, so once the protocol is gone each +// result is moved to directly follow its call. +export function hideApprovals( + plan: ApprovalPlan, + segments: ModelMessage[][], +): ModelMessage[][] { + const answers = new Map(); + return segments + .map((messages) => + strip(plan, messages).flatMap((message): ModelMessage[] => { + if (message.role === "tool") { + const content = message.content.filter((part) => { + const answer = + part.type === "tool-result" && answers.get(part.toolCallId); + if (!answer) return true; + answer.providerOptions ??= message.providerOptions; + answer.content.push(part); + return false; + }); + return content.length > 0 ? [{ ...message, content }] : []; } - return ( - part.type === "tool-call" || isProviderCall(stored, part.toolCallId) + const answer: ToolModelMessage = { role: "tool", content: [] }; + const callIds = partsOf(message).flatMap((part) => + part.type === "tool-call" && !answers.has(part.toolCallId) + ? [part.toolCallId] + : [], ); - } - if (part.type === "tool-approval-response") { - const { toolCallId } = storedRequests.get(part.approvalId)!; - if (executable.has(toolCallId)) { - toolOptions = message.providerOptions; - return false; - } - return !deferred.has(toolCallId) && isProviderCall(stored, toolCallId); - } - if (part.type === "tool-result" && relocated.has(part.toolCallId)) { - return false; - } - return true; - }); - if (content.length === 0) continue; - projected.push({ ...message, content } as ModelMessage); - if (message.role !== "assistant") continue; - const results = content.flatMap((part) => - part.type === "tool-call" && relocated.has(part.toolCallId) - ? [relocated.get(part.toolCallId)!] - : [], + callIds.forEach((id) => answers.set(id, answer)); + return callIds.length > 0 ? [message, answer] : [message]; + }), + ) + .map((messages) => + messages.filter((m) => m.role !== "tool" || m.content.length > 0), ); - if (results.length > 0) { - projected.push({ - ...results[0], - content: results.flatMap((result) => result.content), - }); - } +} + +// The plan's parts are cloned on the way out, since the SDK hands the prompt +// to prepareStep, which may edit it in place. +export function finishApprovals( + plan: ApprovalPlan, + messages: ModelMessage[], +): ModelMessage[] { + // A handler can return history the plan never fetched, so approvals still + // pending there are found in what it returns. + const returned = scan(messages); + const hidden = new Set(plan.hidden); + for (const request of [...returned.requests.values()].flat()) { + if (!returned.settled(request)) hidden.add(request.toolCallId); } - if (executable.size > 0) { - const assistant: AssistantPart[] = []; - const tool: Response[] = []; - for (const approvalId of approvalIds) { - const { toolCallId } = storedRequests.get(approvalId)!; - if (!executable.has(toolCallId)) continue; - assistant.push( - supplied.get(`tool-call:${toolCallId}`) as Call, - supplied.get(`tool-approval-request:${approvalId}`) as Request, + const context = strip({ ...plan, hidden }, messages).flatMap( + (message): ModelMessage[] => { + const replayed = partsOf(message).flatMap((part) => + part.type === "tool-call" + ? structuredClone(plan.replay.get(part.toolCallId) ?? []) + : [], ); - tool.push( - supplied.get(`tool-approval-response:${approvalId}`) as Response, - ); - } - projected.push( - { - role: "assistant", - content: assistant, - ...withOptions(assistantOptions), - }, - { role: "tool", content: tool, ...withOptions(toolOptions) }, + return message.role === "assistant" && replayed.length > 0 + ? [message, { role: "tool", content: replayed }] + : [message]; + }, + ); + // The SDK runs the continued approvals before it validates the prompt, so an + // unanswered call would fail the generation after they ran, and a retry + // would run them again. + const parts = context.flatMap(partsOf); + const answered = new Set( + parts.map((part) => part.type === "tool-result" && part.toolCallId), + ); + const unanswered = parts.find( + (part) => + part.type === "tool-call" && + !part.providerExecuted && + !answered.has(part.toolCallId), + ); + if (unanswered?.type === "tool-call") { + throw new Error( + `Context has tool call ${unanswered.toolCallId} without its result`, ); } - return projected; + return [...context, ...structuredClone(plan.unit)]; } -function partsOf(messages: ModelMessage[]): (AssistantPart | ToolPart)[] { - return messages.flatMap((message) => - Array.isArray(message.content) - ? (message.content as (AssistantPart | ToolPart)[]) - : [], - ); -} - -function approvalParts( - message: ModelMessage, - approvalCallIds: ReadonlySet, -): [string, ApprovalPart][] { - return partsOf([message]).flatMap((part): [string, ApprovalPart][] => { - switch (part.type) { - case "tool-approval-request": - case "tool-approval-response": - return [[`${part.type}:${part.approvalId}`, part]]; - case "tool-call": - case "tool-result": - return approvalCallIds.has(part.toolCallId) - ? [[`${part.type}:${part.toolCallId}`, part]] - : []; - default: - return []; +function scan(messages: ModelMessage[]) { + const requests = new Map(); + const decisions = new Map(); + const calls = new Map(); + const providerCalls = new Set(); + const results = new Set(); + const consumed = new Set(); + let waiting: string[] = []; + for (const message of messages) { + if (message.role === "assistant") { + // Provider-owned approvals do not always produce a tool-result. A later + // assistant response is the durable evidence that they were consumed. + waiting.forEach((approvalId) => consumed.add(approvalId)); + waiting = []; } - }); -} - -function unitKeys({ approvalId, toolCallId }: Request): string[] { - return [ - `tool-call:${toolCallId}`, - `tool-approval-request:${approvalId}`, - `tool-approval-response:${approvalId}`, - `tool-result:${toolCallId}`, - ]; -} - -function isProviderCall( - stored: ReadonlyMap, - toolCallId: string, -): boolean { - const call = stored.get(`tool-call:${toolCallId}`); - return call?.type === "tool-call" && call.providerExecuted === true; -} - -function withOptions(providerOptions: ModelMessage["providerOptions"]) { - return providerOptions ? { providerOptions } : {}; + for (const part of partsOf(message)) { + if (part.type === "tool-approval-request") { + add(requests, part.approvalId, part); + } else if (part.type === "tool-approval-response") { + add(decisions, part.approvalId, part); + waiting.push(part.approvalId); + } else if (part.type === "tool-call") { + add(calls, part.toolCallId, part); + if (part.providerExecuted) providerCalls.add(part.toolCallId); + } else if (part.type === "tool-result") { + results.add(part.toolCallId); + } + } + } + const settled = ({ approvalId, toolCallId }: Request) => + results.has(toolCallId) || + (providerCalls.has(toolCallId) && consumed.has(approvalId)); + return { requests, decisions, calls, providerCalls, settled }; } -function partId(part: ApprovalPart): string { - return part.type === "tool-approval-request" || - part.type === "tool-approval-response" - ? `approval ${part.approvalId}` - : part.toolCallId; +function strip(plan: ApprovalPlan, messages: ModelMessage[]): ModelMessage[] { + return messages.flatMap((message): ModelMessage[] => { + if (!Array.isArray(message.content)) return [message]; + const content = (message.content as Part[]).filter( + (part) => + part.type !== "tool-approval-request" && + part.type !== "tool-approval-response" && + !( + (part.type === "tool-call" || part.type === "tool-result") && + plan.hidden.has(part.toolCallId) + ), + ); + if (content.length === message.content.length) return [message]; + return content.length > 0 ? [{ ...message, content } as ModelMessage] : []; + }); } -// Part-level provider options are the handler's to add, and a result's output -// is the handler's to redact or trim since it cannot cause execution; -// everything else must match what was stored. -function sameApprovalPart(a: ApprovalPart, b: ApprovalPart): boolean { - return approvalPartSignature(a) === approvalPartSignature(b); +function only(byId: Map, type: string, approvalId: string) { + const parts = byId.get(approvalId) ?? []; + if (parts.length > 1) { + throw new Error( + `Stored context has more than one ${type} for approval ${approvalId}`, + ); + } + return parts; } -function approvalPartSignature(part: ApprovalPart): string { - switch (part.type) { - case "tool-approval-request": - return signature({ - type: part.type, - approvalId: part.approvalId, - toolCallId: part.toolCallId, - isAutomatic: part.isAutomatic ?? false, - signature: part.signature, - }); - case "tool-call": - return signature({ - type: part.type, - toolCallId: part.toolCallId, - toolName: part.toolName, - input: part.input, - providerExecuted: part.providerExecuted ?? false, - }); - case "tool-approval-response": - return signature({ - type: part.type, - approvalId: part.approvalId, - approved: part.approved, - reason: part.reason, - providerExecuted: part.providerExecuted ?? false, - }); - case "tool-result": - return signature({ - type: part.type, - toolCallId: part.toolCallId, - toolName: part.toolName, - }); - } +function add(byId: Map, id: K, part: T) { + byId.set(id, [...(byId.get(id) ?? []), part]); } -function signature(value: unknown): string { - return JSON.stringify(convexToJson((value ?? null) as Value)); +function partsOf(message: ModelMessage): Part[] { + return Array.isArray(message.content) ? (message.content as Part[]) : []; } diff --git a/src/vercel/client/search.test.ts b/src/vercel/client/search.test.ts index 7f74fc26..892e6b1f 100644 --- a/src/vercel/client/search.test.ts +++ b/src/vercel/client/search.test.ts @@ -312,7 +312,7 @@ describe("search.ts", () => { expect(result[1]._id).toBe("2"); }); - it("should keep tool calls with approval request but NO approval response (auto-deny handles them)", () => { + it("should drop tool calls whose approval request has no response", () => { const messages: MessageDoc[] = [ { _id: "1", @@ -336,23 +336,17 @@ describe("search.ts", () => { }, order: 1, } as MessageDoc, - // No approval response provided ]; const result = filterOutOrphanedToolMessages(messages); expect(result).toHaveLength(1); - // The assistant message should keep the tool-call (auto-deny resolves it downstream) const assistantContent = result[0].message?.content; expect(Array.isArray(assistantContent)).toBe(true); if (Array.isArray(assistantContent)) { - expect(assistantContent).toHaveLength(3); - expect(assistantContent.find((p) => p.type === "text")).toBeDefined(); - expect( - assistantContent.find((p) => p.type === "tool-approval-request"), - ).toBeDefined(); - expect( - assistantContent.find((p) => p.type === "tool-call"), - ).toBeDefined(); + expect(assistantContent.map((p) => p.type)).toEqual([ + "text", + "tool-approval-request", + ]); } }); diff --git a/src/vercel/client/search.ts b/src/vercel/client/search.ts index 403f8f00..2c04577d 100644 --- a/src/vercel/client/search.ts +++ b/src/vercel/client/search.ts @@ -29,12 +29,12 @@ import type { ActionCtx, } from "./types.js"; import { inlineMessagesFiles } from "./files.js"; +import { docsToModelMessages, toModelMessage } from "../mapping.js"; import { - autoDenyUnresolvedApprovals, - docsToModelMessages, - toModelMessage, -} from "../mapping.js"; -import { prepareApprovalContext } from "./prepareApprovalContext.js"; + finishApprovals, + hideApprovals, + planApprovals, +} from "./prepareApprovalContext.js"; const DEFAULT_VECTOR_SCORE_THRESHOLD = 0.0; // Bound the rare boundary-extension query; incomplete orders are trimmed below. @@ -275,72 +275,41 @@ export async function fetchRecentAndSearchMessages( } /** - * Filter out tool messages that don't have both a tool call and response. - * For the approval workflow, tool calls with approval responses (but no tool-results yet) - * should also be kept. + * Filter out tool calls without a result and results without a call. As for + * the AI SDK, an approval decision answers its call. * @param docs The messages to filter. * @returns The filtered messages. */ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) { const toolCallIds = new Set(); - const toolResultIds = new Set(); - // Track approval workflow: toolCallId → approvalId - const approvalRequestsByToolCallId = new Map(); - // Track which approvalIds have responses - const approvalResponseIds = new Set(); - - const result: MessageDoc[] = []; + const answered = new Set(); + const callIdsByApprovalId = new Map(); + const decided: string[] = []; for (const doc of docs) { - if (doc.message && Array.isArray(doc.message.content)) { - for (const content of doc.message.content) { - if (content.type === "tool-call") { - toolCallIds.add(content.toolCallId); - } else if (content.type === "tool-result") { - toolResultIds.add(content.toolCallId); - } else if (content.type === "tool-approval-request") { - const approvalRequest = content as { - type: "tool-approval-request"; - toolCallId: string; - approvalId: string; - }; - approvalRequestsByToolCallId.set( - approvalRequest.toolCallId, - approvalRequest.approvalId, - ); - } else if (content.type === "tool-approval-response") { - const approvalResponse = content as { - type: "tool-approval-response"; - approvalId: string; - }; - approvalResponseIds.add(approvalResponse.approvalId); - } + if (!Array.isArray(doc.message?.content)) continue; + for (const part of doc.message.content) { + if (part.type === "tool-call") toolCallIds.add(part.toolCallId); + else if (part.type === "tool-result") answered.add(part.toolCallId); + else if (part.type === "tool-approval-request") { + callIdsByApprovalId.set(part.approvalId, part.toolCallId); + } else if (part.type === "tool-approval-response") { + decided.push(part.approvalId); } } } + for (const approvalId of decided) { + const toolCallId = callIdsByApprovalId.get(approvalId); + if (toolCallId) answered.add(toolCallId); + } - // Helper: check if tool call has a corresponding approval response - const hasApprovalResponse = (toolCallId: string) => { - const approvalId = approvalRequestsByToolCallId.get(toolCallId); - return approvalId !== undefined && approvalResponseIds.has(approvalId); - }; - - // Helper: check if tool call has a pending approval request - // (auto-deny handles these downstream, so they must survive the filter) - const hasApprovalRequest = (toolCallId: string) => { - return approvalRequestsByToolCallId.has(toolCallId); - }; - + const result: MessageDoc[] = []; for (const doc of docs) { if ( doc.message?.role === "assistant" && Array.isArray(doc.message.content) ) { const content = doc.message.content.filter( - (p) => - p.type !== "tool-call" || - toolResultIds.has(p.toolCallId) || - hasApprovalResponse(p.toolCallId) || - hasApprovalRequest(p.toolCallId), + (p) => p.type !== "tool-call" || answered.has(p.toolCallId), ); if (content.length) { result.push({ @@ -352,14 +321,9 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) { }); } } else if (doc.message?.role === "tool") { - const content = doc.message.content.filter((c) => { - // tool-result parts have toolCallId - if (c.type === "tool-result") { - return toolCallIds.has(c.toolCallId); - } - // tool-approval-response parts don't have toolCallId, so include them - return true; - }); + const content = doc.message.content.filter( + (p) => p.type !== "tool-result" || toolCallIds.has(p.toolCallId), + ); if (content.length) { result.push({ ...doc, @@ -677,15 +641,25 @@ export async function fetchContextWithPrompt( } } - const continuation = + const approvals = planApprovals( + recentMessages, + searchMessages, args.prompt === undefined - ? approvalContinuation(promptMessage, searchMessages, recentMessages) - : undefined; - const search = docsToModelMessages(searchMessages); - const recent = docsToModelMessages(prePromptDocs); - const inputMessages = messages.map(toModelMessage); - const inputPrompt = promptArray.map(toModelMessage); - const existingResponses = docsToModelMessages(existingResponseDocs); + ? continuedApprovals(promptMessage, recentMessages) + : [], + ); + const [search] = hideApprovals(approvals, [ + docsToModelMessages(searchMessages), + ]); + const [recent, inputMessages, inputPrompt, existingResponses] = hideApprovals( + approvals, + [ + docsToModelMessages(prePromptDocs), + messages.map(toModelMessage), + promptArray.map(toModelMessage), + docsToModelMessages(existingResponseDocs), + ], + ); const allMessages = [ ...search, @@ -694,31 +668,21 @@ export async function fetchContextWithPrompt( ...inputPrompt, ...existingResponses, ]; - let processedMessages = args.contextHandler - ? await args.contextHandler(ctx, { - allMessages, - search, - recent, - inputMessages, - inputPrompt, - existingResponses, - userId, - threadId, - }) - : allMessages; - - if (continuation) { - processedMessages = prepareApprovalContext( - processedMessages, - continuation.approvalIds, - continuation.storedContext, - ); - } - - // Post-process: auto-deny unresolved approvals so the AI SDK sees a - // complete history. Applied after contextHandler so custom handlers - // don't need to handle this. - processedMessages = autoDenyUnresolvedApprovals(processedMessages); + let processedMessages = finishApprovals( + approvals, + args.contextHandler + ? await args.contextHandler(ctx, { + allMessages, + search, + recent, + inputMessages, + inputPrompt, + existingResponses, + userId, + threadId, + }) + : allMessages, + ); // Process messages to inline localhost files (if not, file urls pointing to localhost will be sent to LLM providers) if (process.env.CONVEX_CLOUD_URL?.startsWith("http://127.0.0.1")) { @@ -732,59 +696,36 @@ export async function fetchContextWithPrompt( }; } -// Read before the context handler runs and from copies, since the messages -// it is given share nested objects with these docs. -function approvalContinuation( +// A decision prompt continues every request in the one message on its order +// that it answers. +function continuedApprovals( promptMessage: MessageDoc | undefined, - searchMessages: MessageDoc[], recentMessages: MessageDoc[], -): { approvalIds: Set; storedContext: ModelMessage[] } | undefined { - if ( - promptMessage?.message?.role !== "tool" || - !promptMessage.message.content.some( - (part) => part.type === "tool-approval-response", - ) - ) { - return undefined; - } - const responseIds = new Set( - promptMessage.message.content.flatMap((part) => - part.type === "tool-approval-response" ? [part.approvalId] : [], - ), +): string[] { + if (promptMessage?.message?.role !== "tool") return []; + const answered = promptMessage.message.content.flatMap((part) => + part.type === "tool-approval-response" ? [part.approvalId] : [], ); - const requestContents = recentMessages.flatMap((doc) => - doc.order === promptMessage.order && - doc.message?.role === "assistant" && - Array.isArray(doc.message.content) && - doc.message.content.some( - (part) => - part.type === "tool-approval-request" && - responseIds.has(part.approvalId), + if (answered.length === 0) return []; + const requests = recentMessages + .filter((doc) => doc.order === promptMessage.order) + .map((doc) => + doc.message?.role === "assistant" && Array.isArray(doc.message.content) + ? doc.message.content.flatMap((part) => + part.type === "tool-approval-request" ? [part.approvalId] : [], + ) + : [], ) - ? [doc.message.content] - : [], - ); - if (requestContents.length !== 1) { - throw new Error( - "Approval continuation requires one complete request message", - ); - } - const approvalIds = new Set( - requestContents[0].flatMap((part) => - part.type === "tool-approval-request" ? [part.approvalId] : [], - ), - ); - if ([...responseIds].some((approvalId) => !approvalIds.has(approvalId))) { + .filter((ids) => ids.some((id) => answered.includes(id))); + if ( + requests.length !== 1 || + answered.some((id) => !requests[0].includes(id)) + ) { throw new Error( "Approval continuation requires one complete request message", ); } - return { - approvalIds, - storedContext: docsToModelMessages( - structuredClone([...searchMessages, ...recentMessages]), - ), - }; + return requests[0]; } export function getPromptArray( diff --git a/src/vercel/client/toolApprovals.test.ts b/src/vercel/client/toolApprovals.test.ts index c3e15d64..c3106621 100644 --- a/src/vercel/client/toolApprovals.test.ts +++ b/src/vercel/client/toolApprovals.test.ts @@ -1,8 +1,21 @@ +import type { + LanguageModelV4CallOptions, + LanguageModelV4Content, +} from "@ai-sdk/provider"; import type { ModelMessage, ToolModelMessage } from "ai"; import { afterEach, describe, expect, test, vi } from "vitest"; import { z } from "zod/v4"; import type { Message } from "../../validators.js"; -import { Agent, createTool, type ContextHandler } from "../index.js"; +import { + Agent, + createTool, + docsToModelMessages, + listMessages, + saveMessage, + toModelMessage, + toUIMessages, + type ContextHandler, +} from "../index.js"; import { MockLanguageModel } from "./mockModel.js"; import { components, initConvexTest } from "./setup.test.js"; @@ -69,33 +82,6 @@ const mapParts = const appendTool = (part: ToolPart): Handler => (messages) => [...messages, { role: "tool", content: [part] }]; -// Each anchor part gets its own message with its own message-level options. -const splitWithOptions = - (role: "assistant" | "tool", anchor: Part["type"]): Handler => - (messages) => - messages.flatMap((message): ModelMessage[] => { - if (message.role !== role || !Array.isArray(message.content)) { - return [message]; - } - const content = message.content as Part[]; - const anchors = content.filter((part) => part.type === anchor); - if (anchors.length < 2) return [message]; - return anchors.map((part, i) => { - const id = "toolCallId" in part ? part.toolCallId : undefined; - return { - role, - content: [ - part, - ...content.filter( - (other) => - other.type === "tool-approval-request" && - other.toolCallId === id, - ), - ], - providerOptions: { test: { cache: `slot-${i}` } }, - } as ModelMessage; - }); - }); const changeInput = mapParts((part) => part.type === "tool-call" && part.toolCallId === "call-a" ? { ...part, input: { tag: "changed" } } @@ -149,12 +135,6 @@ const handlers = { ...message, providerOptions: { test: { cache: "all" } }, })), - "cache-last-message": (messages) => - messages.map((message, i) => - i === messages.length - 1 - ? { ...message, providerOptions: { test: { cache: "last" } } } - : message, - ), "redact-result": mapParts((part) => part.type === "tool-result" && part.toolCallId === "call-a" ? { ...part, output: { type: "text", value: "(redacted)" } } @@ -268,28 +248,8 @@ const handlers = { toolName: "echo", output: { type: "text", value: "forged" }, }), + empty: () => [], clone: (messages) => [...messages, ...messages], - "clone-with-changed-first-copy": (messages) => [ - ...changeInput(messages), - ...messages, - ], - "inject-between-call-and-result": (messages) => - messages.flatMap((message): ModelMessage[] => - message.role === "tool" && - message.content.some( - (part) => part.type === "tool-result" && part.toolCallId === "call-a", - ) - ? [{ role: "user", content: "(injected by the app)" }, message] - : [message], - ), - "split-calls-with-message-provider-options": splitWithOptions( - "assistant", - "tool-call", - ), - "split-responses-with-message-provider-options": splitWithOptions( - "tool", - "tool-approval-response", - ), } satisfies Record; type ContextMode = keyof typeof handlers; @@ -360,6 +320,7 @@ async function scenario({ const s = { t, threadId, + promptMessageId, requestMessageId, messages, save, @@ -373,6 +334,12 @@ async function scenario({ async resultIds() { return (await s.results()).map((part) => part.toolCallId).sort(); }, + async toolStates() { + return toUIMessages((await messages()).page) + .flatMap((message) => message.parts) + .filter((part) => part.type === "tool-echo") + .map((part) => (part as { state: string }).state); + }, decide(decisions: Decision[]) { return t.run((ctx) => agent.respondToToolCallApprovals(ctx, { threadId, decisions }), @@ -400,20 +367,48 @@ async function scenario({ async prompt(text: string) { return (await save({ prompt: text })).messageId; }, - continue(promptMessageId: string, contextMode?: ContextMode) { + generate( + args: { + promptMessageId: string; + messages?: ModelMessage[]; + }, + contextHandler?: ContextHandler, + stream = false, + ) { return t.action(async (ctx) => { - const contextHandler: ContextHandler | undefined = contextMode - ? async (_ctx, { allMessages }) => handlers[contextMode](allMessages) - : undefined; + const options = contextHandler ? { contextHandler } : undefined; + if (stream) { + const streamed = await agent.streamText( + ctx, + { threadId }, + args, + options, + ); + await streamed.consumeStream(); + return await streamed.text; + } const result = await agent.generateText( ctx, { threadId }, - { promptMessageId }, - contextHandler ? { contextHandler } : undefined, + args, + options, ); return result.text; }); }, + continue( + promptMessageId: string, + contextMode?: ContextMode, + stream = false, + ) { + return s.generate( + { promptMessageId }, + contextMode + ? async (_ctx, { allMessages }) => handlers[contextMode](allMessages) + : undefined, + stream, + ); + }, // `a` approved and continued, `b` still open. async completeA() { const first = await s.approve("a"); @@ -433,7 +428,8 @@ async function scenario({ type Scenario = Awaited>; function lastModelPrompt() { - const prompt = model.doGenerateCalls.at(-1)?.prompt; + const prompt = + model.doStreamCalls.at(-1)?.prompt ?? model.doGenerateCalls.at(-1)?.prompt; expect(prompt).toBeDefined(); return prompt ?? []; } @@ -450,6 +446,12 @@ function partsFor(id: string) { ); } +function callInputs(id: string) { + return partsFor(id).flatMap((part) => + part.type === "tool-call" ? [part.input] : [], + ); +} + function expectNoDanglingLocalCalls() { const calls = new Set(); const results = new Set(); @@ -468,13 +470,14 @@ function expectNoDanglingLocalCalls() { // Providers with strict turn ordering need each local call answered by the // very next message, and a result must not be shown twice per call. -function expectEveryCallAnsweredNext() { - const prompt = lastModelPrompt(); +function expectEveryCallAnsweredNext( + prompt: { role: string; content: unknown }[] = lastModelPrompt(), +) { for (const [i, message] of prompt.entries()) { if (message.role !== "assistant" || !Array.isArray(message.content)) { continue; } - const callIds = message.content.flatMap((part) => + const callIds = (message.content as Part[]).flatMap((part) => part.type === "tool-call" && !part.providerExecuted ? [part.toolCallId] : [], @@ -482,12 +485,11 @@ function expectEveryCallAnsweredNext() { if (callIds.length === 0) continue; const next = prompt[i + 1]; expect(next?.role).toBe("tool"); - const answered = - next?.role === "tool" - ? next.content.flatMap((part) => - part.type === "tool-result" ? [part.toolCallId] : [], - ) - : []; + const answered = Array.isArray(next?.content) + ? (next.content as Part[]).flatMap((part) => + part.type === "tool-result" ? [part.toolCallId] : [], + ) + : []; for (const id of callIds) { expect(answered.filter((answeredId) => answeredId === id)).toHaveLength( 1, @@ -499,6 +501,7 @@ function expectEveryCallAnsweredNext() { afterEach(() => { executed.length = 0; model.doGenerateCalls.length = 0; + model.doStreamCalls.length = 0; vi.restoreAllMocks(); }); @@ -597,6 +600,18 @@ describe("tool approval semantics", () => { expectEveryCallAnsweredNext(); }); + test("siblings continued separately each answer their call in later context", async () => { + const s = await scenario(); + await s.completeA(); + await s.continue(await s.approve("b")); + + await s.generate({ promptMessageId: await s.prompt("What happened?") }); + + expect(executed).toEqual(["a", "b"]); + expectNoDanglingLocalCalls(); + expectEveryCallAnsweredNext(); + }); + test("a consumed provider-owned decision on another request stays in history", async () => { const s = await scenario({ ids: ["provider"], providerExecuted: true }); const first = await s.approve("provider"); @@ -780,80 +795,43 @@ describe("tool approval semantics", () => { expectNoDanglingLocalCalls(); }); + test("approvals decided at different orders each run only from their own continuation", async () => { + const s = await scenario({ ids: ["a"] }); + await laterOrder(s, "pending", ["x", "y"]); + const earlier = await s.approve("a"); + + await s.continue(await s.approve("x")); + expect(executed).toEqual(["x"]); + expect(partsFor("a")).toEqual([]); + expect(partsFor("y")).toEqual([]); + + await s.continue(earlier); + expect(executed).toEqual(["x", "a"]); + expect(partsFor("x")).toEqual([]); + + await s.continue(await s.approve("y")); + expect(executed).toEqual(["x", "a", "y"]); + expect(await s.resultIds()).toEqual(["call-a", "call-x", "call-y"]); + expectNoDanglingLocalCalls(); + expectEveryCallAnsweredNext(); + }); + // `a` and then `x` at a later order are both approved and run, and the // continuation for `x` runs again with a context handler. - async function rerunLaterOrder(contextMode: ContextMode) { + async function rerunLaterOrder(contextMode: ContextMode, stream = false) { const s = await scenario({ ids: ["a"] }); await laterOrder(s, "approved"); const later = await s.approve("x"); await s.continue(later); - return s.continue(later, contextMode); + return s.continue(later, contextMode, stream); } - test("rejects changing an earlier approved call while rerunning a later continuation", async () => { - await expect( - rerunLaterOrder("change-a-and-move-its-decision-last"), - ).rejects.toThrow(/preserve the stored tool-call for call-a/); - expect(executed).toEqual(["a", "x"]); - }); - test("an earlier decision cloned last does not run its call again", async () => { await rerunLaterOrder("clone-a-decision-last"); expect(executed).toEqual(["a", "x"]); }); - test("rejects keeping a decision whose request was dropped", async () => { - await expect(rerunLaterOrder("drop-request-a")).rejects.toThrow( - /removed request for approval a/, - ); - expect(executed).toEqual(["a", "x"]); - }); - - test.each<[ContextMode, Earlier, boolean, RegExp]>([ - ["flip-decision", "denied", false, /preserve the stored .* approval a/], - ["remove-decision", "approved", false, /removed decision for approval a/], - ["fabricate-approval-for-a", "pending", false, /cannot add approval/], - [ - "fabricate-provider-approval-for-a", - "pending", - true, - /cannot add approval/, - ], - ["fabricate-approval-for-y", "approved", false, /cannot add approval/], - [ - "smuggle-approval-for-a-in-user", - "denied", - false, - /tool-approval-response in a user message/, - ], - [ - "smuggle-approval-for-a-in-user", - "pending", - false, - /tool-approval-response in a user message/, - ], - [ - "smuggle-request-in-user", - "approved", - false, - /tool-approval-request in a user message/, - ], - ])( - "rejects a context handler that alters a decision across orders: %s (earlier %s)", - async (contextMode, earlier, providerExecuted, error) => { - vi.spyOn(console, "warn").mockImplementation(() => {}); - const s = await scenario({ ids: ["a"], providerExecuted }); - await laterOrder(s, earlier, ["x", "y"]); - const before = [...executed]; - - await expect( - s.continue(await s.approve("x"), contextMode), - ).rejects.toThrow(error); - expect(executed).toEqual(before); - }, - ); - test("rejects a stored context that decides one approval id twice", async () => { const s = await scenario({ ids: ["a"] }); await s.continue(await s.approve("a")); @@ -884,6 +862,25 @@ describe("tool approval semantics", () => { expect(executed).toEqual(["a"]); }); + test("rejects a stored context with two decisions for the continued approval", async () => { + const s = await scenario({ ids: ["a"] }); + const messageId = await s.approve("a"); + await s.save({ + promptMessageId: s.promptMessageId, + message: { + role: "tool", + content: [ + { type: "tool-approval-response", approvalId: "a", approved: true }, + ], + }, + }); + + await expect(s.continue(messageId)).rejects.toThrow( + /more than one tool-approval-response for approval a/, + ); + expect(executed).toEqual([]); + }); + test("rejects a response message that answers another order's request", async () => { const s = await scenario(); const later = await s.prompt("Next"); @@ -905,79 +902,6 @@ describe("tool approval semantics", () => { expect(executed).toEqual([]); }); - // A context handler may reorder, wrap, or annotate the stored record, but it - // cannot change what was decided or executed. Shapes: "pending" has `a` - // approved but not run; "completed" has `a` run and `b` approved; "with - // another request" has `a` approved and a second request `c` at the order. - type Shape = "pending" | "completed" | "with-another-request"; - async function shaped(s: Scenario, shape: Shape) { - if (shape === "pending") return s.approve("a"); - if (shape === "completed") { - await s.completeA(); - return s.approve("b"); - } - const first = await s.approve("a"); - await s.addRequest(first, ["c"]); - return first; - } - - test.each<[ContextMode, Shape, RegExp]>([ - ["change-input", "pending", /call-a/], - ["flip-decision", "pending", /approval a/], - ["flip-provider-executed", "pending", /approval a/], - [ - "sign-request", - "pending", - /preserve the stored tool-approval-request for approval a/, - ], - [ - "mark-request-automatic", - "pending", - /preserve the stored tool-approval-request for approval a/, - ], - ["fabricate-result", "pending", /cannot add results/], - ["fabricate-approval", "with-another-request", /cannot add approval/], - [ - "fabricate-approval-for-another-request", - "with-another-request", - /cannot add approval/, - ], - [ - "clone-with-changed-first-copy", - "completed", - /preserve the stored tool-call for call-a/, - ], - ["remove-result", "completed", /removed result for call-a/], - ])( - "rejects a context handler that alters the record: %s (%s)", - async (contextMode, shape, error) => { - const s = await scenario(); - const messageId = await shaped(s, shape); - const before = [...executed]; - - await expect(s.continue(messageId, contextMode)).rejects.toThrow(error); - expect(executed).toEqual(before); - }, - ); - - test.each([ - "clone", - "inject-between-call-and-result", - "fabricate-result", - ])( - "a reshaped context still answers each call once: %s", - async (contextMode) => { - const s = await scenario(); - await s.completeA(); - - await s.continue(await s.approve("b"), contextMode); - - expect(executed).toEqual(["a", "b"]); - expect(await s.resultIds()).toEqual(["call-a", "call-b"]); - expectEveryCallAnsweredNext(); - }, - ); - test.each([true, false])( "a handler may redact an approved call's result (at an earlier order: %s)", async (earlier) => { @@ -1004,37 +928,7 @@ describe("tool approval semantics", () => { }, ); - test.each(["a", "x"])( - "rejects a handler that edits call-%s's input in place", - async (id) => { - const s = await scenario({ ids: ["a"] }); - await laterOrder(s, "approved"); - - await expect( - s.continue( - await s.approve("x"), - id === "a" ? "mutate-a-in-place" : "mutate-x-in-place", - ), - ).rejects.toThrow(`preserve the stored tool-call for call-${id}`); - expect(executed).toEqual(["a"]); - }, - ); - - test("allows a context handler to add call provider options", async () => { - const s = await scenario({ ids: ["a"] }); - - await s.continue(await s.approve("a"), "add-call-provider-options"); - - expect(executed).toEqual(["a"]); - expect(partsFor("a")).toContainEqual( - expect.objectContaining({ - type: "tool-call", - providerOptions: { test: { trace: "allowed" } }, - }), - ); - }); - - test("message provider options a handler applies survive result relocation", async () => { + test("the continued approvals follow the handler's messages", async () => { const s = await scenario(); await s.completeA(); @@ -1044,47 +938,483 @@ describe("tool approval semantics", () => { ); expect(executed).toEqual(["a", "b"]); - for (const message of lastModelPrompt()) { + const prompt = lastModelPrompt(); + for (const message of prompt.slice(0, -2)) { expect(message.providerOptions).toEqual({ test: { cache: "all" } }); } + for (const message of prompt.slice(-2)) { + expect(message.providerOptions).toBeUndefined(); + } + expect(partsFor("b").map((part) => part.type)).toEqual([ + "tool-call", + "tool-result", + ]); }); - // A handler that targets one message (cache control on the last message, - // say) must keep working when the pair is synthesized from several. - test.each<[ContextMode, boolean, "assistant" | "tool", string]>([ - ["split-calls-with-message-provider-options", false, "assistant", "slot-1"], - ["split-responses-with-message-provider-options", true, "tool", "slot-1"], - ["cache-last-message", true, "tool", "last"], - ])( - "the synthesized pair takes the last absorbed message's provider options: %s", - async (contextMode, providerExecuted, role, cache) => { - const s = await scenario({ providerExecuted }); - const messageId = - contextMode === "cache-last-message" - ? (await s.approve("a"), await s.approve("b")) - : await s.approve("a", "b"); - - await s.continue(messageId, contextMode); - - const final = lastModelPrompt() - .filter((message) => message.role === role) - .at(-1); - expect(final?.providerOptions).toEqual({ test: { cache } }); - expect(partsFor("a")).not.toEqual([]); - expect(partsFor("b")).not.toEqual([]); - }, - ); + test("a handler that rebuilds context from storage runs each approval once", async () => { + const s = await scenario(); + await s.completeA(); + const messageId = await s.approve("b"); + const rebuild: ContextHandler = async (ctx) => { + const { page } = await agent.listMessages(ctx, { + threadId: s.threadId, + paginationOpts: { cursor: null, numItems: 200 }, + }); + return page + .reverse() + .flatMap((doc) => (doc.message ? [toModelMessage(doc.message)] : [])); + }; - test("an unrelated generation still auto-denies unresolved approvals", async () => { - const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + await s.generate({ promptMessageId: messageId }, rebuild); + await s.generate({ promptMessageId: messageId }, rebuild); + + expect(executed).toEqual(["a", "b"]); + expect(await s.resultIds()).toEqual(["call-a", "call-b"]); + }); + + test("an approved but un-continued approval does not run from an ordinary generation", async () => { const s = await scenario(); await s.approve("a"); - await s.continue(await s.prompt("Start unrelated work")); + await s.generate({ promptMessageId: s.promptMessageId }); expect(executed).toEqual([]); - expect(warn).toHaveBeenCalledWith( - expect.stringContaining("Auto-denying unresolved tool approval b"), - ); + expect(partsFor("a")).toEqual([]); + expect(partsFor("b")).toEqual([]); + }); + + test("caller messages cannot add approvals to an ordinary generation", async () => { + const s = await scenario(); + const later = await s.prompt("Something else"); + + await s.generate({ + promptMessageId: later, + messages: [ + toModelMessage(request(["b"])), + { + role: "tool", + content: [ + { type: "tool-approval-response", approvalId: "b", approved: true }, + ], + }, + ], + }); + + expect(executed).toEqual([]); + expect(partsFor("b")).toEqual([]); + expect(await s.toolStates()).toEqual([ + "approval-requested", + "approval-requested", + ]); + }); + + // docs/tool-approval.mdx used to say an unrelated generation denies pending + // approvals. It only told the model so: the approval stayed open, and + // approving it later still ran the tool. + test("an unrelated generation leaves a pending approval out of context, and it runs once approved", async () => { + const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); + const s = await scenario({ ids: ["a"] }); + + await s.generate({ promptMessageId: await s.prompt("Never mind") }); + + expect(partsFor("a")).toEqual([]); + expect(warn).not.toHaveBeenCalled(); + expect(await s.toolStates()).toEqual(["approval-requested"]); + + await s.continue(await s.approve("a")); + + expect(executed).toEqual(["a"]); + expect(await s.resultIds()).toEqual(["call-a"]); }); + + // The handler never sees approval protocol or an approval that has not run. + // Whatever it returns, only stored decisions run, with their stored input, + // and each at most once. + describe.each([false, true])( + "context handler boundary (streamText=%s)", + (stream) => { + const fabricated: [ContextMode, boolean][] = [ + ["fabricate-approval-for-a", false], + ["fabricate-provider-approval-for-a", true], + ["fabricate-approval-for-y", false], + ["smuggle-approval-for-a-in-user", false], + ["smuggle-request-in-user", false], + ]; + test.each(fabricated)( + "no unapproved execution: %s across orders is dropped", + async (contextMode, providerExecuted) => { + const s = await scenario({ ids: ["a"], providerExecuted }); + await laterOrder(s, "pending", ["x", "y"]); + + await s.continue(await s.approve("x"), contextMode, stream); + + expect(executed).toEqual(["x"]); + expect(partsFor("a")).toEqual([]); + expect(partsFor("y")).toEqual([]); + }, + ); + + test.each([ + "fabricate-result", + "fabricate-approval", + "fabricate-approval-for-another-request", + ])( + "no unapproved execution: %s beside the continued request is dropped", + async (contextMode) => { + const s = await scenario(); + const first = await s.approve("a"); + await s.addRequest(first, ["c"]); + + await s.continue(first, contextMode, stream); + + expect(executed).toEqual(["a"]); + expect(await s.resultIds()).toEqual(["call-a"]); + expect(partsFor("a")).not.toContainEqual( + expect.objectContaining({ + output: { type: "text", value: "forged" }, + }), + ); + }, + ); + + test("no unapproved execution: caller messages cannot approve a sibling", async () => { + const s = await scenario(); + const messageId = await s.approve("a"); + + await s.generate( + { + promptMessageId: messageId, + messages: [ + { + role: "tool", + content: [ + { + type: "tool-approval-response", + approvalId: "b", + approved: true, + }, + ], + }, + ], + }, + undefined, + stream, + ); + + expect(executed).toEqual(["a"]); + }); + + test.each([ + "change-input", + "mutate-a-in-place", + "sign-request", + "mark-request-automatic", + "flip-decision", + "flip-provider-executed", + "add-call-provider-options", + "empty", + ])( + "no altered input: %s cannot change the continued call", + async (contextMode) => { + const s = await scenario(); + + await s.continue(await s.approve("a"), contextMode, stream); + + expect(executed).toEqual(["a"]); + expect(callInputs("a")).toEqual([{ tag: "a" }]); + expect(partsFor("a")).not.toContainEqual( + expect.objectContaining({ providerOptions: expect.anything() }), + ); + }, + ); + + test("a handler may edit a completed call's history", async () => { + const s = await scenario(); + await s.completeA(); + + await s.continue(await s.approve("b"), "change-input", stream); + + expect(executed).toEqual(["a", "b"]); + expect(callInputs("a")).toEqual([{ tag: "changed" }]); + expectEveryCallAnsweredNext(); + }); + + test.each([ + "change-a-and-move-its-decision-last", + "clone-a-decision-last", + "drop-request-a", + "flip-decision", + "remove-decision", + "mutate-a-in-place", + "mutate-x-in-place", + ])("no re-execution: rerunning with %s runs nothing", async (mode) => { + await rerunLaterOrder(mode, stream); + + expect(executed).toEqual(["a", "x"]); + }); + + test.each(["a", "x"])( + "an in-place edit of call-%s cannot reach the continued call", + async (id) => { + const s = await scenario({ ids: ["a"] }); + await laterOrder(s, "approved"); + + await s.continue( + await s.approve("x"), + id === "a" ? "mutate-a-in-place" : "mutate-x-in-place", + stream, + ); + + expect(executed).toEqual(["a", "x"]); + expect(callInputs("x")).toEqual([{ tag: "x" }]); + }, + ); + + test.each(["clone", "fabricate-result"])( + "a duplicated completed result runs nothing again: %s", + async (contextMode) => { + const s = await scenario(); + await s.completeA(); + + await s.continue(await s.approve("b"), contextMode, stream); + + expect(executed).toEqual(["a", "b"]); + expect(await s.resultIds()).toEqual(["call-a", "call-b"]); + }, + ); + + test("results redactable: the model sees the redacted output", async () => { + const s = await scenario(); + await s.completeA(); + + await s.continue(await s.approve("b"), "redact-result", stream); + + expect(partsFor("a")).toContainEqual( + expect.objectContaining({ + type: "tool-result", + output: { type: "text", value: "(redacted)" }, + }), + ); + expectEveryCallAnsweredNext(); + }); + + test("a completed result the handler drops fails before anything runs", async () => { + const s = await scenario(); + await s.completeA(); + + await expect( + s.continue(await s.approve("b"), "remove-result", stream), + ).rejects.toThrow(/call-a without its result/); + + expect(executed).toEqual(["a"]); + expect(await s.resultIds()).toEqual(["call-a"]); + }); + + test("the handler sees no approval protocol or pending call", async () => { + const s = await scenario({ ids: ["a", "b", "c"] }); + await s.completeA(); + const messageId = await s.approve("b"); + let seen: ModelMessage[] = []; + + await s.generate( + { promptMessageId: messageId }, + async (_ctx, { allMessages }) => { + seen = structuredClone(allMessages); + return allMessages; + }, + stream, + ); + + const parts = seen.flatMap((message) => + Array.isArray(message.content) ? (message.content as Part[]) : [], + ); + expect( + parts.filter( + (part) => + part.type === "tool-approval-request" || + part.type === "tool-approval-response", + ), + ).toEqual([]); + expect( + parts.flatMap((part) => + "toolCallId" in part ? [part.toolCallId] : [], + ), + ).toEqual(["call-a", "call-a"]); + expectEveryCallAnsweredNext(seen); + expect(executed).toEqual(["a", "b"]); + }); + }, + ); +}); + +// https://github.com/get-convex/agent/issues/411 +describe("a context handler that returns the whole thread", () => { + type Prompt = LanguageModelV4CallOptions["prompt"]; + const render = (prompt: Prompt) => + prompt.map( + (message) => + `${message.role}: ${ + typeof message.content === "string" + ? message.content + : message.content + .map((part) => + part.type === "text" + ? part.text + : `${part.type}:${"toolCallId" in part ? part.toolCallId : ""}`, + ) + .join(" ") + }`, + ); + const turn = (n: number, done: boolean) => [ + `user: Delete file-${n}.txt`, + `assistant: tool-call:call-${n}`, + `tool: tool-result:call-${n}`, + ...(done ? [`assistant: Deleted file-${n}.txt`] : []), + ]; + + async function approvals({ + turns, + recentMessages, + handler = true, + first = "approve", + }: { + turns: number; + recentMessages: number; + handler?: boolean; + first?: "approve" | "decline" | "leave pending"; + }) { + const reply = (content: LanguageModelV4Content) => ({ + content: [content], + finishReason: { unified: "stop" as const, raw: undefined }, + usage: { + inputTokens: { total: 1, noCache: 1, cacheRead: 0, cacheWrite: 0 }, + outputTokens: { total: 1, text: 1, reasoning: 0 }, + }, + warnings: [], + }); + const steps = Array.from({ length: turns }, (_, i) => i + 1).flatMap( + (n) => [ + reply({ + type: "tool-call", + toolCallId: `call-${n}`, + toolName: "deleteFile", + input: JSON.stringify({ name: `file-${n}.txt` }), + }), + ...(n === 1 && first === "leave pending" + ? [] + : [reply({ type: "text", text: `Deleted file-${n}.txt` })]), + ], + ); + const prompts: Prompt[] = []; + const deleted: string[] = []; + const component = components.agent; + const wholeThread: ContextHandler = async (ctx, { threadId }) => { + const { page } = await listMessages(ctx, component, { + threadId: threadId!, + paginationOpts: { cursor: null, numItems: 100 }, + }); + return docsToModelMessages(page.reverse()); + }; + const repro = new Agent(component, { + name: "repro", + languageModel: new MockLanguageModel({ + doGenerate: async ({ prompt }) => { + prompts.push(prompt); + return steps.shift()!; + }, + doStream: async () => { + throw new Error("unused"); + }, + }), + tools: { + deleteFile: createTool({ + inputSchema: z.object({ name: z.string() }), + needsApproval: true, + execute: async (_ctx, { name }) => { + deleted.push(name); + return `deleted ${name}`; + }, + }), + }, + contextOptions: { recentMessages }, + ...(handler ? { contextHandler: wholeThread } : {}), + }); + const t = initConvexTest(); + const { threadId } = await t.run((ctx) => repro.createThread(ctx)); + for (let n = 1; n <= turns; n++) { + const declined = first === "decline" && n === 1; + const approvalId = await t.action(async (ctx) => { + const { messageId } = await saveMessage(ctx, component, { + threadId, + prompt: `Delete file-${n}.txt`, + }); + const result = await repro.generateText( + ctx, + { threadId }, + { promptMessageId: messageId }, + ); + return result.content.find( + (part) => part.type === "tool-approval-request", + )!.approvalId; + }); + if (first === "leave pending" && n === 1) continue; + const { messageId } = await t.run((ctx) => + declined + ? repro.denyToolCall(ctx, { threadId, approvalId }) + : repro.approveToolCall(ctx, { threadId, approvalId }), + ); + const text = await t.action(async (ctx) => { + const result = await repro.generateText( + ctx, + { threadId }, + { promptMessageId: messageId }, + ); + return result.text; + }); + expect(text).toBe(`Deleted file-${n}.txt`); + expect(deleted.filter((name) => name === `file-${n}.txt`)).toHaveLength( + declined ? 0 : 1, + ); + const sent = render(prompts.at(-1)!); + expect(sent.slice(-3)).toEqual(turn(n, false)); + if (handler) { + expect(sent).toEqual([ + ...(first === "leave pending" ? ["user: Delete file-1.txt"] : []), + ...Array.from({ length: n - 1 }, (_, i) => turn(i + 1, true)) + .flat() + .slice(first === "leave pending" ? 4 : 0), + ...turn(n, false), + ]); + } + expect(sent.filter((line) => line.includes("approval"))).toEqual([]); + expect(new Set(sent).size).toBe(sent.length); + } + expect(deleted).toHaveLength(first === "approve" ? turns : turns - 1); + } + + test.each([ + { turns: 2, recentMessages: 1 }, + { turns: 2, recentMessages: 3 }, + { turns: 2, recentMessages: 4 }, + { turns: 3, recentMessages: 1 }, + { turns: 3, recentMessages: 3 }, + { turns: 5, recentMessages: 1 }, + { turns: 5, recentMessages: 3 }, + ])( + "continues every approval over $turns turns with recentMessages: $recentMessages", + (args) => approvals(args), + ); + + test.each([2, 3])( + "continues later approvals after the first was declined over %i turns", + (turns) => approvals({ turns, recentMessages: 1, first: "decline" }), + ); + + test.each([2, 3])( + "continues later approvals while the first is still pending over %i turns", + (turns) => approvals({ turns, recentMessages: 1, first: "leave pending" }), + ); + + test("continues every approval without a handler", () => + approvals({ turns: 2, recentMessages: 1, handler: false })); }); diff --git a/src/vercel/mapping.test.ts b/src/vercel/mapping.test.ts index f8a11fdd..894af5d6 100644 --- a/src/vercel/mapping.test.ts +++ b/src/vercel/mapping.test.ts @@ -1,4 +1,4 @@ -import { describe, test, expect, vi } from "vitest"; +import { describe, test, expect } from "vitest"; import { guessMimeType, serializeDataOrUrl, @@ -9,7 +9,6 @@ import { serializeContent, toModelMessageContent, toUIFilePart, - autoDenyUnresolvedApprovals, serializeWarnings, } from "./mapping.js"; import { api } from "../component/_generated/api.js"; @@ -900,183 +899,4 @@ describe("mapping", () => { expect(contentTypes(res.messages[2].message)).toEqual(["tool-result"]); }); }); - - describe("autoDenyUnresolvedApprovals", () => { - test("returns messages unchanged when no unresolved approvals", () => { - const messages = [ - { role: "user" as const, content: "hello" }, - { - role: "assistant" as const, - content: [ - { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} }, - { - type: "tool-approval-request", - approvalId: "ap1", - toolCallId: "tc1", - }, - ], - }, - { - role: "tool" as const, - content: [ - { - type: "tool-approval-response", - approvalId: "ap1", - approved: true, - }, - ], - }, - ] as any; - - const result = autoDenyUnresolvedApprovals(messages); - expect(result).toBe(messages); // same reference, no changes - }); - - test("injects synthetic denial for a single unresolved approval", () => { - const messages = [ - { role: "user" as const, content: "hello" }, - { - role: "assistant" as const, - content: [ - { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} }, - { - type: "tool-approval-request", - approvalId: "ap1", - toolCallId: "tc1", - }, - ], - }, - { role: "user" as const, content: "new message" }, - ] as any; - - const result = autoDenyUnresolvedApprovals(messages); - expect(result).toHaveLength(4); // original 3 + 1 synthetic tool message - // Synthetic denial should be inserted right after the assistant message (index 1) - expect(result[2].role).toBe("tool"); - const denialContent = result[2].content as any[]; - expect(denialContent).toHaveLength(1); - expect(denialContent[0].type).toBe("tool-approval-response"); - expect(denialContent[0].approvalId).toBe("ap1"); - expect(denialContent[0].approved).toBe(false); - expect(denialContent[0].reason).toBe( - "auto-denied: new generation started", - ); - // The new user message should follow - expect(result[3].role).toBe("user"); - expect(result[3].content).toBe("new message"); - }); - - test("groups multiple unresolved approvals from the same step into a single synthetic message", () => { - const messages = [ - { - role: "assistant" as const, - content: [ - { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} }, - { type: "tool-call", toolCallId: "tc2", toolName: "b", input: {} }, - { - type: "tool-approval-request", - approvalId: "ap1", - toolCallId: "tc1", - }, - { - type: "tool-approval-request", - approvalId: "ap2", - toolCallId: "tc2", - }, - ], - }, - ] as any; - - const result = autoDenyUnresolvedApprovals(messages); - expect(result).toHaveLength(2); // assistant + 1 synthetic tool message - expect(result[1].role).toBe("tool"); - const denialContent = result[1].content as any[]; - expect(denialContent).toHaveLength(2); - expect(denialContent[0].approvalId).toBe("ap1"); - expect(denialContent[0].approved).toBe(false); - expect(denialContent[1].approvalId).toBe("ap2"); - expect(denialContent[1].approved).toBe(false); - }); - - test("only auto-denies unresolved approvals, leaves resolved ones alone", () => { - const messages = [ - { - role: "assistant" as const, - content: [ - { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} }, - { type: "tool-call", toolCallId: "tc2", toolName: "b", input: {} }, - { - type: "tool-approval-request", - approvalId: "ap1", - toolCallId: "tc1", - }, - { - type: "tool-approval-request", - approvalId: "ap2", - toolCallId: "tc2", - }, - ], - }, - { - role: "tool" as const, - content: [ - { - type: "tool-approval-response", - approvalId: "ap1", - approved: true, - }, - ], - }, - { role: "user" as const, content: "next question" }, - ] as any; - - const result = autoDenyUnresolvedApprovals(messages); - // Should inject a denial for ap2 (unresolved) after the assistant message - expect(result).toHaveLength(4); // assistant + existing tool + synthetic denial + user - // The synthetic denial is inserted after the assistant (index 0) - expect(result[0].role).toBe("assistant"); - expect(result[1].role).toBe("tool"); // synthetic denial for ap2 - const denialContent = result[1].content as any[]; - expect(denialContent).toHaveLength(1); - expect(denialContent[0].approvalId).toBe("ap2"); - expect(denialContent[0].approved).toBe(false); - // Original tool message (ap1 response) follows - expect(result[2].role).toBe("tool"); - const originalToolContent = result[2].content as any[]; - expect(originalToolContent[0].approvalId).toBe("ap1"); - expect(originalToolContent[0].approved).toBe(true); - // User message last - expect(result[3].role).toBe("user"); - }); - - test("emits console.warn for each auto-denied approval", () => { - const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); - const messages = [ - { - role: "assistant" as const, - content: [ - { type: "tool-call", toolCallId: "tc1", toolName: "a", input: {} }, - { type: "tool-call", toolCallId: "tc2", toolName: "b", input: {} }, - { - type: "tool-approval-request", - approvalId: "ap1", - toolCallId: "tc1", - }, - { - type: "tool-approval-request", - approvalId: "ap2", - toolCallId: "tc2", - }, - ], - }, - ] as any; - - autoDenyUnresolvedApprovals(messages); - - expect(warnSpy).toHaveBeenCalledTimes(2); - expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("ap1")); - expect(warnSpy).toHaveBeenCalledWith(expect.stringContaining("ap2")); - warnSpy.mockRestore(); - }); - }); }); diff --git a/src/vercel/mapping.ts b/src/vercel/mapping.ts index 8cb79121..d68f6ab4 100644 --- a/src/vercel/mapping.ts +++ b/src/vercel/mapping.ts @@ -147,97 +147,6 @@ export function docsToModelMessages(messages: MessageDoc[]): ModelMessage[] { .map(toModelMessage); } -/** - * Scan messages for unresolved `tool-approval-request` parts and inject - * synthetic `tool-approval-response` denials so that the AI SDK receives - * a complete history (every tool-call has a corresponding result or denial). - * - * This handles the case where a user sends a new message instead of - * resolving pending approvals — the old approvals are auto-denied rather - * than silently dropped. - */ -export function autoDenyUnresolvedApprovals( - messages: ModelMessage[], -): ModelMessage[] { - // Collect all approval requests: approvalId → { toolCallId, messageIndex } - const requests = new Map< - string, - { toolCallId: string; messageIndex: number } - >(); - // Collect all resolved approval IDs - const resolvedIds = new Set(); - - for (let i = 0; i < messages.length; i++) { - const msg = messages[i]; - if (!Array.isArray(msg.content)) continue; - for (const part of msg.content as any[]) { - if (part.type === "tool-approval-request") { - requests.set(part.approvalId, { - toolCallId: part.toolCallId, - messageIndex: i, - }); - } else if (part.type === "tool-approval-response") { - resolvedIds.add(part.approvalId); - } - } - } - - // Find unresolved approvals - const unresolved: Array<{ - approvalId: string; - toolCallId: string; - messageIndex: number; - }> = []; - for (const [approvalId, info] of requests) { - if (!resolvedIds.has(approvalId)) { - unresolved.push({ approvalId, ...info }); - } - } - - if (unresolved.length === 0) { - return messages; - } - - // Group unresolved approvals by the assistant message index they came from - const byMessageIndex = new Map< - number, - Array<{ approvalId: string; toolCallId: string }> - >(); - for (const entry of unresolved) { - console.warn( - `Auto-denying unresolved tool approval ${entry.approvalId} ` + - `(toolCallId: ${entry.toolCallId}): new generation started`, - ); - let group = byMessageIndex.get(entry.messageIndex); - if (!group) { - group = []; - byMessageIndex.set(entry.messageIndex, group); - } - group.push(entry); - } - - // Build result by inserting synthetic denial messages after each relevant - // assistant message - const result: ModelMessage[] = []; - for (let i = 0; i < messages.length; i++) { - result.push(messages[i]); - const group = byMessageIndex.get(i); - if (group) { - result.push({ - role: "tool", - content: group.map((entry) => ({ - type: "tool-approval-response" as const, - approvalId: entry.approvalId, - approved: false, - reason: "auto-denied: new generation started", - })), - }); - } - } - - return result; -} - export function serializeUsage(usage: LanguageModelUsage): Usage { return { promptTokens: usage.inputTokens,