@wolfx/opencode-magic-context 0.28.0 → 0.30.3

package/src/shared/rpc-notifications.test.ts

@@ -1,5 +1,11 @@
 import { describe, expect, test } from "bun:test";
-import { drainNotifications, isTuiConnected, pushNotification } from "./rpc-notifications";
+import {
+    drainNotifications,
+    isTuiConnected,
+    type NotificationSink,
+    pushNotification,
+    registerNotificationSink,
+} from "./rpc-notifications";
 
 describe("rpc notifications", () => {
     test("keeps messages queued until the client acks their id", () => {
@@ -45,17 +51,68 @@ describe("rpc notifications", () => {
         expect(poll.map((m) => m.type).sort()).toEqual(["x", "y"]);
     });
 
-    test("isTuiConnected is per-session: a TUI on session A does not mark session B connected", () => {
-        // A TUI draining for tuiA must not make tuiB's producers think a TUI is
-        // polling for tuiB (which would route tuiB's /ctx-status, upgrade
-        // reminder, etc. to the dialog path and lose them in the unrelated TUI).
-        // Use ids no other test drains so the per-session window is unambiguous.
-        drainNotifications(0, "ses_tuiA_only");
-        expect(isTuiConnected("ses_tuiA_only")).toBe(true);
-        expect(isTuiConnected("ses_tuiB_never_drained")).toBe(false);
-        // The session-less (global) query still reports recent activity for the
-        // legacy callers that have no session context.
+    test("isTuiConnected reflects live WS sinks per-session", () => {
+        // No sinks → nothing connected.
+        expect(isTuiConnected("ses_anything")).toBe(false);
+        expect(isTuiConnected()).toBe(false);
+
+        // A live sink scoped to session A marks ONLY A connected (so B's producers
+        // don't misroute B's /ctx-status / upgrade reminder to the dialog path and
+        // lose it in an unrelated TUI), and the global query is also "connected".
+        const unregister = registerNotificationSink({ sessionId: "ses_A", send: () => {} });
+        expect(isTuiConnected("ses_A")).toBe(true);
+        expect(isTuiConnected("ses_B")).toBe(false);
         expect(isTuiConnected()).toBe(true);
+
+        // Closing the socket removes the sink → disconnected again.
+        unregister();
+        expect(isTuiConnected("ses_A")).toBe(false);
+        expect(isTuiConnected()).toBe(false);
+    });
+
+    test("a session-less sink counts as connected for any session query", () => {
+        const unregister = registerNotificationSink({ sessionId: undefined, send: () => {} });
+        expect(isTuiConnected("ses_whatever")).toBe(true);
+        expect(isTuiConnected()).toBe(true);
+        unregister();
+    });
+
+    test("pushNotification fans out live to a matching sink and skips a foreign session", () => {
+        drainNotifications(Number.MAX_SAFE_INTEGER);
+        const received: string[] = [];
+        const sink: NotificationSink = {
+            sessionId: "ses_live",
+            send: (n) => received.push(n.type),
+        };
+        const unregister = registerNotificationSink(sink);
+
+        pushNotification("for-live", { action: "show-status-dialog" }, "ses_live");
+        pushNotification("for-other", { action: "show-status-dialog" }, "ses_other");
+        pushNotification("global", { action: "show-status-dialog" });
+
+        // The sink sees its own session + global, never the foreign session.
+        expect(received.sort()).toEqual(["for-live", "global"]);
+        unregister();
+    });
+
+    test("a dead sink (throwing send) does not block delivery to other sinks", () => {
+        drainNotifications(Number.MAX_SAFE_INTEGER);
+        const live: string[] = [];
+        const unregDead = registerNotificationSink({
+            sessionId: undefined,
+            send: () => {
+                throw new Error("socket dead");
+            },
+        });
+        const unregLive = registerNotificationSink({
+            sessionId: undefined,
+            send: (n) => live.push(n.type),
+        });
+        // Must not throw, and the live sink still receives it.
+        expect(() => pushNotification("resilient", { ok: true })).not.toThrow();
+        expect(live).toEqual(["resilient"]);
+        unregDead();
+        unregLive();
     });
 
     test("queue-cap eviction is session-fair: a noisy session cannot evict another session's newest unseen item", () => {

package/src/shared/rpc-notifications.ts

@@ -16,31 +16,73 @@ export interface RpcNotification {
 
 let queue: RpcNotification[] = [];
 let nextNotificationId = 1;
-// Timestamp of last drain — used to detect if a TUI is actively polling.
-// The TUI polls every 500ms; we consider it connected if it polled within
-// the last 3 seconds (6× the poll interval, tolerates transient delays).
-//
-// PER-SESSION: a single server process can serve MANY sessions (e.g. a TUI on
-// session A plus an OpenCode Desktop opened on session B for the same project,
-// whose newer RPC server this TUI's port discovery then selects). The TUI
-// poller drains with ITS active session id, so a session is "TUI-connected"
-// only if a TUI recently drained FOR THAT session. A process-global timestamp
-// would make session B's producers (`/ctx-status`, upgrade reminder) take the
-// TUI-dialog path because session A's TUI is polling — queuing a B-scoped
-// dialog action that A's poller correctly refuses to show, so B's notice is
-// lost (it also suppressed B's non-TUI fallback). Tracking drains per session
-// routes each producer to the right delivery path.
-const lastDrainAtBySession = new Map<string, number>();
-let lastDrainAtAny = 0;
-const TUI_CONNECTED_WINDOW_MS = 3_000;
 
-/** Push a notification for TUI to pick up via polling. */
+/**
+ * A connected TUI notification sink — one per authenticated WebSocket. The RPC
+ * server registers a sink when a TUI socket authenticates (hello) and removes
+ * it on close. `send` is sink-agnostic (the server owns the actual WS socket)
+ * so this module stays free of Bun/WS types.
+ */
+export interface NotificationSink {
+    /** The TUI's active session at connect time (its hello scope). */
+    sessionId?: string;
+    /** Deliver one notification over this sink's live socket. */
+    send: (notification: RpcNotification) => void;
+}
+
+// Live sinks replace the old poll-drain-timestamp inference. "TUI connected for
+// a session" is now exact socket liveness — accurate and immediate — instead of
+// "did a 500ms poll drain within the last 3s". Per-session scoping still matters:
+// one process can serve MANY sessions (a TUI on session A plus an OpenCode
+// Desktop opened on session B for the same project, whose newer RPC server this
+// TUI's port discovery then selects). Each sink carries ITS session, so a
+// B-scoped producer (`/ctx-status`, upgrade reminder) only sees B's TUI as
+// connected and routes its dialog there, never to A.
+const sinks = new Set<NotificationSink>();
+
+/** Register a live TUI sink. Returns an unregister fn (call on socket close). */
+export function registerNotificationSink(sink: NotificationSink): () => void {
+    sinks.add(sink);
+    return () => {
+        sinks.delete(sink);
+    };
+}
+
+/** Whether a given notification may be delivered to a given sink. A global
+ *  notification (no sessionId) reaches every sink; a session-scoped one reaches
+ *  only sinks for that session (or session-less sinks). Mirrors the drain filter
+ *  from the sink's perspective. */
+function notificationMatchesSink(notification: RpcNotification, sink: NotificationSink): boolean {
+    return (
+        notification.sessionId === undefined ||
+        sink.sessionId === undefined ||
+        notification.sessionId === sink.sessionId
+    );
+}
+
+/** Push a notification to the TUI. Fans out to any live WS sink immediately and
+ *  also enqueues it so a TUI that is momentarily disconnected (reconnecting, or
+ *  not yet connected) still receives it on its next hello via the backlog drain.
+ *  At-least-once: a live push that the socket drops is re-delivered from the
+ *  queue on reconnect (pruned only when the client acks via `lastReceivedId`). */
 export function pushNotification(
     type: string,
     payload: Record<string, unknown>,
     sessionId?: string,
 ): void {
-    queue.push({ id: nextNotificationId++, type, payload, sessionId });
+    const notification: RpcNotification = { id: nextNotificationId++, type, payload, sessionId };
+    queue.push(notification);
+    // Fan out to every live sink this notification is scoped to. A delivery throw
+    // (dead socket mid-send) must not block other sinks or the caller.
+    for (const sink of sinks) {
+        if (!notificationMatchesSink(notification, sink)) continue;
+        try {
+            sink.send(notification);
+        } catch {
+            // Socket died between liveness check and send; the close handler will
+            // unregister it, and the queue backlog re-delivers on reconnect.
+        }
+    }
     // Cap queue size to prevent unbounded growth if a TUI is not draining.
     // Session-FAIR eviction: a naive `slice(-50)` drops the globally-oldest
     // items, so a noisy session could evict ANOTHER session's single unseen
@@ -66,13 +108,12 @@ export function pushNotification(
 }
 
 /** Return pending notifications after acking the client's last received id.
- *  Updates lastDrainAt so isTuiConnected() reflects recent activity.
  *
  *  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 draining a process that also serves
+ *  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
@@ -82,11 +123,9 @@ export function pushNotification(
  *
  *  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 lost poll response re-delivers on the next poll. */
+ *  `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[] {
-    const now = Date.now();
-    lastDrainAtAny = now;
-    if (sessionId !== undefined) lastDrainAtBySession.set(sessionId, now);
     const matchesClient = (notification: RpcNotification): boolean =>
         sessionId === undefined ||
         notification.sessionId === undefined ||
@@ -103,20 +142,20 @@ export function drainNotifications(lastReceivedId = 0, sessionId?: string): RpcN
     );
 }
 
-/** Whether a TUI client is actively polling for notifications.
- *  Returns true only if a TUI has drained within the last 3 seconds.
+/** Whether a TUI client is connected via a live notification socket.
+ *  Now exact socket liveness (a registered WS sink), not a poll-drain timestamp.
  *
- *  Pass `sessionId` (preferred) to ask whether a TUI is polling FOR THAT
+ *  Pass `sessionId` (preferred) to ask whether a TUI is connected FOR THAT
  *  SESSION — this is what producers (`/ctx-status`, `/ctx-recomp`, the upgrade
  *  reminder) must use to decide dialog-vs-message, so a TUI on a different
- *  session in the same process does not misroute their delivery. Omit it only
- *  for legacy/global callers that genuinely have no session context; they fall
- *  back to "any session recently drained" (the pre-per-session behavior). */
+ *  session in the same process does not misroute their delivery. A session-less
+ *  sink (legacy/global) counts for any session query. Omit `sessionId` only for
+ *  callers with no session context; they get "any sink connected". */
 export function isTuiConnected(sessionId?: string): boolean {
-    const now = Date.now();
-    if (sessionId !== undefined) {
-        const at = lastDrainAtBySession.get(sessionId) ?? 0;
-        return at > 0 && now - at < TUI_CONNECTED_WINDOW_MS;
+    if (sinks.size === 0) return false;
+    if (sessionId === undefined) return true;
+    for (const sink of sinks) {
+        if (sink.sessionId === undefined || sink.sessionId === sessionId) return true;
     }
-    return lastDrainAtAny > 0 && now - lastDrainAtAny < TUI_CONNECTED_WINDOW_MS;
+    return false;
 }