From dd53eb56eed8842cb03754bcc14bf5569de64540 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 17 Mar 2026 17:11:18 +0100 Subject: [PATCH] fix(ai): expose provider response ids on assistant messages fixes #2245 --- packages/ai/CHANGELOG.md | 4 + packages/ai/README.md | 2 + packages/ai/src/providers/anthropic.ts | 1 + .../ai/src/providers/google-gemini-cli.ts | 3 + packages/ai/src/providers/google-vertex.ts | 3 + packages/ai/src/providers/google.ts | 3 + packages/ai/src/providers/mistral.ts | 3 + .../ai/src/providers/openai-completions.ts | 3 + .../src/providers/openai-responses-shared.ts | 7 +- packages/ai/src/types.ts | 1 + packages/ai/test/responseid.test.ts | 145 ++++++++++++++++++ 11 files changed, 174 insertions(+), 1 deletion(-) create mode 100644 packages/ai/test/responseid.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index a85e9cde..4f6dd1db 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Added provider-specific `responseId` support on `AssistantMessage` for providers that expose upstream response or message identifiers, including Anthropic, OpenAI, Google, Gemini CLI, and Mistral, and added end-to-end coverage for supported OAuth and API key providers ([#2245](https://github.com/badlogic/pi-mono/issues/2245)) + ## [0.58.4] - 2026-03-16 ## [0.58.3] - 2026-03-15 diff --git a/packages/ai/README.md b/packages/ai/README.md index f3ca38bd..c5793ca7 100644 --- a/packages/ai/README.md +++ b/packages/ai/README.md @@ -519,6 +519,8 @@ Every `AssistantMessage` includes a `stopReason` field that indicates how the ge - `"error"` - An error occurred during generation - `"aborted"` - Request was cancelled via abort signal +`AssistantMessage` may also include `responseId`, a provider-specific upstream response or message identifier when the underlying API exposes one. Do not assume it is always present across providers. + ## Error Handling When a request ends with an error (including aborts and tool call validation errors), the streaming API emits an error event: diff --git a/packages/ai/src/providers/anthropic.ts b/packages/ai/src/providers/anthropic.ts index 8f9748f2..4af0c6ff 100644 --- a/packages/ai/src/providers/anthropic.ts +++ b/packages/ai/src/providers/anthropic.ts @@ -264,6 +264,7 @@ export const streamAnthropic: StreamFunction<"anthropic-messages", AnthropicOpti for await (const event of anthropicStream) { if (event.type === "message_start") { + output.responseId = event.message.id; // Capture initial token usage from message_start event // This ensures we have input token counts even if the stream is aborted early output.usage.input = event.message.usage.input_tokens || 0; diff --git a/packages/ai/src/providers/google-gemini-cli.ts b/packages/ai/src/providers/google-gemini-cli.ts index 67eccdbc..a7ece809 100644 --- a/packages/ai/src/providers/google-gemini-cli.ts +++ b/packages/ai/src/providers/google-gemini-cli.ts @@ -548,6 +548,9 @@ export const streamGoogleGeminiCli: StreamFunction<"google-gemini-cli", GoogleGe // Unwrap the response const responseData = chunk.response; if (!responseData) continue; + // Cloud Code Assist mirrors Gemini's responseId field. Keep the first non-empty one. + // A single streamed response should retain the same ID across chunks. + output.responseId ||= responseData.responseId; const candidate = responseData.candidates?.[0]; if (candidate?.content?.parts) { diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index 3f04b33f..a56aae36 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -101,6 +101,9 @@ export const streamGoogleVertex: StreamFunction<"google-vertex", GoogleVertexOpt const blocks = output.content; const blockIndex = () => blocks.length - 1; for await (const chunk of googleStream) { + // Vertex uses the same @google/genai GenerateContentResponse type as Gemini. + // responseId is documented there as an output-only identifier for each response. + output.responseId ||= chunk.responseId; const candidate = chunk.candidates?.[0]; if (candidate?.content?.parts) { for (const part of candidate.content.parts) { diff --git a/packages/ai/src/providers/google.ts b/packages/ai/src/providers/google.ts index 991d5c90..46ce2299 100644 --- a/packages/ai/src/providers/google.ts +++ b/packages/ai/src/providers/google.ts @@ -86,6 +86,9 @@ export const streamGoogle: StreamFunction<"google-generative-ai", GoogleOptions> const blocks = output.content; const blockIndex = () => blocks.length - 1; for await (const chunk of googleStream) { + // @google/genai documents GenerateContentResponse.responseId as an output-only field + // used to identify each response. Keep the first non-empty one from the stream. + output.responseId ||= chunk.responseId; const candidate = chunk.candidates?.[0]; if (candidate?.content?.parts) { for (const part of candidate.content.parts) { diff --git a/packages/ai/src/providers/mistral.ts b/packages/ai/src/providers/mistral.ts index 95d3f839..c70e72cf 100644 --- a/packages/ai/src/providers/mistral.ts +++ b/packages/ai/src/providers/mistral.ts @@ -285,6 +285,9 @@ async function consumeChatStream( for await (const event of mistralStream) { const chunk = event.data; + // Mistral's streamed CompletionChunk carries an id field. Keep the first non-empty one, + // mirroring how OpenAI-style streaming exposes a stable response identifier per stream. + output.responseId ||= chunk.id; if (chunk.usage) { output.usage.input = chunk.usage.promptTokens || 0; diff --git a/packages/ai/src/providers/openai-completions.ts b/packages/ai/src/providers/openai-completions.ts index 343d127f..bcb7d7ef 100644 --- a/packages/ai/src/providers/openai-completions.ts +++ b/packages/ai/src/providers/openai-completions.ts @@ -127,6 +127,9 @@ export const streamOpenAICompletions: StreamFunction<"openai-completions", OpenA }; for await (const chunk of openaiStream) { + // OpenAI documents ChatCompletionChunk.id as the unique chat completion identifier, + // and each chunk in a streamed completion carries the same id. + output.responseId ||= chunk.id; if (chunk.usage) { output.usage = parseChunkUsage(chunk.usage, model); } diff --git a/packages/ai/src/providers/openai-responses-shared.ts b/packages/ai/src/providers/openai-responses-shared.ts index 51f654cc..ba15ee75 100644 --- a/packages/ai/src/providers/openai-responses-shared.ts +++ b/packages/ai/src/providers/openai-responses-shared.ts @@ -285,7 +285,9 @@ export async function processResponsesStream( const blockIndex = () => blocks.length - 1; for await (const event of openaiStream) { - if (event.type === "response.output_item.added") { + if (event.type === "response.created") { + output.responseId = event.response.id; + } else if (event.type === "response.output_item.added") { const item = event.item; if (item.type === "reasoning") { currentItem = item; @@ -442,6 +444,9 @@ export async function processResponsesStream( } } else if (event.type === "response.completed") { const response = event.response; + if (response?.id) { + output.responseId = response.id; + } if (response?.usage) { const cachedTokens = response.usage.input_tokens_details?.cached_tokens || 0; output.usage = { diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 24699e63..86318571 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -193,6 +193,7 @@ export interface AssistantMessage { api: Api; provider: Provider; model: string; + responseId?: string; // Provider-specific response/message identifier when the upstream API exposes one usage: Usage; stopReason: StopReason; errorMessage?: string; diff --git a/packages/ai/test/responseid.test.ts b/packages/ai/test/responseid.test.ts new file mode 100644 index 00000000..d03ce5bd --- /dev/null +++ b/packages/ai/test/responseid.test.ts @@ -0,0 +1,145 @@ +import { describe, expect, it } from "vitest"; +import { getModel } from "../src/models.js"; +import { complete } from "../src/stream.js"; +import type { Api, Context, Model, StreamOptions } from "../src/types.js"; +import { hasAzureOpenAICredentials, resolveAzureDeploymentName } from "./azure-utils.js"; +import { resolveApiKey } from "./oauth.js"; + +type StreamOptionsWithExtras = StreamOptions & Record; + +const oauthTokens = await Promise.all([ + resolveApiKey("github-copilot"), + resolveApiKey("google-gemini-cli"), + resolveApiKey("google-antigravity"), + resolveApiKey("openai-codex"), +]); +const [githubCopilotToken, geminiCliToken, antigravityToken, openaiCodexToken] = oauthTokens; + +async function expectResponseId(model: Model, options: StreamOptionsWithExtras = {}) { + const context: Context = { + systemPrompt: "You are a helpful assistant. Be concise.", + messages: [{ role: "user", content: "Reply with exactly: response id test", timestamp: Date.now() }], + }; + + const response = await complete(model, context, options); + + expect(response.stopReason, response.errorMessage).not.toBe("error"); + expect(response.responseId).toBeTruthy(); + expect(typeof response.responseId).toBe("string"); +} + +describe("responseId E2E Tests", () => { + describe.skipIf(!process.env.GEMINI_API_KEY)("Google Provider", () => { + const llm = getModel("google", "gemini-2.5-flash"); + + it("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm); + }); + }); + + describe("Google Vertex Provider", () => { + const vertexProject = process.env.GOOGLE_CLOUD_PROJECT || process.env.GCLOUD_PROJECT; + const vertexLocation = process.env.GOOGLE_CLOUD_LOCATION; + const vertexApiKey = process.env.GOOGLE_CLOUD_API_KEY; + const isVertexConfigured = Boolean(vertexProject && vertexLocation); + const vertexOptions = { project: vertexProject, location: vertexLocation } as const; + const llm = getModel("google-vertex", "gemini-3-flash-preview"); + + it.skipIf(!isVertexConfigured)("should expose responseId with ADC", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm, vertexOptions); + }); + + it.skipIf(!vertexApiKey)("should expose responseId with API key", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm, { apiKey: vertexApiKey! }); + }); + }); + + describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Completions Provider", () => { + const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini"); + void _compat; + const llm: Model<"openai-completions"> = { + ...baseModel, + api: "openai-completions", + }; + + it("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm); + }); + }); + + describe.skipIf(!process.env.OPENAI_API_KEY)("OpenAI Responses Provider", () => { + const llm = getModel("openai", "gpt-5-mini"); + + it("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm); + }); + }); + + describe.skipIf(!process.env.ANTHROPIC_API_KEY)("Anthropic Provider", () => { + const llm = getModel("anthropic", "claude-sonnet-4-5"); + + it("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm); + }); + }); + + describe.skipIf(!hasAzureOpenAICredentials())("Azure OpenAI Responses Provider", () => { + const llm = getModel("azure-openai-responses", "gpt-4o-mini"); + const azureDeploymentName = resolveAzureDeploymentName(llm.id); + const azureOptions = azureDeploymentName ? { azureDeploymentName } : {}; + + it("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm, azureOptions); + }); + }); + + describe.skipIf(!process.env.MISTRAL_API_KEY)("Mistral Provider", () => { + const llm = getModel("mistral", "devstral-medium-latest"); + + it("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + await expectResponseId(llm); + }); + }); + + describe("GitHub Copilot Provider", () => { + it.skipIf(!githubCopilotToken)("OpenAI path should expose responseId", { retry: 3, timeout: 30000 }, async () => { + const llm = getModel("github-copilot", "gpt-5.3-codex"); + await expectResponseId(llm, { apiKey: githubCopilotToken }); + }); + + it.skipIf(!githubCopilotToken)( + "Anthropic path should expose responseId", + { retry: 3, timeout: 30000 }, + async () => { + const llm = getModel("github-copilot", "claude-sonnet-4"); + await expectResponseId(llm, { apiKey: githubCopilotToken }); + }, + ); + }); + + describe("Google Gemini CLI Provider", () => { + it.skipIf(!geminiCliToken)("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + const llm = getModel("google-gemini-cli", "gemini-2.5-flash"); + await expectResponseId(llm, { apiKey: geminiCliToken }); + }); + }); + + describe("Google Antigravity Provider", () => { + it.skipIf(!antigravityToken)("Gemini path should expose responseId", { retry: 3, timeout: 30000 }, async () => { + const llm = getModel("google-antigravity", "gemini-3.1-pro-high"); + await expectResponseId(llm, { apiKey: antigravityToken }); + }); + + it.skipIf(!antigravityToken)("Claude path should expose responseId", { retry: 3, timeout: 30000 }, async () => { + const llm = getModel("google-antigravity", "claude-sonnet-4-6"); + await expectResponseId(llm, { apiKey: antigravityToken }); + }); + }); + + describe("OpenAI Codex Provider", () => { + it.skipIf(!openaiCodexToken)("should expose responseId", { retry: 3, timeout: 30000 }, async () => { + const llm = getModel("openai-codex", "gpt-5.2-codex"); + await expectResponseId(llm, { apiKey: openaiCodexToken }); + }); + }); +});