pi-webui/tests/session-broker.test.ts

import assert from "node:assert/strict";
import { readFileSync } from "node:fs";
import { test } from "node:test";
import { BrokerRuntime, SessionBroker, type BrokerMessageHooks, type BrokerSession, type SocketSink } from "../server/session-broker.ts";

type Message = { type: string };
const messagesOf = (socket: FakeSocket): Message[] => socket.messages as Message[];

class FakeClock {
  time = 0;
  now = (): number => this.time;
}

class FakeSocket implements SocketSink {
  open = true;
  messages: unknown[] = [];
  isOpen = (): boolean => this.open;
  send = (message: unknown): void => { this.messages.push(message); };
}

class FakeSession implements BrokerSession {
  readonly sessionId: string;
  isStreaming = false;
  isIdle = true;
  abortCount = 0;
  disposeCount = 0;
  promptCalls: Array<{ text: string; preflight: (accepted: boolean) => void }> = [];
  private listeners = new Set<(event: unknown) => void>();
  constructor(sessionId: string) { this.sessionId = sessionId; }
  subscribe = (listener: (event: unknown) => void): (() => void) => { this.listeners.add(listener); return () => this.listeners.delete(listener); };
  prompt = async (text: string, options?: { preflightResult?: (accepted: boolean) => void }): Promise<void> => {
    this.isStreaming = true;
    this.isIdle = false;
    this.promptCalls.push({ text, preflight: options?.preflightResult ?? (() => {}) });
  };
  steer = async (): Promise<void> => {};
  followUp = async (): Promise<void> => {};
  abort = async (): Promise<void> => { this.abortCount++; this.isStreaming = false; this.isIdle = true; };
  dispose = (): void => { this.disposeCount++; };
  emit(event: unknown): void { for (const listener of this.listeners) listener(event); }
  finish(): void { this.isStreaming = false; this.isIdle = true; this.emit({ type: "agent_settled" }); }
}

const testMessages: BrokerMessageHooks = {
  ready: (runtime) => ({ type: "ready", runtimeId: runtime.runtimeId, sessionId: runtime.session.sessionId }),
  stateSync: (_runtime, snapshot) => ({ type: "stateSync", state: snapshot }),
  event: (event) => ({ type: "event", event }),
  choiceRequest: (request) => ({ type: "choiceRequest", ...request }),
  choiceResolved: (requestId, value) => ({ type: "choiceResolved", requestId, value }),
  promptAccepted: (requestId) => ({ type: "promptAccepted", requestId }),
  promptRejected: (requestId) => ({ type: "promptRejected", requestId }),
};

function setup(options: Partial<ConstructorParameters<typeof SessionBroker>[0]> = {}) {
  const clock = new FakeClock();
  const sessions = new Map<string, FakeSession>();
  const broker = new SessionBroker({
    clock,
    factory: { keyForContext: (context) => String(context.owner), create: async (key) => { const session = new FakeSession(key); sessions.set(key, session); return session; } },
    serializer: { snapshot: (runtime) => ({ sessionId: runtime.session.sessionId }) },
    onDelete: async () => {},
    onError: () => {},
    messages: testMessages,
    ...options,
  });
  return { clock, sessions, broker };
}

async function runtimeAndSocket(broker: SessionBroker, key = "session-1") {
  const runtime = await broker.getOrCreate({ owner: key });
  const socket = new FakeSocket();
  const attachment = runtime.attach(socket);
  return { runtime, socket, attachment, session: runtime.session as FakeSession };
}

test("canonical keys share one in-flight creation and stable runtime identity", async () => {
  let creates = 0;
  let resolve!: (session: FakeSession) => void;
  const broker = new SessionBroker({
    factory: { keyForContext: () => "42", create: async () => { creates++; return new Promise<FakeSession>((done) => { resolve = done; }); } },
    serializer: { snapshot: () => ({}) },
    messages: testMessages,
    onDelete: async () => {},
    onError: () => {},
  });
  const first = broker.getOrCreate({ value: 1 });
  const second = broker.getOrCreate({ value: 2 });
  await Promise.resolve();
  assert.equal(creates, 1);
  const session = new FakeSession("42");
  resolve(session);
  const [a, b] = await Promise.all([first, second]);
  assert.equal(a, b);
  assert.equal(a.runtimeId, b.runtimeId);
  assert.equal(a.context.value, 1);
});

test("attach sends a synchronous snapshot barrier, fanout isolates failures, and detach does not abort", async () => {
  const { broker } = setup();
  const { runtime, socket, attachment, session } = await runtimeAndSocket(broker);
  const second = new FakeSocket();
  runtime.attach(second);
  assert.deepEqual(messagesOf(socket).slice(0, 2), [
    { type: "ready", runtimeId: runtime.runtimeId, sessionId: "session-1" },
    { type: "stateSync", state: { sessionId: "session-1" } },
  ]);
  session.emit({ type: "delta", value: 1 });
  assert.equal(messagesOf(socket)[2].type, "event");
  attachment.detach();
  assert.equal(session.abortCount, 0);
  session.emit({ type: "delta", value: 2 });
  assert.equal(messagesOf(second).filter((m) => m.type === "event").length, 2);
  const broken = new FakeSocket();
  broken.send = () => { throw new Error("broken socket"); };
  assert.throws(() => runtime.attach(broken), /Socket closed during attach/);
  session.emit({ type: "safe" });
  assert.equal(messagesOf(second).filter((m) => m.type === "event").length, 3);
});

test("turn and agent settlement refresh lifecycle activity", async () => {
  const { broker, clock } = setup();
  const { runtime, session } = await runtimeAndSocket(broker);
  clock.time = 10;
  session.emit({ type: "turn_end" });
  assert.equal(runtime.lastUsedAt, 10);
  clock.time = 20;
  session.emit({ type: "agent_end" });
  assert.equal(runtime.lastUsedAt, 20);
});

test("explicit abort is ordered and idempotent", async () => {
  const { broker } = setup();
  const { runtime, attachment, session } = await runtimeAndSocket(broker);
  await runtime.abort(attachment);
  await runtime.abort(attachment);
  assert.equal(session.abortCount, 1);
});

test("abort idempotence resets for the next accepted prompt", async () => {
  const { broker } = setup();
  const { runtime, attachment, session } = await runtimeAndSocket(broker);
  await runtime.abort(attachment);
  await runtime.prompt(attachment, { requestId: "next", text: "hello" }, "next");
  await runtime.abort(attachment);
  assert.equal(session.abortCount, 2);
});

test("mutating commands preserve order and reject stale attachments", async () => {
  const { broker } = setup();
  const { runtime, attachment } = await runtimeAndSocket(broker);
  const order: string[] = [];
  const first = runtime.invoke(attachment, async () => { order.push("first"); await new Promise((resolve) => setTimeout(resolve, 2)); order.push("first-done"); });
  const second = runtime.invoke(attachment, async () => { order.push("second"); });
  await Promise.all([first, second]);
  assert.deepEqual(order, ["first", "first-done", "second"]);
  attachment.detach();
  await assert.rejects(runtime.invoke(attachment, async () => {}), /Attachment/);
});

test("prompt records prevent duplicates and detect fingerprint conflicts", async () => {
  const { broker } = setup();
  const { runtime, attachment, session, socket } = await runtimeAndSocket(broker);
  const request = { requestId: "req-1", text: "hello" };
  const first = runtime.prompt(attachment, request, "same");
  await Promise.resolve();
  assert.equal(session.promptCalls.length, 1);
  const duplicate = runtime.prompt(attachment, request, "same");
  session.promptCalls[0].preflight(true);
  await Promise.all([first, duplicate]);
  assert.equal(session.promptCalls.length, 1);
  assert.deepEqual(messagesOf(socket).filter((m) => m.type === "promptAccepted"), [
    { type: "promptAccepted", requestId: "req-1" },
    { type: "promptAccepted", requestId: "req-1" },
  ]);
  await assert.rejects(runtime.prompt(attachment, { ...request, text: "other" }, "different"), /different payload/);
});

test("prompt preflight rejection releases the chain and records the rejection", async () => {
  const { broker } = setup();
  const { runtime, attachment, session, socket } = await runtimeAndSocket(broker);
  const first = runtime.prompt(attachment, { requestId: "reject", text: "no" }, "x");
  await Promise.resolve();
  session.promptCalls[0].preflight(false);
  await assert.rejects(first, /rejected/);
  const duplicate = runtime.prompt(attachment, { requestId: "reject", text: "no" }, "x");
  await assert.rejects(duplicate, /rejected/);
  assert.deepEqual(messagesOf(socket).filter((m) => m.type === "promptRejected"), [
    { type: "promptRejected", requestId: "reject" },
    { type: "promptRejected", requestId: "reject" },
  ]);
});

test("abort rejects and clears pending choices", async () => {
  const { broker } = setup();
  const { runtime, attachment } = await runtimeAndSocket(broker);
  const pending = runtime.requestChoice({ choiceId: "abort-choice", payload: {} });
  await runtime.abort(attachment);
  await assert.rejects(pending, /Runtime aborted/);
  assert.equal(runtime.hasPendingChoices, false);
});

test("pending choices replay on reconnect and resolve once", async () => {
  const { broker } = setup();
  const { runtime, attachment } = await runtimeAndSocket(broker);
  const choice = runtime.requestChoice({ choiceId: "choice-1", payload: { question: "Continue?" } });
  attachment.detach();
  const socket = new FakeSocket();
  const reconnected = runtime.attach(socket);
  assert.equal(messagesOf(socket).filter((m) => m.type === "choiceRequest").length, 1);
  await runtime.resolveChoice(reconnected, "choice-1", "yes");
  assert.equal(await choice, "yes");
  assert.deepEqual(messagesOf(socket).at(-1), { type: "choiceResolved", requestId: "choice-1", value: "yes" });
  await runtime.resolveChoice(reconnected, "choice-1", "no");
  assert.equal(messagesOf(socket).filter((m) => m.type === "choiceResolved").length, 1);
});

test("choice cancellation rejects the deferred call and clears attached clients", async () => {
  const { broker } = setup();
  const { runtime, socket } = await runtimeAndSocket(broker);
  const pending = runtime.requestChoice({ choiceId: "cancel", payload: {} });
  runtime.rejectChoice("cancel");
  await assert.rejects(pending, /cancelled/);
  assert.equal(messagesOf(socket).filter((message) => message.type === "choiceResolved").length, 1);
});

test("pending choices are bounded per runtime", async () => {
  const { broker } = setup();
  const { runtime } = await runtimeAndSocket(broker);
  const pending: Promise<unknown>[] = [];
  for (let index = 0; index < 16; index++) pending.push(runtime.requestChoice({ choiceId: `choice-${index}`, payload: {} }).catch(() => undefined));
  await assert.rejects(runtime.requestChoice({ choiceId: "choice-17", payload: {} }), /too many/);
  await runtime.dispose();
  await Promise.all(pending);
});

test("deletion is conservative and removes an idle runtime once", async () => {
  const deleted: string[] = [];
  const { broker } = setup({ onDelete: async (key) => { deleted.push(key); } });
  const { runtime, attachment, session } = await runtimeAndSocket(broker);
  const other = new FakeSocket();
  const otherAttachment = runtime.attach(other);
  await assert.rejects(runtime.delete(), /active or has multiple/);
  other.open = false;
  otherAttachment.detach();
  await broker.delete("session-1");
  assert.equal(session.disposeCount, 1);
  assert.deepEqual(deleted, ["session-1"]);
  const replacement = await broker.getOrCreate({ owner: "session-1" });
  assert.notEqual(replacement, runtime);
});

test("sweep retains active work, evicts idle runtimes, and enforces capacity", async () => {
  const { broker, clock, sessions } = setup({ idleTtlMs: 10, activeMaxAgeMs: 100, maxRuntimes: 2 });
  const a = await runtimeAndSocket(broker, "a");
  a.attachment.detach();
  clock.time = 20;
  await broker.sweep();
  assert.equal(a.session.disposeCount, 1);
  const b = await runtimeAndSocket(broker, "b");
  b.attachment.detach();
  const c = await runtimeAndSocket(broker, "c");
  assert.equal(sessions.has("b"), true);
  assert.equal(sessions.has("c"), true);
  assert.equal(b.session.disposeCount, 0);
  c.attachment.detach();
  const d = await runtimeAndSocket(broker, "d");
  assert.equal(b.session.disposeCount, 1);
  d.session.isIdle = false;
  d.session.isStreaming = true;
  d.attachment.detach();
  clock.time = 1000;
  await broker.sweep();
  assert.equal(d.session.abortCount, 1);
});

test("age cleanup aborts an idle session with a pending choice", async () => {
  const { broker, clock } = setup({ activeMaxAgeMs: 10, sweepAbortGraceMs: 5 });
  const { runtime, session, attachment } = await runtimeAndSocket(broker);
  const pending = runtime.requestChoice({ choiceId: "old-choice", payload: {} });
  attachment.detach();
  clock.time = 10;
  await broker.sweep();
  await assert.rejects(pending, /Runtime aborted/);
  assert.equal(session.abortCount, 1);
  assert.equal(runtime.hasPendingChoices, false);
});

test("a new prompt waits for an age abort to settle", async () => {
  const { broker } = setup();
  const { runtime, session, attachment } = await runtimeAndSocket(broker);
  let releaseAbort!: () => void;
  const blockedAbort = new Promise<void>((resolve) => { releaseAbort = resolve; });
  session.isIdle = false;
  session.abort = async () => {
    session.abortCount++;
    await blockedAbort;
    session.isIdle = true;
  };
  attachment.detach();
  await runtime.requestAgeAbort(10);
  const nextSocket = new FakeSocket();
  const nextAttachment = runtime.attach(nextSocket);
  const prompt = runtime.prompt(nextAttachment, { requestId: "after-age-abort", text: "hello" }, "next");
  await Promise.resolve();
  assert.equal(session.promptCalls.length, 0);
  releaseAbort();
  await prompt;
  assert.equal(session.promptCalls.length, 1);
});

test("client preserves input until the runtime is ready", () => {
  const source = readFileSync(new URL("../client/main.ts", import.meta.url), "utf8");
  assert.match(source, /if \(!connected \|\| !agentReady \|\| ws\?\.readyState !== WebSocket\.OPEN\)/);
  assert.match(source, /errorMessage = "Connection is not ready\."/);
});

test("shutdown is bounded, repeated, and disposes late factory results", async () => {
  let resolve!: (session: FakeSession) => void;
  const broker = new SessionBroker({
    factory: { keyForContext: () => "late", create: async () => new Promise<FakeSession>((done) => { resolve = done; }) },
    serializer: { snapshot: () => ({}) },
  });
  const pending = broker.getOrCreate({ owner: "late" });
  const shutting = broker.shutdown(20);
  await Promise.resolve();
  resolve(new FakeSession("late"));
  await assert.rejects(pending, /shut down/);
  await shutting;
  await broker.shutdown(20);
});

test("deletion reservation stays visible until file removal completes", async () => {
  let release!: () => void;
  const { broker } = setup({ onDelete: async () => new Promise<void>((resolve) => { release = resolve; }) });
  const { runtime } = await runtimeAndSocket(broker, "delete-race");
  const deleting = broker.delete("delete-race");
  await Promise.resolve();
  const replacement = broker.getOrCreate({ owner: "delete-race" });
  let created = false;
  void replacement.then(() => { created = true; });
  await Promise.resolve();
  assert.equal(created, false);
  assert.equal(runtime.state, "disposed");
  release();
  assert.notEqual(await replacement, runtime);
  await deleting;
});

test("capacity limits are isolated by an injected group", async () => {
  const broker = new SessionBroker({
    maxRuntimes: 1,
    capacityGroup: (context) => String(context.account),
    factory: { keyForContext: (context) => String(context.key), create: async (key) => new FakeSession(key) },
    serializer: { snapshot: () => ({}) },
    messages: testMessages,
    onError: () => {},
    onDelete: async () => {},
  });
  const a1 = await broker.getOrCreate({ account: "a", key: "a-1" });
  a1.attach(new FakeSocket());
  await broker.getOrCreate({ account: "b", key: "b-1" });
  await assert.rejects(broker.getOrCreate({ account: "a", key: "a-2" }), /capacity/);
  const b2 = await broker.getOrCreate({ account: "b", key: "b-2" });
  assert.equal(b2.context.account, "b");
});

test("capacity reservation rejects concurrent over-limit creation before factory await", async () => {
  let release!: () => void;
  let creates = 0;
  const broker = new SessionBroker({
    maxRuntimes: 1,
    factory: { keyForContext: (context) => String(context.key), create: async (key) => { creates++; return new Promise<FakeSession>((resolve) => { release = () => resolve(new FakeSession(key)); }); } },
    serializer: { snapshot: () => ({}) },
  });
  const first = broker.getOrCreate({ key: "first" });
  await assert.rejects(broker.getOrCreate({ key: "second" }), /capacity/);
  assert.equal(creates, 1);
  release();
  await first;
});

test("synchronous prompt starter failure settles the request record", async () => {
  const { broker } = setup({ promptStarter: () => { throw new Error("starter failed"); } });
  const { runtime, attachment } = await runtimeAndSocket(broker);
  const request = { requestId: "sync-failure", text: "hello" };
  await assert.rejects(runtime.prompt(attachment, request, "throw"), /starter failed/);
  await assert.rejects(runtime.prompt(attachment, request, "throw"), /starter failed/);
});

test("late preflight cannot accept a terminal runtime", async () => {
  let preflight!: (accepted: boolean) => void;
  const { broker } = setup({ promptStarter: (_session, _request, callback) => {
    preflight = callback;
    return new Promise<void>(() => {});
  } });
  const { runtime, attachment } = await runtimeAndSocket(broker);
  const pending = runtime.prompt(attachment, { requestId: "late", text: "hello" }, "late");
  await Promise.resolve();
  await runtime.dispose();
  preflight(true);
  await assert.rejects(pending, /disposed/);
});

test("closed attachments detach and delayed responses require the current generation", async () => {
  const { broker } = setup();
  const { runtime, attachment, session } = await runtimeAndSocket(broker);
  const generation = attachment.generation;
  attachment.attachment.sink.isOpen = () => false;
  session.emit({ type: "delta" });
  assert.equal(runtime.attachmentCount, 0);
  const socket = new FakeSocket();
  const next = runtime.attach(socket);
  assert.equal(runtime.isAttachmentValid(next, generation), false);
  assert.equal(runtime.isAttachmentValid(next, next.generation), true);
  next.detach();
});

test("shutdown disposes a runtime even when abort never settles", async () => {
  let session!: FakeSession;
  const broker = new SessionBroker({
    factory: { keyForContext: () => "hang", create: async (key) => { session = new FakeSession(key); session.isIdle = false; session.isStreaming = true; session.abort = async () => new Promise<void>(() => {}); return session; } },
    serializer: { snapshot: () => ({}) },
  });
  await broker.getOrCreate({ owner: "hang" });
  const keepAlive = setTimeout(() => {}, 50);
  await broker.shutdown(10);
  clearTimeout(keepAlive);
  assert.equal(session.abortCount, 0);
  assert.equal(session.disposeCount, 1);
  await broker.shutdown(10);
  assert.equal(session.disposeCount, 1);
});

test("a throwing event hook does not stop fanout", async () => {
  const errors: unknown[] = [];
  const { broker } = setup({
    onError: (error) => errors.push(error),
    messages: {
      ready: () => ({ type: "ready" }),
      stateSync: () => ({ type: "stateSync" }),
      event: () => { throw new Error("event hook failed"); },
      choiceRequest: (request) => ({ type: "choiceRequest", ...request }),
      choiceResolved: (requestId, value) => ({ type: "choiceResolved", requestId, value }),
      promptAccepted: (requestId) => ({ type: "promptAccepted", requestId }),
      promptRejected: (requestId) => ({ type: "promptRejected", requestId }),
    },
  });
  const { runtime, session } = await runtimeAndSocket(broker);
  const second = new FakeSocket();
  runtime.attach(second);
  session.emit({ type: "event" });
  assert.equal(runtime.attachmentCount, 2);
  assert.equal(errors.length, 1);
});

test("canonical manager key is selected before factory creation", async () => {
  const calls: string[] = [];
  const broker = new SessionBroker({
    factory: {
      keyForContext: (context: { manager: { getSessionId: () => string } }) => context.manager.getSessionId(),
      create: async (key, context) => { calls.push(`${key}:${context.manager.getSessionId()}`); return new FakeSession(key); },
    },
    serializer: { snapshot: () => ({}) },
  });
  await broker.getOrCreate({ manager: { getSessionId: () => "canonical-id" } });
  assert.deepEqual(calls, ["canonical-id:canonical-id"]);
  assert.ok(broker.get("canonical-id"));
  assert.equal(broker.get("requested-path"), undefined);
});

test("disposeAll resets runtimes without shutting down the broker", async () => {
  const { broker, sessions } = setup();
  await runtimeAndSocket(broker, "reset");
  await broker.disposeAll();
  assert.equal(sessions.get("reset")?.disposeCount, 1);
  const replacement = await broker.getOrCreate({ owner: "reset" });
  assert.equal(replacement.session.sessionId, "reset");
});

test("reset clears stale creation reservations and disposes late results", async () => {
  const releases: Array<(session: FakeSession) => void> = [];
  const broker = new SessionBroker({
    factory: {
      keyForContext: (context: { owner: string }) => context.owner,
      create: async () => new Promise<FakeSession>((resolve) => { releases.push(resolve); }),
    },
    serializer: { snapshot: () => ({}) },
    messages: testMessages,
    onDelete: async () => {},
    onError: () => {},
  });
  const pending = broker.getOrCreate({ owner: "old" });
  const reset = broker.disposeAll(1);
  await new Promise((resolve) => setTimeout(resolve, 5));
  await reset;
  const replacementPending = broker.getOrCreate({ owner: "old" });
  assert.equal(releases.length, 2);
  releases[1](new FakeSession("old"));
  const replacement = await replacementPending;
  const stale = new FakeSession("old");
  releases[0](stale);
  await assert.rejects(pending, /reset/);
  assert.equal(stale.disposeCount, 1);
  assert.equal(broker.get("old"), replacement);
});

test("direct runtime deletion does not leave a disposed registry entry", async () => {
  const { broker } = setup();
  const { runtime } = await runtimeAndSocket(broker, "direct-delete");
  await runtime.delete();
  assert.equal(broker.get("direct-delete"), undefined);
  const replacement = await broker.getOrCreate({ owner: "direct-delete" });
  assert.notEqual(replacement, runtime);
});