✨ Wake Penpot tabs from the MCP server through the extension

The extension's service worker holds one WebSocket per MCP server URL
and user token used by its sessions, on the endpoint the plugin uses.
Each connection registers as a wake agent with a stable agent ID, sends
the full session list whenever it changes, and sends a heartbeat every
20 seconds, which also keeps the worker running. It reconnects with
capped backoff and stops for good when the server closes it with code
1008. A connection whose last session goes is closed. An alarm every 30
seconds restarts the worker after Chrome stops it, so the connections
open again.

On a wake request, when waking is turned on and a tab holds the
session, the extension activates the tab and focuses its window, which
makes the browser resume a frozen page, and answers activated.
Otherwise it answers disabled, unknown-session, or tab-closed for a tab
that no longer exists, whose session it then forgets. Waking stays on
until the user turns it off; the setting is kept in
chrome.storage.local.

mcp/packages/browser-extension/src/SessionMap.test.ts

@@ -34,7 +34,8 @@ let sessions: SessionMap;
 beforeEach(() => {
     store = new MemoryStore();
     changes = [];
-    sessions = new SessionMap(store, { sessionsChanged: (change) => changes.push(change) });
+    sessions = new SessionMap(store);
+    sessions.observe({ sessionsChanged: (change) => changes.push(change) });
 });
 
 test("maps a connected session to the announcing tab and saves it", async () => {
@@ -89,14 +90,14 @@ test("forgets only the sessions of a removed tab", async () => {
 
 test("restores the saved sessions when the worker restarts", async () => {
     await sessions.announce(TAB, announcement("connected"));
-    const restarted = new SessionMap(store, { sessionsChanged: () => {} });
+    const restarted = new SessionMap(store);
     await restarted.ready();
     assert.equal(restarted.get("session")?.tabId, 1);
 });
 
 test("changes made while restoring apply to the restored sessions", async () => {
     const saved = new MemoryStore({ first: { ...TAB, serverUrl: "wss://penpot/mcp/ws", token: null } });
-    const restored = new SessionMap(saved, { sessionsChanged: () => {} });
+    const restored = new SessionMap(saved);
     await restored.announce(OTHER_TAB, announcement("connected", "second"));
     assert.deepEqual(Object.keys(saved.saved!).sort(), ["first", "second"]);
 });
@@ -108,7 +109,8 @@ test("the discard guard protects tabs holding a session and releases them after"
             calls.push([tabId, autoDiscardable]);
         },
     });
-    const guarded = new SessionMap(new MemoryStore(), guard);
+    const guarded = new SessionMap(new MemoryStore());
+    guarded.observe(guard);
     await guarded.announce(TAB, announcement("connected"));
     await guarded.removeTab(1);
     assert.deepEqual(calls, [

mcp/packages/browser-extension/src/SessionMap.ts

@@ -36,12 +36,10 @@ export interface SessionMapObserver {
  */
 export class SessionMap {
     private entries = new Map<string, SessionEntry>();
+    private readonly observers: SessionMapObserver[] = [];
     private readonly loaded: Promise<void>;
 
-    constructor(
-        private readonly store: SessionStore,
-        private readonly observer: SessionMapObserver
-    ) {
+    constructor(private readonly store: SessionStore) {
         this.loaded = store.load().then(
             (saved) => {
                 this.entries = new Map(Object.entries(saved ?? {}));
@@ -50,6 +48,11 @@ export class SessionMap {
         );
     }
 
+    /** Adds a receiver of the changes of this map. */
+    observe(observer: SessionMapObserver): void {
+        this.observers.push(observer);
+    }
+
     /** Resolves once the saved entries have been restored. */
     ready(): Promise<void> {
         return this.loaded;
@@ -105,10 +108,13 @@ export class SessionMap {
         }
         const after = this.tabs();
         await this.store.save(Object.fromEntries(this.entries));
-        this.observer.sessionsChanged({
+        const change = {
             tabsAdded: [...after].filter((tabId) => !before.has(tabId)),
             tabsRemoved: [...before].filter((tabId) => !after.has(tabId)),
-        });
+        };
+        for (const observer of this.observers) {
+            observer.sessionsChanged(change);
+        }
     }
 
     private deleteTab(tabId: number): void {

mcp/packages/browser-extension/src/TabWaker.test.ts

@@ -0,0 +1,63 @@
+import assert from "node:assert/strict";
+import { beforeEach, test } from "node:test";
+import { SessionMap } from "./SessionMap.ts";
+import { TabWaker, type TabControl, type WakeSwitch } from "./TabWaker.ts";
+
+let sessions: SessionMap;
+let activated: number[];
+let openTabs: Set<number>;
+let enabled: boolean;
+let waker: TabWaker;
+
+beforeEach(async () => {
+    sessions = new SessionMap({ load: async () => undefined, save: async () => {} });
+    activated = [];
+    openTabs = new Set([7]);
+    enabled = true;
+    const tabs: TabControl = {
+        activate: async (tabId) => {
+            if (!openTabs.has(tabId)) return false;
+            activated.push(tabId);
+            return true;
+        },
+    };
+    const wakeSwitch: WakeSwitch = {
+        isEnabled: async () => enabled,
+        setEnabled: async (value) => {
+            enabled = value;
+        },
+    };
+    waker = new TabWaker(sessions, tabs, wakeSwitch);
+    await sessions.announce(
+        { tabId: 7, windowId: 1 },
+        {
+            type: "penpot-mcp:session",
+            sessionId: "tab",
+            serverUrl: "wss://penpot/mcp/ws",
+            token: null,
+            status: "connected",
+        }
+    );
+});
+
+test("activates the tab holding the session", async () => {
+    assert.equal((await waker.wake("tab")).outcome, "activated");
+    assert.deepEqual(activated, [7]);
+});
+
+test("does nothing while the user has turned waking off", async () => {
+    enabled = false;
+    assert.equal((await waker.wake("tab")).outcome, "disabled");
+    assert.deepEqual(activated, []);
+});
+
+test("reports a session no tab holds", async () => {
+    assert.equal((await waker.wake("other")).outcome, "unknown-session");
+    assert.deepEqual(activated, []);
+});
+
+test("reports a closed tab and forgets its session", async () => {
+    openTabs.clear();
+    assert.equal((await waker.wake("tab")).outcome, "tab-closed");
+    assert.equal(sessions.get("tab"), undefined);
+});

mcp/packages/browser-extension/src/TabWaker.ts

@@ -0,0 +1,46 @@
+import type { SessionMap } from "./SessionMap.ts";
+import type { WakeHandler, WakeReport } from "./WakeAgentConnection.ts";
+
+/** Browser control over tab activation. */
+export interface TabControl {
+    /**
+     * Activates a tab and focuses its window.
+     *
+     * @returns Whether the tab still exists
+     */
+    activate(tabId: number): Promise<boolean>;
+}
+
+/** The user's switch that pauses waking without removing the extension. */
+export interface WakeSwitch {
+    isEnabled(): Promise<boolean>;
+    setEnabled(enabled: boolean): Promise<void>;
+}
+
+/**
+ * Wakes the tab holding a session by activating it and focusing its window, which makes the
+ * browser resume a frozen page.
+ */
+export class TabWaker implements WakeHandler {
+    constructor(
+        private readonly sessions: SessionMap,
+        private readonly tabs: TabControl,
+        private readonly wakeSwitch: WakeSwitch
+    ) {}
+
+    async wake(sessionId: string): Promise<WakeReport> {
+        if (!(await this.wakeSwitch.isEnabled())) {
+            return { outcome: "disabled" };
+        }
+        await this.sessions.ready();
+        const entry = this.sessions.get(sessionId);
+        if (!entry) {
+            return { outcome: "unknown-session" };
+        }
+        if (!(await this.tabs.activate(entry.tabId))) {
+            await this.sessions.removeTab(entry.tabId);
+            return { outcome: "tab-closed" };
+        }
+        return { outcome: "activated" };
+    }
+}

mcp/packages/browser-extension/src/WakeAgentConnection.test.ts

@@ -0,0 +1,238 @@
+import assert from "node:assert/strict";
+import { afterEach, beforeEach, test } from "node:test";
+import { WakeAgentConnection, type AgentSocket, type WakeHandler, type WakeReport } from "./WakeAgentConnection.ts";
+import { WakeAgentConnections } from "./WakeAgentConnections.ts";
+import { SessionMap } from "./SessionMap.ts";
+
+/** Socket recording what is sent and letting the test play the server. */
+class FakeSocket implements AgentSocket {
+    readyState = 0;
+    readonly sent: any[] = [];
+    closedWith?: number;
+    onopen: ((event: Event) => void) | null = null;
+    onmessage: ((event: MessageEvent) => void) | null = null;
+    onclose: ((event: CloseEvent) => void) | null = null;
+    onerror: ((event: Event) => void) | null = null;
+
+    constructor(readonly url: string) {}
+
+    send(data: string): void {
+        this.sent.push(JSON.parse(data));
+    }
+
+    close(code?: number): void {
+        this.closedWith = code;
+        this.readyState = 3;
+    }
+
+    open(): void {
+        this.readyState = 1;
+        this.onopen?.({} as Event);
+    }
+
+    receive(message: object): void {
+        this.onmessage?.({ data: JSON.stringify(message) } as MessageEvent);
+    }
+
+    drop(code = 1006): void {
+        this.readyState = 3;
+        this.onclose?.({ code } as CloseEvent);
+    }
+}
+
+const TIMING = { heartbeatMs: 10, reconnectBaseMs: 5, reconnectMaxMs: 20 };
+
+let sockets: FakeSocket[];
+let report: WakeReport;
+let woken: string[];
+let handler: WakeHandler;
+let connection: WakeAgentConnection;
+
+beforeEach(() => {
+    sockets = [];
+    report = { outcome: "activated", details: "Brought the tab forward." };
+    woken = [];
+    handler = {
+        wake: async (sessionId) => {
+            woken.push(sessionId);
+            return report;
+        },
+    };
+    connection = create("secret");
+});
+
+afterEach(() => {
+    connection.close();
+});
+
+function create(token: string | null): WakeAgentConnection {
+    return new WakeAgentConnection(
+        "wss://penpot.example/mcp/ws",
+        token,
+        "agent",
+        (url) => {
+            const socket = new FakeSocket(url);
+            sockets.push(socket);
+            return socket;
+        },
+        handler,
+        TIMING
+    );
+}
+
+function until(condition: () => boolean): Promise<void> {
+    return new Promise((resolve) => {
+        const check = () => (condition() ? resolve() : setTimeout(check, 1));
+        check();
+    });
+}
+
+test("connects with the user token and registers its sessions as a wake agent", () => {
+    connection.update(["a", "b"]);
+    assert.equal(sockets[0].url, "wss://penpot.example/mcp/ws?userToken=secret");
+    assert.equal(connection.status, "connecting");
+    sockets[0].open();
+    assert.equal(connection.status, "connected");
+    assert.deepEqual(sockets[0].sent, [{ type: "initialize-wake-agent", agentId: "agent", sessionIds: ["a", "b"] }]);
+});
+
+test("opens a WebSocket URL for an endpoint announced with an http(s) scheme", () => {
+    connection = new WakeAgentConnection(
+        "https://localhost:3449/mcp/ws",
+        null,
+        "agent",
+        (url) => {
+            const socket = new FakeSocket(url);
+            sockets.push(socket);
+            return socket;
+        },
+        handler,
+        TIMING
+    );
+    connection.update(["a"]);
+    assert.equal(sockets[0].url, "wss://localhost:3449/mcp/ws");
+});
+
+test("connects without a token parameter when the server needs none", () => {
+    connection = create(null);
+    connection.update(["a"]);
+    assert.equal(sockets[0].url, "wss://penpot.example/mcp/ws");
+});
+
+test("sends the full session list when it changes", () => {
+    connection.update(["a"]);
+    sockets[0].open();
+    connection.update(["a"]);
+    connection.update(["a", "b"]);
+    assert.deepEqual(sockets[0].sent.slice(1), [{ type: "wake-sessions", sessionIds: ["a", "b"] }]);
+});
+
+test("sends heartbeats while connected", async () => {
+    connection.update(["a"]);
+    sockets[0].open();
+    await until(() => sockets[0].sent.filter((message) => message.type === "heartbeat").length >= 2);
+});
+
+test("answers a wake request with the handler's outcome and details", async () => {
+    connection.update(["a"]);
+    sockets[0].open();
+    report = { outcome: "disabled", details: "Waking is turned off." };
+    sockets[0].receive({ type: "wake", requestId: "request", sessionId: "a" });
+    await until(() => sockets[0].sent.length === 2);
+    assert.deepEqual(woken, ["a"]);
+    assert.deepEqual(sockets[0].sent[1], {
+        type: "wake-result",
+        requestId: "request",
+        outcome: "disabled",
+        details: "Waking is turned off.",
+    });
+});
+
+test("answers tab-closed when the handler fails", async () => {
+    handler.wake = () => Promise.reject(new Error("No tab with id"));
+    connection.update(["a"]);
+    sockets[0].open();
+    sockets[0].receive({ type: "wake", requestId: "request", sessionId: "a" });
+    await until(() => sockets[0].sent.length === 2);
+    assert.equal(sockets[0].sent[1].outcome, "tab-closed");
+    assert.equal(sockets[0].sent[1].details, "No tab with id");
+});
+
+test("ignores other messages from the server", async () => {
+    connection.update(["a"]);
+    sockets[0].open();
+    sockets[0].onmessage?.({ data: "not json" } as MessageEvent);
+    sockets[0].receive({ type: "initialized" });
+    await new Promise((resolve) => setTimeout(resolve, 5));
+    assert.deepEqual(woken, []);
+    assert.equal(sockets[0].sent.length, 1);
+});
+
+test("reconnects after a dropped connection and registers the current sessions", async () => {
+    connection.update(["a"]);
+    sockets[0].open();
+    sockets[0].drop();
+    assert.equal(connection.status, "disconnected");
+    connection.update(["a", "b"]);
+    await until(() => sockets.length === 2);
+    sockets[1].open();
+    assert.deepEqual(sockets[1].sent[0].sessionIds, ["a", "b"]);
+});
+
+test("stops reconnecting when the server rejects the connection", async () => {
+    connection.update(["a"]);
+    sockets[0].drop(1008);
+    assert.equal(connection.status, "rejected");
+    await new Promise((resolve) => setTimeout(resolve, 30));
+    assert.equal(sockets.length, 1);
+});
+
+test("closing stops reconnects and heartbeats", async () => {
+    connection.update(["a"]);
+    sockets[0].open();
+    connection.close();
+    assert.equal(sockets[0].closedWith, 1000);
+    assert.equal(connection.status, "closed");
+    await new Promise((resolve) => setTimeout(resolve, 30));
+    assert.equal(sockets.length, 1);
+    assert.equal(sockets[0].sent.length, 1);
+});
+
+test("the connection set opens one connection per server and token", async () => {
+    const sessions = new SessionMap({ load: async () => undefined, save: async () => {} });
+    const opened: FakeSocket[] = [];
+    const set = new WakeAgentConnections(
+        sessions,
+        Promise.resolve("agent"),
+        (url) => {
+            const socket = new FakeSocket(url);
+            opened.push(socket);
+            return socket;
+        },
+        handler,
+        TIMING
+    );
+    sessions.observe(set);
+    const announce = (tabId: number, sessionId: string, token: string) =>
+        sessions.announce(
+            { tabId, windowId: 1 },
+            { type: "penpot-mcp:session", sessionId, serverUrl: "wss://penpot/mcp/ws", token, status: "connected" }
+        );
+    await announce(1, "a", "alice");
+    await announce(2, "b", "alice");
+    await announce(3, "c", "bob");
+    await set.sync();
+    assert.deepEqual(
+        opened.map((socket) => socket.url),
+        ["wss://penpot/mcp/ws?userToken=alice", "wss://penpot/mcp/ws?userToken=bob"]
+    );
+    assert.deepEqual(
+        set.summaries().map((summary) => summary.sessionIds),
+        [["a", "b"], ["c"]]
+    );
+    await sessions.removeTab(3);
+    await set.sync();
+    assert.equal(opened[1].closedWith, 1000);
+    assert.equal(set.summaries().length, 1);
+    for (const socket of opened) socket.drop(1008);
+});

mcp/packages/browser-extension/src/WakeAgentConnection.ts

@@ -0,0 +1,212 @@
+import type { WakeAgentInit, WakeAgentSessions, WakeRequest, WakeResult } from "@penpot/mcp-common";
+
+/** WebSocket close code the MCP server uses for policy violations, after which retrying is futile. */
+export const WS_CLOSE_POLICY_VIOLATION = 1008;
+
+const WS_OPEN = 1;
+
+/** WebSocket operations the connection uses; the service worker passes browser WebSockets. */
+export interface AgentSocket {
+    readonly readyState: number;
+    onopen: ((event: Event) => void) | null;
+    onmessage: ((event: MessageEvent) => void) | null;
+    onclose: ((event: CloseEvent) => void) | null;
+    onerror: ((event: Event) => void) | null;
+    send(data: string): void;
+    close(code?: number, reason?: string): void;
+}
+
+/** Creates the WebSocket for a server URL. */
+export type SocketFactory = (url: string) => AgentSocket;
+
+/** Outcome of a wake request and what the extension did, as reported to the server. */
+export type WakeReport = Pick<WakeResult, "outcome" | "details">;
+
+/** Handler of the wake requests the server sends. */
+export interface WakeHandler {
+    wake(sessionId: string): Promise<WakeReport>;
+}
+
+/** Timing of a wake agent connection, in ms. */
+export interface ConnectionTiming {
+    heartbeatMs: number;
+    reconnectBaseMs: number;
+    reconnectMaxMs: number;
+}
+
+export const DEFAULT_TIMING: ConnectionTiming = {
+    // Chrome keeps a service worker running while its WebSocket is active within 30s windows
+    heartbeatMs: 20_000,
+    reconnectBaseMs: 1_000,
+    reconnectMaxMs: 30_000,
+};
+
+/**
+ * State of a wake agent connection.
+ *
+ * - `connecting`: opening the socket
+ * - `connected`: the socket is open and the agent has registered its sessions
+ * - `disconnected`: the socket closed; a reconnect is scheduled
+ * - `rejected`: the server refused the connection; no reconnect follows
+ * - `closed`: the connection was closed because it has no sessions left
+ */
+export type ConnectionStatus = "connecting" | "connected" | "disconnected" | "rejected" | "closed";
+
+/**
+ * Connection of this extension, as a wake agent, to one MCP server for one user token.
+ *
+ * It registers the sessions it can wake, keeps the connection and the service worker active with
+ * heartbeats, answers wake requests through its handler, and reconnects with capped backoff unless
+ * the server rejects it.
+ */
+export class WakeAgentConnection {
+    private socket: AgentSocket | null = null;
+    private sessionIds: string[] = [];
+    private statusValue: ConnectionStatus = "closed";
+    private attempts = 0;
+    private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
+    private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
+
+    /**
+     * @param serverUrl - The WebSocket URL of the MCP server
+     * @param token - The user's MCP token, or null when the server needs none
+     * @param agentId - The stable identifier of this extension
+     * @param createSocket - The factory of the WebSocket
+     * @param handler - The handler of wake requests
+     * @param timing - The heartbeat and reconnect timing
+     */
+    constructor(
+        readonly serverUrl: string,
+        readonly token: string | null,
+        private readonly agentId: string,
+        private readonly createSocket: SocketFactory,
+        private readonly handler: WakeHandler,
+        private readonly timing: ConnectionTiming = DEFAULT_TIMING
+    ) {}
+
+    get status(): ConnectionStatus {
+        return this.statusValue;
+    }
+
+    /** The sessions this connection can wake. */
+    get sessions(): readonly string[] {
+        return this.sessionIds;
+    }
+
+    /**
+     * Sets the sessions this connection can wake, connecting if not connected yet.
+     *
+     * @param sessionIds - The full list of sessions using this server and token
+     */
+    update(sessionIds: string[]): void {
+        const changed = JSON.stringify(sessionIds) !== JSON.stringify(this.sessionIds);
+        this.sessionIds = sessionIds;
+        if (this.statusValue === "closed") {
+            this.open();
+        } else if (changed && this.socket?.readyState === WS_OPEN) {
+            this.send({ type: "wake-sessions", sessionIds } satisfies WakeAgentSessions);
+        }
+    }
+
+    /** Closes the connection without reconnecting. */
+    close(): void {
+        this.statusValue = "closed";
+        this.stopTimers();
+        const socket = this.socket;
+        this.socket = null;
+        socket?.close(1000, "No Penpot sessions left to wake.");
+    }
+
+    private open(): void {
+        this.stopTimers();
+        const url = new URL(this.serverUrl);
+        // Penpot may announce the endpoint with an http(s) scheme
+        url.protocol = url.protocol.replace(/^http/, "ws");
+        if (this.token) {
+            url.searchParams.set("userToken", this.token);
+        }
+        this.statusValue = "connecting";
+        let socket: AgentSocket;
+        try {
+            socket = this.createSocket(url.toString());
+        } catch {
+            this.statusValue = "disconnected";
+            this.scheduleReconnect();
+            return;
+        }
+        this.socket = socket;
+        socket.onopen = () => {
+            if (this.socket !== socket) return;
+            this.attempts = 0;
+            this.statusValue = "connected";
+            this.send({
+                type: "initialize-wake-agent",
+                agentId: this.agentId,
+                sessionIds: this.sessionIds,
+            } satisfies WakeAgentInit);
+            this.heartbeatTimer = setInterval(() => this.send({ type: "heartbeat" }), this.timing.heartbeatMs);
+        };
+        socket.onmessage = (event) => {
+            if (this.socket !== socket) return;
+            void this.handleMessage(event.data);
+        };
+        socket.onclose = (event) => {
+            if (this.socket !== socket) return;
+            this.socket = null;
+            this.stopTimers();
+            if (event.code === WS_CLOSE_POLICY_VIOLATION) {
+                this.statusValue = "rejected";
+                return;
+            }
+            this.statusValue = "disconnected";
+            this.scheduleReconnect();
+        };
+        // a failed socket also reports its close, which schedules the reconnect
+        socket.onerror = () => {};
+    }
+
+    private async handleMessage(data: unknown): Promise<void> {
+        let request: WakeRequest;
+        try {
+            request = JSON.parse(String(data));
+        } catch {
+            return;
+        }
+        if (request?.type !== "wake" || typeof request.requestId !== "string") {
+            return;
+        }
+        let report: WakeReport;
+        try {
+            report = await this.handler.wake(request.sessionId);
+        } catch (error) {
+            report = { outcome: "tab-closed", details: (error as Error).message };
+        }
+        this.send({ type: "wake-result", requestId: request.requestId, ...report } satisfies WakeResult);
+    }
+
+    private send(message: object): void {
+        if (this.socket?.readyState === WS_OPEN) {
+            this.socket.send(JSON.stringify(message));
+        }
+    }
+
+    private scheduleReconnect(): void {
+        const delay = Math.min(this.timing.reconnectBaseMs * 2 ** this.attempts, this.timing.reconnectMaxMs);
+        this.attempts++;
+        this.reconnectTimer = setTimeout(() => {
+            this.reconnectTimer = null;
+            this.open();
+        }, delay);
+    }
+
+    private stopTimers(): void {
+        if (this.heartbeatTimer !== null) {
+            clearInterval(this.heartbeatTimer);
+            this.heartbeatTimer = null;
+        }
+        if (this.reconnectTimer !== null) {
+            clearTimeout(this.reconnectTimer);
+            this.reconnectTimer = null;
+        }
+    }
+}

mcp/packages/browser-extension/src/WakeAgentConnections.ts

@@ -0,0 +1,88 @@
+import type { SessionMap, SessionMapObserver } from "./SessionMap.ts";
+import {
+    DEFAULT_TIMING,
+    WakeAgentConnection,
+    type ConnectionStatus,
+    type ConnectionTiming,
+    type SocketFactory,
+    type WakeHandler,
+} from "./WakeAgentConnection.ts";
+
+/** Summary of one server connection, as shown in the popup. */
+export interface ConnectionSummary {
+    serverUrl: string;
+    status: ConnectionStatus;
+    sessionIds: string[];
+}
+
+/**
+ * The wake agent connections of this extension: one per MCP server URL and user token used by the
+ * mapped sessions.
+ *
+ * It opens a connection when a server and token pair gains its first session, sends each
+ * connection its current sessions, and closes a connection whose last session goes.
+ */
+export class WakeAgentConnections implements SessionMapObserver {
+    private readonly connections = new Map<string, WakeAgentConnection>();
+
+    /**
+     * @param sessions - The map of sessions to tabs
+     * @param agentId - The stable identifier of this extension, once loaded
+     * @param createSocket - The factory of WebSockets
+     * @param handler - The handler of wake requests
+     * @param timing - The heartbeat and reconnect timing
+     */
+    constructor(
+        private readonly sessions: SessionMap,
+        private readonly agentId: Promise<string>,
+        private readonly createSocket: SocketFactory,
+        private readonly handler: WakeHandler,
+        private readonly timing: ConnectionTiming = DEFAULT_TIMING
+    ) {}
+
+    sessionsChanged(): void {
+        void this.sync();
+    }
+
+    /** Brings the connections in line with the current sessions. */
+    async sync(): Promise<void> {
+        const [agentId] = await Promise.all([this.agentId, this.sessions.ready()]);
+        const wanted = new Map<string, { serverUrl: string; token: string | null; sessionIds: string[] }>();
+        for (const [sessionId, entry] of this.sessions.all()) {
+            const key = JSON.stringify([entry.serverUrl, entry.token]);
+            const group = wanted.get(key) ?? { serverUrl: entry.serverUrl, token: entry.token, sessionIds: [] };
+            group.sessionIds.push(sessionId);
+            wanted.set(key, group);
+        }
+        for (const [key, connection] of [...this.connections]) {
+            if (!wanted.has(key)) {
+                connection.close();
+                this.connections.delete(key);
+            }
+        }
+        for (const [key, group] of wanted) {
+            let connection = this.connections.get(key);
+            if (!connection) {
+                connection = new WakeAgentConnection(
+                    group.serverUrl,
+                    group.token,
+                    agentId,
+                    this.createSocket,
+                    this.handler,
+                    this.timing
+                );
+                this.connections.set(key, connection);
+            }
+            connection.update(group.sessionIds.sort());
+        }
+    }
+
+    /** Summarizes the connections. */
+    summaries(): ConnectionSummary[] {
+        return [...this.connections.values()].map((connection) => ({
+            serverUrl: connection.serverUrl,
+            status: connection.status,
+            sessionIds: [...connection.sessions],
+        }));
+    }
+}

mcp/packages/browser-extension/src/background.ts

@@ -7,12 +7,20 @@
  * Chrome may stop the worker at any time, so listeners are registered synchronously at start
  * and state lives in session storage.
  */
-import { ChromeDiscardControl, ChromeSessionStore } from "./chrome.ts";
+import { ChromeDiscardControl, ChromeSessionStore, ChromeTabControl, ChromeWakeSwitch, getAgentId } from "./chrome.ts";
 import { DiscardGuard } from "./DiscardGuard.ts";
 import { parseAnnouncement, type ContentMessage } from "./messages.ts";
 import { SessionMap } from "./SessionMap.ts";
+import { TabWaker } from "./TabWaker.ts";
+import { WakeAgentConnections } from "./WakeAgentConnections.ts";
 
-const sessions = new SessionMap(new ChromeSessionStore(), new DiscardGuard(new ChromeDiscardControl()));
+const RESTART_ALARM = "restart-connections";
+
+const sessions = new SessionMap(new ChromeSessionStore());
+const waker = new TabWaker(sessions, new ChromeTabControl(), new ChromeWakeSwitch());
+const connections = new WakeAgentConnections(sessions, getAgentId(), (url) => new WebSocket(url), waker);
+sessions.observe(new DiscardGuard(new ChromeDiscardControl()));
+sessions.observe(connections);
 
 chrome.runtime.onMessage.addListener((message: ContentMessage, sender) => {
     const tab = sender.tab;
@@ -32,3 +40,13 @@ chrome.runtime.onMessage.addListener((message: ContentMessage, sender) => {
 chrome.tabs.onRemoved.addListener((tabId) => void sessions.removeTab(tabId));
 
 chrome.tabs.onReplaced.addListener((_addedTabId, removedTabId) => void sessions.removeTab(removedTabId));
+
+// the alarm starts the worker again after Chrome has stopped it, which reopens the connections
+chrome.alarms.onAlarm.addListener((alarm) => {
+    if (alarm.name === RESTART_ALARM) {
+        void connections.sync();
+    }
+});
+void chrome.alarms.create(RESTART_ALARM, { periodInMinutes: 0.5 });
+
+void connections.sync();

mcp/packages/browser-extension/src/chrome.ts

@@ -1,7 +1,27 @@
 import type { DiscardControl } from "./DiscardGuard.ts";
 import type { SessionEntry, SessionStore } from "./SessionMap.ts";
+import type { TabControl, WakeSwitch } from "./TabWaker.ts";
 
 const SESSIONS_KEY = "sessions";
+const AGENT_ID_KEY = "agentId";
+const WAKE_ENABLED_KEY = "wakeEnabled";
+
+/**
+ * Returns the identifier of this extension installation, creating it on first use.
+ *
+ * The MCP server replaces an agent's earlier connection when the same agent reconnects, so the
+ * identifier must survive service worker restarts.
+ */
+export async function getAgentId(): Promise<string> {
+    const stored = await chrome.storage.local.get(AGENT_ID_KEY);
+    const existing = stored[AGENT_ID_KEY];
+    if (typeof existing === "string") {
+        return existing;
+    }
+    const agentId = crypto.randomUUID();
+    await chrome.storage.local.set({ [AGENT_ID_KEY]: agentId });
+    return agentId;
+}
 
 /**
  * Session store in `chrome.storage.session`, which lives in memory, survives service worker
@@ -24,3 +44,33 @@ export class ChromeDiscardControl implements DiscardControl {
         await chrome.tabs.update(tabId, { autoDiscardable });
     }
 }
+
+/** Tab activation through the Chrome tabs and windows APIs. */
+export class ChromeTabControl implements TabControl {
+    async activate(tabId: number): Promise<boolean> {
+        let tab: chrome.tabs.Tab | undefined;
+        try {
+            tab = await chrome.tabs.update(tabId, { active: true });
+        } catch {
+            return false;
+        }
+        if (!tab) {
+            return false;
+        }
+        // the tab may have moved to another window since it announced its session
+        await chrome.windows.update(tab.windowId, { focused: true });
+        return true;
+    }
+}
+
+/** Wake switch kept in `chrome.storage.local`; waking is enabled until the user turns it off. */
+export class ChromeWakeSwitch implements WakeSwitch {
+    async isEnabled(): Promise<boolean> {
+        const stored = await chrome.storage.local.get(WAKE_ENABLED_KEY);
+        return stored[WAKE_ENABLED_KEY] !== false;
+    }
+
+    async setEnabled(enabled: boolean): Promise<void> {
+        await chrome.storage.local.set({ [WAKE_ENABLED_KEY]: enabled });
+    }
+}