Skip to content

fix(agent): consume Gemini replay metadata once #3824

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
kojiwakayama wants to merge 1 commit into from
Closed

fix(agent): consume Gemini replay metadata once #3824

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 179 additions & 0 deletions src/agent/runtime/provider-metadata-continuation.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Record<string, unknown>> }).prompt ?? [];
return prompt
.filter((message) => message.role === "assistant" && Object.hasOwn(message, "providerMetadata"))
.map((message) => message.providerMetadata);
}

function streamFrom(parts: unknown[]): ReadableStream<unknown> {
return new ReadableStream({
start(controller) {
Expand Down Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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 () => {
Expand Down
9 changes: 9 additions & 0 deletions src/agent/runtime/provider-metadata.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,3 +19,12 @@ export function readAttachedProviderMetadata(
): Record<string, unknown> | undefined {
return providerMetadataByMessage.get(message);
}

/** Read and remove provider replay metadata after it enters one provider request. */
export function consumeAttachedProviderMetadata(
message: Message,
): Record<string, unknown> | undefined {
const providerMetadata = providerMetadataByMessage.get(message);
providerMetadataByMessage.delete(message);
return providerMetadata;
}
84 changes: 84 additions & 0 deletions src/agent/runtime/text-generation-runtime-message-converter.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import {
convertToTextGenerationRuntimeMessages,
convertToTextGenerationRuntimeRequestMessages,
} from "./text-generation-runtime-message-converter.ts";
import { attachProviderMetadata } from "./provider-metadata.ts";
import type {
TextGenerationRuntimeAssistantMessage,
TextGenerationRuntimeMessage,
Expand Down Expand Up @@ -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[] = [
{
Expand Down
Loading
Loading