diff --git a/electron/remote/server-ws.test.ts b/electron/remote/server-ws.test.ts index e3e0abcc..30083dc7 100644 --- a/electron/remote/server-ws.test.ts +++ b/electron/remote/server-ws.test.ts @@ -11,6 +11,9 @@ import { chmodSync, mkdirSync, mkdtempSync, writeFileSync, rmSync } from 'node:f import { tmpdir } from 'node:os'; import { join } from 'node:path'; +// The server's PTY event listeners, so a test can respawn an agent. +const ptyListeners = vi.hoisted(() => new Map void>()); + vi.mock('../ipc/pty.js', () => ({ writeToAgent: vi.fn(), resizeAgent: vi.fn(), @@ -25,7 +28,10 @@ vi.mock('../ipc/pty.js', () => ({ getAgentMeta: vi.fn(() => null), getAgentCols: vi.fn(() => 80), getAgentRows: vi.fn(() => 24), - onPtyEvent: vi.fn(() => vi.fn()), // returns an unsubscribe fn + onPtyEvent: vi.fn((event: string, listener: (agentId: string, data?: unknown) => void) => { + ptyListeners.set(event, listener); + return vi.fn(); // the unsubscribe fn + }), })); const pty = await import('../ipc/pty.js'); @@ -229,6 +235,76 @@ describe('subscribe snapshot', () => { }); }); +describe('agent respawned under the same ID', () => { + // Switching a task back from chat to terminal starts a new PTY with the old + // agent ID. The phone never unsubscribed, so it does not subscribe again. + beforeEach(() => { + vi.mocked(pty.subscribeToAgentRendered).mockImplementation(() => null); + }); + + function snapshots(ws: WebSocket): unknown[] { + const seen: unknown[] = []; + ws.on('message', (raw) => { + const msg = JSON.parse(String(raw)) as { type: string }; + if (msg.type === 'scrollback') seen.push(msg); + }); + return seen; + } + + async function subscribeWithoutPty(ws: WebSocket, agentId: string): Promise { + ws.send(JSON.stringify({ type: 'subscribe', agentId })); + await vi.waitFor(() => + expect(pty.subscribeToAgentRendered).toHaveBeenCalledWith( + agentId, + expect.any(Function), + expect.any(Function), + ), + ); + } + + function respawnWithSnapshot(): void { + vi.mocked(pty.subscribeToAgentRendered).mockImplementation((_id, onSnapshot) => { + onSnapshot({ data: 'bmV3', cols: 80, rows: 24 }); + return vi.fn(); + }); + } + + it('streams the new PTY to a phone that was watching the agent', async () => { + const ws = await connectAndAuth(mobileToken); + const seen = snapshots(ws); + await subscribeWithoutPty(ws, 'agent-1'); + + respawnWithSnapshot(); + ptyListeners.get('spawn')?.('agent-1'); + + await vi.waitFor(() => + expect(seen).toEqual([ + { type: 'scrollback', agentId: 'agent-1', data: 'bmV3', cols: 80, rows: 24 }, + ]), + ); + ws.close(); + }); + + it('leaves phones that unsubscribed, and reattached PTYs, alone', async () => { + const ws = await connectAndAuth(mobileToken); + const seen = snapshots(ws); + await subscribeWithoutPty(ws, 'agent-1'); + await subscribeWithoutPty(ws, 'agent-2'); + ws.send(JSON.stringify({ type: 'unsubscribe', agentId: 'agent-1' })); + // Messages are handled in order, so this round trip settles the unsubscribe. + await subscribeWithoutPty(ws, 'agent-3'); + + respawnWithSnapshot(); + vi.mocked(pty.subscribeToAgentRendered).mockClear(); + ptyListeners.get('spawn')?.('agent-1'); + ptyListeners.get('spawn')?.('agent-2', { reattached: true }); + + expect(pty.subscribeToAgentRendered).not.toHaveBeenCalled(); + expect(seen).toEqual([]); + ws.close(); + }); +}); + describe('paired phone view size', () => { it('sizes the PTY and hands it back when the phone disconnects', async () => { const ws = await connectAndAuth(await pair()); diff --git a/electron/remote/server.ts b/electron/remote/server.ts index 19ffdf8d..e45d5b26 100644 --- a/electron/remote/server.ts +++ b/electron/remote/server.ts @@ -2042,6 +2042,8 @@ export function startRemoteServer(opts: { }); const clientSubs = new WeakMap void>>(); + // The agents each client asked to stream; unlike clientSubs, this outlives the PTY. + const clientWants = new WeakMap>(); // Which phone currently sizes each agent's PTY (see ViewSizeCommand). const viewSizeOwners = new Map(); const authenticatedClients = new Set(); @@ -2053,6 +2055,65 @@ export function startRemoteServer(opts: { // past this backlog, chat frames wait rather than pile up in the send buffer. const SOCKET_BACKLOG_BYTES = 1024 * 1024; + /** + * Stream an agent's terminal to a phone, starting from its rendered snapshot. + * Also runs again when the agent's PTY respawns under the same ID (switching + * a task back from chat to terminal), since the old PTY's subscribers die with it. + */ + function streamAgent(ws: WebSocket, agentId: string): void { + const subs = clientSubs.get(ws); + const sendScrollback = (data: string, cols: number, rows: number) => { + if (ws.readyState !== WebSocket.OPEN) return; + ws.send( + JSON.stringify({ + type: 'scrollback', + agentId, + data, + cols, + rows, + } satisfies ServerMessage), + ); + }; + const sendRawScrollback = () => { + const scrollback = getAgentScrollback(agentId); + if (scrollback) sendScrollback(scrollback, getAgentCols(agentId), getAgentRows(agentId)); + }; + const cb = (encoded: string) => { + if (ws.readyState === WebSocket.OPEN) { + if (ws.bufferedAmount > SOCKET_BACKLOG_BYTES) { + // Terminal chunks cannot be skipped safely. Release the queued bytes; + // the phone reconnects and resets its terminal from scrollback replay. + ws.terminate(); + return; + } + ws.send( + JSON.stringify({ + type: 'output', + agentId, + data: encoded, + } satisfies ServerMessage), + ); + } + }; + // Phones get the rendered screen and history; a live agent's raw + // replay can be nothing but repaints of its last screen. + const subscriber = subscribeToAgentRendered( + agentId, + (snapshot) => { + if (snapshot) sendScrollback(snapshot.data, snapshot.cols, snapshot.rows); + else sendRawScrollback(); + }, + cb, + ); + if (subscriber) { + subs?.set(agentId, subscriber); + } else { + // Not a live PTY session: replay what is left, as before. + sendRawScrollback(); + if (subscribeToAgent(agentId, cb)) subs?.set(agentId, cb); + } + } + function broadcast(msg: ServerMessage): void { const json = JSON.stringify(msg); for (const client of wss.clients) { @@ -2069,7 +2130,19 @@ export function startRemoteServer(opts: { push.snapshot(list); broadcast({ type: 'agents', list }); }; - const unsubSpawn = onPtyEvent('spawn', publishAgents); + const unsubSpawn = onPtyEvent('spawn', (agentId, data) => { + publishAgents(); + if ((data as { reattached?: boolean } | undefined)?.reattached) return; + // A new PTY under a known ID (a task switched back from chat to terminal) + // starts with no subscribers, and the phone still believes it is streaming. + for (const client of wss.clients) { + if (!clientWants.get(client)?.has(agentId)) continue; + const stale = clientSubs.get(client)?.get(agentId); + if (stale) unsubscribeFromAgent(agentId, stale); + clientSubs.get(client)?.delete(agentId); + streamAgent(client, agentId); + } + }); const unsubChats = opts.chats?.onChange(() => { @@ -2111,6 +2184,7 @@ export function startRemoteServer(opts: { wss.on('connection', (ws, req) => { clientSubs.set(ws, new Map()); + clientWants.set(ws, new Set()); ws.on('pong', () => awaitingPong.delete(ws)); if (opts.chats) clientChats.set( @@ -2341,64 +2415,14 @@ export function startRemoteServer(opts: { break; case 'subscribe': { - const subs = clientSubs.get(ws); - if (subs?.has(msg.agentId)) break; - - const sendScrollback = (data: string, cols: number, rows: number) => { - if (ws.readyState !== WebSocket.OPEN) return; - ws.send( - JSON.stringify({ - type: 'scrollback', - agentId: msg.agentId, - data, - cols, - rows, - } satisfies ServerMessage), - ); - }; - const sendRawScrollback = () => { - const scrollback = getAgentScrollback(msg.agentId); - if (scrollback) - sendScrollback(scrollback, getAgentCols(msg.agentId), getAgentRows(msg.agentId)); - }; - const cb = (encoded: string) => { - if (ws.readyState === WebSocket.OPEN) { - if (ws.bufferedAmount > SOCKET_BACKLOG_BYTES) { - // Terminal chunks cannot be skipped safely. Release the queued bytes; - // the phone reconnects and resets its terminal from scrollback replay. - ws.terminate(); - return; - } - ws.send( - JSON.stringify({ - type: 'output', - agentId: msg.agentId, - data: encoded, - } satisfies ServerMessage), - ); - } - }; - // Phones get the rendered screen and history; a live agent's raw - // replay can be nothing but repaints of its last screen. - const subscriber = subscribeToAgentRendered( - msg.agentId, - (snapshot) => { - if (snapshot) sendScrollback(snapshot.data, snapshot.cols, snapshot.rows); - else sendRawScrollback(); - }, - cb, - ); - if (subscriber) { - subs?.set(msg.agentId, subscriber); - } else { - // Not a live PTY session: replay what is left, as before. - sendRawScrollback(); - if (subscribeToAgent(msg.agentId, cb)) subs?.set(msg.agentId, cb); - } + clientWants.get(ws)?.add(msg.agentId); + if (clientSubs.get(ws)?.has(msg.agentId)) break; + streamAgent(ws, msg.agentId); break; } case 'unsubscribe': { + clientWants.get(ws)?.delete(msg.agentId); const subs = clientSubs.get(ws); const cb = subs?.get(msg.agentId); if (cb) {