Skip to content

fix(remote): stream a terminal again after its PTY respawns #312

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
johannesjo merged 1 commit into from
Oct 11, 2026
Merged

fix(remote): stream a terminal again after its PTY respawns #312

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 77 additions & 1 deletion electron/remote/server-ws.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, (agentId: string, data?: unknown) => void>());

vi.mock('../ipc/pty.js', () => ({
writeToAgent: vi.fn(),
resizeAgent: vi.fn(),
Expand All @@ -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');
Expand Down Expand Up @@ -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<void> {
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());
Expand Down
134 changes: 79 additions & 55 deletions electron/remote/server.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2042,6 +2042,8 @@ export function startRemoteServer(opts: {
});

const clientSubs = new WeakMap<WebSocket, Map<string, (data: string) => void>>();
// The agents each client asked to stream; unlike clientSubs, this outlives the PTY.
const clientWants = new WeakMap<WebSocket, Set<string>>();
// Which phone currently sizes each agent's PTY (see ViewSizeCommand).
const viewSizeOwners = new Map<string, WebSocket>();
const authenticatedClients = new Set<WebSocket>();
Expand All @@ -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) {
Expand All @@ -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(() => {
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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) {
Expand Down
Loading