From cd95c2749f491ecbf59ccd80975ae1206f704805 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 23 Jun 2026 16:35:45 +0200 Subject: [PATCH] fix(ai): require OpenAI Responses terminal events --- packages/agent/CHANGELOG.md | 4 + .../src/harness/compaction/compaction.ts | 9 +- .../agent/test/harness/compaction.test.ts | 17 ++ packages/ai/CHANGELOG.md | 1 + .../ai/src/api/openai-responses-shared.ts | 68 ++--- ...nai-responses-partial-json-cleanup.test.ts | 5 + .../openai-responses-terminal-event.test.ts | 233 ++++++++++++++++++ packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/agent-session.ts | 12 +- .../src/core/compaction/compaction.ts | 11 +- .../test/agent-session-stats.test.ts | 24 ++ packages/coding-agent/test/compaction.test.ts | 32 +++ 12 files changed, 378 insertions(+), 39 deletions(-) create mode 100644 packages/ai/test/openai-responses-terminal-event.test.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index a48ede06..edbed28a 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -8,6 +8,10 @@ - `compact()`, `generateSummary()`, and `generateBranchSummary()` take a `Models` parameter and no longer accept explicit `apiKey`/`headers`. - `StreamFn` is defined structurally (`(model, context, options?) => AssistantMessageEventStream | Promise<...>`); `Models.streamSimple` satisfies it. +### Fixed + +- Fixed harness compaction estimates to ignore malformed all-zero assistant usage after truncated responses ([#5526](https://github.com/earendil-works/pi/pull/5526) by [@dmmulroy](https://github.com/dmmulroy)). + ## [0.79.10] - 2026-06-22 ## [0.79.9] - 2026-06-20 diff --git a/packages/agent/src/harness/compaction/compaction.ts b/packages/agent/src/harness/compaction/compaction.ts index d6874c33..2d6d5583 100644 --- a/packages/agent/src/harness/compaction/compaction.ts +++ b/packages/agent/src/harness/compaction/compaction.ts @@ -121,14 +121,19 @@ export function calculateContextTokens(usage: Usage): number { function getAssistantUsage(msg: AgentMessage): Usage | undefined { if (msg.role === "assistant" && "usage" in msg) { const assistantMsg = msg as AssistantMessage; - if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { + if ( + assistantMsg.stopReason !== "aborted" && + assistantMsg.stopReason !== "error" && + assistantMsg.usage && + calculateContextTokens(assistantMsg.usage) > 0 + ) { return assistantMsg.usage; } } return undefined; } -/** Return usage from the last successful assistant message in session entries. */ +/** Return usage from the last valid assistant message in session entries. */ export function getLastAssistantUsage(entries: SessionTreeEntry[]): Usage | undefined { for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; diff --git a/packages/agent/test/harness/compaction.test.ts b/packages/agent/test/harness/compaction.test.ts index 95148c3b..8a228eeb 100644 --- a/packages/agent/test/harness/compaction.test.ts +++ b/packages/agent/test/harness/compaction.test.ts @@ -301,11 +301,28 @@ describe("harness compaction", () => { createMessageEntry({ ...assistant, stopReason: "error" }), ]), ).toBeUndefined(); + expect( + getLastAssistantUsage([ + createMessageEntry(createUserMessage("user")), + createMessageEntry(assistant), + createMessageEntry(createAssistantMessage("partial", createMockUsage(0, 0))), + ]), + ).toBe(usage); expect(estimateContextTokens([createUserMessage("no usage")]).lastUsageIndex).toBeNull(); expect(estimateContextTokens([assistant, createUserMessage("tail")])).toMatchObject({ usageTokens: 20, lastUsageIndex: 0, }); + const estimate = estimateContextTokens([ + createUserMessage("Hello"), + assistant, + createUserMessage("continue"), + createAssistantMessage("Partial thinking", createMockUsage(0, 0)), + ]); + expect(estimate.usageTokens).toBe(20); + expect(estimate.lastUsageIndex).toBe(1); + expect(estimate.trailingTokens).toBeGreaterThan(0); + expect(estimate.tokens).toBe(20 + estimate.trailingTokens); }); it("builds session context with a compaction entry", () => { diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index db6631d3..cdcae32c 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -67,6 +67,7 @@ Migration guide: ### Fixed +- Fixed OpenAI Responses streams to fail when they end before a terminal response event and to treat `response.incomplete` as a length stop ([#5526](https://github.com/earendil-works/pi/pull/5526) by [@dmmulroy](https://github.com/dmmulroy)). - Fixed Amazon Bedrock endpoint resolution to honor scoped `AWS_PROFILE` values. - Fixed Cloudflare providers to require account/gateway configuration and route built-in `/compat` requests through provider auth. - Fixed OpenAI Codex Responses WebSocket sessions to reconnect once when OpenAI's connection limit is reached before output starts ([#5973](https://github.com/earendil-works/pi/issues/5973)). diff --git a/packages/ai/src/api/openai-responses-shared.ts b/packages/ai/src/api/openai-responses-shared.ts index 6fd59a44..72b4b8ae 100644 --- a/packages/ai/src/api/openai-responses-shared.ts +++ b/packages/ai/src/api/openai-responses-shared.ts @@ -294,8 +294,41 @@ export async function processResponsesStream( ): Promise { let currentItem: ResponseReasoningItem | ResponseOutputMessage | ResponseFunctionToolCall | null = null; let currentBlock: ThinkingContent | TextContent | (ToolCall & { partialJson: string }) | null = null; + let sawTerminalResponseEvent = false; const blocks = output.content; const blockIndex = () => blocks.length - 1; + const finalizeResponse = ( + response: Extract["response"], + ): void => { + sawTerminalResponseEvent = true; + if (response?.id) { + output.responseId = response.id; + } + if (response?.usage) { + const cachedTokens = response.usage.input_tokens_details?.cached_tokens || 0; + output.usage = { + // OpenAI includes cached tokens in input_tokens, so subtract to get non-cached input + input: (response.usage.input_tokens || 0) - cachedTokens, + output: response.usage.output_tokens || 0, + cacheRead: cachedTokens, + cacheWrite: 0, + totalTokens: response.usage.total_tokens || 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }; + } + calculateCost(model, output.usage); + if (options?.applyServiceTierPricing) { + const serviceTier = options.resolveServiceTier + ? options.resolveServiceTier(response?.service_tier, options.serviceTier) + : (response?.service_tier ?? options.serviceTier); + options.applyServiceTierPricing(output.usage, serviceTier); + } + // Map status to stop reason + output.stopReason = mapStopReason(response?.status); + if (output.content.some((b) => b.type === "toolCall") && output.stopReason === "stop") { + output.stopReason = "toolUse"; + } + }; for await (const event of openaiStream) { if (event.type === "response.created") { @@ -491,38 +524,12 @@ export async function processResponsesStream( currentBlock = null; stream.push({ type: "toolcall_end", contentIndex: blockIndex(), toolCall, partial: output }); } - } 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 = { - // OpenAI includes cached tokens in input_tokens, so subtract to get non-cached input - input: (response.usage.input_tokens || 0) - cachedTokens, - output: response.usage.output_tokens || 0, - cacheRead: cachedTokens, - cacheWrite: 0, - totalTokens: response.usage.total_tokens || 0, - cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, - }; - } - calculateCost(model, output.usage); - if (options?.applyServiceTierPricing) { - const serviceTier = options.resolveServiceTier - ? options.resolveServiceTier(response?.service_tier, options.serviceTier) - : (response?.service_tier ?? options.serviceTier); - options.applyServiceTierPricing(output.usage, serviceTier); - } - // Map status to stop reason - output.stopReason = mapStopReason(response?.status); - if (output.content.some((b) => b.type === "toolCall") && output.stopReason === "stop") { - output.stopReason = "toolUse"; - } + } else if (event.type === "response.completed" || event.type === "response.incomplete") { + finalizeResponse(event.response); } else if (event.type === "error") { throw new Error(`Error Code ${event.code}: ${event.message}` || "Unknown error"); } else if (event.type === "response.failed") { + sawTerminalResponseEvent = true; const error = event.response?.error; const details = event.response?.incomplete_details; const msg = error @@ -533,6 +540,9 @@ export async function processResponsesStream( throw new Error(msg); } } + if (!sawTerminalResponseEvent) { + throw new Error("OpenAI Responses stream ended before a terminal response event"); + } } function mapStopReason(status: OpenAI.Responses.ResponseStatus | undefined): StopReason { diff --git a/packages/ai/test/openai-responses-partial-json-cleanup.test.ts b/packages/ai/test/openai-responses-partial-json-cleanup.test.ts index 26a6704a..e4f4de4d 100644 --- a/packages/ai/test/openai-responses-partial-json-cleanup.test.ts +++ b/packages/ai/test/openai-responses-partial-json-cleanup.test.ts @@ -57,6 +57,11 @@ async function* createFunctionCallEvents(argumentsJson: string): AsyncIterable { diff --git a/packages/ai/test/openai-responses-terminal-event.test.ts b/packages/ai/test/openai-responses-terminal-event.test.ts new file mode 100644 index 00000000..c37896a3 --- /dev/null +++ b/packages/ai/test/openai-responses-terminal-event.test.ts @@ -0,0 +1,233 @@ +import type { ResponseStreamEvent } from "openai/resources/responses/responses.js"; +import { describe, expect, it, vi } from "vitest"; +import { stream as streamOpenAIResponses } from "../src/api/openai-responses.ts"; +import { processResponsesStream } from "../src/api/openai-responses-shared.ts"; +import type { AssistantMessage, AssistantMessageEvent, Context, Model } from "../src/types.ts"; +import { AssistantMessageEventStream } from "../src/utils/event-stream.ts"; + +vi.mock("openai", () => { + async function* createMockResponsesStream(): AsyncIterable { + yield { + type: "response.created", + sequence_number: 0, + response: { id: "resp_wrapper_early_eof" }, + } as ResponseStreamEvent; + yield { + type: "response.output_item.added", + sequence_number: 1, + output_index: 0, + item: { type: "reasoning", id: "rs_wrapper_early_eof", summary: [] }, + } as ResponseStreamEvent; + yield { + type: "response.reasoning_text.delta", + sequence_number: 2, + output_index: 0, + content_index: 0, + item_id: "rs_wrapper_early_eof", + delta: "partial reasoning before the wrapper stream ends", + } as ResponseStreamEvent; + } + + class FakeOpenAI { + responses = { + create: () => { + const responseStream = createMockResponsesStream(); + const promise = Promise.resolve(responseStream) as Promise> & { + withResponse: () => Promise<{ + data: AsyncIterable; + response: { status: number; headers: Headers }; + }>; + }; + promise.withResponse = async () => ({ + data: responseStream, + response: { status: 200, headers: new Headers() }, + }); + return promise; + }, + }; + } + + return { default: FakeOpenAI }; +}); + +function createModel(): Model<"openai-responses"> { + return { + id: "gpt-5-mini", + name: "GPT-5 Mini", + api: "openai-responses", + provider: "openai", + baseUrl: "https://api.openai.com/v1", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 400000, + maxTokens: 128000, + }; +} + +function createOutput(model: Model<"openai-responses">): AssistantMessage { + return { + role: "assistant", + content: [], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; +} + +async function* createEarlyEofEvents(): AsyncIterable { + yield { + type: "response.created", + sequence_number: 0, + response: { id: "resp_early_eof" }, + } as ResponseStreamEvent; + yield { + type: "response.output_item.added", + sequence_number: 1, + output_index: 0, + item: { type: "reasoning", id: "rs_early_eof", summary: [] }, + } as ResponseStreamEvent; + yield { + type: "response.reasoning_text.delta", + sequence_number: 2, + output_index: 0, + content_index: 0, + item_id: "rs_early_eof", + delta: "partial reasoning before the stream ends", + } as ResponseStreamEvent; +} + +async function* createCompletedEvents(): AsyncIterable { + yield { + type: "response.completed", + sequence_number: 0, + response: { + id: "resp_completed", + status: "completed", + usage: { + input_tokens: 20, + output_tokens: 7, + total_tokens: 27, + input_tokens_details: { cached_tokens: 2 }, + }, + }, + } as ResponseStreamEvent; +} + +async function* createIncompleteEvents(): AsyncIterable { + yield { + type: "response.incomplete", + sequence_number: 0, + response: { + id: "resp_incomplete", + status: "incomplete", + usage: { + input_tokens: 30, + output_tokens: 12, + total_tokens: 42, + input_tokens_details: { cached_tokens: 5 }, + }, + }, + } as ResponseStreamEvent; +} + +async function* createFailedEvents(): AsyncIterable { + yield { + type: "response.failed", + sequence_number: 0, + response: { + id: "resp_failed", + status: "failed", + error: { code: "server_error", message: "boom" }, + }, + } as ResponseStreamEvent; +} + +describe("OpenAI Responses terminal event handling", () => { + it("rejects streams that end before a terminal response event", async () => { + const model = createModel(); + const output = createOutput(model); + const stream = new AssistantMessageEventStream(); + + await expect(processResponsesStream(createEarlyEofEvents(), output, stream, model)).rejects.toThrow( + "OpenAI Responses stream ended before a terminal response event", + ); + }); + + it("emits an error final result when the wrapper stream ends before a terminal response event", async () => { + const model = createModel(); + const context: Context = { + systemPrompt: "", + messages: [{ role: "user", content: [{ type: "text", text: "hi" }], timestamp: 0 }], + tools: [], + }; + const stream = streamOpenAIResponses(model, context, { apiKey: "test" }); + const events: AssistantMessageEvent[] = []; + + for await (const event of stream) { + events.push(event); + } + + const result = await stream.result(); + const lastEvent = events.at(-1); + expect(lastEvent?.type).toBe("error"); + expect(result.stopReason).toBe("error"); + expect(result.errorMessage).toBe("OpenAI Responses stream ended before a terminal response event"); + }); + + it("finalizes completed terminal events as stop", async () => { + const model = createModel(); + const output = createOutput(model); + const stream = new AssistantMessageEventStream(); + + await processResponsesStream(createCompletedEvents(), output, stream, model); + + expect(output.responseId).toBe("resp_completed"); + expect(output.stopReason).toBe("stop"); + expect(output.usage).toMatchObject({ + input: 18, + output: 7, + cacheRead: 2, + cacheWrite: 0, + totalTokens: 27, + }); + }); + + it("finalizes incomplete terminal events as length stops", async () => { + const model = createModel(); + const output = createOutput(model); + const stream = new AssistantMessageEventStream(); + + await processResponsesStream(createIncompleteEvents(), output, stream, model); + + expect(output.responseId).toBe("resp_incomplete"); + expect(output.stopReason).toBe("length"); + expect(output.usage).toMatchObject({ + input: 25, + output: 12, + cacheRead: 5, + cacheWrite: 0, + totalTokens: 42, + }); + }); + + it("rejects failed terminal events with the provider error", async () => { + const model = createModel(); + const output = createOutput(model); + const stream = new AssistantMessageEventStream(); + + await expect(processResponsesStream(createFailedEvents(), output, stream, model)).rejects.toThrow( + "server_error: boom", + ); + }); +}); diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index b71110ca..324822ad 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -14,6 +14,7 @@ ### Fixed +- Fixed context usage and compaction estimates to ignore malformed all-zero assistant usage after truncated responses ([#5526](https://github.com/earendil-works/pi/pull/5526) by [@dmmulroy](https://github.com/dmmulroy)). - Fixed `pi --resume` to load user package themes and resolve automatic light/dark theme settings. - Fixed `models.json` custom providers so stored credentials can satisfy auth without a redundant provider-level `apiKey` ([#5953](https://github.com/earendil-works/pi/issues/5953)). diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 26022f63..65a6ebd8 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1874,10 +1874,12 @@ export class AgentSession { } // Case 2: Threshold - context is getting large - // For error messages (no usage data), estimate from last successful response. - // This ensures sessions that hit persistent API errors (e.g. 529) can still compact. + // For error messages or all-zero usage messages, estimate from the last valid response. + // This ensures sessions that hit persistent API errors (e.g. 529) or malformed zero-usage + // responses can still compact and do not reset context accounting. let contextTokens: number; - if (assistantMessage.stopReason === "error") { + const directContextTokens = assistantMessage.usage ? calculateContextTokens(assistantMessage.usage) : 0; + if (assistantMessage.stopReason === "error" || directContextTokens === 0) { const messages = this.agent.state.messages; const estimate = estimateContextTokens(messages); if (estimate.lastUsageIndex === null) return false; // No usage data at all @@ -1894,7 +1896,7 @@ export class AgentSession { } contextTokens = estimate.tokens; } else { - contextTokens = calculateContextTokens(assistantMessage.usage); + contextTokens = directContextTokens; } if (shouldCompact(contextTokens, contextWindow, settings)) { return await this._runAutoCompaction("threshold", false); @@ -3011,8 +3013,8 @@ export class AgentSession { const contextTokens = calculateContextTokens(assistant.usage); if (contextTokens > 0) { hasPostCompactionUsage = true; + break; } - break; } } } diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 83b0db57..f369601c 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -139,12 +139,17 @@ export function calculateContextTokens(usage: Usage): number { /** * Get usage from an assistant message if available. - * Skips aborted and error messages as they don't have valid usage data. + * Skips aborted, error, and all-zero usage messages as they don't have valid usage data. */ function getAssistantUsage(msg: AgentMessage): Usage | undefined { if (msg.role === "assistant" && "usage" in msg) { const assistantMsg = msg as AssistantMessage; - if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { + if ( + assistantMsg.stopReason !== "aborted" && + assistantMsg.stopReason !== "error" && + assistantMsg.usage && + calculateContextTokens(assistantMsg.usage) > 0 + ) { return assistantMsg.usage; } } @@ -152,7 +157,7 @@ function getAssistantUsage(msg: AgentMessage): Usage | undefined { } /** - * Find the last non-aborted assistant message usage from session entries. + * Find the last valid assistant message usage from session entries. */ export function getLastAssistantUsage(entries: SessionEntry[]): Usage | undefined { for (let i = entries.length - 1; i >= 0; i--) { diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index 9ba30a94..7e2a5f3f 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -140,4 +140,28 @@ describe("AgentSession.getSessionStats", () => { session.dispose(); } }); + + it("ignores zero-usage messages when checking for post-compaction context usage", () => { + const { session, sessionManager } = createSession(); + + try { + sessionManager.appendMessage(createUserMessage("first", 1)); + sessionManager.appendMessage(createAssistantMessage("response1", 180_000, 2)); + const keptUserId = sessionManager.appendMessage(createUserMessage("second", 3)); + sessionManager.appendMessage(createAssistantMessage("response2", 195_000, 4)); + sessionManager.appendCompaction("summary", keptUserId, 195_000); + sessionManager.appendMessage(createUserMessage("third", 5)); + sessionManager.appendMessage(createAssistantMessage("response3", 25_000, 6)); + sessionManager.appendMessage(createUserMessage("continue", 7)); + sessionManager.appendMessage(createAssistantMessage("partial", 0, 8)); + syncAgentMessages(session, sessionManager); + + const stats = session.getSessionStats(); + expect(stats.contextUsage).toBeDefined(); + expect(stats.contextUsage?.tokens).not.toBeNull(); + expect(stats.contextUsage?.tokens ?? 0).toBeGreaterThan(25_000); + } finally { + session.dispose(); + } + }); }); diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index a245a3fc..ab09c3bf 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -9,6 +9,7 @@ import { calculateContextTokens, compact, DEFAULT_COMPACTION_SETTINGS, + estimateContextTokens, findCutPoint, getLastAssistantUsage, prepareCompaction, @@ -217,12 +218,43 @@ describe("getLastAssistantUsage", () => { expect(usage!.input).toBe(100); }); + it("should skip all-zero assistant usage", () => { + const entries: SessionEntry[] = [ + createMessageEntry(createUserMessage("Hello")), + createMessageEntry(createAssistantMessage("Hi", createMockUsage(100, 50))), + createMessageEntry(createUserMessage("continue")), + createMessageEntry(createAssistantMessage("Partial", createMockUsage(0, 0))), + ]; + + const usage = getLastAssistantUsage(entries); + expect(usage).not.toBeNull(); + expect(usage!.input).toBe(100); + }); + it("should return undefined if no assistant messages", () => { const entries: SessionEntry[] = [createMessageEntry(createUserMessage("Hello"))]; expect(getLastAssistantUsage(entries)).toBeUndefined(); }); }); +describe("estimateContextTokens", () => { + it("uses the last non-zero assistant usage as the context anchor", () => { + const messages: AgentMessage[] = [ + createUserMessage("Hello"), + createAssistantMessage("Hi", createMockUsage(100, 50)), + createUserMessage("continue"), + createAssistantMessage("Partial thinking", createMockUsage(0, 0)), + ]; + + const estimate = estimateContextTokens(messages); + + expect(estimate.usageTokens).toBe(150); + expect(estimate.lastUsageIndex).toBe(1); + expect(estimate.trailingTokens).toBeGreaterThan(0); + expect(estimate.tokens).toBe(150 + estimate.trailingTokens); + }); +}); + describe("shouldCompact", () => { it("should return true when context exceeds threshold", () => { const settings: CompactionSettings = {