mason: fix TUI RPC lifecycle issues

Co-authored-by: Alfonso [Magic Context] <288211368+alfonso-magic-context@users.noreply.github.com>

Author: ualtinok

packages/plugin/.gitignore

@@ -30,8 +30,9 @@
     "README.md"
   ],
   "scripts": {
-    "build": "bun build src/index.ts --outdir dist --target node --format esm --external @opencode-ai/plugin --external @huggingface/transformers --external onnxruntime-web --external bun:sqlite --external node:sqlite && tsc --emitDeclarationOnly",
+    "build": "bun run build:tui && bun build src/index.ts --outdir dist --target node --format esm --external @opencode-ai/plugin --external @huggingface/transformers --external onnxruntime-web --external bun:sqlite --external node:sqlite && tsc --emitDeclarationOnly",
     "build:tui": "bun scripts/build-tui.ts",
+    "check:tui-compiled": "bun run build:tui && git diff --exit-code -- src/tui-compiled && test -z \"$(git status --porcelain -- src/tui-compiled)\"",
     "typecheck": "tsc --noEmit && tsc -p tsconfig.scripts.json",
     "test": "bun test",
     "lint": "biome check .",

packages/plugin/package.json

@@ -1,11 +1,12 @@
 import { afterEach, describe, expect, test } from "bun:test";
-import { mkdirSync, mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs";
+import { mkdirSync, mkdtempSync, readdirSync, readFileSync, rmSync, writeFileSync } from "node:fs";
 import { createServer } from "node:http";
 import { tmpdir } from "node:os";
 import { dirname, join } from "node:path";
 import { MagicContextRpcClient } from "./rpc-client";
+import { drainNotifications, isTuiConnected, pushNotification } from "./rpc-notifications";
 import { MagicContextRpcServer } from "./rpc-server";
-import { parseRpcPortFile, rpcPortFilePath } from "./rpc-utils";
+import { parseRpcPortFile, type RpcPortFileRecord, rpcPortDir, rpcPortFilePath } from "./rpc-utils";
 
 interface TestServer {
     port: number;
@@ -56,6 +57,78 @@ function writePortFileForPid(
     writeFileSync(portFile, JSON.stringify({ port, pid, started_at: startedAt }), "utf-8");
 }
 
+function readNewestPortRecord(storageDir: string, directory: string): RpcPortFileRecord | null {
+    const records: RpcPortFileRecord[] = [];
+    for (const entry of readdirSync(rpcPortDir(storageDir, directory))) {
+        if (!entry.startsWith("port-") || !entry.endsWith(".json")) continue;
+        const record = parseRpcPortFile(
+            readFileSync(join(rpcPortDir(storageDir, directory), entry), "utf-8"),
+        );
+        if (record) records.push(record);
+    }
+    records.sort((a, b) => b.started_at - a.started_at);
+    return records[0] ?? null;
+}
+
+async function waitFor(condition: () => boolean, label: string, timeoutMs = 2_000): Promise<void> {
+    const start = Date.now();
+    while (Date.now() - start < timeoutMs) {
+        if (condition()) return;
+        await new Promise((resolve) => setTimeout(resolve, 25));
+    }
+    throw new Error(`Timed out waiting for ${label}`);
+}
+
+async function openSocket(port: number, token: string): Promise<WebSocket> {
+    const ws = new WebSocket(`ws://127.0.0.1:${port}/ws?token=${encodeURIComponent(token)}`);
+    await new Promise<void>((resolve, reject) => {
+        const timeout = setTimeout(() => reject(new Error("socket open timed out")), 2_000);
+        ws.addEventListener(
+            "open",
+            () => {
+                clearTimeout(timeout);
+                resolve();
+            },
+            { once: true },
+        );
+        ws.addEventListener(
+            "error",
+            () => {
+                clearTimeout(timeout);
+                reject(new Error("socket open failed"));
+            },
+            { once: true },
+        );
+    });
+    return ws;
+}
+
+function waitForJsonMessage<T extends { type?: string }>(
+    ws: WebSocket,
+    predicate: (message: T) => boolean,
+    timeoutMs = 2_000,
+): Promise<T> {
+    return new Promise((resolve, reject) => {
+        const timeout = setTimeout(() => {
+            ws.removeEventListener("message", onMessage);
+            reject(new Error("socket message timed out"));
+        }, timeoutMs);
+        const onMessage = (event: MessageEvent) => {
+            let message: T;
+            try {
+                message = JSON.parse(String(event.data)) as T;
+            } catch {
+                return;
+            }
+            if (!predicate(message)) return;
+            clearTimeout(timeout);
+            ws.removeEventListener("message", onMessage);
+            resolve(message);
+        };
+        ws.addEventListener("message", onMessage);
+    });
+}
+
 async function startRpcServer(handler: (method: string) => Response | object): Promise<TestServer> {
     const server = createServer(async (req, res) => {
         if (req.method === "GET" && req.url === "/health") {
@@ -145,9 +218,7 @@ describe("MagicContextRpcClient", () => {
         const port = await server.start();
         try {
             // Sanity: the port file carries a non-empty token.
-            const record = parseRpcPortFile(
-                readFileSync(rpcPortFilePath(storageDir, directory), "utf-8"),
-            );
+            const record = readNewestPortRecord(storageDir, directory);
             expect(typeof record?.token).toBe("string");
             expect((record?.token ?? "").length).toBeGreaterThan(0);
 
@@ -167,6 +238,94 @@ describe("MagicContextRpcClient", () => {
         }
     });
 
+    test("websocket upgrade rejects missing bearer token before a socket is created", async () => {
+        const storageDir = makeTempDir();
+        const directory = "/repo-ws-auth";
+        const server = new MagicContextRpcServer(storageDir, directory);
+        const port = await server.start();
+        try {
+            const res = await fetch(`http://127.0.0.1:${port}/ws`);
+            expect(res.status).toBe(401);
+            expect(isTuiConnected()).toBe(false);
+        } finally {
+            server.stop();
+        }
+    });
+
+    test("re-hello replaces the previous websocket notification sink", async () => {
+        drainNotifications(Number.MAX_SAFE_INTEGER);
+        const storageDir = makeTempDir();
+        const directory = "/repo-ws-rehello";
+        const server = new MagicContextRpcServer(storageDir, directory);
+        const port = await server.start();
+        const record = readNewestPortRecord(storageDir, directory);
+        expect(typeof record?.token).toBe("string");
+
+        const ws = await openSocket(port, record?.token ?? "");
+        const notifications: unknown[] = [];
+        ws.addEventListener("message", (event) => {
+            const message = JSON.parse(String(event.data)) as {
+                type?: string;
+                notification?: unknown;
+            };
+            if (message.type === "notification") notifications.push(message.notification);
+        });
+
+        try {
+            ws.send(JSON.stringify({ type: "hello", token: record?.token, sessionId: "ses_A" }));
+            await waitForJsonMessage(ws, (message) => message.type === "hello-ack");
+            expect(isTuiConnected("ses_A")).toBe(true);
+
+            ws.send(JSON.stringify({ type: "hello", token: record?.token, sessionId: "ses_B" }));
+            await waitForJsonMessage(ws, (message) => message.type === "hello-ack");
+            expect(isTuiConnected("ses_A")).toBe(false);
+            expect(isTuiConnected("ses_B")).toBe(true);
+
+            ws.send(JSON.stringify({ type: "hello", token: record?.token, sessionId: "ses_B" }));
+            await waitForJsonMessage(ws, (message) => message.type === "hello-ack");
+            pushNotification("live", { ok: true }, "ses_B");
+            await waitFor(() => notifications.length >= 1, "one live notification");
+            await new Promise((resolve) => setTimeout(resolve, 50));
+            expect(notifications).toHaveLength(1);
+
+            ws.close();
+            await waitFor(() => !isTuiConnected(), "socket sink cleanup");
+        } finally {
+            try {
+                ws.close();
+            } catch {
+                // best-effort
+            }
+            server.stop();
+        }
+    });
+
+    test("same-process servers keep distinct port files during overlap", async () => {
+        const storageDir = makeTempDir();
+        const directory = "/repo-port-collision";
+        const first = new MagicContextRpcServer(storageDir, directory);
+        const second = new MagicContextRpcServer(storageDir, directory);
+        await first.start();
+        const secondPort = await second.start();
+
+        try {
+            const files = readdirSync(rpcPortDir(storageDir, directory)).filter(
+                (entry) => entry.startsWith("port-") && entry.endsWith(".json"),
+            );
+            expect(files.length).toBeGreaterThanOrEqual(2);
+
+            first.stop();
+            const remaining = readNewestPortRecord(storageDir, directory);
+            expect(remaining?.port).toBe(secondPort);
+
+            const client = new MagicContextRpcClient(storageDir, directory);
+            expect((await client.resolveEndpoint())?.port).toBe(secondPort);
+        } finally {
+            first.stop();
+            second.stop();
+        }
+    });
+
     test("gives up when the port file points at a dead server", async () => {
         const storageDir = makeTempDir();
         const directory = "/repo";

packages/plugin/src/shared/rpc-client.test.ts

@@ -107,38 +107,90 @@ export function pushNotification(
     }
 }
 
-/** Return pending notifications after acking the client's last received id.
+export interface DrainNotificationsOptions {
+    /**
+     * Cursor for global notifications when a session-scoped client sends separate
+     * session and global watermarks.
+     */
+    globalLastReceivedId?: number;
+    /** Ack/drain only the named session, not global notifications. */
+    sessionOnly?: boolean;
+    /** Ack/drain only session-less global notifications. */
+    globalOnly?: boolean;
+}
+
+function cursor(value: number | undefined): number {
+    return typeof value === "number" && Number.isFinite(value) && value > 0 ? value : 0;
+}
+
+/** Return pending notifications after pruning only the scopes the client acked.
  *
- *  Session scoping: when `sessionId` is provided, only notifications tagged for
- *  that session (or session-less/global ones) are returned and pruned — a
- *  notification tagged for a DIFFERENT session is never handed to this client
- *  and is never pruned by this client's ack. This matters because the in-memory
- *  queue is per-process but a TUI can end up bound to a process that also serves
- *  OTHER sessions: e.g. opening OpenCode Desktop on the same project starts a
- *  newer RPC server that the TUI's port discovery (newest-pid-wins) then selects,
- *  so a Desktop-session upgrade-dialog action would otherwise surface in an
- *  unrelated TUI session. Each client also tracks its own `lastReceivedId`, so a
- *  global watermark prune would let session A's ack drop session B's still-unseen
- *  notification — scoping the prune to the acking session prevents that too.
+ *  Session-scoped and global notifications have independent cursors. A TUI can
+ *  switch from session A to session B after handling a high id in A; that high
+ *  watermark must never prune B's lower, still-unseen ids. Global notifications
+ *  are also tracked separately so a global dialog does not become a session
+ *  watermark. Legacy callers that omit options keep the original single-cursor
+ *  behavior.
  *
  *  Delivery is at-least-once (non-destructive return + prune-on-ack): a returned
- *  notification stays queued until a later call acks it via a higher
- *  `lastReceivedId`, so a dropped WS socket re-delivers the backlog on reconnect
- *  (the client sends its `lastReceivedId` in the hello). */
-export function drainNotifications(lastReceivedId = 0, sessionId?: string): RpcNotification[] {
+ *  notification stays queued until a later call acks it via the matching scope's
+ *  cursor, so a dropped WS socket re-delivers unhandled backlog on reconnect. */
+export function drainNotifications(
+    lastReceivedId = 0,
+    sessionId?: string,
+    options: DrainNotificationsOptions = {},
+): RpcNotification[] {
+    const sessionCursor = cursor(lastReceivedId);
+
+    if (options.globalOnly) {
+        queue = queue.filter(
+            (notification) =>
+                notification.sessionId !== undefined || notification.id > sessionCursor,
+        );
+        return queue.filter(
+            (notification) =>
+                notification.sessionId === undefined && notification.id > sessionCursor,
+        );
+    }
+
+    if (options.sessionOnly) {
+        if (sessionId === undefined) return [];
+        queue = queue.filter(
+            (notification) =>
+                notification.sessionId !== sessionId || notification.id > sessionCursor,
+        );
+        return queue.filter(
+            (notification) =>
+                notification.sessionId === sessionId && notification.id > sessionCursor,
+        );
+    }
+
+    if (sessionId !== undefined && options.globalLastReceivedId !== undefined) {
+        const globalCursor = cursor(options.globalLastReceivedId);
+        queue = queue.filter((notification) => {
+            if (notification.sessionId === undefined) return notification.id > globalCursor;
+            if (notification.sessionId === sessionId) return notification.id > sessionCursor;
+            return true;
+        });
+        return queue.filter((notification) => {
+            if (notification.sessionId === undefined) return notification.id > globalCursor;
+            return notification.sessionId === sessionId && notification.id > sessionCursor;
+        });
+    }
+
     const matchesClient = (notification: RpcNotification): boolean =>
         sessionId === undefined ||
         notification.sessionId === undefined ||
         notification.sessionId === sessionId;
-    if (lastReceivedId > 0) {
-        // Prune only notifications THIS client both owns (session-matched) and has
-        // acked (id <= lastReceivedId). Other sessions' notifications survive.
+    if (sessionCursor > 0) {
+        // Legacy single-cursor mode prunes the scopes this client can see. New WS
+        // clients pass dual cursors above so cross-session watermarks stay isolated.
         queue = queue.filter(
-            (notification) => !(notification.id <= lastReceivedId && matchesClient(notification)),
+            (notification) => !(notification.id <= sessionCursor && matchesClient(notification)),
         );
     }
     return queue.filter(
-        (notification) => notification.id > lastReceivedId && matchesClient(notification),
+        (notification) => notification.id > sessionCursor && matchesClient(notification),
     );
 }