diff --git a/mcp/packages/browser-extension/src/SessionMap.test.ts b/mcp/packages/browser-extension/src/SessionMap.test.ts index 94dd636300..b39c893244 100644 --- a/mcp/packages/browser-extension/src/SessionMap.test.ts +++ b/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, [ diff --git a/mcp/packages/browser-extension/src/SessionMap.ts b/mcp/packages/browser-extension/src/SessionMap.ts index 65482627c6..440eb7167f 100644 --- a/mcp/packages/browser-extension/src/SessionMap.ts +++ b/mcp/packages/browser-extension/src/SessionMap.ts @@ -36,12 +36,10 @@ export interface SessionMapObserver { */ export class SessionMap { private entries = new Map(); + private readonly observers: SessionMapObserver[] = []; private readonly loaded: Promise; - 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 { 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 { diff --git a/mcp/packages/browser-extension/src/TabWaker.test.ts b/mcp/packages/browser-extension/src/TabWaker.test.ts new file mode 100644 index 0000000000..b386ae486f --- /dev/null +++ b/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; +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); +}); diff --git a/mcp/packages/browser-extension/src/TabWaker.ts b/mcp/packages/browser-extension/src/TabWaker.ts new file mode 100644 index 0000000000..f3e74f1846 --- /dev/null +++ b/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; +} + +/** The user's switch that pauses waking without removing the extension. */ +export interface WakeSwitch { + isEnabled(): Promise; + setEnabled(enabled: boolean): Promise; +} + +/** + * 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 { + 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" }; + } +} diff --git a/mcp/packages/browser-extension/src/WakeAgentConnection.test.ts b/mcp/packages/browser-extension/src/WakeAgentConnection.test.ts new file mode 100644 index 0000000000..1a456be7b8 --- /dev/null +++ b/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 { + 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); +}); diff --git a/mcp/packages/browser-extension/src/WakeAgentConnection.ts b/mcp/packages/browser-extension/src/WakeAgentConnection.ts new file mode 100644 index 0000000000..8bbdc13554 --- /dev/null +++ b/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; + +/** Handler of the wake requests the server sends. */ +export interface WakeHandler { + wake(sessionId: string): Promise; +} + +/** 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 | null = null; + private reconnectTimer: ReturnType | 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 { + 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; + } + } +} diff --git a/mcp/packages/browser-extension/src/WakeAgentConnections.ts b/mcp/packages/browser-extension/src/WakeAgentConnections.ts new file mode 100644 index 0000000000..b55a01c5c2 --- /dev/null +++ b/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(); + + /** + * @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, + 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 { + const [agentId] = await Promise.all([this.agentId, this.sessions.ready()]); + const wanted = new Map(); + 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], + })); + } +} diff --git a/mcp/packages/browser-extension/src/background.ts b/mcp/packages/browser-extension/src/background.ts index 143a04f0bb..2db7251e70 100644 --- a/mcp/packages/browser-extension/src/background.ts +++ b/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(); diff --git a/mcp/packages/browser-extension/src/chrome.ts b/mcp/packages/browser-extension/src/chrome.ts index 38954c7ced..0bed1f1777 100644 --- a/mcp/packages/browser-extension/src/chrome.ts +++ b/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 { + 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 { + 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 { + const stored = await chrome.storage.local.get(WAKE_ENABLED_KEY); + return stored[WAKE_ENABLED_KEY] !== false; + } + + async setEnabled(enabled: boolean): Promise { + await chrome.storage.local.set({ [WAKE_ENABLED_KEY]: enabled }); + } +}