diff --git a/src/agent/runtime/provider-metadata-continuation.test.ts b/src/agent/runtime/provider-metadata-continuation.test.ts index 286162b2d6..56bc8bb313 100644 --- a/src/agent/runtime/provider-metadata-continuation.test.ts +++ b/src/agent/runtime/provider-metadata-continuation.test.ts @@ -36,6 +36,13 @@ function readAssistantProviderMetadata(options: unknown): unknown { return prompt.find((message) => message.role === "assistant")?.providerMetadata; } +function readAssistantProviderMetadataList(options: unknown): unknown[] { + const prompt = (options as { prompt?: Array> }).prompt ?? []; + return prompt + .filter((message) => message.role === "assistant" && Object.hasOwn(message, "providerMetadata")) + .map((message) => message.providerMetadata); +} + function streamFrom(parts: unknown[]): ReadableStream { return new ReadableStream({ start(controller) { @@ -116,6 +123,83 @@ describe("agent provider metadata continuation", () => { assertEquals(JSON.stringify(modelCallEvents).includes("test-thought-signature"), false); }); + it("consumes each provider metadata attachment after one model request", async () => { + const nextProviderMetadata = { + google: { + rawAssistantParts: [{ + functionCall: { + id: "lookup-2", + name: "lookup", + args: { query: "Gemini" }, + }, + thoughtSignature: "next-thought-signature", + }], + }, + }; + let callCount = 0; + const replayedMetadata: unknown[][] = []; + const model: ModelRuntime = { + provider: "google", + modelId: "gemini-3.5-flash", + async doGenerate(options: unknown) { + callCount++; + if (callCount === 1) { + return { + content: [{ + type: "tool-call", + toolCallId: "lookup-1", + toolName: "lookup", + input: '{"query":"Veryfront"}', + }], + finishReason: "tool-calls", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + providerMetadata, + }; + } + + replayedMetadata.push(readAssistantProviderMetadataList(options)); + if (callCount === 2) { + return { + content: [{ + type: "tool-call", + toolCallId: "lookup-2", + toolName: "lookup", + input: '{"query":"Gemini"}', + }], + finishReason: "tool-calls", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + providerMetadata: nextProviderMetadata, + }; + } + + return { + content: [{ type: "text", text: "Done" }], + finishReason: "stop", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + }; + }, + async doStream() { + throw new Error("not used"); + }, + }; + const assistant = agent({ + model: "google/gemini-3.5-flash", + system: "Use the lookup tool.", + tools: { lookup: createLookupTool() }, + maxSteps: 3, + resolveModelTransport: () => ({ model }), + }); + + const result = await assistant.generate({ input: "Look up Veryfront and Gemini" }); + + assertEquals(result.text, "Done"); + assertEquals(callCount, 3); + assertEquals(replayedMetadata, [ + [providerMetadata], + [nextProviderMetadata], + ]); + }); + for (const lifecycleMode of ["legacy", "active"] as const) { it(`replays streamed provider metadata through the ${lifecycleMode} lifecycle`, () => withStreamLifecycleMode(lifecycleMode, async () => { @@ -177,6 +261,101 @@ describe("agent provider metadata continuation", () => { assertEquals(continuedProviderMetadata, providerMetadata); assertEquals(body.includes("test-thought-signature"), false); })); + + it(`consumes streamed metadata between ${lifecycleMode} lifecycle requests`, () => + withStreamLifecycleMode(lifecycleMode, async () => { + const nextProviderMetadata = { + google: { + rawAssistantParts: [{ + functionCall: { + id: "lookup-2", + name: "lookup", + args: { query: "Gemini" }, + }, + thoughtSignature: "next-thought-signature", + }], + }, + }; + let callCount = 0; + const replayedMetadata: unknown[][] = []; + const model: ModelRuntime = { + provider: "google", + modelId: "gemini-3.5-flash", + async doGenerate() { + throw new Error("not used"); + }, + async doStream(options: unknown) { + callCount++; + if (callCount === 1) { + return { + stream: streamFrom([ + { + type: "tool-call", + toolCallId: "lookup-1", + toolName: "lookup", + input: '{"query":"Veryfront"}', + }, + { + type: "finish", + finishReason: "tool-calls", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + providerMetadata, + }, + ]), + }; + } + + replayedMetadata.push(readAssistantProviderMetadataList(options)); + if (callCount === 2) { + return { + stream: streamFrom([ + { + type: "tool-call", + toolCallId: "lookup-2", + toolName: "lookup", + input: '{"query":"Gemini"}', + }, + { + type: "finish", + finishReason: "tool-calls", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + providerMetadata: nextProviderMetadata, + }, + ]), + }; + } + + return { + stream: streamFrom([ + { type: "text-delta", delta: "Done" }, + { + type: "finish", + finishReason: "stop", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + }, + ]), + }; + }, + }; + const assistant = agent({ + model: "google/gemini-3.5-flash", + system: "Use the lookup tool.", + tools: { lookup: createLookupTool() }, + maxSteps: 3, + resolveModelTransport: () => ({ model }), + }); + + const body = await (await assistant.stream({ input: "Look up Veryfront and Gemini" })) + .toDataStreamResponse() + .text(); + + assertStringIncludes(body, "Done"); + assertEquals(callCount, 3); + assertEquals(replayedMetadata, [ + [providerMetadata], + [nextProviderMetadata], + ]); + })); } it("replays the exact signed Gemini tool call before its function response", async () => { diff --git a/src/agent/runtime/provider-metadata.ts b/src/agent/runtime/provider-metadata.ts index 33f7260533..dd320dfac5 100644 --- a/src/agent/runtime/provider-metadata.ts +++ b/src/agent/runtime/provider-metadata.ts @@ -19,3 +19,12 @@ export function readAttachedProviderMetadata( ): Record | undefined { return providerMetadataByMessage.get(message); } + +/** Read and remove provider replay metadata after it enters one provider request. */ +export function consumeAttachedProviderMetadata( + message: Message, +): Record | undefined { + const providerMetadata = providerMetadataByMessage.get(message); + providerMetadataByMessage.delete(message); + return providerMetadata; +} diff --git a/src/agent/runtime/text-generation-runtime-message-converter.test.ts b/src/agent/runtime/text-generation-runtime-message-converter.test.ts index e686038370..67cbf5ae25 100644 --- a/src/agent/runtime/text-generation-runtime-message-converter.test.ts +++ b/src/agent/runtime/text-generation-runtime-message-converter.test.ts @@ -6,6 +6,7 @@ import { convertToTextGenerationRuntimeMessages, convertToTextGenerationRuntimeRequestMessages, } from "./text-generation-runtime-message-converter.ts"; +import { attachProviderMetadata } from "./provider-metadata.ts"; import type { TextGenerationRuntimeAssistantMessage, TextGenerationRuntimeMessage, @@ -855,6 +856,89 @@ describe("text-generation-runtime-message-converter", () => { }); describe("convertToTextGenerationRuntimeRequestMessages", () => { + it("consumes replay metadata while leaving history conversion read-only", () => { + const providerMetadata = { google: { thoughtSignature: "signature" } }; + const assistantMessage = attachProviderMetadata({ + id: "assistant_tool", + role: "assistant", + parts: [{ + type: "tool-call", + toolCallId: "tool_1", + toolName: "lookup", + input: {}, + }], + }, providerMetadata); + const messages: Message[] = [ + assistantMessage, + { + id: "tool_1", + role: "tool", + parts: [{ + type: "tool-result", + toolCallId: "tool_1", + toolName: "lookup", + result: { value: "done" }, + }], + }, + ]; + + const firstHistory = convertToTextGenerationRuntimeMessages(messages); + const secondHistory = convertToTextGenerationRuntimeMessages(messages); + const firstRequest = convertToTextGenerationRuntimeRequestMessages(messages); + const secondRequest = convertToTextGenerationRuntimeRequestMessages(messages); + + assertEquals( + (firstHistory[0] as TextGenerationRuntimeAssistantMessage).providerMetadata, + providerMetadata, + ); + assertEquals( + (secondHistory[0] as TextGenerationRuntimeAssistantMessage).providerMetadata, + providerMetadata, + ); + assertEquals( + (firstRequest[0] as TextGenerationRuntimeAssistantMessage).providerMetadata, + providerMetadata, + ); + assertEquals( + (secondRequest[0] as TextGenerationRuntimeAssistantMessage).providerMetadata, + undefined, + ); + }); + + it("does not consume metadata from a trailing assistant message that is trimmed", () => { + const providerMetadata = { google: { thoughtSignature: "signature" } }; + const assistantMessage = attachProviderMetadata({ + id: "assistant_tool", + role: "assistant", + parts: [{ + type: "tool-call", + toolCallId: "tool_1", + toolName: "lookup", + input: {}, + }], + }, providerMetadata); + const messages: Message[] = [assistantMessage]; + + assertEquals(convertToTextGenerationRuntimeRequestMessages(messages), []); + + messages.push({ + id: "tool_1", + role: "tool", + parts: [{ + type: "tool-result", + toolCallId: "tool_1", + toolName: "lookup", + result: { value: "done" }, + }], + }); + const completedRequest = convertToTextGenerationRuntimeRequestMessages(messages); + + assertEquals( + (completedRequest[0] as TextGenerationRuntimeAssistantMessage).providerMetadata, + providerMetadata, + ); + }); + it("drops trailing assistant-only continuation text before provider requests", () => { const messages: Message[] = [ { diff --git a/src/agent/runtime/text-generation-runtime-message-converter.ts b/src/agent/runtime/text-generation-runtime-message-converter.ts index 32951bf298..a5d3bd892d 100644 --- a/src/agent/runtime/text-generation-runtime-message-converter.ts +++ b/src/agent/runtime/text-generation-runtime-message-converter.ts @@ -19,7 +19,19 @@ import type { import { assertProviderReachableAttachment } from "./attachment-reachability.ts"; import { buildDataFileAnnotation } from "#veryfront/chat/types.ts"; import { getTextFromParts, getToolArguments, type Message, type ToolCallPart } from "../types.ts"; -import { readAttachedProviderMetadata } from "./provider-metadata.ts"; +import { + consumeAttachedProviderMetadata, + readAttachedProviderMetadata, +} from "./provider-metadata.ts"; + +interface ProviderMetadataAttachment { + sourceMessage: Message; + runtimeMessage: TextGenerationRuntimeAssistantMessage; +} + +type ProviderMetadataAttachmentObserver = ( + attachment: ProviderMetadataAttachment, +) => void; function getStringPartField(part: unknown, key: string): string | undefined { if (!part || typeof part !== "object" || Array.isArray(part)) return undefined; @@ -251,14 +263,12 @@ export interface TextGenerationRuntimeConversionOptions { requireInternetReachableAttachments?: boolean; } -/** - * Convert a veryfront Message to the current text-generation runtime message format. - */ -export function convertToTextGenerationRuntimeMessage( +function convertToTextGenerationRuntimeMessageWithMetadataObserver( msg: Message, options: & { providerExecutedToolCallIds?: Set } - & TextGenerationRuntimeConversionOptions = {}, + & TextGenerationRuntimeConversionOptions, + onProviderMetadataAttached?: ProviderMetadataAttachmentObserver, ): TextGenerationRuntimeMessage { const providerExecutedToolCallIds = options.providerExecutedToolCallIds ?? new Set(); const requireInternetReachableAttachments = options.requireInternetReachableAttachments ?? true; @@ -325,6 +335,9 @@ export function convertToTextGenerationRuntimeMessage( content, ...(providerMetadata === undefined ? {} : { providerMetadata }), }; + if (providerMetadata !== undefined) { + onProviderMetadataAttached?.({ sourceMessage: msg, runtimeMessage: assistantMessage }); + } return assistantMessage; } @@ -357,6 +370,21 @@ export function convertToTextGenerationRuntimeMessage( } } +/** + * Convert a veryfront Message to the current text-generation runtime message format. + */ +export function convertToTextGenerationRuntimeMessage( + msg: Message, + options: + & { providerExecutedToolCallIds?: Set } + & TextGenerationRuntimeConversionOptions = {}, +): TextGenerationRuntimeMessage { + return convertToTextGenerationRuntimeMessageWithMetadataObserver( + msg, + options, + ); +} + function hasProviderSendableAssistantContent(message: Message): boolean { if (message.role !== "assistant") return true; @@ -373,6 +401,7 @@ function hasProviderSendableAssistantContent(message: Message): boolean { function convertAssistantMessageToTextGenerationRuntimeMessages( message: Message, providerExecutedToolCallIds: Set, + onProviderMetadataAttached?: ProviderMetadataAttachmentObserver, ): TextGenerationRuntimeMessage[] { const assistantContent: TextGenerationRuntimeAssistantMessage["content"] = []; const deferredAssistantContent: TextGenerationRuntimeAssistantMessage["content"] = []; @@ -477,17 +506,19 @@ function convertAssistantMessageToTextGenerationRuntimeMessages( // conversion split that response would pair it with an incomplete projection. if (providerMetadata !== undefined && assistantMessages.length === 1) { assistantMessages[0]!.providerMetadata = providerMetadata; + onProviderMetadataAttached?.({ + sourceMessage: message, + runtimeMessage: assistantMessages[0]!, + }); } return messages; } -/** - * Convert an array of veryfront Messages to the current text-generation runtime message format. - */ -export function convertToTextGenerationRuntimeMessages( +function convertToTextGenerationRuntimeMessagesWithMetadataObserver( messages: Message[], - options: TextGenerationRuntimeConversionOptions = {}, + options: TextGenerationRuntimeConversionOptions, + onProviderMetadataAttached?: ProviderMetadataAttachmentObserver, ): TextGenerationRuntimeMessage[] { const textGenerationRuntimeMessages: TextGenerationRuntimeMessage[] = []; const providerExecutedToolCallIds = new Set(); @@ -509,11 +540,15 @@ export function convertToTextGenerationRuntimeMessages( } const convertedMessages = message.role === "assistant" - ? convertAssistantMessageToTextGenerationRuntimeMessages(message, providerExecutedToolCallIds) - : [convertToTextGenerationRuntimeMessage(message, { + ? convertAssistantMessageToTextGenerationRuntimeMessages( + message, + providerExecutedToolCallIds, + onProviderMetadataAttached, + ) + : [convertToTextGenerationRuntimeMessageWithMetadataObserver(message, { providerExecutedToolCallIds, ...options, - })]; + }, onProviderMetadataAttached)]; for (const convertedMessage of convertedMessages) { if (convertedMessage.role === "tool" && convertedMessage.content.length === 0) { @@ -534,6 +569,19 @@ export function convertToTextGenerationRuntimeMessages( return textGenerationRuntimeMessages; } +/** + * Convert an array of veryfront Messages to the current text-generation runtime message format. + */ +export function convertToTextGenerationRuntimeMessages( + messages: Message[], + options: TextGenerationRuntimeConversionOptions = {}, +): TextGenerationRuntimeMessage[] { + return convertToTextGenerationRuntimeMessagesWithMetadataObserver( + messages, + options, + ); +} + /** * Convert messages for a provider request. * @@ -546,11 +594,23 @@ export function convertToTextGenerationRuntimeRequestMessages( messages: Message[], options: TextGenerationRuntimeConversionOptions = {}, ): TextGenerationRuntimeMessage[] { - const requestMessages = convertToTextGenerationRuntimeMessages(messages, options); + const providerMetadataAttachments: ProviderMetadataAttachment[] = []; + const requestMessages = convertToTextGenerationRuntimeMessagesWithMetadataObserver( + messages, + options, + (attachment) => providerMetadataAttachments.push(attachment), + ); while (requestMessages.at(-1)?.role === "assistant") { requestMessages.pop(); } + const retainedRequestMessages = new Set(requestMessages); + for (const attachment of providerMetadataAttachments) { + if (retainedRequestMessages.has(attachment.runtimeMessage)) { + consumeAttachedProviderMetadata(attachment.sourceMessage); + } + } + return requestMessages; }