From 3ea18cef60e6f32e123963cb0a7f14239d94d7a9 Mon Sep 17 00:00:00 2001 From: Ali Mohammad Date: Tue, 3 Nov 2026 06:27:53 +0300 Subject: [PATCH] feat(studio): AA04 reattach ring buffer and chat.resume A client that loses its socket mid-turn reconnects, sends chat.resume, and gets the missed events of a running turn; a settled turn is never replayed (it could resurrect a card saved later). - shared/authorWire: chat.resume gains optional afterSeq; GET /assistant/sessions/:id now returns AuthorSessionRecord (AuthorSession + runningRequestId) - server/author/replayBuffer.ts: per-session ring of the last turn's events (REPLAY_CAPACITY 4000, evicting text.delta first); settled entries answer only the terminal event and expire after 120 s - server/author/authorStream.ts: records every runner event and answers chat.resume (replay tail / terminal event / UNKNOWN_REQUEST) - server/author/routes.ts: GET /:id returns the record; runtimeClient getAuthorSession typed accordingly - src/author/reattach.ts: streamEndRecovery, acceptEvent seq dedupe, pending-request localStorage helpers, reattach constants Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../server/__tests__/authorSessions.test.ts | 5 +- .../server/__tests__/authorStream.test.ts | 160 +++++++++++++++++- .../author/__tests__/replayBuffer.test.ts | 150 ++++++++++++++++ .../agent-forge/server/author/authorStream.ts | 79 ++++++++- .../agent-forge/server/author/replayBuffer.ts | 130 ++++++++++++++ apps/agent-forge/server/author/routes.ts | 15 +- apps/agent-forge/shared/authorWire.ts | 21 ++- .../src/author/__tests__/reattach.test.ts | 109 ++++++++++++ apps/agent-forge/src/author/authorStream.ts | 4 - apps/agent-forge/src/author/reattach.ts | 118 +++++++++++++ apps/agent-forge/src/runtime/runtimeClient.ts | 9 +- 11 files changed, 769 insertions(+), 31 deletions(-) create mode 100644 apps/agent-forge/server/author/__tests__/replayBuffer.test.ts create mode 100644 apps/agent-forge/server/author/replayBuffer.ts create mode 100644 apps/agent-forge/src/author/__tests__/reattach.test.ts create mode 100644 apps/agent-forge/src/author/reattach.ts diff --git a/apps/agent-forge/server/__tests__/authorSessions.test.ts b/apps/agent-forge/server/__tests__/authorSessions.test.ts index 7295e8ff..12bc1fa5 100644 --- a/apps/agent-forge/server/__tests__/authorSessions.test.ts +++ b/apps/agent-forge/server/__tests__/authorSessions.test.ts @@ -116,11 +116,12 @@ describe('author session routes', () => { expect('transcript' in list[0]).toBe(false); }); - it('GET /assistant/sessions/:id returns the session; an unknown id -> 404 NOT_FOUND', async () => { + it('GET /assistant/sessions/:id returns the session record; an unknown id -> 404 NOT_FOUND', async () => { const session = await createSession(); const res = await authed(request(app).get(`/assistant/sessions/${session.id}`)); expect(res.status).toBe(200); - expect(res.body).toEqual(session); + // The record is the stored session plus the running turn's requestId (none here). + expect(res.body).toEqual({ ...session, runningRequestId: null }); const missing = await authed(request(app).get('/assistant/sessions/as_0000000000000000')); expect(missing.status).toBe(404); diff --git a/apps/agent-forge/server/__tests__/authorStream.test.ts b/apps/agent-forge/server/__tests__/authorStream.test.ts index b1a44553..2e9b4cf3 100644 --- a/apps/agent-forge/server/__tests__/authorStream.test.ts +++ b/apps/agent-forge/server/__tests__/authorStream.test.ts @@ -3,7 +3,8 @@ import * as fs from 'node:fs'; import * as os from 'node:os'; import * as path from 'node:path'; import { WebSocket } from 'ws'; -import { mockModel, type MockModel } from 'lousho/testing'; +import { mockModel } from 'lousho/testing'; +import type { LLMProvider, StreamChunk } from 'lousho'; import { startStudioServer, type StudioServerHandle } from '../index'; import { STUDIO_TOKEN_HEADER } from '../auth'; import { AuthorSessionStore } from '../author/sessionStore'; @@ -16,6 +17,67 @@ import type { AuthorServerEvent, AuthorSession } from '../../shared/authorWire'; const TOKEN = 'author-stream-token-0123456'; const REQ_A = 'req-aa03-0001'; const REQ_B = 'req-aa03-0002'; +const REQ_RESUME = 'req-aa04-0001'; + +/** A promise the test resolves by hand (the gated model's stream pauses on it). */ +function deferred(): { promise: Promise; release: () => void } { + let release!: () => void; + const promise = new Promise((resolve) => { + release = resolve; + }); + return { promise, release }; +} + +/** + * A provider whose one stream emits `first` deltas, waits on `gateA`, + * emits `second` deltas (resolving `emittedSecond` once the last was + * yielded), waits on `gateB`, then emits `last` and finishes. Lets the + * test park a turn mid-flight, drop the client, and resume against the + * replay buffer while the model is still "typing". + */ +function gatedModel(script: { + first: string[]; + second: string[]; + last: string[]; + gateA: Promise; + gateB: Promise; +}): { provider: LLMProvider; emittedSecond: Promise } { + let markSecond!: () => void; + const emittedSecond = new Promise((resolve) => { + markSecond = resolve; + }); + const full = [...script.first, ...script.second, ...script.last].join(''); + const provider: LLMProvider = { + name: 'gated-mock', + defaultModel: 'gated-model', + generate: () => Promise.reject(new Error('generate() is not used by the Author')), + stream: () => { + const fullStream = (async function* (): AsyncGenerator { + for (const text of script.first) yield { type: 'text-delta', textDelta: text }; + await script.gateA; + for (const text of script.second) yield { type: 'text-delta', textDelta: text }; + markSecond(); + await script.gateB; + for (const text of script.last) yield { type: 'text-delta', textDelta: text }; + yield { type: 'finish', finishReason: 'stop' }; + })(); + return Promise.resolve({ + fullStream, + textStream: (async function* () { + yield full; + })(), + text: Promise.resolve(full), + usage: Promise.resolve(undefined), + finishReason: Promise.resolve('stop'), + toolCalls: Promise.resolve([]), + }); + }, + supportsTools: () => true, + supportsStreaming: () => true, + getModels: () => Promise.resolve(['gated-model']), + }; + return { provider, emittedSecond }; +} /** * `WS /assistant/stream` end to end: a real `startStudioServer` with the @@ -26,16 +88,17 @@ const REQ_B = 'req-aa03-0002'; describe('WS /assistant/stream (startStudioServer)', () => { let baseDir: string; let handle: StudioServerHandle; - let model: MockModel; + let model: LLMProvider; let origin: string; let session: AuthorSession; let otherSession: AuthorSession; + let runner: AuthorTurnRunner; beforeAll(async () => { baseDir = fs.mkdtempSync(path.join(os.tmpdir(), 'studio-aa03-')); const home = path.join(baseDir, '.lousho'); model = mockModel(['Hi']); - const runner = new AuthorTurnRunner({ + runner = new AuthorTurnRunner({ sessions: new AuthorSessionStore(studioPathsFor(home).sessions), settings: new SettingsStore(home), secrets: new SecretsStore(home), @@ -230,11 +293,17 @@ describe('WS /assistant/stream (startStudioServer)', () => { message: 'Studio could not read that message.', code: 'BAD_MESSAGE', }); - // chat.resume is AA04's reattach: AA03 answers it like any unknown frame, - // echoing the requestId it carried. + // chat.resume for a requestId the replay buffer never saw answers + // UNKNOWN_REQUEST on this socket only (seq 0, per-socket reply). a.ws.send(JSON.stringify({ type: 'chat.resume', requestId: 'req-resume-01' })); const resume = await waitForEvent(a.events, (e) => e.type === 'turn.error' && e.requestId === 'req-resume-01'); - expect(resume).toMatchObject({ type: 'turn.error', requestId: 'req-resume-01', seq: 0, code: 'BAD_MESSAGE' }); + expect(resume).toEqual({ + type: 'turn.error', + requestId: 'req-resume-01', + seq: 0, + message: 'This reply is no longer available. Showing the saved chat.', + code: 'UNKNOWN_REQUEST', + }); // A structurally wrong message with a requestId echoes it too. a.ws.send(JSON.stringify({ type: 'chat.send', requestId: 'req-nope-0001' })); const noText = await waitForEvent(a.events, (e) => e.type === 'turn.error' && e.requestId === 'req-nope-0001'); @@ -246,4 +315,83 @@ describe('WS /assistant/stream (startStudioServer)', () => { b.ws.close(); } }); + + it('chat.resume mid-turn replays seq > afterSeq, then the live tail, with no duplicates', async () => { + const gateA = deferred(); + const gateB = deferred(); + const gated = gatedModel({ first: ['aa', 'bb'], second: ['cc', 'dd'], last: ['ee'], gateA: gateA.promise, gateB: gateB.promise }); + model = gated.provider; + const a = await connect(session.id); + let b: StreamClient | undefined; + try { + chatSend(a.ws, REQ_RESUME, 'talk'); + // seq 1 turn.started, seqs 2-3 the two first-stage deltas. + await waitForEvent(a.events, (e) => e.seq === 3); + a.ws.close(); + // The turn keeps running with no subscribers: two more deltas land + // in the replay buffer, then the model parks on gateB. Wait on the + // runner's own emit for seq 5 (recorded before the buffer replays). + gateA.release(); + await new Promise((resolve, reject) => { + const deadline = setTimeout(() => reject(new Error('seq 5 never emitted')), 5000); + const check = (_s: string, e: AuthorServerEvent) => { + if (e.requestId === REQ_RESUME && e.seq === 5) { + runner.off('event', check); + clearTimeout(deadline); + resolve(); + } + }; + runner.on('event', check); + }); + await gated.emittedSecond; + // While it runs, the session record reports the in-flight requestId. + const record = (await ( + await fetch(`${origin}/assistant/sessions/${session.id}`, { headers: { [STUDIO_TOKEN_HEADER]: TOKEN } }) + ).json()) as { runningRequestId: string | null }; + expect(record.runningRequestId).toBe(REQ_RESUME); + + // The dropped client reconnects and resumes from its last seq. + b = await connect(session.id); + b.ws.send(JSON.stringify({ type: 'chat.resume', requestId: REQ_RESUME, afterSeq: 3 })); + // The replayed buffer tail: seqs 4-5 only. + await waitForEvent(b.events, (e) => e.seq === 5); + expect(b.events.map((e) => e.seq)).toEqual([4, 5]); + // gateB releases the live tail: delta seq 6, turn.finished seq 7. + gateB.release(); + await waitForEvent(b.events, (e) => e.type === 'turn.finished'); + expect(b.events.map((e) => e.seq)).toEqual([4, 5, 6, 7]); + expect(b.events.slice(0, 3).every((e) => e.type === 'text.delta')).toBe(true); + expect(b.events[3]).toMatchObject({ type: 'turn.finished', requestId: REQ_RESUME, finishReason: 'stop' }); + } finally { + gateA.release(); + gateB.release(); + a.ws.close(); + b?.ws.close(); + } + }); + + it('chat.resume after the turn settled sends only the terminal event', async () => { + model = mockModel(['quick reply']); + const a = await connect(session.id); + let b: StreamClient | undefined; + try { + chatSend(a.ws, REQ_A, 'hi'); + const end = await waitForEvent(a.events, (e) => e.type === 'turn.finished'); + // The reconnect happens after the turn settled: the resume answer + // is the terminal event alone - never the turn's deltas again. + b = await connect(session.id); + b.ws.send(JSON.stringify({ type: 'chat.resume', requestId: REQ_A, afterSeq: 0 })); + await waitForEvent(b.events, (e) => e.type === 'turn.finished' || e.type === 'turn.error'); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(b.events).toEqual([end]); + // And the session record reports no running reply. + const record = (await ( + await fetch(`${origin}/assistant/sessions/${session.id}`, { headers: { [STUDIO_TOKEN_HEADER]: TOKEN } }) + ).json()) as { runningRequestId: string | null }; + expect(record.runningRequestId).toBeNull(); + } finally { + a.ws.close(); + b?.ws.close(); + } + }); }); diff --git a/apps/agent-forge/server/author/__tests__/replayBuffer.test.ts b/apps/agent-forge/server/author/__tests__/replayBuffer.test.ts new file mode 100644 index 00000000..6165704d --- /dev/null +++ b/apps/agent-forge/server/author/__tests__/replayBuffer.test.ts @@ -0,0 +1,150 @@ +import { describe, it, expect, beforeEach } from 'vitest'; +import { + ReplayBuffer, + REPLAY_CAPACITY, + REPLAY_SETTLED_TTL_MS, +} from '../replayBuffer'; +import type { AuthorServerEvent } from '../../../shared/authorWire'; + +/** + * ReplayBuffer: the per-session ring of one turn's `AuthorServerEvent`s a + * dropped client re-fetches with `chat.resume`. A running turn answers + * `replay` (events after the client's cursor), a settled turn answers only + * its terminal event (never a replay - it could re-show a card saved + * later), and anything else is `unknown`. The `now` seam stands in for a + * fake clock so the settled-TTL test doesn't wait two minutes. + */ +describe('ReplayBuffer', () => { + const SESSION = 'sess-1'; + const REQ = 'req-aa04-0001'; + let now: number; + let buffer: ReplayBuffer; + + beforeEach(() => { + now = 1_000_000; + buffer = new ReplayBuffer({ now: () => now }); + }); + + const started = (seq: number, requestId = REQ): AuthorServerEvent => ({ + type: 'turn.started', + requestId, + seq, + }); + const delta = (seq: number): AuthorServerEvent => ({ + type: 'text.delta', + requestId: REQ, + seq, + text: `d${seq}`, + }); + const toolCall = (seq: number): AuthorServerEvent => ({ + type: 'tool.call', + requestId: REQ, + seq, + toolCallId: `call_${seq}`, + name: 'listAgents', + input: {}, + }); + const toolResult = (seq: number): AuthorServerEvent => ({ + type: 'tool.result', + requestId: REQ, + seq, + toolCallId: `call_${seq}`, + name: 'listAgents', + output: { ok: true }, + state: 'done', + }); + const finished = (seq: number): AuthorServerEvent => ({ + type: 'turn.finished', + requestId: REQ, + seq, + finishReason: 'stop', + }); + + /** Records one running turn: turn.started seq 1, then `deltas` text.delta events (seqs 2..n). */ + function recordRunning(deltas: number, sessionId = SESSION, requestId = REQ): AuthorServerEvent[] { + const events: AuthorServerEvent[] = [started(1, requestId)]; + for (let seq = 2; seq <= deltas + 1; seq += 1) events.push(delta(seq)); + for (const event of events) buffer.record(sessionId, event); + return events; + } + + it('replays only the events after the resume cursor while the turn runs', () => { + const events = recordRunning(4); // seqs 1..5 + expect(buffer.resume(SESSION, REQ, 2)).toEqual({ kind: 'replay', events: events.slice(2) }); + expect(buffer.resume(SESSION, REQ, 0)).toEqual({ kind: 'replay', events }); + expect(buffer.resume(SESSION, REQ, 99)).toEqual({ kind: 'replay', events: [] }); + }); + + it('answers unknown for a requestId it never saw, another session, or a replaced turn', () => { + recordRunning(2); + expect(buffer.resume(SESSION, 'req-never-seen', 0)).toEqual({ kind: 'unknown' }); + expect(buffer.resume('other-session', REQ, 0)).toEqual({ kind: 'unknown' }); + // A new turn.started replaces the session's entry entirely. + buffer.record(SESSION, started(1, 'req-aa04-0002')); + expect(buffer.resume(SESSION, REQ, 0)).toEqual({ kind: 'unknown' }); + expect(buffer.resume(SESSION, 'req-aa04-0002', 0)).toEqual({ + kind: 'replay', + events: [started(1, 'req-aa04-0002')], + }); + }); + + it('answers a settled turn with only its terminal event, never a replay', () => { + recordRunning(3); + const end = finished(5); + buffer.record(SESSION, end); + expect(buffer.resume(SESSION, REQ, 0)).toEqual({ kind: 'settled', final: end }); + expect(buffer.resume(SESSION, REQ, 4)).toEqual({ kind: 'settled', final: end }); + }); + + it('treats a runner turn.error as the settled terminal event', () => { + recordRunning(1); + const end: AuthorServerEvent = { type: 'turn.error', requestId: REQ, seq: 3, message: 'boom' }; + buffer.record(SESSION, end); + expect(buffer.resume(SESSION, REQ, 0)).toEqual({ kind: 'settled', final: end }); + }); + + it('answers unknown once a settled entry outlives REPLAY_SETTLED_TTL_MS', () => { + recordRunning(1); + buffer.record(SESSION, finished(3)); + now += REPLAY_SETTLED_TTL_MS; // exactly at the TTL it still answers + expect(buffer.resume(SESSION, REQ, 0).kind).toBe('settled'); + now += 1; + expect(buffer.resume(SESSION, REQ, 0)).toEqual({ kind: 'unknown' }); + // An expired entry is also evicted lazily by the next record on any session. + recordRunning(0, 'sess-2', 'req-aa04-0003'); + expect(buffer.resume(SESSION, REQ, 0)).toEqual({ kind: 'unknown' }); + }); + + it('keeps a running turn past the TTL - only settled entries expire', () => { + recordRunning(2); + now += REPLAY_SETTLED_TTL_MS * 10; + expect(buffer.resume(SESSION, REQ, 0).kind).toBe('replay'); + }); + + it('on overflow drops the oldest text.delta events first and keeps every tool.* event', () => { + // seq 1 = turn.started, seqs 2..2001 deltas, 2002..3501 tool.call, + // 3502..4501 tool.result, 4502..5101 deltas: 5 101 events, 1 101 over. + const events: AuthorServerEvent[] = [started(1)]; + for (let seq = 2; seq <= 2001; seq += 1) events.push(delta(seq)); + for (let seq = 2002; seq <= 3501; seq += 1) events.push(toolCall(seq)); + for (let seq = 3502; seq <= 4501; seq += 1) events.push(toolResult(seq)); + for (let seq = 4502; seq <= 5101; seq += 1) events.push(delta(seq)); + for (const event of events) buffer.record(SESSION, event); + + const answer = buffer.resume(SESSION, REQ, 0); + if (answer.kind !== 'replay') throw new Error(`expected replay, got ${answer.kind}`); + expect(answer.events.length).toBe(REPLAY_CAPACITY); + // turn.started survived (it is not a text.delta). + expect(answer.events[0]).toEqual(started(1)); + // Every tool.* event survived; only text.delta events were evicted. + const tools = answer.events.filter((e) => e.type === 'tool.call' || e.type === 'tool.result'); + expect(tools.length).toBe(2500); + const deltas = answer.events.filter((e) => e.type === 'text.delta'); + // 2600 recorded deltas - 1101 evictions = 1499 left, oldest gone first. + expect(deltas.length).toBe(1499); + expect(deltas[0].seq).toBe(1103); + // Emission order is preserved after eviction (still sorted by seq). + const seqs = answer.events.map((e) => e.seq); + expect(seqs).toEqual([...seqs].sort((a, b) => a - b)); + }); +}); diff --git a/apps/agent-forge/server/author/authorStream.ts b/apps/agent-forge/server/author/authorStream.ts index 737cd385..8709b21a 100644 --- a/apps/agent-forge/server/author/authorStream.ts +++ b/apps/agent-forge/server/author/authorStream.ts @@ -4,9 +4,12 @@ * One socket subscribes to one session's turns: the server forwards every * `AuthorServerEvent` the `AuthorTurnRunner` emits for that session (all * subscribers of the same session see the same frames), and accepts - * `chat.send`/`chat.stop` frames back. `chat.resume` (AA04's reattach) and - * any unreadable or unknown frame get a `BAD_MESSAGE` `turn.error` on the - * sending socket only - one client's garbage never reaches the others. + * `chat.send`/`chat.stop`/`chat.resume` frames back. `chat.resume` is + * answered from the session's `ReplayBuffer` entry - the missed events of + * a running turn, the terminal event of a settled one, or an + * `UNKNOWN_REQUEST` `turn.error` - while any unreadable or unknown frame + * gets a `BAD_MESSAGE` `turn.error`. Both land on the sending socket only: + * one client's garbage never reaches the others. * * Registered on the server's one `UpgradeRouter` under * {@link AUTHOR_STREAM_PATH}, so the shared Host/Origin/token gate has @@ -27,6 +30,7 @@ import { isValidAgentId } from '../../shared/agentId'; import type { UpgradeHandler } from '../ws/upgradeRouter'; import type { AuthorSessionStore } from './sessionStore'; import type { AuthorTurnRunner } from './turnRunner'; +import { ReplayBuffer } from './replayBuffer'; /** * The close code an unknown or malformed `session` id gets after the @@ -39,11 +43,19 @@ const SESSION_NOT_FOUND_REASON = 'Chat not found'; /** The BAD_MESSAGE `turn.error`'s text: one fixed sentence, never an echo of the frame. */ const BAD_MESSAGE = 'Studio could not read that message.'; -/** What `chat.resume` (and every unknown frame) is answered with until AA04 lands it. */ +/** What `chat.resume` gets when the replay buffer holds no entry for the `requestId`. */ +const UNKNOWN_REQUEST = 'This reply is no longer available. Showing the saved chat.'; + +/** What every unreadable or unknown frame is answered with (seq 0: a per-socket reply, not a turn event). */ function badMessage(requestId: string): AuthorServerEvent { return { type: 'turn.error', requestId, seq: 0, message: BAD_MESSAGE, code: 'BAD_MESSAGE' }; } +/** The `chat.resume` answer for a `requestId` the buffer never saw (or whose settled entry expired). */ +function unknownRequest(requestId: string): AuthorServerEvent { + return { type: 'turn.error', requestId, seq: 0, message: UNKNOWN_REQUEST, code: 'UNKNOWN_REQUEST' }; +} + function send(ws: WebSocket, event: AuthorServerEvent): void { if (ws.readyState === WebSocket.OPEN) { ws.send(JSON.stringify(event)); @@ -65,8 +77,7 @@ function parseFocus(value: unknown): { ok: true; focus?: AuthorFocus } | { ok: f /** * One raw text frame as a client message, or `undefined` when it isn't * strict JSON of a known shape (wrong `type`, missing `requestId`, - * non-string `text`, malformed `focus`). `chat.resume` parses fine - it is - * a known type - and the dispatch answers it BAD_MESSAGE until AA04. + * non-string `text`, malformed `focus`, malformed `afterSeq`). */ function parseClientMessage(raw: string): AuthorClientMessage | undefined { let data: unknown; @@ -82,7 +93,7 @@ function parseClientMessage(raw: string): AuthorClientMessage | undefined { case 'chat.stop': return { type: 'chat.stop', requestId: message.requestId }; case 'chat.resume': - return { type: 'chat.resume', requestId: message.requestId }; + return parseResume(message as { requestId: string; afterSeq?: unknown }); case 'chat.send': { if (typeof message.text !== 'string') return undefined; const parsed = parseFocus(message.focus); @@ -96,6 +107,20 @@ function parseClientMessage(raw: string): AuthorClientMessage | undefined { } } +/** + * A `chat.resume` frame with `requestId` already checked: `afterSeq` is + * optional but, when present, must be a non-negative integer (the last + * seq the client applied) - anything else is a BAD_MESSAGE like every + * other malformed field. + */ +function parseResume(message: { requestId: string; afterSeq?: unknown }): AuthorClientMessage | undefined { + if (message.afterSeq === undefined) return { type: 'chat.resume', requestId: message.requestId }; + if (typeof message.afterSeq !== 'number' || !Number.isInteger(message.afterSeq) || message.afterSeq < 0) { + return undefined; + } + return { type: 'chat.resume', requestId: message.requestId, afterSeq: message.afterSeq }; +} + /** Best-effort `requestId` out of a frame that failed to parse, so its BAD_MESSAGE still lands in the right chat bubble. */ function requestIdIn(raw: string): string { try { @@ -126,6 +151,9 @@ export function createAuthorUpgradeHandler(deps: { sessions: AuthorSessionStore; }): UpgradeHandler { const wss = new WebSocketServer({ noServer: true, maxPayload: 64 * 1024 }); + // The server's one replay buffer: every runner event lands here before + // it broadcasts, so a `chat.resume` always sees the turn's full tail. + const replay = new ReplayBuffer(); // sessionId -> connected sockets: a turn's events fan out to every tab // watching that chat and to no one else. const subscribers = new Map>(); @@ -136,7 +164,10 @@ export function createAuthorUpgradeHandler(deps: { for (const ws of sockets) send(ws, event); } - deps.runner.on('event', (sessionId: string, event: AuthorServerEvent) => broadcast(sessionId, event)); + deps.runner.on('event', (sessionId: string, event: AuthorServerEvent) => { + replay.record(sessionId, event); + broadcast(sessionId, event); + }); /** `chat.send`: run the turn; a refused start is a `turn.error` (seq 0) on the sending socket only. */ async function onChatSend( @@ -170,6 +201,31 @@ export function createAuthorUpgradeHandler(deps: { } } + /** + * `chat.resume`: answer from the replay buffer on the asking socket + * only. `replay` sends the missed events in order (the socket is + * already subscribed, so live events simply follow); `settled` sends + * only the terminal event - a settled turn is never replayed (its card + * may have been saved since); `unknown` is the UNKNOWN_REQUEST + * `turn.error` the client reconciles from. + */ + function onChatResume( + ws: WebSocket, + sessionId: string, + message: Extract + ): void { + const answer = replay.resume(sessionId, message.requestId, message.afterSeq ?? 0); + if (answer.kind === 'replay') { + for (const event of answer.events) send(ws, event); + return; + } + if (answer.kind === 'settled') { + send(ws, answer.final); + return; + } + send(ws, unknownRequest(message.requestId)); + } + /** One incoming text frame of an already-subscribed socket. */ function onFrame(ws: WebSocket, sessionId: string, raw: string): void { const message = parseClientMessage(raw); @@ -184,8 +240,13 @@ export function createAuthorUpgradeHandler(deps: { case 'chat.stop': deps.runner.stop(sessionId, message.requestId); return; + case 'chat.resume': + onChatResume(ws, sessionId, message); + return; default: - send(ws, badMessage(message.requestId)); + // Unreachable over the current union; kept so a widened + // AuthorClientMessage type is a BAD_MESSAGE, not a silent drop. + send(ws, badMessage(requestIdIn(raw))); } } diff --git a/apps/agent-forge/server/author/replayBuffer.ts b/apps/agent-forge/server/author/replayBuffer.ts new file mode 100644 index 00000000..87ada8b7 --- /dev/null +++ b/apps/agent-forge/server/author/replayBuffer.ts @@ -0,0 +1,130 @@ +/** + * The reattach ring buffer (PLAN 8.1): one entry per session holding the + * last turn's `AuthorServerEvent`s in emission order, so a client that + * lost its socket mid-turn can send `chat.resume` and get exactly the + * events it missed (`seq` is per turn, so `afterSeq` is the cursor). + * + * A settled turn is never replayed: its entry answers `settled` with only + * the terminal `turn.finished`/`turn.error` event (replaying deltas could + * resurrect a proposal card that was saved after the turn ended - the + * client reconciles from the stored transcript instead). Settled entries + * linger for `REPLAY_SETTLED_TTL_MS` so a client that dropped just before + * the last event can still learn the turn ended, then they are evicted + * lazily on the next `record`/`resume`. + */ +import type { AuthorServerEvent } from '../../shared/authorWire'; + +/** The most events one turn's entry keeps; overflow evicts `text.delta`s first. */ +export const REPLAY_CAPACITY = 4000; + +/** How long a settled turn's terminal event stays answerable after the turn ended. */ +export const REPLAY_SETTLED_TTL_MS = 120_000; + +/** What `resume` says the reconnecting socket should be sent. */ +export type ResumeAnswer = + /** The turn still runs: every buffered event with `seq > afterSeq` (live events follow on the same socket). */ + | { kind: 'replay'; events: AuthorServerEvent[] } + /** The turn ended: only its `turn.finished`/`turn.error` - the deltas are deliberately withheld. */ + | { kind: 'settled'; final: AuthorServerEvent } + /** No such turn: unknown requestId, another turn owns the session now, or the entry expired. */ + | { kind: 'unknown' }; + +/** One session's replay state: the running (or just-settled) turn's events. */ +interface ReplayEntry { + requestId: string; + events: AuthorServerEvent[]; + settled: boolean; + /** `now()` at the terminal event (the TTL counts from here). */ + settledAt: number; + /** Overflow evicted old `text.delta`s: the replayed tail is incomplete, so the client reconciles at the terminal event. */ + truncated: boolean; +} + +function isTerminal(event: AuthorServerEvent): boolean { + return event.type === 'turn.finished' || event.type === 'turn.error'; +} + +/** + * Records every event `AuthorTurnRunner` emits (`authorStream.ts` feeds it + * on the runner's `'event'` channel, one buffer per server) and answers + * `chat.resume` from it. `deps.now` is the clock seam for tests. + */ +export class ReplayBuffer { + private readonly entries = new Map(); + private readonly now: () => number; + + constructor(deps?: { now?: () => number }) { + this.now = deps?.now ?? (() => Date.now()); + } + + /** + * Adds `event` to the session's entry. `turn.started` always begins a + * fresh entry (one running turn per session, so the previous entry - + * running or settled - is dead the moment a new turn starts); any later + * event is appended only while it names the entry's own `requestId`. + * The terminal event flips the entry to settled and starts its TTL. + */ + record(sessionId: string, event: AuthorServerEvent): void { + this.evictExpired(); + let entry = this.entries.get(sessionId); + if (event.type === 'turn.started') { + entry = { requestId: event.requestId, events: [], settled: false, settledAt: 0, truncated: false }; + this.entries.set(sessionId, entry); + } + if (entry === undefined || entry.requestId !== event.requestId) return; + entry.events.push(event); + if (isTerminal(event)) { + entry.settled = true; + entry.settledAt = this.now(); + } + this.enforceCapacity(entry); + } + + /** + * What a `chat.resume { requestId, afterSeq }` gets back for `sessionId`: + * the buffered tail of a running turn, the terminal event of a settled + * one, or `unknown` when the buffer holds no entry for that request. + */ + resume(sessionId: string, requestId: string, afterSeq: number): ResumeAnswer { + this.evictExpired(); + const entry = this.entries.get(sessionId); + if (entry === undefined || entry.requestId !== requestId) return { kind: 'unknown' }; + if (entry.settled) { + const final = entry.events.at(-1); + // A settled entry always holds its terminal event (the last push); + // `undefined` is only reachable through pathological truncation. + return final === undefined ? { kind: 'unknown' } : { kind: 'settled', final }; + } + return { kind: 'replay', events: entry.events.filter((event) => event.seq > afterSeq) }; + } + + /** Lazily deletes settled entries past the TTL (called from `record`/`resume`, not a timer). */ + private evictExpired(): void { + const now = this.now(); + for (const [sessionId, entry] of this.entries) { + if (entry.settled && now - entry.settledAt > REPLAY_SETTLED_TTL_MS) { + this.entries.delete(sessionId); + } + } + } + + /** + * Caps the entry at `REPLAY_CAPACITY`: evict the oldest `text.delta` + * first (deltas are re-renderable text - losing some leaves a gap the + * client repairs on reconcile), and only when no delta remains evict + * the oldest event at all, so every `tool.*` and the terminal event + * survive any realistic turn. `truncated` marks that the replay tail + * has a hole. + */ + private enforceCapacity(entry: ReplayEntry): void { + while (entry.events.length > REPLAY_CAPACITY) { + const deltaIndex = entry.events.findIndex((event) => event.type === 'text.delta'); + if (deltaIndex === -1) { + entry.events.shift(); + } else { + entry.events.splice(deltaIndex, 1); + } + entry.truncated = true; + } + } +} diff --git a/apps/agent-forge/server/author/routes.ts b/apps/agent-forge/server/author/routes.ts index 2112992f..851ebce0 100644 --- a/apps/agent-forge/server/author/routes.ts +++ b/apps/agent-forge/server/author/routes.ts @@ -14,7 +14,7 @@ * is the persisted-session CRUD they all build on. */ import type { Express, Request, Response } from 'express'; -import { AUTHOR_SESSION_ID_RE, type AuthorFocus } from '../../shared/authorWire'; +import { AUTHOR_SESSION_ID_RE, type AuthorFocus, type AuthorSessionRecord } from '../../shared/authorWire'; import { isValidAgentId } from '../../shared/agentId'; import { asyncRoute, paramId, sendError } from '../http/routeUtils'; import { workspaceOf } from '../workspace/middleware'; @@ -60,7 +60,8 @@ function sessionIdParam(req: Request, res: Response): string | undefined { * session belongs to still comes from `workspaceOf(req)` on create. * `deps.runner` (the app's AuthorTurnRunner) makes DELETE refuse while a * reply is running in that chat - deleting the file mid-turn would lose - * the in-flight assistant message. + * the in-flight assistant message - and gives GET `/:id` its + * `runningRequestId` for the reattach decision. */ export function registerAuthorRoutes( app: Express, @@ -88,7 +89,9 @@ export function registerAuthorRoutes( }) ); - // One session with its transcript: 200 AuthorSession, 404 when unknown. + // One session with its transcript: 200 AuthorSessionRecord (the stored + // session plus the running turn's requestId, so a client that reopens + // the chat knows whether to send chat.resume), 404 when unknown. app.get( '/assistant/sessions/:id', asyncRoute(async (req, res) => { @@ -99,7 +102,11 @@ export function registerAuthorRoutes( sendError(res, 404, 'Chat not found.', 'NOT_FOUND'); return; } - res.json(session); + const record: AuthorSessionRecord = { + ...session, + runningRequestId: deps.runner?.running(id)?.requestId ?? null, + }; + res.json(record); }) ); diff --git a/apps/agent-forge/shared/authorWire.ts b/apps/agent-forge/shared/authorWire.ts index 6d8dc6d0..5b558a9c 100644 --- a/apps/agent-forge/shared/authorWire.ts +++ b/apps/agent-forge/shared/authorWire.ts @@ -81,6 +81,16 @@ export interface AuthorSession { createdAgentIds: string[]; } +/** + * `GET /assistant/sessions/:id`'s response: the stored session plus the + * requestId of its running turn, when one is in flight (from + * `AuthorTurnRunner.running`). A client that (re)opens a chat uses it to + * decide whether to send `chat.resume` for the in-flight reply. + */ +export interface AuthorSessionRecord extends AuthorSession { + runningRequestId: string | null; +} + /** The `GET /assistant/sessions` row: a session without its transcript. */ export interface AuthorSessionSummary { id: string; @@ -146,14 +156,17 @@ export const AUTHOR_STREAM_PATH = '/assistant/stream'; /** * Everything a client may send on `WS /assistant/stream`. `chat.send` * starts a turn (the `requestId` ties its events back to this send), - * `chat.stop` ends the running turn early, and `chat.resume` is AA04's - * reattach - the AA03 transport answers it (like any unreadable or - * unknown frame) with a `BAD_MESSAGE` `turn.error`. + * `chat.stop` ends the running turn early, and `chat.resume` reattaches a + * dropped client to its turn: `afterSeq` is the last seq the client + * applied (omitted = replay everything still buffered). The server + * answers from the replay buffer - the missed events while the turn + * still runs, only the terminal event once it settled, or an + * `UNKNOWN_REQUEST` `turn.error` for a requestId it never saw. */ export type AuthorClientMessage = | { type: 'chat.send'; requestId: string; text: string; focus?: AuthorFocus } | { type: 'chat.stop'; requestId: string } - | { type: 'chat.resume'; requestId: string }; + | { type: 'chat.resume'; requestId: string; afterSeq?: number }; /** * Every event one Author turn produces, in order. `AuthorTurnRunner` diff --git a/apps/agent-forge/src/author/__tests__/reattach.test.ts b/apps/agent-forge/src/author/__tests__/reattach.test.ts new file mode 100644 index 00000000..2983626a --- /dev/null +++ b/apps/agent-forge/src/author/__tests__/reattach.test.ts @@ -0,0 +1,109 @@ +import { describe, it, expect, beforeEach, afterEach } from 'vitest'; +import { breakStorage, installMemoryStorage } from '../../ui/__tests__/memoryStorage'; +import { + acceptEvent, + loadPendingRequest, + PENDING_REQUEST_KEY, + REATTACH_GAVE_UP, + REATTACH_MAX_ATTEMPTS, + REATTACH_TIMEOUT_MS, + savePendingRequest, + streamEndRecovery, +} from '../reattach'; +import type { AuthorServerEvent } from '../../../shared/authorWire'; + +/** + * The reattach helpers AA15a's chat hook drives: the stream-end verdict + * (`none`/`drop`/`busy`), the seq dedupe a resumed stream applies, and the + * pending-request marker in localStorage that survives a reload. Every + * storage access is try/catch - a blocked or throwing localStorage must + * never break the chat. + */ +describe('streamEndRecovery', () => { + it.each([ + ['stoppedByUser wins over everything', { wasStreaming: true, socketClosed: true, errorCode: 'BUSY', stoppedByUser: true }, 'none'], + ['a BUSY turn.error means another reply is running', { wasStreaming: false, socketClosed: true, errorCode: 'BUSY', stoppedByUser: false }, 'busy'], + ['a drop mid-stream reattaches', { wasStreaming: true, socketClosed: true, stoppedByUser: false }, 'drop'], + ['anything else needs no recovery', { wasStreaming: true, socketClosed: false, stoppedByUser: false }, 'none'], + ['a clean socket close without streaming needs no recovery', { wasStreaming: false, socketClosed: true, stoppedByUser: false }, 'none'], + ] as const)('%s', (_label, input, expected) => { + expect(streamEndRecovery(input)).toBe(expected); + }); +}); + +describe('acceptEvent', () => { + const REQ = 'req-aa04-0001'; + const delta = (seq: number, requestId = REQ): AuthorServerEvent => ({ + type: 'text.delta', + requestId, + seq, + text: 'x', + }); + + it('accepts a new seq and drops an old or equal one for the same requestId', () => { + expect(acceptEvent(3, REQ, delta(4))).toBe(true); + expect(acceptEvent(3, REQ, delta(3))).toBe(false); + expect(acceptEvent(3, REQ, delta(2))).toBe(false); + }); + + it('accepts any seq from a different requestId, and anything when no turn is active', () => { + expect(acceptEvent(3, REQ, delta(1, 'req-aa04-0002'))).toBe(true); + expect(acceptEvent(3, null, delta(1))).toBe(true); + }); + + it('always accepts seq 0 (per-socket replies bypass dedupe)', () => { + const error: AuthorServerEvent = { + type: 'turn.error', + requestId: REQ, + seq: 0, + message: 'nope', + code: 'UNKNOWN_REQUEST', + }; + expect(acceptEvent(3, REQ, error)).toBe(true); + expect(acceptEvent(3, null, error)).toBe(true); + }); +}); + +describe('reattach constants', () => { + it('pins the values the hook retries and gives up with', () => { + expect(REATTACH_TIMEOUT_MS).toBe(5000); + expect(REATTACH_MAX_ATTEMPTS).toBe(3); + expect(REATTACH_GAVE_UP).toContain('Could not reconnect'); + expect(PENDING_REQUEST_KEY).toBe('forge.author.pendingRequest'); + }); +}); + +describe('pending-request storage', () => { + beforeEach(() => { + installMemoryStorage(); + }); + + afterEach(() => { + window.localStorage.removeItem(PENDING_REQUEST_KEY); + }); + + it('round-trips a PendingRequest and clears it on null', () => { + const pending = { sessionId: 'sess-1', requestId: 'req-aa04-0001', lastSeq: 7 }; + expect(loadPendingRequest()).toBeNull(); + savePendingRequest(pending); + expect(loadPendingRequest()).toEqual(pending); + savePendingRequest(null); + expect(loadPendingRequest()).toBeNull(); + }); + + it('returns null for malformed stored JSON or a wrong shape', () => { + window.localStorage.setItem(PENDING_REQUEST_KEY, '{not json'); + expect(loadPendingRequest()).toBeNull(); + window.localStorage.setItem(PENDING_REQUEST_KEY, JSON.stringify({ sessionId: 's' })); + expect(loadPendingRequest()).toBeNull(); + window.localStorage.setItem(PENDING_REQUEST_KEY, JSON.stringify('a string')); + expect(loadPendingRequest()).toBeNull(); + }); + + it('survives a throwing localStorage (save does not throw, load returns null)', () => { + breakStorage(); // every getItem/setItem throws, as blocked site data does + expect(() => savePendingRequest({ sessionId: 's', requestId: 'r', lastSeq: 1 })).not.toThrow(); + expect(loadPendingRequest()).toBeNull(); + expect(() => savePendingRequest(null)).not.toThrow(); + }); +}); diff --git a/apps/agent-forge/src/author/authorStream.ts b/apps/agent-forge/src/author/authorStream.ts index 6fe6ec5e..9b988dd9 100644 --- a/apps/agent-forge/src/author/authorStream.ts +++ b/apps/agent-forge/src/author/authorStream.ts @@ -17,7 +17,6 @@ import { runtimeClient } from '../runtime/runtimeClient'; export type AuthorStreamState = 'connecting' | 'open' | 'reconnecting' | 'closed'; /** The stream's callbacks: events verbatim, state on transitions, `onOpen` after each successful (re)connect. */ -// fallow-ignore-next-line unused-export -- the Author chat hook (AA15a) implements it export interface AuthorStreamHandlers { onEvent(event: AuthorServerEvent): void; onState(state: AuthorStreamState): void; @@ -26,7 +25,6 @@ export interface AuthorStreamHandlers { } /** The handle `connectAuthorStream` returns: write client messages, or end the stream for good. */ -// fallow-ignore-next-line unused-export -- the Author chat hook (AA15a) holds it export interface AuthorStream { /** JSON-sends `message` when the socket is open; `false` (nothing sent) otherwise. */ send(message: AuthorClientMessage): boolean; @@ -46,7 +44,6 @@ const SESSION_NOT_FOUND_CLOSE = 4404; * `options.open` is the socket seam (tests pass a fake; the default is * `runtimeClient.openAuthorSocket`, which adds the `?token=`). */ -// fallow-ignore-next-line unused-export -- the Author chat hook (AA15a) connects through it export function connectAuthorStream( sessionId: string, handlers: AuthorStreamHandlers, @@ -129,7 +126,6 @@ export function connectAuthorStream( } /** A fresh `chat.send`/`chat.stop` `requestId` (the runner's charset, no dashes). */ -// fallow-ignore-next-line unused-export -- the Author chat hook (AA15a) stamps requestIds with it export function newRequestId(): string { return `r_${crypto.randomUUID().replace(/-/g, '')}`; } diff --git a/apps/agent-forge/src/author/reattach.ts b/apps/agent-forge/src/author/reattach.ts new file mode 100644 index 00000000..6ef06c91 --- /dev/null +++ b/apps/agent-forge/src/author/reattach.ts @@ -0,0 +1,118 @@ +/** + * The reattach helpers the Author chat hook drives when its stream ends + * or a chat (re)opens: a verdict on whether a lost socket is worth + * reattaching (`streamEndRecovery`), a dedupe that drops replayed events + * the client already rendered (`acceptEvent`), and the pending-request + * marker in `localStorage` that survives a reload. + * + * The pass itself lives in the hook: on `drop`/`busy` (or a session + * record whose `runningRequestId` is set) it waits for the socket to + * open, sends `chat.resume { requestId, afterSeq: lastSeq }`, applies + * events through `acceptEvent`, and on a terminal event - or after + * `REATTACH_TIMEOUT_MS` of silence - refetches the session record and + * replaces local messages with the stored transcript (server wins). If + * the record still shows a running reply it retries up to + * `REATTACH_MAX_ATTEMPTS`, then shows `REATTACH_GAVE_UP`. + */ +import type { AuthorServerEvent } from '../../shared/authorWire'; + +/** What the end of a stream means for reattachment. */ +export type StreamEnd = 'none' | 'drop' | 'busy'; + +/** + * Whether a finished stream needs the reattach pass - the Author port of + * the "stream ends" state machine, simplified to WS drops: a user stop + * needs nothing, a `BUSY` `turn.error` means another reply is running + * (reattach to it, not to ours), and a socket that closed while a reply + * streamed is a `drop`. Everything else is `none`. + */ +export function streamEndRecovery(input: { + wasStreaming: boolean; + socketClosed: boolean; + errorCode?: string; + stoppedByUser: boolean; +}): StreamEnd { + if (input.stoppedByUser) return 'none'; + if (input.errorCode === 'BUSY') return 'busy'; + if (input.wasStreaming && input.socketClosed) return 'drop'; + return 'none'; +} + +/** The one in-flight `chat.send` a reload can still reattach to. */ +export interface PendingRequest { + sessionId: string; + requestId: string; + lastSeq: number; +} + +/** The localStorage key the pending request is written under (JSON `PendingRequest`). */ +export const PENDING_REQUEST_KEY = 'forge.author.pendingRequest'; + +/** + * Writes (or, for `null`, removes) the pending-request marker. Every + * storage access is wrapped in try/catch: a denied or full localStorage + * must never break the chat - losing the marker only costs one resume. + */ +export function savePendingRequest(p: PendingRequest | null): void { + try { + if (p === null) { + window.localStorage.removeItem(PENDING_REQUEST_KEY); + return; + } + window.localStorage.setItem(PENDING_REQUEST_KEY, JSON.stringify(p)); + } catch { + // Storage denied or full: the chat still works, reattach just can't. + } +} + +/** + * Reads the pending-request marker back (post-reload reattach), `null` + * when absent, malformed or unreadable. Malformed entries are treated + * like absent ones - a wrong shape is never trusted. + */ +export function loadPendingRequest(): PendingRequest | null { + try { + const raw = window.localStorage.getItem(PENDING_REQUEST_KEY); + if (raw === null) return null; + const parsed = JSON.parse(raw) as Partial | null; + if ( + parsed !== null && + typeof parsed === 'object' && + typeof parsed.sessionId === 'string' && + typeof parsed.requestId === 'string' && + typeof parsed.lastSeq === 'number' + ) { + return { sessionId: parsed.sessionId, requestId: parsed.requestId, lastSeq: parsed.lastSeq }; + } + return null; + } catch { + return null; + } +} + +/** + * The dedupe a resumed stream applies to each incoming event: `seq 0` + * (per-socket replies like `UNKNOWN_REQUEST`) always applies; an event + * from another `requestId` than the active turn always applies (a new + * turn); an event of the active turn applies only past `lastSeq`, so + * replayed frames the client already rendered are dropped. + */ +export function acceptEvent( + lastSeq: number, + activeRequestId: string | null, + event: AuthorServerEvent +): boolean { + if (event.seq === 0) return true; + if (activeRequestId === null || event.requestId !== activeRequestId) return true; + return event.seq > lastSeq; +} + +/** Silence after a `chat.resume` at which the hook gives up waiting and reconciles. */ +export const REATTACH_TIMEOUT_MS = 5000; + +/** How many times the hook resends `chat.resume` while the record still shows a running reply. */ +export const REATTACH_MAX_ATTEMPTS = 3; + +/** The ErrorNote shown when every reattach attempt timed out. */ +export const REATTACH_GAVE_UP = + 'Could not reconnect to this reply. Your chat is saved; reopen it from History to see the latest.'; diff --git a/apps/agent-forge/src/runtime/runtimeClient.ts b/apps/agent-forge/src/runtime/runtimeClient.ts index bfed035b..fceb2323 100644 --- a/apps/agent-forge/src/runtime/runtimeClient.ts +++ b/apps/agent-forge/src/runtime/runtimeClient.ts @@ -60,6 +60,7 @@ import { AUTHOR_STREAM_PATH, type AuthorFocus, type AuthorSession, + type AuthorSessionRecord, type AuthorSessionSummary, type AuthorSettingsPayload, } from '../../shared/authorWire'; @@ -635,9 +636,13 @@ class RuntimeClient implements AgentApi { return this.request('/assistant/sessions', { method: 'GET' }); } - /** One Author chat with its transcript (`GET /assistant/sessions/:id`); rejects on 404. */ + /** + * One Author chat with its transcript (`GET /assistant/sessions/:id`); + * rejects on 404. `runningRequestId` names the in-flight reply's + * requestId, when one is running (the reattach decision). + */ // fallow-ignore-next-line unused-class-member -- the Author panel calls it once AA15 lands - async getAuthorSession(id: string): Promise { + async getAuthorSession(id: string): Promise { return this.request(`/assistant/sessions/${encodeURIComponent(id)}`, { method: 'GET' }); }