feat(terminal-agent): sessionId-aware grant + scoped restart + eager spawn

Wires the pty-session-lease primitive (3aada48b) into terminal-agent so
the Commit 2 work in server.ts (next commit) can route /pty-restart and
re-attach by session identity rather than by single-use token.

Changes:

  * validTokens: Set<string> → Map<string, string|null>. Each grant carries
    its bound sessionId (or null for legacy single-grant callers). On WS
    upgrade, the agent surfaces the bound sessionId via ws.data so open()
    can register the session in the new reverse index.
  * sessionsById: Map<sessionId, PtySession> — populated in open(),
    cleared in close(). Required so /internal/restart can find and dispose
    one specific session by id rather than enumerating all live sessions.
  * /internal/restart: scoped to one sessionId. Codex T2 of the eng review
    caught the gap — pre-spec the route would have disposed every PTY on
    the agent, breaking pair-agent and any future multi-sidebar setup.
    The body now requires `{sessionId}`; missing or unknown id returns
    `{killed: 0}` and leaves siblings alone.
  * maybeSpawnPty(ws, session): hoisted from the inline binary-frame spawn
    block so both the legacy "spawn on first keystroke" trigger AND the
    new `{type:"start"}` text-frame trigger land in the same code path.
    Idempotent on session.spawned.
  * `{type:"start"}` text frame: explicit spawn trigger. forceRestart
    (extension side, lands in Commit 2C) sends this immediately on every
    fresh WS so claude boots without requiring a keystroke. Pre-v1.44 the
    lazy-binary-spawn pattern made the restart feel stuck.
  * close(ws): drops the sessionsById entry alongside the existing
    sessions WeakMap + validTokens cleanup. Commit 3 will revisit this to
    keep the session alive for a 60s detach window before disposing.

Test (browse/test/terminal-agent-session-routing.test.ts):
  * 8 static-grep tripwires pinning the load-bearing properties: validTokens
    is a Map (not Set), sessionsById exists, /internal/restart is scoped
    (negative-assert against enumerate-all patterns), WS upgrade plumbs
    sessionId, maybeSpawnPty is the single spawn entry, close() drops the
    index. Live spawn cycles belong in the e2e tier.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Garry Tan 2026-05-23 23:15:49 -07:00
parent 3aada48bf9
commit 449144cda5
No known key found for this signature in database
GPG Key ID: C1F69E85C74EFE1D
2 changed files with 233 additions and 48 deletions

View File

@ -41,9 +41,31 @@ const INTERNAL_TOKEN = crypto.randomBytes(32).toString('base64url'); // shared w
*/
const CURRENT_GEN = crypto.randomBytes(16).toString('base64url');
// In-memory cookie token registry. Parent posts /internal/grant after
// /pty-session; we validate WS cookies against this set.
const validTokens = new Set<string>();
// In-memory attach-token registry. Parent posts /internal/grant after
// /pty-session; we validate WS upgrades against this map.
//
// v1.44+: each token is bound to a v1.44 sessionId (the stable, non-secret
// identifier from browse/src/pty-session-lease.ts). The token grants ONE
// attach for ONE session — re-attach within the lease window comes through
// /pty-session/reattach, which mints a fresh token for the same sessionId.
//
// Legacy callers can still pass `{token}` without sessionId (the value
// stays null and the WS upgrade still works); those callers don't get
// re-attach because there's no stable identifier to match against.
const validTokens = new Map<string, string | null>(); // token → sessionId
/**
* Reverse index for re-attach lookups: sessionId live PtySession.
* Populated when a WS first attaches with a known sessionId; cleared when
* the session is disposed or the lease expires. Used by:
* - /ws upgrade: if the incoming attachToken maps to a sessionId that
* already has a live session, REPLACE its ws ref instead of spawning.
* - /internal/restart: enumerate by sessionId, dispose that one session.
*
* Kept separate from the WeakMap<ws,PtySession> so re-attach can find the
* session by id even after the original ws has gone.
*/
const sessionsById = new Map<string, PtySession>();
// Active PTY session per WS. One terminal per connection. Codex finding #4:
// uncaught handlers below catch bugs in framing/cleanup so they don't kill
@ -60,6 +82,12 @@ interface PtySession {
cols: number;
rows: number;
cookie: string;
/**
* v1.44+ stable session identifier (from pty-session-lease). Null for
* legacy /internal/grant callers that didn't pass one. Used for
* targeted /internal/restart and (in Commit 3) re-attach lookups.
*/
sessionId: string | null;
spawned: boolean;
/**
* 25s server-side WS keepalive interval (v1.44+). Set in the WS `open`
@ -285,6 +313,56 @@ async function internalHandler<T>(
}
}
/**
* Spawn the claude PTY for a session if it hasn't been spawned yet.
* Used by both the legacy binary-frame spawn trigger and the v1.44 explicit
* `{type:"start"}` text-frame trigger. Idempotent on `session.spawned`.
*
* Returns true if claude is now running, false if spawn failed (e.g. claude
* binary not on PATH). On failure, the caller is expected to have already
* surfaced the error to the client (or will via the next frame).
*/
function maybeSpawnPty(ws: any, session: PtySession): boolean {
if (session.spawned) return true;
session.spawned = true;
let leftover = Buffer.alloc(0);
const proc = spawnClaude(session.cols, session.rows, (chunk) => {
const combined = Buffer.concat([leftover, Buffer.from(chunk)]);
// UTF-8 boundary detection (issue #1272). Look back at most 3 bytes
// for the start of an incomplete multibyte sequence and defer it.
let safeEnd = combined.length;
for (let i = combined.length - 1; i >= Math.max(0, combined.length - 3); i--) {
const b = combined[i];
if ((b & 0x80) === 0) { safeEnd = i + 1; break; }
if ((b & 0xC0) === 0x80) continue;
const expected = (b & 0xE0) === 0xC0 ? 2 : (b & 0xF0) === 0xE0 ? 3 : 4;
safeEnd = (combined.length - i >= expected) ? combined.length : i;
break;
}
const flush = combined.slice(0, safeEnd);
leftover = combined.slice(safeEnd);
if (flush.length) {
try { ws.sendBinary(flush); } catch {}
}
});
if (!proc) {
try {
ws.send(JSON.stringify({
type: 'error',
code: 'CLAUDE_NOT_FOUND',
message: 'claude CLI not on PATH. Install: https://docs.anthropic.com/en/docs/claude-code',
}));
ws.close(4404, 'claude not found');
} catch {}
return false;
}
session.proc = proc;
proc.exited?.then?.(() => {
try { ws.close(1000, 'pty exited'); } catch {}
});
return true;
}
function buildServer() {
return Bun.serve({
hostname: '127.0.0.1',
@ -295,10 +373,18 @@ function buildServer() {
const url = new URL(req.url);
// /internal/grant — loopback-only handshake from parent server.
// v1.44+: accepts `{token, sessionId?}`. The sessionId binding lets
// the agent route re-attach attempts (same sessionId, fresh token)
// back to the same PtySession. Legacy callers passing just `{token}`
// still work — sessionId becomes null and re-attach is unavailable
// for that grant.
if (url.pathname === '/internal/grant' && req.method === 'POST') {
return internalHandler(req, (body) => {
if (typeof body?.token === 'string' && body.token.length > 16) {
validTokens.add(body.token);
const sid = typeof body?.sessionId === 'string' && body.sessionId.length > 0
? body.sessionId
: null;
validTokens.set(body.token, sid);
}
});
}
@ -310,6 +396,24 @@ function buildServer() {
});
}
// /internal/restart — dispose the PtySession for a specific sessionId.
// Scoped to one caller (not enumerate-all). Server.ts /pty-restart
// posts here with the caller's sessionId; we kill ONLY that PTY,
// leaving any other live sidebar tabs untouched. Codex T2 of the
// eng review caught this gap — pre-spec the route would have
// disposed all sessions.
if (url.pathname === '/internal/restart' && req.method === 'POST') {
return internalHandler(req, (body) => {
const sid = typeof body?.sessionId === 'string' ? body.sessionId : null;
if (!sid) return { killed: 0 };
const session = sessionsById.get(sid);
if (!session) return { killed: 0 };
disposeSession(session);
sessionsById.delete(sid);
return { killed: 1 };
});
}
// /internal/healthz — liveness probe used by the v1.44 watchdog.
// Returns this agent's pid + gen + active session count without
// touching claude binary lookup (which can fail for non-process
@ -394,8 +498,13 @@ function buildServer() {
return new Response('unauthorized', { status: 401 });
}
// v1.44+: surface the token's sessionId binding to the upgraded ws.
// open() reads it via ws.data and registers the session in
// sessionsById so /internal/restart and (Commit 3) re-attach
// lookups can find it.
const sessionId = validTokens.get(token) ?? null;
const upgraded = server.upgrade(req, {
data: { cookie: token },
data: { cookie: token, sessionId },
// Echo the protocol back so the browser accepts the upgrade.
// Required when the client sends Sec-WebSocket-Protocol — the
// server MUST select one of the offered protocols, otherwise
@ -410,17 +519,21 @@ function buildServer() {
websocket: {
/**
* Start the 25s server-side keepalive ping. Stored on the session so
* close() can clear it. We construct the PtySession here (not lazily
* in message()) so the ping interval has somewhere to live before any
* data flows. Lazy claude spawn still happens on first binary frame.
* Spawn the claude PTY for `session` if it hasn't been spawned yet.
* Called from both message paths: the legacy binary-frame trigger
* (any keystroke) AND the v1.44 explicit `{type:"start"}` trigger
* (forceRestart sends this on every fresh WS to get an eager prompt
* without requiring the user to type). Idempotent a second call
* after `spawned: true` is a no-op.
*/
open(ws) {
const sessionId = (ws.data as any)?.sessionId ?? null;
const session: PtySession = {
proc: null,
cols: 80,
rows: 24,
cookie: (ws.data as any)?.cookie || '',
sessionId,
spawned: false,
pingInterval: null,
};
@ -432,6 +545,8 @@ function buildServer() {
}
}, KEEPALIVE_INTERVAL_MS);
sessions.set(ws, session);
// Index by sessionId for /internal/restart + (Commit 3) re-attach.
if (sessionId) sessionsById.set(sessionId, session);
},
message(ws, raw) {
@ -445,10 +560,12 @@ function buildServer() {
cols: 80,
rows: 24,
cookie: (ws.data as any)?.cookie || '',
sessionId: (ws.data as any)?.sessionId ?? null,
spawned: false,
pingInterval: null,
};
sessions.set(ws, session);
if (session.sessionId) sessionsById.set(session.sessionId, session);
}
// Text frames are control messages: {type: "resize", cols, rows},
@ -482,50 +599,24 @@ function buildServer() {
// future agent peer) mirrors our server-side ping shape.
return;
}
if (msg?.type === 'start') {
// v1.44 explicit spawn trigger. forceRestart sends this
// immediately on every fresh WS so claude boots without the
// user having to type a keystroke (pre-v1.44, the lazy-binary
// spawn made restart look stuck until the user typed). No-op
// if already spawned.
maybeSpawnPty(ws, session);
return;
}
// Unknown text frame — ignore.
return;
}
// Binary input. Lazy-spawn claude on the first byte.
// Binary input. Lazy-spawn claude on the first byte if `start`
// wasn't sent first. Both paths land in the same maybeSpawnPty
// helper for behavior parity.
if (!session.spawned) {
session.spawned = true;
// UTF-8 boundary detection to prevent splitting multi-byte characters (issue #1272).
// Buffer incomplete UTF-8 sequences until the next chunk completes them.
let leftover = Buffer.alloc(0);
const proc = spawnClaude(session.cols, session.rows, (chunk) => {
const combined = Buffer.concat([leftover, Buffer.from(chunk)]);
// Find the last index where a UTF-8 codepoint ends. Look back at most 3 bytes.
let safeEnd = combined.length;
for (let i = combined.length - 1; i >= Math.max(0, combined.length - 3); i--) {
const b = combined[i];
if ((b & 0x80) === 0) { safeEnd = i + 1; break; } // ASCII
if ((b & 0xC0) === 0x80) continue; // continuation byte
const expected = (b & 0xE0) === 0xC0 ? 2 : (b & 0xF0) === 0xE0 ? 3 : 4;
safeEnd = (combined.length - i >= expected) ? combined.length : i;
break;
}
const flush = combined.slice(0, safeEnd);
leftover = combined.slice(safeEnd);
if (flush.length) {
try { ws.sendBinary(flush); } catch {}
}
});
if (!proc) {
try {
ws.send(JSON.stringify({
type: 'error',
code: 'CLAUDE_NOT_FOUND',
message: 'claude CLI not on PATH. Install: https://docs.anthropic.com/en/docs/claude-code',
}));
ws.close(4404, 'claude not found');
} catch {}
return;
}
session.proc = proc;
// Watch for child exit so the WS closes cleanly when claude exits.
proc.exited?.then?.(() => {
try { ws.close(1000, 'pty exited'); } catch {}
});
if (!maybeSpawnPty(ws, session)) return;
}
try {
// raw is a Uint8Array; Bun.Terminal.write accepts string|Buffer.
@ -548,6 +639,10 @@ function buildServer() {
// Drop the cookie so it can't be replayed against a new PTY.
validTokens.delete(session.cookie);
}
// Drop the sessionId index too (Commit 3 will revisit this for
// re-attach where we WANT the session to outlive the original ws
// for the 60s detach window; for now, close = full dispose).
if (session.sessionId) sessionsById.delete(session.sessionId);
sessions.delete(ws);
}
},

View File

@ -0,0 +1,90 @@
import { describe, test, expect } from 'bun:test';
import * as fs from 'fs';
import * as path from 'path';
// v1.44 Commit 2 — terminal-agent sessionId routing + eager spawn.
//
// Live spawn tests would require a real claude binary on PATH and a Bun.serve
// listener; both are e2e-tier. These static-grep tripwires defend the load-
// bearing protocol changes:
// - validTokens carries the sessionId binding (Map, not Set)
// - sessionsById index exists for /internal/restart + (Commit 3) re-attach
// - /internal/restart is scoped to one sessionId (codex T2 fix)
// - {type:"start"} triggers spawn for eager UX after forceRestart
// - maybeSpawnPty helper is the single entry point for both spawn paths
const AGENT_TS = path.resolve(new URL(import.meta.url).pathname, '..', '..', 'src', 'terminal-agent.ts');
describe('terminal-agent session routing (v1.44+ Commit 2)', () => {
test('1. validTokens is a Map binding token → sessionId', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
// Pre-Commit 2 was `Set<string>`; the Map carries the sessionId
// binding that /internal/restart and (Commit 3) re-attach depend on.
expect(src).toMatch(/const validTokens = new Map<string, string \| null>\(\)/);
expect(src).not.toMatch(/const validTokens = new Set</);
});
test('2. sessionsById reverse index exists', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
expect(src).toMatch(/const sessionsById = new Map<string, PtySession>\(\)/);
// Populated in open() — required so /internal/restart can find the session.
expect(src).toMatch(/if \(sessionId\) sessionsById\.set\(sessionId, session\)/);
});
test('3. /internal/grant binds an optional sessionId to the token', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
const block = sliceBetween(src, "url.pathname === '/internal/grant'", "url.pathname === '/internal/revoke'");
expect(block).toContain('validTokens.set(body.token, sid)');
expect(block).toContain('body?.sessionId');
});
test('4. /internal/restart is scoped to one sessionId, not dispose-all', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
const block = sliceBetween(src, "url.pathname === '/internal/restart'", "// /claude-available");
expect(block).toContain('sessionsById.get(sid)');
expect(block).toContain('disposeSession(session)');
expect(block).toContain('sessionsById.delete(sid)');
// Negative: must NOT enumerate all live sessions and dispose them
// (codex T2 caught this — pre-spec the route killed every PTY on the
// agent, breaking multi-sidebar / pair-agent setups).
expect(block).not.toMatch(/for\s*\(\s*const\s+\[?ws/);
});
test('5. WS upgrade surfaces sessionId on ws.data', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
expect(src).toContain('validTokens.get(token) ?? null');
expect(src).toMatch(/data:\s*\{\s*cookie:\s*token,\s*sessionId\s*\}/);
});
test('6. eager spawn via {type:"start"} text frame', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
expect(src).toMatch(/msg\?\.type === 'start'/);
// Both spawn paths route through the same helper for parity.
expect(src).toContain('function maybeSpawnPty(');
expect(src).toMatch(/maybeSpawnPty\(ws, session\)/);
});
test('7. close() drops sessionsById entry alongside ws cleanup', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
const block = sliceBetween(src, 'close(ws) {', 'function handleTabState');
expect(block).toContain('sessionsById.delete(session.sessionId)');
});
test('8. PtySession interface carries the sessionId field', () => {
const src = fs.readFileSync(AGENT_TS, 'utf-8');
// Whole interface — close paren is sufficient.
const i = src.indexOf('interface PtySession {');
expect(i).toBeGreaterThan(-1);
const j = src.indexOf('\n}', i);
const block = src.slice(i, j);
expect(block).toContain('sessionId: string | null');
});
});
function sliceBetween(source: string, start: string, end: string): string {
const i = source.indexOf(start);
if (i === -1) throw new Error(`marker not found: ${start}`);
const j = source.indexOf(end, i + start.length);
if (j === -1) throw new Error(`end marker not found: ${end}`);
return source.slice(i, j);
}