OmniRoute/tests/unit/inspector-sse-merger.test.ts

import test from "node:test";
import assert from "node:assert/strict";
import {
  detectApiFormat,
  mergeStream,
  parseSseStream,
  rebuildAnthropic,
  rebuildGemini,
  rebuildOpenAI,
  type SseEvent,
} from "../../src/mitm/inspector/sseMerger.ts";

function jsonChunks(payloads: unknown[]): SseEvent[] {
  return payloads.map((p) => ({ data: JSON.stringify(p), json: p }));
}

test("detectApiFormat — message_start → anthropic", () => {
  const chunks = jsonChunks([{ type: "message_start", message: { id: "msg_1" } }]);
  assert.equal(detectApiFormat(chunks), "anthropic");
});

test("detectApiFormat — content_block_delta → anthropic", () => {
  const chunks = jsonChunks([
    { type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "Hi" } },
  ]);
  assert.equal(detectApiFormat(chunks), "anthropic");
});

test("detectApiFormat — Anthropic SSE event name is authoritative when JSON type is absent", () => {
  const chunks: SseEvent[] = [{ event: "message_start", data: "{}", json: {} }];
  assert.equal(detectApiFormat(chunks), "anthropic");
});

test("detectApiFormat — choices[].delta → openai", () => {
  const chunks = jsonChunks([{ choices: [{ index: 0, delta: { content: "Hi" } }] }]);
  assert.equal(detectApiFormat(chunks), "openai");
});

test("detectApiFormat — candidates → gemini", () => {
  const chunks = jsonChunks([{ candidates: [{ content: { parts: [{ text: "Hello" }] } }] }]);
  assert.equal(detectApiFormat(chunks), "gemini");
});

test("detectApiFormat — Gemini usage-only terminal chunk remains identifiable", () => {
  const chunks = jsonChunks([
    { usageMetadata: { promptTokenCount: 3, candidatesTokenCount: 2 }, modelVersion: "gemini" },
  ]);
  assert.equal(detectApiFormat(chunks), "gemini");
});

test("detectApiFormat — no JSON chunks → unknown", () => {
  assert.equal(detectApiFormat([{ data: "weird non-json" }]), "unknown");
});

test("rebuildAnthropic — concat text_delta by index", () => {
  const chunks = jsonChunks([
    { type: "message_start", message: { id: "msg_1" } },
    { type: "content_block_start", index: 0, content_block: { type: "text" } },
    { type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "Hello" } },
    { type: "content_block_delta", index: 0, delta: { type: "text_delta", text: " world" } },
    { type: "content_block_stop", index: 0 },
    { type: "message_delta", delta: { stop_reason: "end_turn" }, usage: { output_tokens: 7 } },
  ]);
  const merged = rebuildAnthropic(chunks);
  assert.equal(merged.format, "anthropic");
  const msg = merged.message as {
    content: Array<{ text: string }>;
    stop_reason: string;
    usage: { output_tokens: number };
  };
  assert.equal(msg.content[0].text, "Hello world");
  assert.equal(msg.stop_reason, "end_turn");
  assert.equal(msg.usage.output_tokens, 7);
});

test("rebuildAnthropic — input_json_delta merges and JSON.parses", () => {
  const chunks = jsonChunks([
    { type: "message_start", message: {} },
    {
      type: "content_block_start",
      index: 0,
      content_block: { type: "tool_use", id: "tu_1", name: "search" },
    },
    {
      type: "content_block_delta",
      index: 0,
      delta: { type: "input_json_delta", partial_json: '{"q":' },
    },
    {
      type: "content_block_delta",
      index: 0,
      delta: { type: "input_json_delta", partial_json: '"hi"}' },
    },
    { type: "content_block_stop", index: 0 },
  ]);
  const merged = rebuildAnthropic(chunks);
  const msg = merged.message as {
    content: Array<{ type: string; input: { q: string }; name: string }>;
  };
  assert.equal(msg.content[0].type, "tool_use");
  assert.equal(msg.content[0].name, "search");
  assert.deepEqual(msg.content[0].input, { q: "hi" });
});

test("rebuildAnthropic — thinking_delta accumulates", () => {
  const chunks = jsonChunks([
    { type: "content_block_start", index: 0, content_block: { type: "thinking" } },
    { type: "content_block_delta", index: 0, delta: { type: "thinking_delta", thinking: "I am " } },
    {
      type: "content_block_delta",
      index: 0,
      delta: { type: "thinking_delta", thinking: "thinking..." },
    },
  ]);
  const merged = rebuildAnthropic(chunks);
  const msg = merged.message as { content: Array<{ thinking: string }> };
  assert.equal(msg.content[0].thinking, "I am thinking...");
});

test("rebuildOpenAI — concat delta.content per choice index", () => {
  const chunks = jsonChunks([
    {
      id: "c1",
      model: "gpt-4",
      choices: [{ index: 0, delta: { role: "assistant", content: "Hel" } }],
    },
    { choices: [{ index: 0, delta: { content: "lo" }, finish_reason: null }] },
    { choices: [{ index: 0, delta: {}, finish_reason: "stop" }] },
    { usage: { prompt_tokens: 3, completion_tokens: 1 } },
  ]);
  const merged = rebuildOpenAI(chunks);
  assert.equal(merged.format, "openai");
  const msg = merged.message as {
    id: string;
    model: string;
    choices: Array<{ message: { content: string; role: string }; finish_reason: string }>;
    usage: { prompt_tokens: number };
  };
  assert.equal(msg.id, "c1");
  assert.equal(msg.model, "gpt-4");
  assert.equal(msg.choices[0].message.content, "Hello");
  assert.equal(msg.choices[0].message.role, "assistant");
  assert.equal(msg.choices[0].finish_reason, "stop");
  assert.equal(msg.usage.prompt_tokens, 3);
});

test("rebuildOpenAI — merges tool_calls per (choice, tool) index", () => {
  const chunks = jsonChunks([
    {
      choices: [
        {
          index: 0,
          delta: {
            tool_calls: [
              {
                index: 0,
                id: "tc_1",
                type: "function",
                function: { name: "search", arguments: '{"q":' },
              },
            ],
          },
        },
      ],
    },
    {
      choices: [
        {
          index: 0,
          delta: {
            tool_calls: [{ index: 0, function: { arguments: '"hi"}' } }],
          },
        },
      ],
    },
  ]);
  const merged = rebuildOpenAI(chunks);
  const msg = merged.message as {
    choices: Array<{
      message: {
        tool_calls: Array<{ id: string; function: { name: string; arguments: string } }>;
      };
    }>;
  };
  assert.equal(msg.choices[0].message.tool_calls[0].id, "tc_1");
  assert.equal(msg.choices[0].message.tool_calls[0].function.name, "search");
  assert.equal(msg.choices[0].message.tool_calls[0].function.arguments, '{"q":"hi"}');
});

test("rebuildOpenAI — reconstructs Responses API text items and completion metadata", () => {
  const chunks = jsonChunks([
    { type: "response.created", response: { id: "resp_1", object: "response", output: [] } },
    {
      type: "response.output_item.added",
      output_index: 0,
      item: { id: "msg_1", type: "message", role: "assistant", content: [] },
    },
    {
      type: "response.content_part.added",
      output_index: 0,
      content_index: 0,
      part: { type: "output_text", text: "" },
    },
    { type: "response.output_text.delta", output_index: 0, content_index: 0, delta: "Hel" },
    { type: "response.output_text.delta", output_index: 0, content_index: 0, delta: "lo" },
    {
      type: "response.completed",
      response: {
        id: "resp_1",
        object: "response",
        status: "completed",
        output: [],
        usage: { input_tokens: 2, output_tokens: 1, total_tokens: 3 },
      },
    },
  ]);

  const merged = rebuildOpenAI(chunks);
  const response = merged.message as {
    status: string;
    usage: { total_tokens: number };
    output: Array<{ id: string; content: Array<{ text: string }> }>;
  };
  assert.equal(response.status, "completed");
  assert.equal(response.usage.total_tokens, 3);
  assert.equal(response.output[0].id, "msg_1");
  assert.equal(response.output[0].content[0].text, "Hello");
});

test("rebuildGemini — merges parts across candidates", () => {
  const chunks = jsonChunks([
    { candidates: [{ content: { parts: [{ text: "Hello " }] } }] },
    { candidates: [{ content: { parts: [{ text: "world" }] } }] },
    { usageMetadata: { promptTokenCount: 2, candidatesTokenCount: 2 } },
  ]);
  const merged = rebuildGemini(chunks);
  assert.equal(merged.format, "gemini");
  const msg = merged.message as {
    candidates: Array<{ content: { parts: Array<{ text: string }> } }>;
    usageMetadata: { promptTokenCount: number };
  };
  assert.equal(msg.candidates[0].content.parts.length, 2);
  assert.equal(msg.candidates[0].content.parts[0].text, "Hello ");
  assert.equal(msg.candidates[0].content.parts[1].text, "world");
  assert.equal(msg.usageMetadata.promptTokenCount, 2);
});

test("mergeStream — unknown format returns raw fallback (no crash)", () => {
  const chunks: SseEvent[] = [{ data: "garbage" }, { data: "{}", json: { foo: 1 } }];
  const merged = mergeStream(chunks);
  assert.equal(merged.format, "unknown");
  assert.ok(Array.isArray(merged.raw));
  assert.equal(merged.raw!.length, 2);
});

test("parseSseStream — parses event/data blocks separated by blank lines", () => {
  const raw = "event: foo\ndata: 1\n\n" + 'data: {"x":1}\n\n' + "data: [DONE]\n\n";
  const events = parseSseStream(raw);
  assert.equal(events.length, 3);
  assert.equal(events[0].event, "foo");
  assert.deepEqual(events[1].json, { x: 1 });
  assert.equal(events[2].data, "[DONE]");
});

test("parseSseStream — keeps raw data when JSON parse fails", () => {
  const events = parseSseStream("data: {not-json\n\n");
  assert.equal(events.length, 1);
  assert.equal(events[0].data, "{not-json");
  assert.equal(events[0].json, undefined);
});

test("mergeStream — dispatches by detected format", () => {
  const anth = mergeStream(jsonChunks([{ type: "message_start", message: {} }]));
  assert.equal(anth.format, "anthropic");
  const oai = mergeStream(jsonChunks([{ choices: [{ delta: { content: "x" } }] }]));
  assert.equal(oai.format, "openai");
  const gem = mergeStream(jsonChunks([{ candidates: [{ content: { parts: [{ text: "x" }] } }] }]));
  assert.equal(gem.format, "gemini");
});