@wolfx/opencode-magic-context 0.28.0 → 0.30.3

package/src/tui/badge-contrast.ts

@@ -1,39 +1,52 @@
 /**
- * Pick a readable text color (black or white) for text drawn ON TOP of a given
- * background color.
+ * Pick the text color for the sidebar header badge (a bold label drawn on a
+ * `theme.accent` background).
  *
- * The sidebar header badge previously drew its label with `fg={theme.background}`
- * on a `theme.accent` background. That breaks for themes that set
- * `background: "none"` (transparent) to respect terminal transparency: the
- * resolved background is `RGBA(0,0,0,0)`, so the badge text renders fully
- * transparent and disappears (issue #186). The badge background (`accent`) is
- * always opaque, so deriving the text color from it is transparency-proof.
+ * Primary rule (matches AFT's sidebar by construction): paint the theme's own
+ * `background` color as the label, the inverse-of-panel look. Because it is a
+ * fixed theme token rather than an accent-derived computation, MC's badge and
+ * AFT's badge agree on EVERY accent automatically, so the same theme can never
+ * make one badge black and the other white (issue #198).
  *
- * The pick is WHITE-BIASED off the accent's relative luminance: white for any
- * accent in the dark half (luminance < 0.5), black only for genuinely light
- * accents. A strict "higher-contrast-wins" pick (crossover at luminance 0.179)
- * flips ordinary mid-tone accents to black: a typical orange/amber sidebar
- * accent (luminance ~0.3) reads black ~5:1 vs white ~3.7:1, so contrast-wins
- * picks black even though white at ~3.7:1 is perfectly legible for a short bold
- * label. That looks heavy and clashes with the sibling status badges, so we
- * prefer white across the whole dark half and only fall to black once the accent
- * is actually light (pale/pastel/near-white), where white would be unreadable.
+ * Fallback rule (the reason a luminance pick exists at all): `theme.background`
+ * can be unusable as a label color in two degenerate cases, where it would
+ * render the label invisible on the accent:
+ *   1. Transparent background. Themes that respect terminal transparency set
+ *      `background: "none"`, which resolves to `RGBA(0,0,0,0)`; drawing it as
+ *      text renders fully transparent and the label disappears (issue #186).
+ *   2. Background ~= accent. If the theme's background and accent are nearly the
+ *      same color, background-on-accent text has no contrast.
+ * In either case we fall back to a black/white pick that is guaranteed visible
+ * on the always-opaque accent.
  *
- * `RGBA` channels from @opentui/core are normalized 0..1 floats. We accept the
- * minimal `{ r, g, b }` shape so this stays a pure, trivially testable function
- * independent of the native color class.
+ * `RGBA` channels from @opentui/core are normalized 0..1 floats (alpha included).
+ * We accept the minimal `{ r, g, b, a? }` shape so this stays a pure, trivially
+ * testable function independent of the native color class, and we return the
+ * passed-in `background` object unchanged on the primary path so it stays the
+ * exact same theme token AFT uses.
  */
 
-// Luminance midpoint: accents below this keep white text, accents at/above it
-// (light/pastel/near-white) get black. White-biased relative to the strict
-// equal-contrast crossover (~0.179) so saturated mid-tone accents stay white.
+type Color = { r: number; g: number; b: number; a?: number };
+
+// Below this alpha the theme background is too transparent to read as a label on
+// the accent (issue #186: background:"none" resolves to alpha 0).
+const MIN_OPAQUE_ALPHA = 0.5;
+
+// If the theme background and accent are within this per-channel distance they
+// are effectively the same color, so background-on-accent text is unreadable.
+const MIN_CHANNEL_DISTANCE = 0.06;
+
+// Luminance midpoint for the fallback pick: accents below this keep white text,
+// accents at/above it (light/pastel/near-white) get black. White-biased relative
+// to the strict equal-contrast crossover (~0.179) so saturated mid-tone accents
+// stay white. Only consulted on the degenerate fallback path.
 const LIGHT_ACCENT_LUMINANCE = 0.5;
 
 function srgbChannelToLinear(c: number): number {
     return c <= 0.03928 ? c / 12.92 : ((c + 0.055) / 1.055) ** 2.4;
 }
 
-function relativeLuminance(bg: { r: number; g: number; b: number }): number {
+function relativeLuminance(bg: Color): number {
     return (
         0.2126 * srgbChannelToLinear(bg.r) +
         0.7152 * srgbChannelToLinear(bg.g) +
@@ -41,6 +54,31 @@ function relativeLuminance(bg: { r: number; g: number; b: number }): number {
     );
 }
 
-export function readableTextColorOn(bg: { r: number; g: number; b: number }): string {
+function nearlyEqual(a: Color, b: Color): boolean {
+    return (
+        Math.abs(a.r - b.r) < MIN_CHANNEL_DISTANCE &&
+        Math.abs(a.g - b.g) < MIN_CHANNEL_DISTANCE &&
+        Math.abs(a.b - b.b) < MIN_CHANNEL_DISTANCE
+    );
+}
+
+/**
+ * Pure black/white pick by accent luminance. Used as the badge fallback and kept
+ * exported for callers that only have the accent.
+ */
+export function readableTextColorOn(bg: Color): string {
     return relativeLuminance(bg) < LIGHT_ACCENT_LUMINANCE ? "#ffffff" : "#000000";
 }
+
+/**
+ * Badge label color on the accent: the theme background (AFT parity) when it is
+ * usable, else a guaranteed-visible black/white fallback. Returns the passed-in
+ * `background` reference unchanged on the primary path.
+ */
+export function badgeTextColor<T extends Color>(accent: T, background: T): T | string {
+    const alpha = background.a ?? 1;
+    if (alpha >= MIN_OPAQUE_ALPHA && !nearlyEqual(accent, background)) {
+        return background;
+    }
+    return readableTextColorOn(accent);
+}

package/src/tui/data/context-db.ts

@@ -5,18 +5,12 @@
 import os from "node:os";
 import path from "node:path";
 import { MagicContextRpcClient } from "../../shared/rpc-client";
-import type {
-    EmbedDetail,
-    RpcNotificationMessage,
-    SidebarSnapshot,
-    StatusDetail,
-} from "../../shared/rpc-types";
+import type { EmbedDetail, SidebarSnapshot, StatusDetail } from "../../shared/rpc-types";
 
 export type { EmbedDetail, SidebarSnapshot, StatusDetail };
 
 let rpcClient: MagicContextRpcClient | null = null;
 let rpcGeneration = 0;
-const lastReceivedNotificationIdBySession = new Map<string, number>();
 
 function getStorageDir(): string {
     // Plugin v0.16+ uses the shared cortexkit/magic-context path so OpenCode
@@ -31,9 +25,9 @@ function getStorageDir(): string {
 export function initRpcClient(directory: string): void {
     const storageDir = getStorageDir();
     // Bump the generation before replacing the client so late notification
-    // responses from a disposed client cannot repopulate cleared cursors.
+    // responses from a disposed client are ignored (the WS socket observes the
+    // new generation and abandons its in-flight connect).
     rpcGeneration += 1;
-    lastReceivedNotificationIdBySession.clear();
     rpcClient = new MagicContextRpcClient(storageDir, directory);
 }
 
@@ -41,14 +35,19 @@ export function getRpcGeneration(): number {
     return rpcGeneration;
 }
 
+/** The live RPC client (for the WS notification socket's endpoint discovery).
+ *  Null before init / after close. */
+export function getRpcClient(): MagicContextRpcClient | null {
+    return rpcClient;
+}
+
 /** Clean up the RPC client. */
 export function closeRpc(): void {
     // Closing invalidates any already-issued RPC calls; their callbacks must
-    // observe the new generation and avoid advancing stale notification cursors.
+    // observe the new generation and abandon (the WS socket checks it too).
     rpcGeneration += 1;
     rpcClient?.reset();
     rpcClient = null;
-    lastReceivedNotificationIdBySession.clear();
 }
 
 const EMPTY_SNAPSHOT: SidebarSnapshot = {
@@ -311,13 +310,6 @@ export async function loadToastDurationMs(): Promise<number> {
     }
 }
 
-export interface TuiMessage {
-    id: number;
-    type: string;
-    payload: Record<string, unknown>;
-    sessionId?: string;
-}
-
 /**
  * Fetch the current startup announcement from the server, if any.
  * Returns `{show: false}` when there's nothing to announce or when the
@@ -360,49 +352,3 @@ export async function markAnnounced(): Promise<boolean> {
         return false;
     }
 }
-
-/** Poll for pending server→TUI notifications via RPC. */
-export async function consumeTuiMessages(sessionId: string): Promise<TuiMessage[]> {
-    if (!rpcClient) return [];
-    try {
-        const result = await rpcClient.call<{ messages: RpcNotificationMessage[] }>(
-            "pending-notifications",
-            // Pass the TUI's active session so the server only drains
-            // notifications scoped to it (or global ones). Without this, a
-            // notification for another session served by the same process (e.g.
-            // OpenCode Desktop on the same project) could surface here. The
-            // cursor is per-session and is advanced by the poller only after it
-            // has delivered the returned batch.
-            {
-                lastReceivedId: lastReceivedNotificationIdBySession.get(sessionId) ?? 0,
-                sessionId,
-            },
-        );
-        return (result.messages ?? []).map((m) => ({
-            id: m.id,
-            type: m.type,
-            payload: m.payload,
-            sessionId: m.sessionId,
-        }));
-    } catch {
-        return [];
-    }
-}
-
-/**
- * Advance the delivered-message cursor for one active TUI session.
- * Callers must pass only the contiguous handled prefix of the drained batch;
- * this helper remains empty-safe and monotonic for that prefix.
- */
-export function markTuiMessagesHandled(sessionId: string, messages: TuiMessage[]): void {
-    const previous = lastReceivedNotificationIdBySession.get(sessionId) ?? 0;
-    let next = previous;
-    for (const message of messages) {
-        if (message.id > next) {
-            next = message.id;
-        }
-    }
-    if (next > previous) {
-        lastReceivedNotificationIdBySession.set(sessionId, next);
-    }
-}

package/src/tui/data/notification-socket.ts

@@ -0,0 +1,229 @@
+/**
+ * Persistent WebSocket to the server plugin's RPC server, replacing the old
+ * 500ms HTTP notification poll.
+ *
+ * Why this exists: the TUI plugin and the server plugin run in separate Bun
+ * runners in the same process, so they bridge over a localhost socket. The old
+ * bridge polled `pending-notifications` over HTTP every 500ms — and each poll
+ * opened a NEW loopback TCP connection (Bun's fetch isn't pooled to our server),
+ * which was the entire source of idle TUI CPU (#200). A single long-lived WS
+ * carries server→TUI pushes with zero per-event connection cost, and the server
+ * pushes notifications the instant they're queued (no polling latency).
+ *
+ * Session scope: the socket carries the TUI's active session in its `hello` so
+ * the server delivers only that session's (plus global) notifications and its
+ * `isTuiConnected(session)` routing stays correct. The active session is tracked
+ * with a cheap watcher that only reads `api.route.current` (a property access,
+ * no IPC) and re-scopes the socket ONLY when the session actually changes — so
+ * unlike the old poll it does no network work at idle.
+ */
+
+import { getRpcClient, getRpcGeneration } from "./context-db";
+
+export interface SocketNotification {
+    id: number;
+    type: string;
+    payload: Record<string, unknown>;
+    sessionId?: string;
+}
+
+interface NotificationSocketOptions {
+    /** Current active session id (re-read cheaply to follow session switches). */
+    getSessionId: () => string | null;
+    /** Handle one delivered notification. Returns true if it was consumed (so its
+     *  id can advance the ack cursor). Async because dialog handlers await. */
+    onNotification: (notification: SocketNotification) => boolean | Promise<boolean>;
+}
+
+const RECONNECT_BASE_MS = 500;
+const RECONNECT_MAX_MS = 10_000;
+/** Cheap session-watch interval. Reads a property only; no network. The CPU bug
+ *  was the per-tick fetch, not the timer — this tick does zero IPC at idle. */
+const SESSION_WATCH_MS = 1_000;
+
+let socket: WebSocket | null = null;
+let reconnectTimer: ReturnType<typeof setTimeout> | undefined;
+let sessionWatchTimer: ReturnType<typeof setInterval> | undefined;
+let reconnectAttempt = 0;
+let closed = false;
+let helloedSession: string | null = null;
+let opts: NotificationSocketOptions | null = null;
+/** Generation of the rpc client at connect time; a dispose/reinit bumps it and
+ *  invalidates an in-flight socket so its late callbacks are ignored. */
+let connectGeneration = 0;
+/** Highest handled notification id (sent in hello so a reconnect re-delivers the
+ *  unhandled backlog; sent in ack so the server prunes the queue). */
+let lastHandledId = 0;
+
+/** Open the persistent notification socket. Idempotent: a second call while open
+ *  is a no-op. Reconnects on its own after any drop. */
+export function startNotificationSocket(options: NotificationSocketOptions): void {
+    opts = options;
+    closed = false;
+    connectGeneration = getRpcGeneration();
+    connect();
+    if (!sessionWatchTimer) {
+        sessionWatchTimer = setInterval(watchSession, SESSION_WATCH_MS);
+    }
+}
+
+/** Close the socket and stop reconnecting. Call on TUI dispose. */
+export function stopNotificationSocket(): void {
+    closed = true;
+    if (reconnectTimer) {
+        clearTimeout(reconnectTimer);
+        reconnectTimer = undefined;
+    }
+    if (sessionWatchTimer) {
+        clearInterval(sessionWatchTimer);
+        sessionWatchTimer = undefined;
+    }
+    try {
+        socket?.close();
+    } catch {
+        // best-effort
+    }
+    socket = null;
+    helloedSession = null;
+    reconnectAttempt = 0;
+}
+
+function scheduleReconnect(): void {
+    if (closed) return;
+    if (reconnectTimer) return;
+    const delay = Math.min(RECONNECT_BASE_MS * 2 ** reconnectAttempt, RECONNECT_MAX_MS);
+    reconnectAttempt += 1;
+    reconnectTimer = setTimeout(() => {
+        reconnectTimer = undefined;
+        connect();
+    }, delay);
+}
+
+async function connect(): Promise<void> {
+    if (closed) return;
+    if (socket) return; // already connected/connecting
+
+    const client = getRpcClient();
+    if (!client) {
+        scheduleReconnect();
+        return;
+    }
+    const endpoint = await client.resolveEndpoint();
+    // The generation may have bumped (dispose/reinit) while resolving — abandon.
+    if (closed || getRpcGeneration() !== connectGeneration) return;
+    if (!endpoint) {
+        scheduleReconnect();
+        return;
+    }
+
+    let ws: WebSocket;
+    try {
+        ws = new WebSocket(`ws://127.0.0.1:${endpoint.port}/ws`);
+    } catch {
+        scheduleReconnect();
+        return;
+    }
+    socket = ws;
+
+    ws.addEventListener("open", () => {
+        if (socket !== ws) return;
+        reconnectAttempt = 0;
+        sendHello(ws, endpoint.token);
+    });
+
+    ws.addEventListener("message", (event) => {
+        if (socket !== ws) return;
+        void handleSocketMessage(ws, String((event as MessageEvent).data));
+    });
+
+    const onDown = () => {
+        if (socket === ws) {
+            socket = null;
+            helloedSession = null;
+        }
+        scheduleReconnect();
+    };
+    ws.addEventListener("close", onDown);
+    ws.addEventListener("error", onDown);
+}
+
+function sendHello(ws: WebSocket, token: string | null): void {
+    const sessionId = opts?.getSessionId() ?? undefined;
+    helloedSession = sessionId ?? null;
+    ws.send(
+        JSON.stringify({
+            type: "hello",
+            token: token ?? "",
+            sessionId,
+            lastReceivedId: lastHandledId,
+        }),
+    );
+}
+
+async function handleSocketMessage(ws: WebSocket, raw: string): Promise<void> {
+    let msg: { type?: string; notification?: SocketNotification; error?: string };
+    try {
+        msg = JSON.parse(raw);
+    } catch {
+        return;
+    }
+
+    if (msg.type === "notification" && msg.notification) {
+        const notification = msg.notification;
+        // Client-side session filter mirrors the old poller's per-message re-check:
+        // a session-scoped notification is only acted on while the TUI is actually
+        // viewing that session (the active session can change between queueing and
+        // delivery). Global (session-less) notifications always apply.
+        const active = opts?.getSessionId() ?? null;
+        if (notification.sessionId && active && notification.sessionId !== active) {
+            // Not for the session we're viewing — do NOT ack it (a TUI on the right
+            // session, or a later switch back, should still get it). Just skip.
+            return;
+        }
+        let consumed = false;
+        try {
+            consumed = await Promise.resolve(opts?.onNotification(notification) ?? false);
+        } catch {
+            consumed = false;
+        }
+        // A dispose/reinit during an awaited dialog handler invalidates this socket.
+        if (socket !== ws || getRpcGeneration() !== connectGeneration) return;
+        if (consumed && notification.id > lastHandledId) {
+            lastHandledId = notification.id;
+            // Ack so the server prunes the queue during this long-lived connection.
+            // (The id also rides the next reconnect hello via lastReceivedId, so a
+            // dropped ack just re-delivers an already-consumed item, which the
+            // handlers are idempotent against.)
+            try {
+                ws.send(JSON.stringify({ type: "ack", lastReceivedId: lastHandledId }));
+            } catch {
+                // best-effort; reconnect hello re-syncs via lastReceivedId
+            }
+        }
+        return;
+    }
+
+    if (msg.type === "error") {
+        // Server rejected us (bad token, etc.). Close and let backoff retry after
+        // rediscovering the port/token (the server may have been replaced).
+        try {
+            ws.close();
+        } catch {
+            // best-effort
+        }
+    }
+}
+
+/** Cheap session-change watcher: re-scope the socket only when the active session
+ *  actually changes. Reads a property; no network at idle. */
+function watchSession(): void {
+    if (closed || !socket || socket.readyState !== WebSocket.OPEN) return;
+    const current = opts?.getSessionId() ?? null;
+    if (current === helloedSession) return;
+    // Re-hello with the new session; the server replaces this socket's sink scope.
+    const client = getRpcClient();
+    void client?.resolveEndpoint().then((endpoint) => {
+        if (!socket || socket.readyState !== WebSocket.OPEN) return;
+        sendHello(socket, endpoint?.token ?? null);
+    });
+}

package/src/tui/index.tsx

@@ -6,7 +6,8 @@ import { createMemo } from "solid-js"
 import type { TuiPlugin, TuiPluginApi, TuiThemeCurrent } from "@opencode-ai/plugin/tui"
 import { createSidebarContentSlot, kickRecompProgressRefresh } from "./slots/sidebar-content"
 import packageJson from "../../package.json"
-import { closeRpc, consumeTuiMessages, dismissUpgradeReminder, getAnnouncement, getCompartmentCount, getRpcGeneration, initRpcClient, loadEmbedDetail, loadStatusDetail, loadToastDurationMs, markAnnounced, markTuiMessagesHandled, requestRecomp, requestUpgrade, type EmbedDetail, type TuiMessage, type StatusDetail } from "./data/context-db"
+import { closeRpc, dismissUpgradeReminder, getAnnouncement, getCompartmentCount, getRpcGeneration, initRpcClient, loadEmbedDetail, loadStatusDetail, loadToastDurationMs, markAnnounced, requestRecomp, requestUpgrade, type EmbedDetail, type StatusDetail } from "./data/context-db"
+import { startNotificationSocket, stopNotificationSocket, type SocketNotification } from "./data/notification-socket"
 import { formatThresholdPercent } from "../shared/format-threshold"
 import { detectConflicts } from "../shared/conflict-detector"
 import { fixConflicts } from "../shared/conflict-fixer"
@@ -869,128 +870,77 @@ const tui: TuiPlugin = async (api, _options, meta) => {
     // only when keymap is missing.
     registerCommandPaletteEntries(api)
 
-    // Poll for server→TUI messages: toasts and dialog requests.
-    // The poller owns cursor advancement so notifications are acked only after
-    // they are accepted for the still-active session and delivered to the UI.
-    let pollInFlight = false
-    const messagePoller = setInterval(() => {
-        // Scope the drain to the TUI's active session so notifications tagged for
-        // a different session (served by the same RPC process) are not consumed
-        // here. Do not poll on non-session routes: a session-scoped action fetched
-        // while sessionless could otherwise be acked without being shown.
-        // Avoid overlapping read-only drains: the server re-delivers until acked,
-        // so a second in-flight poll can fetch and dispatch the same batch twice.
-        if (pollInFlight) return
-
+    // Receive server→TUI notifications (toasts + dialog requests) over a single
+    // persistent WebSocket, pushed the instant the server queues them. This
+    // replaces the old 500ms HTTP poll whose new-connection-per-tick cost was the
+    // source of idle TUI CPU (#200). The socket carries the active session in its
+    // hello so the server scopes delivery; here we re-check the active session per
+    // notification (it can change between queue and delivery) before acting.
+    const handleNotification = async (n: SocketNotification): Promise<boolean> => {
         const requestedSessionId = getSessionId(api)
-        if (!requestedSessionId) return
+        const generation = getRpcGeneration()
+        // A session-scoped notification only applies while we're viewing that
+        // session; global (session-less) ones always apply. Returning false leaves
+        // it unacked so a TUI on the right session (or a later switch back) still
+        // gets it.
+        if (n.sessionId && requestedSessionId && n.sessionId !== requestedSessionId) {
+            return false
+        }
+        if (n.type === "toast") {
+            const p = n.payload
+            showToast(api, {
+                message: String(p.message ?? ""),
+                variant: (p.variant as "info" | "warning" | "error" | "success") ?? "info",
+                durationOverrideMs:
+                    typeof p.duration === "number" && Number.isFinite(p.duration)
+                        ? p.duration
+                        : undefined,
+            })
+            return true
+        }
+        if (n.type !== "action") return false
+        const action = n.payload?.action
+        const stillActive = () =>
+            getRpcGeneration() === generation && getSessionId(api) === requestedSessionId
+        if (action === "show-status-dialog") {
+            return stillActive() && (await showStatusDialog(api, requestedSessionId))
+        }
+        if (action === "show-recomp-dialog") {
+            return stillActive() && (await showRecompDialog(api, requestedSessionId))
+        }
+        if (action === "show-upgrade-dialog") {
+            const resume =
+                n.payload?.resume === true
+                    ? {
+                          stagedCount: Number(n.payload?.stagedCount ?? 0),
+                          stagedThrough: Number(n.payload?.stagedThrough ?? 0),
+                      }
+                    : undefined
+            return stillActive() && showUpgradeDialog(api, resume, requestedSessionId)
+        }
+        if (action === "show-embed-dialog") {
+            return stillActive() && (await showEmbedDialog(api, requestedSessionId))
+        }
+        if (action === "show-flush-dialog") {
+            const flushMsg = String(n.payload?.message ?? "Flushed.")
+            return stillActive() && showResultDialog(api, "Flush", flushMsg)
+        }
+        if (action === "show-result-dialog") {
+            const title = String(n.payload?.title ?? "Magic Context")
+            const body = String(n.payload?.message ?? "")
+            return stillActive() && showResultDialog(api, title, body)
+        }
+        return false
+    }
 
-        pollInFlight = true
-        const pollGeneration = getRpcGeneration()
-        void consumeTuiMessages(requestedSessionId).then(async (messages) => {
-            if (unifiedToastDurationMs === DEFAULT_TOAST_DURATION_MS) {
-                void refreshToastDurationMs()
-            }
-            // The dialog handlers read the current session when they run. If the
-            // user switched routes while the RPC was in flight, drop this whole
-            // batch without advancing the cursor; the next poll for the new
-            // session will fetch the right notifications.
-            // Ignore late responses from an older RPC client generation; close/init
-            // clears cursors and stale callbacks must not recreate them.
-            if (getRpcGeneration() !== pollGeneration) return
-
-            if (getSessionId(api) !== requestedSessionId) return
-
-            const orderedMessages = [...messages].sort((a, b) => a.id - b.id)
-            const handledMessageIds = new Set<number>()
-            for (const msg of orderedMessages) {
-                // A dialog helper earlier in this batch may have awaited; re-check
-                // the route before EACH message so a later action/toast in the same
-                // batch can't paint into a session the user switched to mid-batch
-                // (the pre-batch + pre-ack guards alone don't cover mid-batch awaits).
-                if (getRpcGeneration() !== pollGeneration) return
-                if (getSessionId(api) !== requestedSessionId) return
-                // Drop any action/dialog whose sessionId doesn't match this TUI's
-                // active session (session-less/global notifications still apply).
-                if (
-                    msg.type === "action" &&
-                    msg.sessionId &&
-                    msg.sessionId !== requestedSessionId
-                ) {
-                    continue
-                }
-                if (msg.type === "toast") {
-                    const p = msg.payload
-                    showToast(api, {
-                        message: String(p.message ?? ""),
-                        variant: (p.variant as "info" | "warning" | "error" | "success") ?? "info",
-                        durationOverrideMs:
-                            typeof p.duration === "number" && Number.isFinite(p.duration)
-                                ? p.duration
-                                : undefined,
-                    })
-                    handledMessageIds.add(msg.id)
-                } else if (msg.type === "action") {
-                    const action = msg.payload?.action
-                    if (action === "show-status-dialog") {
-                        if (await showStatusDialog(api, requestedSessionId)) {
-                            handledMessageIds.add(msg.id)
-                        }
-                    } else if (action === "show-recomp-dialog") {
-                        if (await showRecompDialog(api, requestedSessionId)) {
-                            handledMessageIds.add(msg.id)
-                        }
-                    } else if (action === "show-upgrade-dialog") {
-                        const resume =
-                            msg.payload?.resume === true
-                                ? {
-                                      stagedCount: Number(msg.payload?.stagedCount ?? 0),
-                                      stagedThrough: Number(msg.payload?.stagedThrough ?? 0),
-                                  }
-                                : undefined
-                        if (showUpgradeDialog(api, resume, requestedSessionId)) {
-                            handledMessageIds.add(msg.id)
-                        }
-                    } else if (action === "show-embed-dialog") {
-                        if (await showEmbedDialog(api, requestedSessionId)) {
-                            handledMessageIds.add(msg.id)
-                        }
-                    } else if (action === "show-flush-dialog") {
-                        const flushMsg = String(msg.payload?.message ?? "Flushed.")
-                        if (showResultDialog(api, "Flush", flushMsg)) {
-                            handledMessageIds.add(msg.id)
-                        }
-                    } else if (action === "show-result-dialog") {
-                        const title = String(msg.payload?.title ?? "Magic Context")
-                        const body = String(msg.payload?.message ?? "")
-                        if (showResultDialog(api, title, body)) {
-                            handledMessageIds.add(msg.id)
-                        }
-                    }
-                }
-            }
-            const handledPrefixMessages: TuiMessage[] = []
-            for (const msg of orderedMessages) {
-                if (!handledMessageIds.has(msg.id)) break
-                handledPrefixMessages.push(msg)
-            }
-            // A dialog helper may have awaited more RPC work; re-check before
-            // acking so a dispose/reinit or route switch during that await cannot
-            // advance a stale cursor.
-            if (getRpcGeneration() !== pollGeneration) return
-            if (getSessionId(api) !== requestedSessionId) return
-
-            markTuiMessagesHandled(requestedSessionId, handledPrefixMessages)
-        }).catch(() => {
-            // Intentional: message polling should never crash the TUI
-        }).finally(() => {
-            pollInFlight = false
-        })
-    }, 500)
+    startNotificationSocket({
+        getSessionId: () => getSessionId(api),
+        onNotification: handleNotification,
+    })
 
     // Clean up on dispose
     api.lifecycle.onDispose(() => {
-        clearInterval(messagePoller)
+        stopNotificationSocket()
         closeRpc()
     })