236 lines · 8.1 KB · TypeScript
RawDownload
import assert from "node:assert/strict";import test from "node:test";import {  createSessionRuntimeController,  type SessionRuntimeBinding,} from "../src/features/agent/runtime/session-runtime-controller";import type {  RuntimeEventPayload,  RuntimeEventSubscription,  RuntimeSessionSummary,} from "../src/features/agent/runtime/api";import type { Session } from "../src/features/agent/runtime/types"; type OpenedSubscription = {  runtime: string;  after: number;  send: (payload: RuntimeEventPayload) => void;  closed: boolean;}; function makeSession(overrides: Partial<Session> = {}): Session {  return {    id: "tab-1",    piSessionId: "pi-1",    title: "Adoption test",    messages: [      { id: "user-1", role: "user", text: "hello" },      { id: "assistant-1", role: "assistant", text: "", blocks: [] },    ],    status: "running",    error: "",    input: "",    activeAssistantId: "assistant-1",    ...overrides,  };} function createAdoptionHarness(initial: Session) {  let liveSession = initial;  const opened: OpenedSubscription[] = [];  const pendingFetches: ((entries: RuntimeSessionSummary[]) => void)[] = [];   const controller = createSessionRuntimeController({    idleReconnectMs: 0,    pollIntervalMs: 1_000_000,    api: {      listRuntimeSessions: () =>        new Promise<RuntimeSessionSummary[]>((resolve) => {          pendingFetches.push(resolve);        }),      loadRuntimeStatus: async () => null,      subscribeRuntimeEvents: (runtime, after, _piSessionId, handlers): RuntimeEventSubscription => {        const entry: OpenedSubscription = {          runtime,          after,          send: handlers.onPayload,          closed: false,        };        opened.push(entry);        return {          close: () => {            entry.closed = true;          },        };      },    },  });   const binding: SessionRuntimeBinding = {    commit: (sessionId, patch) => {      if (sessionId === liveSession.id) liveSession = patch(liveSession);    },    getSession: (sessionId) => (sessionId === liveSession.id ? liveSession : undefined),    getSessions: () => [liveSession],  };  controller.bind(binding);   return {    controller,    opened,    session: () => liveSession,    setSession: (next: Session) => {      liveSession = next;    },    resolveFetch: async (index: number, entries: RuntimeSessionSummary[]) => {      pendingFetches[index]?.(entries);      await new Promise((resolve) => setTimeout(resolve, 0));    },  };} // Server restart mid-session: the poll's pi match finds the session's runtime// under a NEW server key. Adoption is a controller-internal connection-key// override — session state is untouched, the cursor resets, the SSE reopens on// the new key from seq 0, and events keep flowing into the transcript.test("restart adoption reopens the SSE under the override key with a reset cursor", async () => {  const harness = createAdoptionHarness(makeSession());   try {    harness.controller.reconcile([harness.session()]);    assert.equal(harness.opened.length, 1, "one attachment on first reconcile");    assert.equal(harness.opened[0].runtime, "tab-1", "first attach uses the session id as key");     // Stream a few events so the live cursor climbs.    for (let seq = 1; seq <= 5; seq += 1) {      harness.opened[0].send({        type: "pi",        seq,        event: {          type: "message_update",          message: { role: "assistant", content: [{ type: "text", text: `chunk ${seq}` }] },        },      });    }    harness.controller.flush("tab-1");    assert.equal(harness.session().lastEventSeq, 5, "cursor committed before the restart");     // The server restarts; the runtime now lives under "srv-2" (same pi id).    harness.controller.pollNow();    await harness.resolveFetch(0, [      { sessionId: "srv-2", status: { active: true, piSessionId: "pi-1" } },    ]);     assert.equal(harness.controller.connectionKey("tab-1"), "srv-2", "override recorded");    assert.equal(harness.opened[0].closed, true, "stale attachment closed");    assert.equal(harness.opened.length, 2, "attachment reopened by the controller itself");    assert.equal(harness.opened[1].runtime, "srv-2", "reopened under the adopted server key");    assert.equal(      harness.opened[1].after,      0,      "cursor reset on adoption — the fresh runtime's seq restarts from 0",    );     // Events from the adopted runtime keep flowing into the same session.    harness.opened[1].send({      type: "pi",      seq: 1,      event: {        type: "message_update",        message: { role: "assistant", content: [{ type: "text", text: "after restart" }] },      },    });    harness.controller.flush("tab-1");    const assistant = harness.session().messages.find((m) => m.id === "assistant-1");    assert.ok(      assistant?.blocks?.some((b) => b.kind === "text" && b.text.includes("after restart")),      "post-adoption events land in the transcript",    );    assert.equal(harness.session().lastEventSeq, 1, "cursor tracks the new runtime's seq");  } finally {    harness.controller.closeAll();    harness.controller.unbind();  }  await new Promise((resolve) => setTimeout(resolve, 0));}); // Upgrade path: a session persisted as RUNNING under a legacy pre-alias rt-*// key is reattached via the one-shot connection-key seed read from old// localStorage state; the SSE opens on the legacy key from the persisted// cursor.test("legacy rt-* seed reattaches a restored running session to its old runtime key", async () => {  const harness = createAdoptionHarness(makeSession({ lastEventSeq: 7 }));   try {    harness.controller.seedConnectionKey("tab-1", "rt-legacy");    harness.controller.reconcile([harness.session()]);     assert.equal(harness.opened.length, 1);    assert.equal(harness.opened[0].runtime, "rt-legacy", "attached under the legacy runtime key");    assert.equal(harness.opened[0].after, 7, "resumes from the persisted cursor");    assert.equal(harness.controller.connectionKey("tab-1"), "rt-legacy");     // A seed never clobbers itself or downgrades an existing override.    harness.controller.seedConnectionKey("tab-1", "rt-other");    assert.equal(harness.controller.connectionKey("tab-1"), "rt-legacy");  } finally {    harness.controller.closeAll();    harness.controller.unbind();  }  await new Promise((resolve) => setTimeout(resolve, 0));}); // Crosstalk guard: a piSessionId shared by two sessions must not let one// runtime entry set a connection-key override on the session that does NOT own// the runtime (the pi reverse-index is ambiguous there).test("shared piSessionId never records an override for the non-owning session", async () => {  let sessionA = makeSession({    id: "tab-A",    piSessionId: "pi-shared",    status: "idle",  });  let sessionB = makeSession({    id: "tab-B",    piSessionId: "pi-shared",    status: "idle",  });  const sessions = () => [sessionA, sessionB];   const controller = createSessionRuntimeController({    idleReconnectMs: 0,    pollIntervalMs: 1_000_000,    api: {      listRuntimeSessions: async () => [        { sessionId: "tab-A", status: { active: true, piSessionId: "pi-shared" } },      ],      loadRuntimeStatus: async () => null,      subscribeRuntimeEvents: () => ({ close: () => undefined }),    },  });  controller.bind({    commit: (id, patch) => {      if (id === "tab-A") sessionA = patch(sessionA);      else if (id === "tab-B") sessionB = patch(sessionB);    },    getSession: (id) => sessions().find((s) => s.id === id),    getSessions: sessions,  });   controller.pollNow();  await new Promise((resolve) => setTimeout(resolve, 0));   try {    assert.equal(sessionA.status, "running", "owner promoted via the direct match");    assert.equal(sessionB.status, "idle", "the colliding session is untouched");    assert.equal(      controller.connectionKey("tab-B"),      "tab-B",      "no override crosstalk onto the colliding session",    );    assert.equal(controller.connectionKey("tab-A"), "tab-A", "direct match needs no override");  } finally {    controller.closeAll();    controller.unbind();  }  await new Promise((resolve) => setTimeout(resolve, 0));});