Jyotisha/frontend/tests/rectification-v9-stream.test.ts

import assert from "node:assert/strict";
import test from "node:test";

import {
  mapStreamChunkToPhase,
  safePublicEvent,
  streamToolNames,
} from "../src/lib/rectification-agentic/v9/stream-mapping.ts";
import { runV9AgentTurn, type V9AgentRunOptions } from "../src/lib/rectification-agentic/v9/agent-run.ts";
import {
  CASE_ID,
  SESSION_ID,
  TURN_ID,
  USER_ID,
  dossierFixture,
  fakeAccounting,
  receiptHandlers,
} from "./rectification-v9-test-support.ts";
import { RECTIFICATION_SKILL_NAME } from "../src/lib/rectification-agentic/v9/case-status.ts";

type StreamChunk = {
  type: string;
  payload?: {
    toolName?: unknown;
    text?: unknown;
    args?: unknown;
    error?: unknown;
    result?: unknown;
    output?: unknown;
  };
  object?: unknown;
};

function chunk(type: string, payload?: Record<string, unknown>): StreamChunk {
  return { type, ...(payload ? { payload } : {}) };
}

test("fullStream chunks map to the allowlisted NDJSON phases only", () => {
  assert.equal(mapStreamChunkToPhase(chunk("start") as never), null);
  assert.deepEqual(
    mapStreamChunkToPhase(chunk("tool-call", { toolName: "skill", args: { name: RECTIFICATION_SKILL_NAME } }) as never),
    { type: "skill.started" },
  );
  assert.deepEqual(
    mapStreamChunkToPhase(chunk("tool-result", { toolName: "skill" }) as never),
    { type: "skill.loaded" },
  );
  assert.deepEqual(
    mapStreamChunkToPhase(chunk("tool-call", { toolName: "rectification-compare-candidates" }) as never),
    { type: "candidates.comparing", tool: "rectification-compare-candidates" },
  );
  assert.deepEqual(
    mapStreamChunkToPhase(chunk("tool-result", {
      toolName: "rectification-compare-candidates",
      result: {
        executed_methods: [
          "d1-rashi",
          "vimshottari-dasha",
          "internal-secret-technique",
          "d1-rashi",
        ],
        birth_context: { birth_date: "1997-08-08" },
        event_contribution_matrix: { secret: true },
        rule_ids: ["internal-rule"],
      },
    }) as never),
    {
      type: "candidates.updated",
      tool: "rectification-compare-candidates",
      methods: ["d1-rashi", "vimshottari-dasha"],
    },
  );
  assert.equal(
    mapStreamChunkToPhase(chunk("tool-call", { toolName: "rectification-read-case" }) as never),
    null,
  );
  assert.deepEqual(
    mapStreamChunkToPhase(chunk("tool-result", {
      toolName: "rectification-read-case",
      result: { birth_context: { birth_date: "1997-08-08", latitude: 36.4 } },
    }) as never),
    { type: "case.loaded", tool: "rectification-read-case" },
  );
  assert.deepEqual(
    mapStreamChunkToPhase(chunk("text-delta", { text: "你好" }) as never),
    { type: "answer.delta", text: "你好" },
  );
  assert.equal(mapStreamChunkToPhase(chunk("finish") as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("error", { error: new Error("boom") }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("abort") as never), null);
});

test("reasoning, raw payloads, provider metadata and step internals never map", () => {
  assert.equal(mapStreamChunkToPhase(chunk("reasoning-start", { id: "r1" }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("reasoning-delta", { text: "内部推理" }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("reasoning-end") as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("raw", { payload: { secret: true } }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("step-start", { messageId: "m1" }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("step-finish", { output: { text: "x" } }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("response-metadata", { signature: "s" }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("source", { title: "t" }) as never), null);
  assert.equal(mapStreamChunkToPhase(chunk("file", { mimeType: "text/plain" }) as never), null);
});

test("streamToolNames exposes only allowlisted rectification tools", () => {
  assert.deepEqual(streamToolNames(chunk("tool-call", { toolName: "rectification-read-case" }) as never), ["rectification-read-case"]);
  assert.deepEqual(streamToolNames(chunk("tool-call", { toolName: "skill" }) as never), []);
  assert.deepEqual(streamToolNames(chunk("tool-call", { toolName: "rectification-gate" }) as never), []);
  assert.deepEqual(streamToolNames(chunk("text-delta", { text: "x" }) as never), []);
});

test("safePublicEvent drops anything outside the allowlist", () => {
  assert.deepEqual(safePublicEvent({ type: "answer.delta", text: "你好" }), { type: "answer.delta", text: "你好" });
  assert.deepEqual(safePublicEvent({ type: "skill.loaded" }), { type: "skill.loaded" });
  assert.deepEqual(
    safePublicEvent({
      type: "case.loaded",
      tool: "rectification-read-case",
      text: "1997-08-08 河北省邯郸市",
      methods: ["d1-rashi"],
      birth_context: { latitude: 36.4 },
    }),
    { type: "case.loaded", tool: "rectification-read-case" },
  );
  assert.deepEqual(
    safePublicEvent({
      type: "candidates.updated",
      tool: "rectification-compare-candidates",
      methods: ["d10-dashamsa", "private-method", "d10-dashamsa"],
      rule_ids: ["private-rule"],
    }),
    {
      type: "candidates.updated",
      tool: "rectification-compare-candidates",
      methods: ["d10-dashamsa"],
    },
  );
  assert.equal(safePublicEvent({ type: "provider.reasoning", text: "内部" }), null);
  assert.equal(safePublicEvent({ type: "tool.payload", text: "秘密" }), null);
  assert.equal(safePublicEvent({ type: "raw" }), null);
  assert.equal(safePublicEvent(null), null);
});

type FakeStreamResult = {
  fullStream: AsyncIterable<{ type: string; payload?: Record<string, unknown> }>;
  totalUsage?: Promise<{ inputTokens?: number; outputTokens?: number }>;
};

function fakeAgentStream(chunks: Array<{ type: string; payload?: Record<string, unknown> }>) {
  return {
    stream: async () => ({
      fullStream: (async function* () {
        for (const item of chunks) yield item;
      })(),
      totalUsage: Promise.resolve({ inputTokens: 10, outputTokens: 20 }),
    }) as FakeStreamResult,
    getSkill: async () => ({ name: RECTIFICATION_SKILL_NAME, instructions: "skill" }),
  };
}

function runOptions(overrides: Partial<V9AgentRunOptions> = {}) {
  const emitted: Array<{ type: string; text?: string }> = [];
  const billing = { reserved: 0, completed: 0, released: 0 };
  const accounting = fakeAccounting({
    ...receiptHandlers,
    get_agentic_rectification_case_dossier: () => dossierFixture(),
    append_agentic_rectification_turn: () => ({ turn_id: TURN_ID }),
    finalize_agentic_rectification_turn: () => ({ turn_id: TURN_ID, status: "completed", idempotent: false }),
  });
  const optionsValue: V9AgentRunOptions = {
    userId: USER_ID,
    caseId: CASE_ID,
    sessionId: SESSION_ID,
    requestId: "aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee",
    action: "evidence",
    message: "2016年9月离开家去北京工作",
    modelName: "gpt-4o-mini",
    accounting: accounting.client,
    billing: {
      reserve: async () => { billing.reserved += 1; return { success: true, status: 200 }; },
      complete: async () => { billing.completed += 1; return true; },
      release: async () => { billing.released += 1; return true; },
    },
    emit: (event) => { emitted.push(event); },
    buildAgent: async () => fakeAgentStream([]) as never,
    ...overrides,
  };
  return { options: optionsValue, emitted, billing };
}

test("answer deltas stream in order and reasoning is never forwarded", async () => {
  const { options, emitted } = runOptions({
    buildAgent: async () => fakeAgentStream([
      chunk("start"),
      chunk("tool-call", { toolName: "skill", args: { name: RECTIFICATION_SKILL_NAME } }),
      chunk("tool-result", { toolName: "skill" }),
      chunk("reasoning-start", { id: "r1" }),
      chunk("reasoning-delta", { text: "我应该先……" }),
      chunk("reasoning-end"),
      chunk("text-delta", { text: "好的，" }),
      chunk("text-delta", { text: "先确认一下：" }),
      chunk("raw", { payload: { tool_args: { secret: true } } }),
      chunk("finish"),
    ]) as never,
  });
  const result = await runV9AgentTurn(options);
  assert.equal(result.ok, true);
  const deltas = emitted.filter((event) => event.type === "answer.delta");
  assert.deepEqual(deltas, [
    { type: "answer.delta", text: "好的，" },
    { type: "answer.delta", text: "先确认一下：" },
  ]);
  assert.equal(emitted.some((event) => event.type === "run.completed"), true);
  assert.equal(emitted.some((event) => String(event.type).includes("reasoning")), false);
  assert.equal(emitted.some((event) => String(event.type).includes("raw")), false);
});

test("half-failure never becomes settled history and releases usage", async () => {
  const { options, emitted, billing } = runOptions({
    buildAgent: async () => fakeAgentStream([
      chunk("start"),
      chunk("tool-call", { toolName: "skill", args: { name: RECTIFICATION_SKILL_NAME } }),
      chunk("tool-result", { toolName: "skill" }),
      chunk("text-delta", { text: "正在计算，" }),
      chunk("error", { error: new Error("provider failure") }),
    ]) as never,
  });
  const result = await runV9AgentTurn(options);
  assert.equal(result.ok, false);
  assert.equal(result.turnStatus, "retryable");
  assert.equal(billing.released, 1);
  assert.equal(billing.completed, 0);
  assert.equal(emitted.some((event) => event.type === "run.failed"), true);
  assert.equal(emitted.some((event) => event.type === "run.completed"), false);
});

test("browser disconnect aborts the run, finalizes retryable and releases usage", async () => {
  const controller = new AbortController();
  const aborted = new Promise<void>((resolve) => {
    controller.signal.addEventListener("abort", () => resolve(), { once: true });
  });
  const hanging = (async function* () {
    yield chunk("start");
    yield chunk("tool-call", { toolName: "skill", args: { name: RECTIFICATION_SKILL_NAME } });
    yield chunk("tool-result", { toolName: "skill" });
    yield chunk("text-delta", { text: "你好，" });
    // The provider stream hangs until the client disconnects.
    await aborted;
  })();
  const { options, billing } = runOptions({
    signal: controller.signal,
    buildAgent: async () => ({
      stream: async () => ({
        fullStream: hanging,
        totalUsage: Promise.resolve({ inputTokens: 10, outputTokens: 20 }),
      }),
      getSkill: async () => ({ name: RECTIFICATION_SKILL_NAME, instructions: "skill" }),
    }) as never,
  });
  const pending = runV9AgentTurn(options);
  setTimeout(() => controller.abort(), 30);
  const result = await pending;
  assert.equal(result.ok, false);
  assert.equal(result.errorCode, "stream_aborted");
  assert.equal(billing.released, 1);
});

test("empty stream fails closed without completing billing", async () => {
  const { options, emitted, billing } = runOptions({
    buildAgent: async () => fakeAgentStream([
      chunk("start"),
      chunk("tool-call", { toolName: "skill", args: { name: RECTIFICATION_SKILL_NAME } }),
      chunk("tool-result", { toolName: "skill" }),
      chunk("finish"),
    ]) as never,
  });
  const result = await runV9AgentTurn(options);
  assert.equal(result.ok, false);
  assert.equal(result.errorCode, "empty_stream");
  assert.equal(billing.released, 1);
  assert.equal(emitted.some((event) => event.type === "run.completed"), false);
});

test("execution receipts are persisted per turn (phases + tools)", async () => {
  const accounting = fakeAccounting({
    ...receiptHandlers,
    get_agentic_rectification_case_dossier: () => dossierFixture(),
    append_agentic_rectification_turn: () => ({ turn_id: TURN_ID }),
    finalize_agentic_rectification_turn: () => ({ turn_id: TURN_ID, status: "completed", idempotent: false }),
  });
  const { options } = runOptions({
    accounting: accounting.client,
    buildAgent: async () => fakeAgentStream([
      chunk("start"),
      chunk("tool-call", { toolName: "skill", args: { name: RECTIFICATION_SKILL_NAME } }),
      chunk("tool-result", { toolName: "skill" }),
      chunk("tool-call", { toolName: "rectification-read-case", args: { caseId: CASE_ID } }),
      chunk("tool-result", { toolName: "rectification-read-case" }),
      chunk("text-delta", { text: "你好" }),
      chunk("finish"),
    ]) as never,
  });
  const result = await runV9AgentTurn(options);
  assert.equal(result.ok, true);
  const phases = accounting.calls
    .filter((call) => call.fn === "insert_agentic_rectification_run_phase")
    .map((call) => call.args.p_phase);
  assert.ok(phases.includes("run.started"));
  assert.ok(phases.includes("skill.started"));
  assert.ok(phases.includes("skill.loaded"));
  assert.ok(phases.includes("case.loaded"));
  assert.ok(phases.includes("run.completed"));
  // answer.delta is never persisted per-delta.
  assert.ok(!phases.includes("answer.delta"));
  assert.deepEqual(result.toolsUsed, ["rectification-read-case"]);
});