diff --git a/open-sse/handlers/chatCore/attemptLogging.ts b/open-sse/handlers/chatCore/attemptLogging.ts index bbe049cd63..63a12c3038 100644 --- a/open-sse/handlers/chatCore/attemptLogging.ts +++ b/open-sse/handlers/chatCore/attemptLogging.ts @@ -16,6 +16,7 @@ import { emit } from "@/lib/events/eventBus"; import type { RequestCompletedPayload, RequestFailedPayload } from "@/lib/events/types"; import { saveCallLog } from "@/lib/usageDb"; import { FORMATS } from "../../translator/formats.ts"; +import { takeEarlyKeepaliveBytes } from "../../utils/earlyKeepaliveByteBuffer.ts"; import { cloneBoundedChatLogPayload, truncateForLog } from "./logTruncation.ts"; import { attachLogMeta } from "./cacheUsageMeta.ts"; @@ -244,6 +245,22 @@ export function persistAttemptLogs(args: PersistAttemptLogsArgs, ctx: PersistAtt message: error, }; } + // withEarlyStreamKeepalive writes keepalive/startup/error frames directly + // to the client from OUTSIDE this handler's own reqLogger, so they never + // reach reqLogger.appendConvertedChunk. correlationId is the only thing + // both sides share (see earlyKeepaliveByteBuffer.ts's file doc for why); + // merge here, once, right before persistence, prepended in send order. + if (detailedLoggingEnabled && correlationId) { + const earlyClientBytes = takeEarlyKeepaliveBytes(correlationId); + if (earlyClientBytes.length > 0) { + const existingStreamChunks = + (pipelinePayloads.streamChunks as { client?: string[] } | undefined) ?? {}; + pipelinePayloads.streamChunks = { + ...existingStreamChunks, + client: [...earlyClientBytes, ...(existingStreamChunks.client ?? [])], + }; + } + } } saveCallLog({ diff --git a/open-sse/utils/earlyKeepaliveByteBuffer.ts b/open-sse/utils/earlyKeepaliveByteBuffer.ts new file mode 100644 index 0000000000..510ffcf5b2 --- /dev/null +++ b/open-sse/utils/earlyKeepaliveByteBuffer.ts @@ -0,0 +1,58 @@ +/** + * @file earlyKeepaliveByteBuffer.ts + * @description Bridges bytes withEarlyStreamKeepalive writes directly to the + * client (outside the request handler's own reqLogger) back into that same + * request's persisted call-log artifact. + * + * withEarlyStreamKeepalive wraps a route's handler Promise from OUTSIDE the + * handler's own call tree — it has no reference to the reqLogger the handler + * creates deep inside chatCore.ts, and by the time that reqLogger exists the + * keepalive/startup frames may already be written. A shared correlationId + * (threaded by the route as handleChat's 4th positional arg, and separately + * into withEarlyStreamKeepalive's options) is the only thing both sides + * share, so recordEarlyKeepaliveBytes/takeEarlyKeepaliveBytes key on that + * instead of trying to pass a live object reference across the boundary. + * + * Entries are consumed once (chatCore/attemptLogging.ts calls + * takeEarlyKeepaliveBytes exactly when it assembles the final call-log + * payload) and swept on a TTL so a request that never reaches that point + * (aborted, detailed logging disabled, a route that never wires this up) + * cannot leak buffered bytes forever. + */ + +const MAX_ITEMS_PER_CORRELATION = 200; +const ENTRY_TTL_MS = 10 * 60 * 1000; + +type BufferEntry = { chunks: string[]; createdAt: number }; + +const buffers = new Map(); + +function sweepExpired(): void { + const cutoff = Date.now() - ENTRY_TTL_MS; + for (const [correlationId, entry] of buffers) { + if (entry.createdAt < cutoff) { + buffers.delete(correlationId); + } + } +} + +export function recordEarlyKeepaliveBytes(correlationId: string, chunk: string): void { + if (!correlationId || !chunk) return; + sweepExpired(); + let entry = buffers.get(correlationId); + if (!entry) { + entry = { chunks: [], createdAt: Date.now() }; + buffers.set(correlationId, entry); + } + if (entry.chunks.length < MAX_ITEMS_PER_CORRELATION) { + entry.chunks.push(chunk); + } +} + +export function takeEarlyKeepaliveBytes(correlationId: string): string[] { + sweepExpired(); + const entry = buffers.get(correlationId); + if (!entry) return []; + buffers.delete(correlationId); + return entry.chunks; +} diff --git a/open-sse/utils/earlyStreamKeepalive.ts b/open-sse/utils/earlyStreamKeepalive.ts index fe9964e477..aa7fb63594 100644 --- a/open-sse/utils/earlyStreamKeepalive.ts +++ b/open-sse/utils/earlyStreamKeepalive.ts @@ -32,6 +32,7 @@ */ import { ResponsesOutputIndexStack } from "./responsesOutputIndexStack.ts"; +import { recordEarlyKeepaliveBytes } from "./earlyKeepaliveByteBuffer.ts"; const ENCODER = new TextEncoder(); const KEEPALIVE_FRAME = ENCODER.encode(": keepalive\n\n"); @@ -207,6 +208,19 @@ export type EarlyStreamKeepaliveOptions = { * instead — see the doc comment on the default ERROR_FRAME above for why. */ errorFrame?: Uint8Array; + /** + * Request correlation id, threaded from the route's own handleChat(..., + * correlationId) call. When set, every byte this wrapper writes to the + * client directly (startup frame, periodic keepalive ticks, and any + * in-band error frame) — everything except the verbatim-forwarded real + * response body, which the handler's own reqLogger already captures — is + * recorded via earlyKeepaliveByteBuffer and merged into this same + * request's call-log streamChunks.client by + * chatCore/attemptLogging.ts, so the persisted artifact reflects what + * actually went out on the wire instead of only what the inner handler + * produced. Omit to leave today's behavior unchanged (no recording). + */ + correlationId?: string; }; /** @@ -234,6 +248,15 @@ export async function withEarlyStreamKeepalive( // Responses) — derived from errorFrame itself so the dynamic real-upstream-body case // below stays consistent with the static default-message case without a second option. const errorFrameUsesNamedEvent = new TextDecoder().decode(errorFrame).startsWith("event:"); + const correlationId = options.correlationId; + const frameDecoder = correlationId ? new TextDecoder() : null; + // Records every direct-to-client write EXCEPT the forwarded real response + // body — that one is already captured by the handler's own reqLogger, so + // recording it again here would duplicate it in the persisted artifact. + const recordClientBytes = (chunk: Uint8Array): void => { + if (!correlationId || !frameDecoder) return; + recordEarlyKeepaliveBytes(correlationId, frameDecoder.decode(chunk)); + }; // Settle into a tagged result so neither race branch leaves an unhandled // rejection when the threshold timer wins. @@ -272,6 +295,7 @@ export async function withEarlyStreamKeepalive( if (stopped) return; try { controller.enqueue(keepaliveFrame); + recordClientBytes(keepaliveFrame); } catch { stopped = true; clearInterval(interval); @@ -286,6 +310,7 @@ export async function withEarlyStreamKeepalive( // sub-interval gap, defeating the keepalive for exactly the case it targets. try { controller.enqueue(startupFrame); + recordClientBytes(startupFrame); } catch { /* consumer already gone */ } @@ -326,6 +351,7 @@ export async function withEarlyStreamKeepalive( if (result.status === "rejected") { // Handler rejected — emit a generic error frame (never the raw error/stack). controller.enqueue(errorFrame); + recordClientBytes(errorFrame); } else { const response = result.response; const contentType = (response.headers.get("content-type") || "").toLowerCase(); @@ -352,6 +378,7 @@ export async function withEarlyStreamKeepalive( // the stream end naturally. if (bytesForwarded === 0) { controller.enqueue(errorFrame); + recordClientBytes(errorFrame); } } } else { @@ -366,7 +393,9 @@ export async function withEarlyStreamKeepalive( const framed = errorFrameUsesNamedEvent ? `event: error\ndata: ${dataLine}\n\n` : `data: ${dataLine}\n\n`; - controller.enqueue(ENCODER.encode(framed)); + const framedBytes = ENCODER.encode(framed); + controller.enqueue(framedBytes); + recordClientBytes(framedBytes); } } } catch { @@ -374,6 +403,7 @@ export async function withEarlyStreamKeepalive( if (!aborted) { try { controller.enqueue(errorFrame); + recordClientBytes(errorFrame); } catch { /* consumer gone */ } diff --git a/src/app/api/v1/responses/route.ts b/src/app/api/v1/responses/route.ts index 66270896e5..fe2fcbe8aa 100644 --- a/src/app/api/v1/responses/route.ts +++ b/src/app/api/v1/responses/route.ts @@ -9,6 +9,7 @@ import { resolveResponsesApiModel } from "@/app/api/internal/codex-responses-ws/ import { getModelInfo, getComboForModel } from "@/sse/services/model"; import { resolveKeepaliveThreshold } from "@omniroute/open-sse/utils/keepaliveThreshold"; import { resolveStreamFlag } from "@omniroute/open-sse/utils/aiSdkCompat"; +import { generateRequestId } from "@/shared/utils/requestId"; // NOTE: We do NOT call initTranslators() here — the translator registry is // bootstrapped at module level inside open-sse/translator/index.ts when it @@ -100,11 +101,18 @@ async function postHandler(request: any, context: any, preParsedBody: any = null // Reuse resolvedBody.model — no extra clone/parse needed (#4041). const model = resolvedBody?.model; const thresholdMs = resolveKeepaliveThreshold(model); - return await withEarlyStreamKeepalive(handleChat(resolved, null, resolvedBody), { + // Generated here (rather than left to handleChatImplementation's own + // fallback) so withEarlyStreamKeepalive can tag its own direct-to-client + // writes with the same id chatCore.ts ends up persisting the call log + // under — see earlyKeepaliveByteBuffer.ts for why this is the only way + // the two sides of that boundary can agree on "which request." + const correlationId = generateRequestId(); + return await withEarlyStreamKeepalive(handleChat(resolved, null, resolvedBody, correlationId), { signal: request.signal, thresholdMs, startupFrame: RESPONSES_STARTUP_THINKING_FRAME, errorFrame: OPENAI_RESPONSES_ERROR_FRAME, + correlationId, }); } return await handleChat(resolved, null, resolvedBody); diff --git a/tests/unit/attempt-logging-early-keepalive-merge.test.ts b/tests/unit/attempt-logging-early-keepalive-merge.test.ts new file mode 100644 index 0000000000..da9ef6fbfe --- /dev/null +++ b/tests/unit/attempt-logging-early-keepalive-merge.test.ts @@ -0,0 +1,156 @@ +// tests/unit/attempt-logging-early-keepalive-merge.test.ts +// End-to-end proof that bytes withEarlyStreamKeepalive writes directly to +// the client (outside the handler's own reqLogger) actually reach the +// persisted call-log row's pipeline.streamChunks.client, prepended in the +// order they were sent — the gap flagged against the real 2026-08-13 +// incident: OmniRoute's own call-log artifact never showed the keepalive +// frames that were actually on the wire, only what the inner handler +// produced. Uses a real temp DB + persisted-row polling, same pattern as +// tests/unit/chatcore-attempt-logging.test.ts. +import { test, before, after } from "node:test"; +import assert from "node:assert/strict"; +import fs from "node:fs"; +import os from "node:os"; +import path from "node:path"; + +const testDataDir = fs.mkdtempSync(path.join(os.tmpdir(), "omni-keepalive-merge-test-")); +process.env.DATA_DIR = testDataDir; + +const coreDb = await import("../../src/lib/db/core.ts"); +const { getCallLogById } = await import("../../src/lib/usage/callLogs.ts"); +const { persistAttemptLogs } = await import("../../open-sse/handlers/chatCore/attemptLogging.ts"); +const { recordEarlyKeepaliveBytes, takeEarlyKeepaliveBytes } = + await import("../../open-sse/utils/earlyKeepaliveByteBuffer.ts"); + +function baseCtx(overrides: Record = {}) { + return { + provider: "openai", + connectionId: "conn-1", + model: "gpt-x", + skillRequestId: "skill-1", + detailedLoggingEnabled: true, + reqLogger: null, + pendingRequestId: "REPLACE", + clientRawRequest: { endpoint: "/v1/responses" }, + requestedModel: "gpt-x-requested", + credentials: { connectionId: "conn-1" }, + startTime: Date.now(), + body: { input: [{ role: "user", content: "hi" }] }, + sourceFormat: "openai-responses", + targetFormat: "openai-responses", + comboName: null, + comboStepId: null, + comboExecutionKey: null, + tokensCompressed: 0, + apiKeyInfo: { id: "key-1", name: "Key One" }, + noLogEnabled: false, + ...overrides, + } as Parameters[1]; +} + +async function pollForCallLog(id: string, tries = 120) { + for (let i = 0; i < tries; i++) { + const row = await getCallLogById(id); + if (row) return row as Record; + await new Promise((r) => setTimeout(r, 20)); + } + return null; +} + +before(async () => { + await coreDb.ensureDbInitialized(); +}); + +after(() => { + coreDb.resetDbInstance(); + fs.rmSync(testDataDir, { recursive: true, force: true }); +}); + +test("bytes recorded before persistAttemptLogs are prepended into pipeline.streamChunks.client", async () => { + const id = "attempt-keepalive-merge-1"; + const correlationId = "corr-keepalive-merge-1"; + recordEarlyKeepaliveBytes(correlationId, "[00:00:00.100] : keepalive\n\n"); + recordEarlyKeepaliveBytes( + correlationId, + '[00:00:00.200] event: response.output_item.added\ndata: {"item":{"id":"rs_keepalive"}}\n\n' + ); + + persistAttemptLogs( + { status: 200 }, + baseCtx({ + pendingRequestId: id, + correlationId, + reqLogger: { + getPipelinePayloads: () => ({ + streamChunks: { client: ["[00:00:05.000] real body chunk"] }, + }), + }, + }) + ); + + const row = await pollForCallLog(id); + assert.ok(row, "call log row should be persisted"); + const pipeline = row.pipelinePayloads as { streamChunks?: { client?: string[] } }; + assert.deepEqual(pipeline.streamChunks?.client, [ + "[00:00:00.100] : keepalive\n\n", + '[00:00:00.200] event: response.output_item.added\ndata: {"item":{"id":"rs_keepalive"}}\n\n', + "[00:00:05.000] real body chunk", + ]); +}); + +test("bytes are consumed exactly once — a repeat lookup for the same correlationId finds nothing left to merge", async () => { + const correlationId = "corr-keepalive-merge-2"; + recordEarlyKeepaliveBytes(correlationId, "[00:00:00.100] : keepalive\n\n"); + + const first = takeEarlyKeepaliveBytes(correlationId); + assert.equal(first.length, 1); + + const second = takeEarlyKeepaliveBytes(correlationId); + assert.deepEqual(second, []); +}); + +test("no correlationId on the attempt means no merge is attempted (existing streamChunks untouched)", async () => { + const id = "attempt-keepalive-merge-3"; + // No recordEarlyKeepaliveBytes call at all for this id/correlationId — proves + // the merge path is a strict no-op, not a silent create-empty-array side effect. + persistAttemptLogs( + { status: 200 }, + baseCtx({ + pendingRequestId: id, + correlationId: null, + reqLogger: { + getPipelinePayloads: () => ({ + streamChunks: { client: ["[00:00:05.000] real body chunk"] }, + }), + }, + }) + ); + + const row = await pollForCallLog(id); + assert.ok(row); + const pipeline = row.pipelinePayloads as { streamChunks?: { client?: string[] } }; + assert.deepEqual(pipeline.streamChunks?.client, ["[00:00:05.000] real body chunk"]); +}); + +test("detailedLoggingEnabled=false skips the merge even when early bytes are buffered (matches the existing streamChunks capture gate)", async () => { + const id = "attempt-keepalive-merge-4"; + const correlationId = "corr-keepalive-merge-4"; + recordEarlyKeepaliveBytes(correlationId, "[00:00:00.100] : keepalive\n\n"); + + persistAttemptLogs( + { status: 200 }, + baseCtx({ + pendingRequestId: id, + correlationId, + detailedLoggingEnabled: false, + reqLogger: null, + }) + ); + + const row = await pollForCallLog(id); + assert.ok(row); + // Buffer must still hold the entry — a disabled-detailed-logging attempt + // must not silently drain another (later, detailed-logging-enabled) attempt's + // buffered bytes out from under it. + assert.equal(takeEarlyKeepaliveBytes(correlationId).length, 1); +}); diff --git a/tests/unit/early-keepalive-byte-buffer.test.ts b/tests/unit/early-keepalive-byte-buffer.test.ts new file mode 100644 index 0000000000..43b0a85184 --- /dev/null +++ b/tests/unit/early-keepalive-byte-buffer.test.ts @@ -0,0 +1,51 @@ +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + recordEarlyKeepaliveBytes, + takeEarlyKeepaliveBytes, +} from "../../open-sse/utils/earlyKeepaliveByteBuffer.ts"; + +test("records are returned in the order they were written", () => { + const id = "buf-order-1"; + recordEarlyKeepaliveBytes(id, "a"); + recordEarlyKeepaliveBytes(id, "b"); + recordEarlyKeepaliveBytes(id, "c"); + assert.deepEqual(takeEarlyKeepaliveBytes(id), ["a", "b", "c"]); +}); + +test("take clears the entry — a second take for the same id is empty", () => { + const id = "buf-once-1"; + recordEarlyKeepaliveBytes(id, "a"); + takeEarlyKeepaliveBytes(id); + assert.deepEqual(takeEarlyKeepaliveBytes(id), []); +}); + +test("take on an id that was never recorded returns an empty array, not undefined", () => { + assert.deepEqual(takeEarlyKeepaliveBytes("buf-never-recorded"), []); +}); + +test("different correlation ids do not leak into each other", () => { + recordEarlyKeepaliveBytes("buf-iso-a", "only-a"); + recordEarlyKeepaliveBytes("buf-iso-b", "only-b"); + assert.deepEqual(takeEarlyKeepaliveBytes("buf-iso-a"), ["only-a"]); + assert.deepEqual(takeEarlyKeepaliveBytes("buf-iso-b"), ["only-b"]); +}); + +test("empty chunk and empty correlationId are both no-ops", () => { + const id = "buf-noop-1"; + recordEarlyKeepaliveBytes(id, ""); + recordEarlyKeepaliveBytes("", "should-not-record"); + assert.deepEqual(takeEarlyKeepaliveBytes(id), []); + assert.deepEqual(takeEarlyKeepaliveBytes(""), []); +}); + +test("a pathologically long-lived slow path does not grow the buffer unbounded", () => { + const id = "buf-cap-1"; + for (let i = 0; i < 1000; i++) { + recordEarlyKeepaliveBytes(id, `tick-${i}`); + } + const chunks = takeEarlyKeepaliveBytes(id); + assert.ok(chunks.length < 1000, "buffer must cap well below an unbounded 1000 items"); + assert.ok(chunks.length > 0, "capping must not silently drop everything"); +}); diff --git a/tests/unit/early-stream-keepalive.test.ts b/tests/unit/early-stream-keepalive.test.ts index c52514886f..59fdd5f9c0 100644 --- a/tests/unit/early-stream-keepalive.test.ts +++ b/tests/unit/early-stream-keepalive.test.ts @@ -18,6 +18,7 @@ import { OPENAI_RESPONSES_ERROR_FRAME, } from "../../open-sse/utils/earlyStreamKeepalive.ts"; import { assertResponsesOutputIndexLifecycle } from "../helpers/assertResponsesOutputIndexLifecycle.ts"; +import { takeEarlyKeepaliveBytes } from "../../open-sse/utils/earlyKeepaliveByteBuffer.ts"; async function readAll(response: Response): Promise { const reader = response.body!.getReader(); @@ -301,6 +302,48 @@ test("slow handler emits the Responses API startup frame before the real body", assert.match(body, /data: \[DONE\]/); }); +test("a correlationId records the startup frame and keepalive ticks, but not the forwarded body", async () => { + const correlationId = "corr-record-test-1"; + const slow = new Promise((resolve) => { + setTimeout( + () => resolve(sseResponse("event: response.created\ndata: {}\n\ndata: [DONE]\n\n")), + 65 + ); + }); + + const result = await withEarlyStreamKeepalive(slow, { + thresholdMs: 25, + intervalMs: 20, + startupFrame: RESPONSES_STARTUP_THINKING_FRAME, + correlationId, + }); + await readAll(result); + + const recorded = takeEarlyKeepaliveBytes(correlationId).join(""); + assert.match(recorded, /event: response\.output_item\.added/, "startup frame must be recorded"); + assert.doesNotMatch( + recorded, + /event: response\.created/, + "the verbatim-forwarded real body must NOT be recorded here — the handler's own reqLogger already captures it, and double-recording would duplicate it in the persisted artifact" + ); +}); + +test("omitting correlationId leaves the buffer untouched (today's behavior, unchanged)", async () => { + const correlationId = "corr-record-test-omitted"; + const slow = new Promise((resolve) => { + setTimeout(() => resolve(sseResponse("event: response.created\ndata: {}\n\n")), 65); + }); + + const result = await withEarlyStreamKeepalive(slow, { + thresholdMs: 25, + intervalMs: 20, + startupFrame: RESPONSES_STARTUP_THINKING_FRAME, + }); + await readAll(result); + + assert.deepEqual(takeEarlyKeepaliveBytes(correlationId), []); +}); + test("slow handler emits the custom keepaliveFrame (Anthropic ping) before the body", async () => { const slow = new Promise((resolve) => { setTimeout(