diff --git a/src/server/relay.ts b/src/server/relay.ts index a456c6124e..f3f60d6534 100644 --- a/src/server/relay.ts +++ b/src/server/relay.ts @@ -366,6 +366,7 @@ export function trackSseForRequestLog( ): ReadableStream { const reader = body.getReader(); let terminalReported = false; + let cancelled = false; const reportTerminal = (status: ResponsesTerminalStatus) => { if (terminalReported) return; @@ -394,12 +395,19 @@ export function trackSseForRequestLog( inspector.feed(value); controller.enqueue(value); } catch (err) { + // The upstream read rejected: the 200 body died mid-flight. Client + // cancellation is the caller's separate 499 path, so a cancel-drained + // pending read (cancelled=true) must not carry the truncation marker. + if (!cancelled && !terminalReported && logCtx?.activeAttempt) { + logCtx.activeAttempt.streamAborted = true; + } if (!terminalReported) reportTerminal("incomplete"); inspector.dispose(); try { controller.error(err); } catch { /* already torn down */ } } }, cancel(reason) { + cancelled = true; inspector.dispose(); onCancel(); reader.cancel(reason).catch(() => {}); @@ -1122,6 +1130,10 @@ export function consumeForInspection( if (logCtx) { logCtx.transportPhase = "mid_stream"; logCtx.terminalSource = "synthetic"; + // A truncated 200 body must not meter as a success the client never + // received; the router's equivalent turn carries 502 + streamAborted + // (codex-router #139). + if (logCtx.activeAttempt) logCtx.activeAttempt.streamAborted = true; } onTerminal("failed", 502); } diff --git a/src/usage/log.ts b/src/usage/log.ts index 2067d1e506..7cb6160b11 100644 --- a/src/usage/log.ts +++ b/src/usage/log.ts @@ -36,6 +36,12 @@ export interface PersistedUsageAttempt { adapter: string; status: number; durationMs: number; + /** + * True only when the upstream stream died after its 200 head was committed, + * so the row must not meter as a success the client never received. + * Absent on ordinary attempts so old rows keep their exact shape. + */ + streamAborted?: boolean; /** TTFT relative to THIS attempt's start (WP4); unset for non-streaming/tool-only. */ firstOutputMs?: number; sendCount: number; @@ -277,6 +283,8 @@ function normalizeUsageAttempt(raw: unknown): PersistedUsageAttempt | null { adapter: attempt.adapter, status: attempt.status, durationMs: attempt.durationMs, + // Absent by default; only the literal `true` marker survives the round trip. + ...(attempt.streamAborted === true ? { streamAborted: true } : {}), ...(isNonNegativeFiniteNumber(attempt.firstOutputMs) ? { firstOutputMs: attempt.firstOutputMs } : {}), diff --git a/tests/stream-aborted-marker.test.ts b/tests/stream-aborted-marker.test.ts new file mode 100644 index 0000000000..d4943dd553 --- /dev/null +++ b/tests/stream-aborted-marker.test.ts @@ -0,0 +1,196 @@ +import { afterEach, beforeEach, describe, expect, test } from "bun:test"; +import { mkdtempSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { consumeForInspection, trackSseForRequestLog } from "../src/server/relay"; +import { + addFinalRequestLog, + addRequestLog, + beginRequestAttempt, + httpStatusForRequestLogTerminal, + type RequestLogContext, +} from "../src/server/request-log"; +import { + readUsageEntries, + resetUsageReadCacheForTests, + type PersistedUsageAttempt, +} from "../src/usage/log"; + +// Port of codex-router #139's streamAborted metering marker: an upstream stream +// that dies after its 200 head was committed must meter as a truncated turn +// (synthetic 502 + streamAborted), while a client cancellation keeps opencodex's +// own 499 client_cancel semantics and never carries the marker. + +const encoder = new TextEncoder(); + +let testDir = ""; +let previousHome: string | undefined; + +beforeEach(() => { + previousHome = process.env.OPENCODEX_HOME; + testDir = mkdtempSync(join(tmpdir(), "ocx-stream-aborted-")); + process.env.OPENCODEX_HOME = testDir; + resetUsageReadCacheForTests(); +}); + +afterEach(() => { + if (previousHome === undefined) delete process.env.OPENCODEX_HOME; + else process.env.OPENCODEX_HOME = previousHome; + if (testDir) rmSync(testDir, { recursive: true, force: true }); +}); + +function makeLogCtx(): { logCtx: RequestLogContext; attempt: PersistedUsageAttempt } { + const attempt = beginRequestAttempt(1, "openai", "gpt-test", "openai-responses"); + const logCtx: RequestLogContext = { + provider: "openai", + model: "gpt-test", + activeAttempt: attempt, + activeAttemptStartedAt: Date.now(), + attempts: [attempt], + }; + return { logCtx, attempt }; +} + +/** Enqueues one event (a 200 head is committed), then the upstream read dies. */ +function streamThatFailsMidStream(): ReadableStream { + let reads = 0; + return new ReadableStream({ + pull(controller) { + reads += 1; + if (reads === 1) { + controller.enqueue(encoder.encode('data: {"type":"response.output_text.delta","delta":"hel"}\n\n')); + } else { + controller.error(new Error("socket reset")); + } + }, + }); +} + +/** A stream whose read never resolves on its own; only cancel or abort ends it. */ +function pendingStream(): ReadableStream { + return new ReadableStream({ start() {}, pull() { /* producer is test-controlled */ } }); +} + +describe("streamAborted marker (codex-router #139)", () => { + test("mid-stream death after a 200 head meters as 502 + streamAborted", async () => { + const { logCtx, attempt } = makeLogCtx(); + const terminalReported = Promise.withResolvers(); + const terminals: Array<[string, number | undefined]> = []; + let cancels = 0; + consumeForInspection( + streamThatFailsMidStream(), + (status, httpStatusOverride) => { + terminals.push([status, httpStatusOverride]); + terminalReported.resolve(); + }, + undefined, + () => {}, + logCtx, + () => { cancels += 1; }, + ); + await terminalReported.promise; + expect(terminals).toEqual([["failed", 502]]); + expect(cancels).toBe(0); + expect(attempt.streamAborted).toBe(true); + + // Finalize exactly like the native-passthrough terminal path in index.ts. + addFinalRequestLog( + "ocx-stream-aborted-e2e", + Date.now(), + logCtx, + httpStatusForRequestLogTerminal("failed", logCtx), + { terminalStatus: "failed", closeReason: "terminal" }, + addRequestLog, + ); + const [row] = readUsageEntries(); + expect(row?.status).toBe(502); + expect(row?.attempts?.[0]?.status).toBe(502); + expect(row?.attempts?.[0]?.streamAborted).toBe(true); + }); + + test("client cancellation keeps 499 and never sets streamAborted", async () => { + const { logCtx, attempt } = makeLogCtx(); + const ac = new AbortController(); + const cancelFired = Promise.withResolvers(); + const doneFired = Promise.withResolvers(); + const terminals: string[] = []; + let cancels = 0; + consumeForInspection( + pendingStream(), + status => { terminals.push(status); }, + ac.signal, + () => doneFired.resolve(), + logCtx, + () => { + cancels += 1; + cancelFired.resolve(); + }, + ); + ac.abort(); + await Promise.all([cancelFired.promise, doneFired.promise]); + expect(terminals).toEqual([]); + expect(cancels).toBe(1); + expect(attempt.streamAborted).toBeUndefined(); + + addFinalRequestLog( + "ocx-cancel-e2e", + Date.now(), + logCtx, + 499, + { closeReason: "client_cancel" }, + addRequestLog, + ); + const [row] = readUsageEntries(); + expect(row?.status).toBe(499); + expect(row?.attempts?.[0]?.streamAborted).toBeUndefined(); + }); + + test("translated SSE read failure marks the attempt streamAborted", async () => { + const { logCtx, attempt } = makeLogCtx(); + const terminals: string[] = []; + let cancels = 0; + const relayed = trackSseForRequestLog( + streamThatFailsMidStream(), + status => { terminals.push(status); }, + () => { cancels += 1; }, + logCtx, + ); + await expect(new Response(relayed).text()).rejects.toThrow("socket reset"); + expect(terminals).toEqual(["incomplete"]); + expect(cancels).toBe(0); + expect(attempt.streamAborted).toBe(true); + }); + + test("translated SSE client cancel never sets the truncation marker", async () => { + const { logCtx, attempt } = makeLogCtx(); + const cancelFired = Promise.withResolvers(); + const terminals: string[] = []; + let cancels = 0; + const relayed = trackSseForRequestLog( + pendingStream(), + status => { terminals.push(status); }, + () => { + cancels += 1; + cancelFired.resolve(); + }, + logCtx, + ); + await relayed.getReader().cancel(new DOMException("client closed", "AbortError")); + await cancelFired.promise; + expect(cancels).toBe(1); + expect(terminals).toEqual([]); + expect(attempt.streamAborted).toBeUndefined(); + + addFinalRequestLog( + "ocx-cancel-translated", + Date.now(), + logCtx, + 499, + { closeReason: "client_cancel" }, + addRequestLog, + ); + const [row] = readUsageEntries(); + expect(row?.status).toBe(499); + expect(row?.attempts?.[0]?.streamAborted).toBeUndefined(); + }); +}); diff --git a/tests/usage-log.test.ts b/tests/usage-log.test.ts index 5a4b905043..090e02e250 100644 --- a/tests/usage-log.test.ts +++ b/tests/usage-log.test.ts @@ -95,24 +95,115 @@ describe("usage log", () => { provider: "blsc", model: "blsc/DeepSeek-V4-Flash", status: 429, - durationMs: 4, - usageStatus: "reported", + durationMs: 1, + usageStatus: "unreported", attempts: [{ ordinal: 1, provider: "blsc", model: "blsc/DeepSeek-V4-Flash", adapter: "openai-chat", status: 429, - durationMs: 4, - sendCount: 2, - recoveryKinds: ["rate-limit-429", "rate-limit-429"], - usageStatus: "reported", + durationMs: 1, + sendCount: 1, + recoveryKinds: ["rate-limit-429"], + usageStatus: "unreported", }], }; appendUsageEntry(entry); expect(readUsageEntries()[0]?.attempts?.[0]?.recoveryKinds).toEqual(["rate-limit-429"]); }); + test("persists the streamAborted marker on truncated-stream attempts", () => { + appendUsageEntry({ + requestId: "ocx-stream-aborted", + timestamp: 1, + provider: "openai", + model: "gpt-test", + status: 502, + durationMs: 3, + usageStatus: "unreported", + terminalStatus: "failed", + closeReason: "terminal", + attempts: [{ + ordinal: 1, + provider: "openai", + model: "gpt-test", + adapter: "openai-responses", + status: 502, + durationMs: 3, + sendCount: 1, + recoveryKinds: [], + usageStatus: "unreported", + streamAborted: true, + }], + }); + const raw = readFileSync(usageLogPath(), "utf-8"); + expect(raw).toContain('"streamAborted":true'); + expect(readUsageEntries()[0]?.attempts?.[0]?.streamAborted).toBe(true); + }); + + test("ordinary attempts omit the streamAborted marker (backward compatible)", () => { + const attempt: PersistedUsageAttempt = { + ordinal: 1, + provider: "openai", + model: "gpt-test", + adapter: "openai-responses", + status: 200, + durationMs: 3, + sendCount: 1, + recoveryKinds: [], + usageStatus: "reported", + }; + const base: PersistedUsageEntry = { + requestId: "ocx-stream-aborted-absent", + timestamp: 1, + provider: "openai", + model: "gpt-test", + status: 200, + durationMs: 3, + usageStatus: "reported", + attempts: [attempt], + }; + for (const marker of [undefined, false]) { + const normalized = normalizeUsageEntryForTest({ + ...base, + attempts: [{ ...attempt, ...(marker === undefined ? {} : { streamAborted: marker }) }], + }); + expect(normalized.attempts?.[0]).not.toHaveProperty("streamAborted"); + } + const marked = normalizeUsageEntryForTest({ + ...base, + attempts: [{ ...attempt, streamAborted: true }], + }); + expect(marked.attempts?.[0]?.streamAborted).toBe(true); + }); + + test("legacy attempts without streamAborted stay readable and unset", () => { + writeFileSync(usageLogPath(), `${JSON.stringify({ + requestId: "legacy-no-marker", + timestamp: 1, + provider: "openai", + model: "gpt-test", + status: 200, + durationMs: 3, + usageStatus: "reported", + attempts: [{ + ordinal: 1, + provider: "openai", + model: "gpt-test", + adapter: "openai-responses", + status: 200, + durationMs: 3, + sendCount: 1, + recoveryKinds: [], + usageStatus: "reported", + }], + })}\n`); + const attempt = readUsageEntries()[0]?.attempts?.[0]; + expect(attempt?.ordinal).toBe(1); + expect(attempt).not.toHaveProperty("streamAborted"); + }); + /** Build one minimal persisted-usage JSONL line for the given request id. */ const persistedLine = (requestId: string) => JSON.stringify({ requestId,