ThothII/backend/test/session-bridge.test.ts

import { test, expect, vi } from "vitest";
import { SessionBridge } from "../src/bridge/session-bridge.js";

function fakeRpc() {
  const sent: any[] = [];
  let evcb: any;
  return {
    rpc: {
      send: (c: any) => sent.push(c),
      on: (_: any, cb: any) => { evcb = cb; },
      request: vi.fn(),
    } as any,
    sent,
    fire: (m: any) => evcb(m),
  };
}

test("real Pi message_update (assistantMessageEvent text_delta) becomes a text_delta to the FE", () => {
  const { rpc, fire } = fakeRpc();
  const b = new SessionBridge(rpc);
  const seen: any[] = [];
  b.onClientEvent((e) => seen.push(e));
  // Shape emitted by real `pi --mode rpc`: text streams inside message_update.assistantMessageEvent.delta
  fire({ type: "message_update", assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: "Avvio" } });
  expect(seen[0]).toEqual({ type: "text_delta", text: "Avvio" });
  // Non-text message_update events (e.g. text_start) are ignored.
  fire({ type: "message_update", assistantMessageEvent: { type: "text_start", contentIndex: 0 } });
  expect(seen).toHaveLength(1);
});

test("real Pi thinking_delta becomes a dedicated activity_delta to the FE", () => {
  const { rpc, fire } = fakeRpc();
  const b = new SessionBridge(rpc);
  const seen: any[] = [];
  b.onClientEvent((e) => seen.push(e));

  fire({
    type: "message_update",
    assistantMessageEvent: { type: "thinking_delta", contentIndex: 0, delta: "Valuto le ambiguità" },
  });

  expect(seen).toEqual([{ type: "activity_delta", text: "Valuto le ambiguità" }]);
});

test("assistant message_end exposes sanitized token usage with the configured context window", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));
  bridge.setContextWindow(200_000);

  fire({
    type: "message_end",
    message: {
      role: "assistant",
      stopReason: "stop",
      usage: {
        input: 110_348,
        cacheRead: 633_344,
        cacheWrite: 12,
        output: 33_652,
        totalTokens: 777_356,
        cost: { total: 99 },
      },
      content: "DO_NOT_FORWARD",
    },
  });

  expect(seen).toEqual([{
    type: "usage",
    usage: {
      input: 110_348,
      cacheRead: 633_344,
      output: 33_652,
      totalTokens: 777_356,
      contextWindow: 200_000,
    },
  }]);
  expect(JSON.stringify(seen)).not.toMatch(/DO_NOT_FORWARD|cost|cacheWrite/);
});

test("assistant provider errors are sanitized and leave the turn failed", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));
  bridge.beginTurn();

  fire({
    type: "message_end",
    message: {
      role: "assistant",
      stopReason: "error",
      errorMessage: "Connection failed for https://secret.invalid/?api_key=DO_NOT_LEAK",
    },
  });
  fire({ type: "agent_end", messages: [] });

  expect(bridge.turnState()).toBe("failed");
  expect(seen).toContainEqual({
    type: "info",
    level: "error",
    text: "Model request failed. Check provider connectivity, then Resume the session.",
  });
  expect(JSON.stringify(seen)).not.toContain("DO_NOT_LEAK");
});

test("subscription model errors become a safe actionable client message", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));

  fire({
    type: "message_end",
    message: {
      role: "assistant",
      stopReason: "error",
      errorMessage: "429: {\"code\":\"1311\",\"message\":\"Your current subscription plan does not yet include access to GLM-5.3\",\"secret\":\"DO_NOT_LEAK\"}",
    },
  });

  expect(seen).toContainEqual({
    type: "info",
    level: "error",
    text: "The selected model is unavailable for the current subscription. Choose another model and start a new session.",
  });
  expect(JSON.stringify(seen)).not.toContain("GLM-5.3");
  expect(JSON.stringify(seen)).not.toContain("DO_NOT_LEAK");
});

test("markFailed records backend-detected failure without emitting raw detail", () => {
  const { rpc } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));
  bridge.beginTurn();

  bridge.markFailed();

  expect(bridge.turnState()).toBe("failed");
  expect(seen).toEqual([]);
});

test("reviewer wait and response transition waiting back to running", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  bridge.beginTurn();
  fire({
    type: "extension_ui_request",
    id: "pi-1",
    method: "input",
    title: JSON.stringify({ id: "gate-1", widget: "select" }),
  });
  expect(bridge.turnState()).toBe("waiting");
  bridge.respond({ id: "gate-1", choices: ["approve"] });
  expect(bridge.turnState()).toBe("running");
  fire({ type: "agent_end", messages: [] });
  expect(bridge.turnState()).toBe("idle");
});

test("agent_start transitions an idle turn to running", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));

  expect(bridge.turnState()).toBe("idle");
  fire({ type: "agent_start" });

  expect(bridge.turnState()).toBe("running");
  expect(seen).toContainEqual({ type: "system_event", event: "agent_start" });
});

test("agent_end leaves a pending reviewer wait intact", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  bridge.beginTurn();
  fire({
    type: "extension_ui_request",
    id: "pi-1",
    method: "input",
    title: JSON.stringify({ id: "gate-1", widget: "select" }),
  });

  fire({ type: "agent_end", messages: [] });

  expect(bridge.turnState()).toBe("waiting");
});

test("extension_ui_request nativo (method:input, title=json) diventa ui_request col descriptor ed è il pendente", () => {
  const { rpc, fire } = fakeRpc();
  const b = new SessionBridge(rpc);
  const seen: any[] = [];
  b.onClientEvent((e) => seen.push(e));
  const descriptor = { id: "u1", widget: "select" };
  // Pi assegna a ctx.ui.input un proprio id RPC (crypto.randomUUID), distinto dall'id
  // interno del descriptor che il gate mette nel `title`. Al frontend va il descriptor.
  fire({ type: "extension_ui_request", id: "pi-req-1", method: "input", title: JSON.stringify(descriptor) });
  expect(seen[0]).toEqual({ type: "ui_request", ui_request: descriptor });
  expect(b.pendingWidget()).toEqual(descriptor);
});

test("respond correla sull'id RPC di Pi (non sull'id del descriptor) e azzera il pendente", () => {
  const { rpc, sent, fire } = fakeRpc();
  const b = new SessionBridge(rpc);
  // Pi emette la richiesta con il SUO id RPC ("pi-req-1"); il descriptor nel title ha id "u1".
  fire({ type: "extension_ui_request", id: "pi-req-1", method: "input", title: JSON.stringify({ id: "u1", widget: "select" }) });
  // Il frontend rimanda l'id del descriptor ("u1").
  expect(b.respond({ id: "u1", choices: ["a"] })).toBe(true);
  // Pi correla la risposta sul SUO id ("pi-req-1") per risolvere ctx.ui.input; il value
  // continua a portare l'id del descriptor, cosi' il check interno del gate regge.
  expect(sent.at(-1)).toEqual({ type: "extension_ui_response", id: "pi-req-1", value: JSON.stringify({ id: "u1", choices: ["a"] }) });
  expect(b.pendingWidget()).toBeNull();
});

test("respond rifiuta risposte senza gate pendente o con id non corrispondente", () => {
  const { rpc, sent, fire } = fakeRpc();
  const b = new SessionBridge(rpc);
  // Nessun gate pendente: la risposta non parte e lo stato non cambia.
  expect(b.respond({ id: "u0", choices: ["a"] })).toBe(false);
  expect(sent).toEqual([]);

  fire({ type: "extension_ui_request", id: "pi-req-1", method: "input", title: JSON.stringify({ id: "u1", widget: "select" }) });
  // Risposta stantia per un ALTRO gate: rifiutata, il gate vero resta pendente in waiting.
  expect(b.respond({ id: "u0", choices: ["a"] })).toBe(false);
  expect(sent).toEqual([]);
  expect(b.turnState()).toBe("waiting");
  expect(b.pendingWidget()).toEqual({ id: "u1", widget: "select" });
  // Doppio submit: il primo passa, il secondo (pendente ormai nullo) viene rifiutato.
  expect(b.respond({ id: "u1", choices: ["a"] })).toBe(true);
  expect(b.respond({ id: "u1", choices: ["a"] })).toBe(false);
  expect(sent).toHaveLength(1);
});

test("agent_end di Pi diventa un system_event agent_end per il FE", () => {
  const { rpc, fire } = fakeRpc();
  const b = new SessionBridge(rpc);
  const seen: any[] = [];
  b.onClientEvent((e) => seen.push(e));
  // Pi emette agent_end alla fine di ogni prompt (pi-agent-core agent-loop); e' il solo
  // segnale di fine turno: senza mapparlo il FE non puo' mai uscire dallo stato "working".
  fire({ type: "agent_end", messages: [] });
  expect(seen).toEqual([{ type: "system_event", event: "agent_end" }]);
});

test("generic Pi system events expose only an allowlisted non-empty event name", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));

  fire({
    type: "system_event",
    event: "session_exit",
    command: "curl https://secret.invalid/?token=DO_NOT_LEAK",
    result: { path: "/srv/private/key" },
  });
  fire({ type: "system_event", event: "   ", stderr: "DO_NOT_LEAK_STDERR" });
  fire({ type: "system_event", event: 42, args: "DO_NOT_LEAK_ARGS" });

  expect(seen).toEqual([{ type: "system_event", event: "session_exit" }]);
  expect(Object.keys(seen[0])).toEqual(["type", "event"]);
  expect(JSON.stringify(seen)).not.toMatch(/DO_NOT_LEAK|command|result|path|stderr|args/);
});

test("steer invia un comando steer e riattiva il turno", () => {
  const { rpc, sent, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  bridge.beginTurn();
  fire({
    type: "extension_ui_request",
    id: "pi-1",
    method: "input",
    title: JSON.stringify({ id: "gate-1", widget: "select" }),
  });
  expect(bridge.turnState()).toBe("waiting");

  bridge.steer("considera solo il 2024");

  expect(sent.at(-1)).toEqual({ type: "steer", message: "considera solo il 2024" });
  expect(bridge.turnState()).toBe("running");
});

test("tool execution emits only a sanitized correlated lifecycle", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));

  fire({
    type: "tool_execution_start",
    toolCallId: "tool-1",
    toolName: "bash",
    args: { command: "curl https://secret.invalid/?token=DO_NOT_LEAK_ARGS" },
  });
  fire({
    type: "tool_execution_update",
    toolCallId: "tool-1",
    toolName: "bash",
    partialResult: { content: "DO_NOT_LEAK_PARTIAL" },
  });
  fire({
    type: "tool_execution_end",
    toolCallId: "tool-1",
    toolName: "bash",
    result: { content: "DO_NOT_LEAK_RESULT" },
    isError: false,
  });

  expect(seen).toEqual([
    {
      type: "activity_event",
      activity: { kind: "tool", toolCallId: "tool-1", toolName: "bash", status: "running" },
    },
    {
      type: "activity_event",
      activity: { kind: "tool", toolCallId: "tool-1", toolName: "bash", status: "completed" },
    },
  ]);
  expect(JSON.stringify(seen)).not.toMatch(/DO_NOT_LEAK|args|partialResult|result|command/);
});

test("failed tool completion is sanitized and invalid starts are ignored", () => {
  const { rpc, fire } = fakeRpc();
  const bridge = new SessionBridge(rpc);
  const seen: any[] = [];
  bridge.onClientEvent((event) => seen.push(event));

  fire({ type: "tool_execution_start", toolName: "read", args: { path: "DO_NOT_LEAK_PATH" } });
  fire({
    type: "tool_execution_end",
    toolCallId: "tool-2",
    toolName: 42,
    result: { error: "DO_NOT_LEAK_ERROR" },
    isError: true,
  });

  expect(seen).toEqual([
    {
      type: "activity_event",
      activity: { kind: "tool", toolCallId: "tool-2", toolName: "Tool", status: "failed" },
    },
  ]);
  expect(JSON.stringify(seen)).not.toMatch(/DO_NOT_LEAK|args|result|error|path/);
});