diff --git a/CHANGELOG.md b/CHANGELOG.md index df5d075e..88af61be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,10 @@ `output-error` (#162, #441) - Store a step's reasoning once, in content, so long reasoning fits the document size limit (#418, #423) +- Fix continuing a tool approval when a context handler returns older history. + Context handlers no longer see approval requests or decisions, and pending + approvals are left out of other generations' context instead of auto-denied + (#411, #416) ## 0.7.9 diff --git a/docs/tool-approval.mdx b/docs/tool-approval.mdx index aa6a0ea4..c3526d6c 100644 --- a/docs/tool-approval.mdx +++ b/docs/tool-approval.mdx @@ -163,16 +163,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..cab465e2 100644 --- a/src/vercel/client/prepareApprovalContext.ts +++ b/src/vercel/client/prepareApprovalContext.ts @@ -1,340 +1,232 @@ 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); +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 = scan(docsToModelMessages(stored)); + collect(plan, state); + collect(plan, scan(docsToModelMessages(search))); + 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); } - const approvalCallIds = new Set( - [...storedRequests.values()].map((request) => request.toolCallId), - ); - - 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); - } - } + if (decisions.length > 0) { + plan.unit = [ + { role: "assistant", content: calls }, + { role: "tool", content: decisions }, + ] as ModelMessage[]; } + return plan; +} - // 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)}`, +// What the context handler is given, and what is kept of what it returns. +// 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, once, as providers require. +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; + const answered = answer.content.some( + (other) => + other.type === "tool-result" && + other.toolCallId === part.toolCallId, + ); + if (!answered) { + answer.providerOptions ??= message.providerOptions; + answer.content.push(part); + } + return false; + }); + return content.length > 0 ? [{ ...message, content }] : []; + } + const answer: ToolModelMessage = { role: "tool", content: [] }; + const callIds = partsOf(message).flatMap((part) => + part.type === "tool-call" ? [part.toolCallId] : [], ); - } - if (!supplied.has(key)) supplied.set(key, part); - } - } - // 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; + 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), + ); +} + +// The plan's parts are cloned on the way out, since they are the stored docs' +// own objects and 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. Its pending calls are + // hidden too, but only stored decisions are replayed. + const returned = { ...plan, hidden: new Set(plan.hidden) }; + collect(returned, scan(messages), false); + const [kept] = hideApprovals(returned, [messages]); + const context = kept.flatMap((message, i): ModelMessage[] => { + const replayed = partsOf(message).flatMap((part) => + part.type === "tool-call" + ? structuredClone(plan.replay.get(part.toolCallId) ?? []) + : [], + ); + // The SDK takes decisions in the prompt's last message as fresh ones. + const followed = i < kept.length - 1 || plan.unit.length > 0; + return message.role === "assistant" && replayed.length > 0 && followed + ? [message, { role: "tool", content: replayed }] + : [message]; + }); + // The SDK runs the continued approvals before it validates the prompt, so a + // call unanswered by the next user or system message would fail the + // generation after they ran, and a retry would run them again. + const unanswered = new Set(); + const requireAnswered = () => { + const [toolCallId] = unanswered; + if (toolCallId) { + throw new Error(`Context has tool call ${toolCallId} without its result`); } - 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)}`, - ); + }; + for (const message of context) { + if (message.role === "user" || message.role === "system") requireAnswered(); + for (const part of partsOf(message)) { + if (part.type === "tool-call" && !part.providerExecuted) { + unanswered.add(part.toolCallId); + } else if (part.type === "tool-result") { + unanswered.delete(part.toolCallId); + } } } + requireAnswered(); + return [...context, ...structuredClone(plan.unit)]; +} - 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; +// A pending approval's call is hidden, and the decisions of a provider-executed +// approval that ran are kept to replay. +function collect( + plan: ApprovalPlan, + { requests, decisions, providerCalls, settled }: ReturnType, + replay = true, +) { + for (const request of [...requests.values()].flat()) { + const { approvalId, toolCallId } = request; + if (plan.replay.has(toolCallId)) continue; + if (!settled(request)) plan.hidden.add(toolCallId); + else if (replay && providerCalls.has(toolCallId)) { + plan.replay.set(toolCallId, decisions.get(approvalId) ?? []); } - if (!approvalIds.has(approvalId)) deferred.add(toolCallId); - else if (stored.has(`tool-approval-response:${approvalId}`)) { - executable.add(toolCallId); - } else undecided.add(toolCallId); } +} - // 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[] = []; +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 (!Array.isArray(message.content)) { - projected.push(message); - continue; + // Provider-owned approvals do not always produce a tool-result. Their + // continuation's response, saved before any later prompt, is the durable + // evidence that they were consumed. + if (message.role === "assistant") { + waiting.forEach((approvalId) => consumed.add(approvalId)); } - 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; - } - return ( - part.type === "tool-call" || isProviderCall(stored, 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; + if (message.role === "assistant" || message.role === "user") waiting = []; + 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); } - 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)!] - : [], - ); - if (results.length > 0) { - projected.push({ - ...results[0], - content: results.flatMap((result) => result.content), - }); - } - } - 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, - ); - 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 projected; -} - -function partsOf(messages: ModelMessage[]): (AssistantPart | ToolPart)[] { - return messages.flatMap((message) => - Array.isArray(message.content) - ? (message.content as (AssistantPart | ToolPart)[]) - : [], - ); + const settled = ({ approvalId, toolCallId }: Request) => + results.has(toolCallId) || + (providerCalls.has(toolCallId) && consumed.has(approvalId)); + return { requests, decisions, calls, providerCalls, settled }; } -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 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] : []; }); } -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 } : {}; -} - -function partId(part: ApprovalPart): string { - return part.type === "tool-approval-request" || - part.type === "tool-approval-response" - ? `approval ${part.approvalId}` - : part.toolCallId; -} - -// 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 a2773410..30b5f8f8 100644 --- a/src/vercel/client/search.test.ts +++ b/src/vercel/client/search.test.ts @@ -347,7 +347,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", @@ -371,23 +371,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 fbbe4059..8a619d1b 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 each page that completes an order split by the recent window. @@ -340,12 +340,6 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) { 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); - }; - for (const doc of docs) { if ( doc.message?.role === "assistant" && @@ -355,8 +349,7 @@ export function filterOutOrphanedToolMessages(docs: MessageDoc[]) { (p) => p.type !== "tool-call" || toolResultIds.has(p.toolCallId) || - hasApprovalResponse(p.toolCallId) || - hasApprovalRequest(p.toolCallId), + hasApprovalResponse(p.toolCallId), ); if (content.length) { result.push({ @@ -693,15 +686,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, @@ -710,31 +713,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")) { @@ -748,59 +741,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 2e736094..30267dfd 100644 --- a/src/vercel/client/toolApprovals.test.ts +++ b/src/vercel/client/toolApprovals.test.ts @@ -1,8 +1,20 @@ +import type { + LanguageModelV4CallOptions, + LanguageModelV4Content, +} from "@ai-sdk/provider"; import { stepCountIs, type ModelMessage, type 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, + toUIMessages, + type ContextHandler, +} from "../index.js"; import { MockLanguageModel, mockModel } from "./mockModel.js"; import { components, initConvexTest } from "./setup.test.js"; @@ -69,79 +81,10 @@ 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" } } - : part, -); -// Edits a call's input on the context's own objects, as a handler that -// mutates instead of copying would. -const mutateInputInPlace = - (toolCallId: string): Handler => - (messages) => { - for (const part of messages.flatMap((message) => - message.role === "assistant" && Array.isArray(message.content) - ? message.content - : [], - )) { - if (part.type === "tool-call" && part.toolCallId === toolCallId) { - (part.input as { tag: string }).tag = "changed"; - } - } - return messages; - }; const handlers = { - "change-input": changeInput, - "flip-decision": mapParts((part) => - part.type === "tool-approval-response" && part.approvalId === "a" - ? { ...part, approved: !part.approved } - : part, - ), - "flip-provider-executed": mapParts((part) => - part.type === "tool-approval-response" && part.approvalId === "a" - ? { ...part, providerExecuted: !part.providerExecuted } - : part, - ), - "sign-request": mapParts((part) => - part.type === "tool-approval-request" && part.approvalId === "a" - ? { ...part, signature: "forged" } - : part, - ), - "mark-request-automatic": mapParts((part) => - part.type === "tool-approval-request" && part.approvalId === "a" - ? { ...part, isAutomatic: true } - : part, - ), - "add-call-provider-options": mapParts((part) => + "change-input": mapParts((part) => part.type === "tool-call" && part.toolCallId === "call-a" - ? { ...part, providerOptions: { test: { trace: "allowed" } } } + ? { ...part, input: { tag: "changed" } } : part, ), "wrap-with-message-provider-options": (messages) => @@ -149,19 +92,11 @@ 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)" } } : part, ), - "mutate-a-in-place": mutateInputInPlace("call-a"), - "mutate-x-in-place": mutateInputInPlace("call-x"), "remove-result": (messages) => messages.flatMap((message): ModelMessage[] => { if (message.role !== "tool") return [message]; @@ -180,15 +115,6 @@ const handlers = { approvalId: "c", approved: true, }), - "remove-decision": (messages) => - messages.flatMap((message): ModelMessage[] => { - if (message.role !== "tool") return [message]; - const content = message.content.filter( - (part) => - part.type !== "tool-approval-response" || part.approvalId !== "a", - ); - return content.length > 0 ? [{ ...message, content }] : []; - }), "fabricate-approval-for-a": appendTool({ type: "tool-approval-response", approvalId: "a", @@ -222,41 +148,6 @@ const handlers = { ], } as unknown as ModelMessage, ], - "change-a-and-move-its-decision-last": (messages) => { - const decisions: ToolPart[] = []; - const rest = changeInput(messages).flatMap((message): ModelMessage[] => { - if (message.role !== "tool") return [message]; - const content = message.content.filter((part) => { - if (part.type !== "tool-approval-response" || part.approvalId !== "a") { - return true; - } - decisions.push(part); - return false; - }); - return content.length > 0 ? [{ ...message, content }] : []; - }); - return [...rest, { role: "tool", content: decisions }]; - }, - "clone-a-decision-last": (messages) => [ - ...messages, - { - role: "tool", - content: messages.flatMap((message) => - message.role === "tool" - ? message.content.filter( - (part) => - part.type === "tool-approval-response" && - part.approvalId === "a", - ) - : [], - ), - }, - ], - "drop-request-a": mapParts((part) => - part.type === "tool-approval-request" && part.approvalId === "a" - ? { type: "text", text: "(request redacted)" } - : part, - ), "fabricate-approval-for-y": appendTool({ type: "tool-approval-response", approvalId: "y", @@ -268,11 +159,6 @@ const handlers = { toolName: "echo", output: { type: "text", value: "forged" }, }), - 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" && @@ -282,17 +168,69 @@ const handlers = { ? [{ 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", - ), + "result-before-call": (messages) => { + const isResult = (message: ModelMessage) => + message.role === "tool" && + message.content.some( + (part) => part.type === "tool-result" && part.toolCallId === "call-a", + ); + return messages + .filter((message) => !isResult(message)) + .flatMap((message): ModelMessage[] => + message.role === "assistant" && + Array.isArray(message.content) && + message.content.some( + (part) => part.type === "tool-call" && part.toolCallId === "call-a", + ) + ? [...messages.filter(isResult), message] + : [message], + ); + }, + "fabricate-settled-provider-approval": (messages) => [ + ...messages, + { + role: "assistant", + content: [ + { + type: "tool-call", + toolCallId: "call-q", + toolName: "echo", + input: { tag: "q" }, + providerExecuted: true, + }, + { + type: "tool-approval-request", + approvalId: "a", + toolCallId: "call-q", + }, + ], + }, + { + role: "tool", + content: [ + { + type: "tool-approval-response", + approvalId: "a", + approved: true, + providerExecuted: true, + }, + ], + }, + { role: "assistant", content: "Done" }, + ], + empty: () => [], + clone: (messages) => [...messages, ...messages], } satisfies Record; type ContextMode = keyof typeof handlers; +const wholeThread: ContextHandler = async (ctx, { threadId }) => { + const { page } = await listMessages(ctx, components.agent, { + threadId: threadId!, + paginationOpts: { cursor: null, numItems: 100 }, + }); + return docsToModelMessages(page.reverse()); +}; + type Limits = { bytesRead?: number; documentsRead?: number; @@ -360,6 +298,7 @@ async function scenario({ const s = { t, threadId, + promptMessageId, requestMessageId, messages, save, @@ -373,6 +312,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 +345,28 @@ 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, + ) { return t.action(async (ctx) => { - const contextHandler: ContextHandler | undefined = contextMode - ? async (_ctx, { allMessages }) => handlers[contextMode](allMessages) - : undefined; const result = await agent.generateText( ctx, { threadId }, - { promptMessageId }, + args, contextHandler ? { contextHandler } : undefined, ); return result.text; }); }, + continue(promptMessageId: string, contextMode?: ContextMode) { + return s.generate( + { promptMessageId }, + contextMode + ? async (_ctx, { allMessages }) => handlers[contextMode](allMessages) + : undefined, + ); + }, // `a` approved and continued, `b` still open. async completeA() { const first = await s.approve("a"); @@ -450,6 +403,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 +427,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 +442,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, @@ -597,6 +556,101 @@ describe("tool approval semantics", () => { expectEveryCallAnsweredNext(); }); + test.each([false, true])( + "siblings continued separately each answer their call in later context (whole-thread handler: %s)", + async (whole) => { + const s = await scenario(); + const handler = whole ? wholeThread : undefined; + await s.completeA(); + await s.generate({ promptMessageId: await s.approve("b") }, handler); + + await s.generate( + { promptMessageId: await s.prompt("What happened?") }, + handler, + ); + + expect(executed).toEqual(["a", "b"]); + expectNoDanglingLocalCalls(); + expectEveryCallAnsweredNext(); + }, + ); + + test.each([false, true])( + "an approved provider-owned decision stays out of unrelated generations until continued (whole-thread handler: %s)", + async (whole) => { + const s = await scenario({ ids: ["a"], providerExecuted: true }); + const handler = whole ? wholeThread : undefined; + const decision = await s.approve("a"); + + await s.generate({ promptMessageId: await s.prompt("First") }, handler); + await s.generate({ promptMessageId: await s.prompt("Second") }, handler); + expect(partsFor("a")).toEqual([]); + + await s.continue(decision); + expect(partsFor("a").map((part) => part.type)).toContain( + "tool-approval-response", + ); + }, + ); + + test("a handler cannot replay a provider-owned decision it fabricates", async () => { + const s = await scenario({ ids: ["a"], providerExecuted: true }); + + await s.continue( + await s.prompt("Next"), + "fabricate-settled-provider-approval", + ); + + expect( + lastModelPrompt().flatMap((message) => + Array.isArray(message.content) + ? message.content.filter( + (part) => part.type === "tool-approval-response", + ) + : [], + ), + ).toEqual([]); + }); + + // The SDK takes a decision in the prompt's last message as a fresh one. + test.each([ + [ + "a newer pending request", + async (s: Scenario, first: string) => { + await s.addRequest(first, ["c"], true); + return s.continue(first); + }, + ], + [ + "a handler that trims its response", + async (s: Scenario, first: string) => { + await s.continue(first); + return s.generate({ promptMessageId: first }, async (_ctx, args) => + args.allMessages.filter( + (message) => + message.role !== "assistant" || + (Array.isArray(message.content) && + message.content.some((part) => part.type !== "text")), + ), + ); + }, + ], + ])( + "re-running a completed provider-owned continuation followed by %s sends no fresh decision", + async (_name, rerun) => { + const s = await scenario({ ids: ["a"], providerExecuted: true }); + + await rerun(s, await s.approve("a")); + + expect(partsFor("c")).toEqual([]); + const last = lastModelPrompt().at(-1); + expect( + Array.isArray(last?.content) && + last.content.some((part) => part.type === "tool-approval-response"), + ).toBe(false); + }, + ); + 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 +834,27 @@ describe("tool approval semantics", () => { expectNoDanglingLocalCalls(); }); - // `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) { + test("approvals decided at different orders each run only from their own continuation", async () => { 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); - } + await laterOrder(s, "pending", ["x", "y"]); + const earlier = await s.approve("a"); - 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"); + await s.continue(await s.approve("x")); + expect(executed).toEqual(["x"]); + expect(partsFor("a")).toEqual([]); + expect(partsFor("y")).toEqual([]); - expect(executed).toEqual(["a", "x"]); - }); + await s.continue(earlier); + expect(executed).toEqual(["x", "a"]); + expect(partsFor("x")).toEqual([]); - 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"]); + 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(); }); - 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 +885,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 +925,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,91 +951,347 @@ 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"); + test("the continued approvals follow the handler's messages", async () => { + const s = await scenario(); + await s.completeA(); + + await s.continue( + await s.approve("b"), + "wrap-with-message-provider-options", + ); + + expect(executed).toEqual(["a", "b"]); + 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", + ]); + }); + + 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.generate({ promptMessageId: s.promptMessageId }); + + expect(executed).toEqual([]); + expect(partsFor("a")).toEqual([]); + expect(partsFor("b")).toEqual([]); + }); + + 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"]); + }); + + // Whatever a context handler returns, only stored decisions run, with their + // stored input, and each at most once. + test.each<[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], + ])( + "a handler cannot approve another order's request: %s", + async (contextMode, providerExecuted) => { + const s = await scenario({ ids: ["a"], providerExecuted }); + await laterOrder(s, "pending", ["x", "y"]); + + await s.continue(await s.approve("x"), contextMode); + + expect(executed).toEqual(["x"]); + expect(partsFor("a")).toEqual([]); + expect(partsFor("y")).toEqual([]); + }, + ); + + test.each([ + "fabricate-result", + "fabricate-approval", + "fabricate-approval-for-another-request", + ])( + "a handler cannot approve or answer a sibling: %s", + async (contextMode) => { + const s = await scenario(); + const first = await s.approve("a"); + await s.addRequest(first, ["c"]); + + await s.continue(first, contextMode); - 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"]); + expect(await s.resultIds()).toEqual(["call-a"]); + expect(partsFor("a")).not.toContainEqual( + expect.objectContaining({ output: { type: "text", value: "forged" } }), + ); }, ); - test("allows a context handler to add call provider options", async () => { - const s = await scenario({ ids: ["a"] }); + test("caller messages cannot approve a sibling", async () => { + const s = await scenario(); - await s.continue(await s.approve("a"), "add-call-provider-options"); + await s.generate({ + promptMessageId: await s.approve("a"), + messages: [ + { + role: "tool", + content: [ + { type: "tool-approval-response", approvalId: "b", approved: true }, + ], + }, + ], + }); 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 call runs with its stored input when the handler returns nothing", async () => { + const s = await scenario(); + + await s.continue(await s.approve("a"), "empty"); + + expect(executed).toEqual(["a"]); + expect(callInputs("a")).toEqual([{ tag: "a" }]); + }); + + 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"), - "wrap-with-message-provider-options", - ); + await s.continue(await s.approve("b"), "change-input"); expect(executed).toEqual(["a", "b"]); - for (const message of lastModelPrompt()) { - expect(message.providerOptions).toEqual({ test: { cache: "all" } }); - } + expect(callInputs("a")).toEqual([{ tag: "changed" }]); + expectEveryCallAnsweredNext(); }); - // 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"], + test.each([ + "clone", + "fabricate-result", + "inject-between-call-and-result", ])( - "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([]); + "a reshaped completed history 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("an unrelated generation still auto-denies unresolved approvals", async () => { - const warn = vi.spyOn(console, "warn").mockImplementation(() => {}); - const s = await scenario(); - await s.approve("a"); + // The SDK runs the continued approvals before it validates the prompt. + test.each(["remove-result", "result-before-call"])( + "a completed call the handler leaves unanswered fails before anything runs: %s", + async (contextMode) => { + const s = await scenario(); + await s.completeA(); - await s.continue(await s.prompt("Start unrelated work")); + await expect( + s.continue(await s.approve("b"), contextMode), + ).rejects.toThrow(/call-a without its result/); - expect(executed).toEqual([]); - expect(warn).toHaveBeenCalledWith( - expect.stringContaining("Auto-denying unresolved tool approval b"), + 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(); + let seen: ModelMessage[] = []; + + await s.generate( + { promptMessageId: await s.approve("b") }, + async (_ctx, { allMessages }) => { + seen = structuredClone(allMessages); + return allMessages; + }, ); + + 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, + first = "approve", + }: { + turns: number; + recentMessages: number; + 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 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 }, + 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).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([1, 3])( + "continues every approval over 3 turns with recentMessages: %i", + (recentMessages) => approvals({ turns: 3, recentMessages }), + ); + + test("continues later approvals after the first was declined", () => + approvals({ turns: 3, recentMessages: 1, first: "decline" })); + + test("continues later approvals while the first is still pending", () => + approvals({ turns: 3, recentMessages: 1, first: "leave pending" })); +}); + describe("approval continuation with excludeToolMessages", () => { function setup() { const ran: string[] = []; 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 66da3eaa..b584d163 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,