diff --git a/open-sse/utils/stream.ts b/open-sse/utils/stream.ts index 0eeccec17d..24999c2eee 100644 --- a/open-sse/utils/stream.ts +++ b/open-sse/utils/stream.ts @@ -74,6 +74,7 @@ import { hasUnsupportedReasoningSignal, } from "./reasoningFields.ts"; import { applyThinkTag, flushThink, initThinkState } from "./thinkTagParser.ts"; +import { sseCommentsEnabled } from "./sseHeartbeat.ts"; import { caseInsensitiveToolNameLookup, restoreOpenAIToolNames, diff --git a/tests/unit/sse-comments-optout-9305.test.ts b/tests/unit/sse-comments-optout-9305.test.ts new file mode 100644 index 0000000000..e8479130f1 --- /dev/null +++ b/tests/unit/sse-comments-optout-9305.test.ts @@ -0,0 +1,185 @@ +import test 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 previousDataDir = process.env.DATA_DIR; +const testDataDir = fs.mkdtempSync(path.join(os.tmpdir(), "omniroute-sse-comments-9305-")); +process.env.DATA_DIR = testDataDir; + +const core = await import("../../src/lib/db/core.ts"); +const usageHistory = await import("../../src/lib/usage/usageHistory.ts"); +const { createSSEStream } = await import("../../open-sse/utils/stream.ts"); +const { FORMATS } = await import("../../open-sse/translator/formats.ts"); + +const encoder = new TextEncoder(); +const EXPECTED_USAGE = { + prompt_tokens: 2, + completion_tokens: 1, +}; + +const EXPECTED_RESPONSE_USAGE = { + ...EXPECTED_USAGE, + total_tokens: 3, +}; + +type CompletionPayload = { + status: number; + usage: unknown; + responseBody?: { + usage?: unknown; + }; +}; + +function contentChunk(): string { + return `data: ${JSON.stringify({ + id: "chatcmpl-sse-comments-9305", + object: "chat.completion.chunk", + created: 1700000000, + model: "test-model", + choices: [ + { + index: 0, + delta: { role: "assistant", content: "ordinary-data" }, + finish_reason: null, + }, + ], + usage: EXPECTED_USAGE, + })}\n\n`; +} + +function countOccurrences(value: string, needle: string): number { + return value.split(needle).length - 1; +} + +async function runFinalizationCase({ + envValue, + upstreamDone, +}: { + envValue: string | undefined; + upstreamDone: boolean; +}) { + const previousComments = process.env.OMNIROUTE_SSE_COMMENTS; + usageHistory.clearPendingRequests(); + + try { + if (envValue === undefined) delete process.env.OMNIROUTE_SSE_COMMENTS; + else process.env.OMNIROUTE_SSE_COMMENTS = envValue; + + const provider = "test-provider"; + const model = "test-model"; + const connectionId = `conn-9305-${envValue ?? "default"}-${upstreamDone ? "done" : "eof"}`; + const requestId = usageHistory.trackPendingRequest(model, provider, connectionId, true); + const convertedChunks: string[] = []; + let completion: CompletionPayload | null = null; + let finalized = false; + + const source = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(contentChunk())); + if (upstreamDone) controller.enqueue(encoder.encode("data: [DONE]\n\n")); + controller.close(); + }, + }); + + const output = await new Response( + source.pipeThrough( + createSSEStream({ + mode: "passthrough", + sourceFormat: FORMATS.OPENAI, + provider, + model, + connectionId, + body: { messages: [{ role: "user", content: "hello" }] }, + reqLogger: { + appendConvertedChunk(value) { + convertedChunks.push(value); + }, + }, + onComplete(payload) { + completion = payload as CompletionPayload; + finalized = usageHistory.finalizePendingRequestById(requestId, { + status: payload.status, + providerResponse: payload.providerPayload, + clientResponse: payload.clientPayload, + }); + }, + }) + ) + ).text(); + + return { + output, + convertedOutput: convertedChunks.join(""), + completion, + finalized, + requestStillPending: usageHistory.getPendingById().has(requestId), + }; + } finally { + usageHistory.clearPendingRequests(); + if (previousComments === undefined) delete process.env.OMNIROUTE_SSE_COMMENTS; + else process.env.OMNIROUTE_SSE_COMMENTS = previousComments; + } +} + +for (const upstreamDone of [true, false]) { + const finalization = upstreamDone ? "upstream [DONE]" : "natural EOF"; + + for (const [label, envValue, commentsExpected] of [ + ["default", undefined, true], + ["explicitly enabled", "yes", true], + ["disabled", "off", false], + ] as const) { + test(`createSSEStream ${finalization} finalization preserves invariants with comments ${label}`, async () => { + const result = await runFinalizationCase({ envValue, upstreamDone }); + const finishMarker = '"finish_reason":"stop"'; + const metadataMarker = ": x-omniroute-response-cost="; + const doneMarker = "data: [DONE]"; + + assert.match(result.output, /"content":"ordinary-data"/); + assert.equal(countOccurrences(result.output, finishMarker), 1); + assert.equal(countOccurrences(result.output, doneMarker), 1); + assert.equal(result.completion?.status, 200); + assert.deepEqual(result.completion?.usage, EXPECTED_USAGE); + assert.deepEqual(result.completion?.responseBody?.usage, EXPECTED_RESPONSE_USAGE); + assert.equal(result.finalized, true, "onComplete should finalize usage accounting"); + assert.equal( + result.requestStillPending, + false, + "successful finalization should clean pending state" + ); + assert.equal( + result.convertedOutput, + result.output, + "logger and client should observe the same order" + ); + + const ordinaryIndex = result.output.indexOf('"content":"ordinary-data"'); + const finishIndex = result.output.indexOf(finishMarker); + const metadataIndex = result.output.indexOf(metadataMarker); + const doneIndex = result.output.indexOf(doneMarker); + assert.ok( + ordinaryIndex < finishIndex, + "ordinary data should precede the synthetic finish chunk" + ); + assert.ok(finishIndex < doneIndex, "synthetic finish chunk should precede [DONE]"); + + if (commentsExpected) { + assert.match(result.output, /: x-omniroute-provider=test-provider/); + assert.ok(metadataIndex > finishIndex, "metadata should follow the finish chunk"); + assert.ok(metadataIndex < doneIndex, "metadata should precede [DONE]"); + } else { + assert.doesNotMatch(result.output, /: x-omniroute-/); + } + }); + } +} + +test.after(() => { + usageHistory.clearPendingRequests(); + core.resetDbInstance(); + fs.rmSync(testDataDir, { recursive: true, force: true }); + if (previousDataDir === undefined) delete process.env.DATA_DIR; + else process.env.DATA_DIR = previousDataDir; +}); diff --git a/tests/unit/sseHeartbeat.test.ts b/tests/unit/sseHeartbeat.test.ts index 5f5bde3908..e83263a933 100644 --- a/tests/unit/sseHeartbeat.test.ts +++ b/tests/unit/sseHeartbeat.test.ts @@ -19,6 +19,38 @@ function withEnv(value: string | undefined, fn: () => void) { } } +async function withEnvAsync(value: string | undefined, fn: () => Promise): Promise { + const prev = process.env.OMNIROUTE_SSE_COMMENTS; + try { + if (value === undefined) delete process.env.OMNIROUTE_SSE_COMMENTS; + else process.env.OMNIROUTE_SSE_COMMENTS = value; + return await fn(); + } finally { + if (prev === undefined) delete process.env.OMNIROUTE_SSE_COMMENTS; + else process.env.OMNIROUTE_SSE_COMMENTS = prev; + } +} + +async function collectHeartbeatOutput( + shape: (typeof HEARTBEAT_SHAPES)[keyof typeof HEARTBEAT_SHAPES] +) { + const enc = new TextEncoder(); + let closeTimer: ReturnType | undefined; + const input = new ReadableStream({ + start(controller) { + controller.enqueue(enc.encode("data: ordinary\n\n")); + closeTimer = setTimeout(() => controller.close(), 35); + }, + cancel() { + if (closeTimer) clearTimeout(closeTimer); + }, + }); + + return new Response( + input.pipeThrough(createSseHeartbeatTransform({ shape, intervalMs: 10 })) + ).text(); +} + test("sseCommentsEnabled defaults to true when the env var is unset", () => { withEnv(undefined, () => assert.equal(sseCommentsEnabled(), true)); }); @@ -39,43 +71,26 @@ test("sseCommentsEnabled is false for 'off', 'false', '0', 'no' (case-insensitiv test("shapeForClientFormat maps known client formats", () => { assert.equal(shapeForClientFormat("claude"), HEARTBEAT_SHAPES.ANTHROPIC_PING); assert.equal(shapeForClientFormat("openai"), HEARTBEAT_SHAPES.OPENAI_CHUNK); - assert.equal(shapeForClientFormat("openai-responses"), HEARTBEAT_SHAPES.OPENAI_RESPONSES_IN_PROGRESS); + assert.equal( + shapeForClientFormat("openai-responses"), + HEARTBEAT_SHAPES.OPENAI_RESPONSES_IN_PROGRESS + ); assert.equal(shapeForClientFormat(undefined), HEARTBEAT_SHAPES.COMMENT); }); -test("createSseHeartbeatTransform suppresses COMMENT heartbeats when OMNIROUTE_SSE_COMMENTS=off", async () => { - const prev = process.env.OMNIROUTE_SSE_COMMENTS; - process.env.OMNIROUTE_SSE_COMMENTS = "off"; - try { - const enc = new TextEncoder(); - const dec = new TextDecoder(); - const input = new ReadableStream({ - start(c) { - c.enqueue(enc.encode("data: hello\n\n")); - c.close(); - }, - }); - const reader = input - .pipeThrough(createSseHeartbeatTransform({ shape: HEARTBEAT_SHAPES.COMMENT, intervalMs: 20 })) - .pipeThrough( - new TransformStream({ - transform(chunk, ctrl) { - ctrl.enqueue(dec.decode(chunk)); - }, - }) - ) - .getReader(); - const chunks: string[] = []; - let res = await reader.read(); - while (!res.done) { - chunks.push(res.value as string); - res = await reader.read(); - } - const out = chunks.join(""); - assert.ok(!out.includes(": keepalive"), "no comment heartbeat should be emitted when disabled"); - assert.ok(out.includes("data: hello"), "original chunk passes through unchanged"); - } finally { - if (prev === undefined) delete process.env.OMNIROUTE_SSE_COMMENTS; - else process.env.OMNIROUTE_SSE_COMMENTS = prev; - } +test("comment opt-out suppresses only COMMENT heartbeats, not data-event heartbeats", async () => { + await withEnvAsync("off", async () => { + const comment = await collectHeartbeatOutput(HEARTBEAT_SHAPES.COMMENT); + assert.doesNotMatch(comment, /: keepalive/, "comment heartbeat should be suppressed"); + assert.match(comment, /data: ordinary/, "ordinary data should pass through unchanged"); + + const openAI = await collectHeartbeatOutput(HEARTBEAT_SHAPES.OPENAI_CHUNK); + assert.match(openAI, /"object":"chat\.completion\.chunk"/); + + const anthropic = await collectHeartbeatOutput(HEARTBEAT_SHAPES.ANTHROPIC_PING); + assert.match(anthropic, /event: ping\ndata: \{"type":"ping"\}/); + + const responses = await collectHeartbeatOutput(HEARTBEAT_SHAPES.OPENAI_RESPONSES_IN_PROGRESS); + assert.match(responses, /data: \{"type":"response\.in_progress"\}/); + }); });