fix(sse): honor comment opt-out for final metadata (#9305) (#9378)

* test(sse): add RED coverage for comment opt-out

* fix(sse): honor comment opt-out for final metadata

---------

Co-authored-by: 千乘妍 (Xiaoyaner) <xiaoyaner0201@users.noreply.github.com>
Co-authored-by: diegosouzapw <8016841+diegosouzapw@users.noreply.github.com>

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,

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<Uint8Array>({
+      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;
+});

tests/unit/sseHeartbeat.test.ts

@@ -19,6 +19,38 @@ function withEnv(value: string | undefined, fn: () => void) {
   }
 }
 
+async function withEnvAsync<T>(value: string | undefined, fn: () => Promise<T>): Promise<T> {
+  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<typeof setTimeout> | undefined;
+  const input = new ReadableStream<Uint8Array>({
+    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<Uint8Array>({
-      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<Uint8Array, string>({
-          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"\}/);
+  });
 });