fix(ai): expose provider response ids on assistant messages fixes #2245

This commit is contained in:
Mario Zechner
2026-03-17 17:11:18 +01:00
parent 18d90b5c48
commit dd53eb56ee
11 changed files with 174 additions and 1 deletions

View File

@@ -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

View File

@@ -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:

View File

@@ -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;

View File

@@ -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) {

View File

@@ -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) {

View File

@@ -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) {

View File

@@ -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;

View File

@@ -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);
}

View File

@@ -285,7 +285,9 @@ export async function processResponsesStream<TApi extends Api>(
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<TApi extends Api>(
}
} 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 = {

View File

@@ -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;

View File

@@ -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<string, unknown>;
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<TApi extends Api>(model: Model<TApi>, 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 });
});
});
});