From 517433454cc07033066be2492ece134d712c30bf Mon Sep 17 00:00:00 2001 From: "haozhe.yang" Date: Thu, 11 Jun 2026 12:31:16 +0800 Subject: [PATCH] feat(sync): add snapshot-based session resync - add session snapshot API and protocol schemas with epoch-aware cursors\n- track in-flight turn state and volatile delta offsets for exact rebuild\n- update web client and tests to seed snapshots before subscribing --- .../src/api/daemon/agentEventProjector.ts | 123 +++- apps/kimi-web/src/api/daemon/client.ts | 68 +- apps/kimi-web/src/api/daemon/eventReducer.ts | 4 +- apps/kimi-web/src/api/daemon/wire.ts | 49 +- apps/kimi-web/src/api/daemon/ws.ts | 89 ++- apps/kimi-web/src/api/types.ts | 52 +- .../src/composables/useKimiWebClient.ts | 83 ++- .../useKimiWebClient-session-cache.test.ts | 35 +- .../daemon-e2e/scenarios/03-refresh-replay.ts | 12 +- .../daemon-e2e/test/refresh-replay.test.ts | 21 +- .../daemon/src/routes/registerApiV1Routes.ts | 2 + packages/daemon/src/routes/snapshot.ts | 133 ++++ .../services/gateway/inFlightTurnTracker.ts | 138 ++++ .../services/gateway/sessionEventJournal.ts | 241 +++++++ .../src/services/gateway/wsBroadcast.ts | 89 ++- .../services/gateway/wsBroadcastService.ts | 180 +++-- packages/daemon/src/ws/connection.ts | 168 +++-- packages/daemon/src/ws/protocol.ts | 40 +- packages/daemon/test/services.test.ts | 224 ++++++- packages/daemon/test/snapshot.e2e.test.ts | 216 ++++++ packages/daemon/test/ws-resync.e2e.test.ts | 120 ++-- .../protocol/src/__tests__/events.test.ts | 53 ++ .../protocol/src/__tests__/snapshot.test.ts | 130 ++++ .../protocol/src/__tests__/ws-control.test.ts | 185 +++++- packages/protocol/src/events.ts | 617 +++++++++++++++++- packages/protocol/src/index.ts | 1 + packages/protocol/src/rest/snapshot.ts | 75 +++ packages/protocol/src/ws-control.ts | 202 +++++- 28 files changed, 3042 insertions(+), 308 deletions(-) create mode 100644 packages/daemon/src/routes/snapshot.ts create mode 100644 packages/daemon/src/services/gateway/inFlightTurnTracker.ts create mode 100644 packages/daemon/src/services/gateway/sessionEventJournal.ts create mode 100644 packages/daemon/test/snapshot.e2e.test.ts create mode 100644 packages/protocol/src/__tests__/snapshot.test.ts create mode 100644 packages/protocol/src/rest/snapshot.ts diff --git a/apps/kimi-web/src/api/daemon/agentEventProjector.ts b/apps/kimi-web/src/api/daemon/agentEventProjector.ts index 6d6c0652b..4810543bd 100644 --- a/apps/kimi-web/src/api/daemon/agentEventProjector.ts +++ b/apps/kimi-web/src/api/daemon/agentEventProjector.ts @@ -16,7 +16,7 @@ // const appEvents = projector.project(rawType, payload, sessionId); // // call reset() when re-subscribing / resyncing a session -import type { AppEvent, AppMessage, AppSessionUsage } from '../types'; +import type { AppEvent, AppInFlightTurn, AppMessage, AppSessionUsage } from '../types'; import { i18n } from '../../i18n'; // --------------------------------------------------------------------------- @@ -60,6 +60,12 @@ interface SessionState { // Assistant message tracking currentAssistantMsgId: string | undefined; + // Per-turn accumulated stream lengths — aligned against the wire `offset` + // on volatile delta frames (v2 sync protocol) to skip duplicates and + // detect gaps after a snapshot seed. + turnTextLen: number; + turnThinkLen: number; + // Tool timing toolStartTimes: Map; @@ -82,6 +88,8 @@ function createSessionState(): SessionState { turnPromptId: new Map(), currentPromptId: undefined, currentAssistantMsgId: undefined, + turnTextLen: 0, + turnThinkLen: 0, toolStartTimes: new Map(), totalInput: 0, totalOutput: 0, @@ -213,14 +221,32 @@ function buildUsageSnapshot(state: SessionState): AppSessionUsage { // AgentProjector // --------------------------------------------------------------------------- +export interface ProjectMeta { + /** + * Wire-level pre-append stream offset on volatile text-delta frames (v2 + * sync protocol). Used to skip duplicate deltas and detect gaps after a + * snapshot seed. + */ + offset?: number; +} + export interface AgentProjector { /** Project a single raw agent-core event into zero or more AppEvents. Never throws. */ - project(rawType: string, payload: unknown, sessionId: string): AppEvent[]; + project(rawType: string, payload: unknown, sessionId: string, meta?: ProjectMeta): AppEvent[]; /** * Bind an externally-known promptId to the next turn.startd for this session. * Call this right after submitPrompt() returns, before the first turn.started arrives. */ bindNextPromptId(sessionId: string, promptId: string): void; + /** + * Seed mid-turn state from a session snapshot's `in_flight_turn` (v2 sync): + * resets per-session state, builds the partially-streamed assistant message + * (thinking + text + running tool_use parts), and returns the AppEvents + * (sessionStatusChanged + messageCreated) to apply to the reducer. Live + * deltas continue appending; their wire `offset` aligns against the seeded + * text so the overlap window around snapshot/subscribe is exact. + */ + seedInFlight(sessionId: string, turn: AppInFlightTurn): AppEvent[]; /** Reset all per-session state (call on re-subscribe / resync). */ reset(sessionId: string): void; } @@ -246,9 +272,54 @@ export function createAgentProjector(): AgentProjector { s.currentPromptId = promptId; } - function project(rawType: string, payload: unknown, sessionId: string): AppEvent[] { + function seedInFlight(sessionId: string, turn: AppInFlightTurn): AppEvent[] { + reset(sessionId); + const s = getOrCreate(sessionId); + + const promptId = ulid('pr_'); + s.currentPromptId = promptId; + s.turnPromptId.set(turn.turnId, promptId); + + const msg = startAssistantMessage(s, sessionId, promptId); + if (turn.thinkingText.length > 0) { + msg.content.push({ type: 'thinking', thinking: turn.thinkingText }); + } + if (turn.assistantText.length > 0) { + msg.content.push({ type: 'text', text: turn.assistantText }); + } + for (const tool of turn.runningTools) { + msg.content.push({ + type: 'toolUse', + toolCallId: tool.toolCallId, + toolName: tool.name, + input: tool.args ?? {}, + }); + s.toolStartTimes.set(tool.toolCallId, Date.now()); + } + s.currentAssistantMsgId = msg.id; + s.turnTextLen = turn.assistantText.length; + s.turnThinkLen = turn.thinkingText.length; + + return [ + { + type: 'sessionStatusChanged', + sessionId, + status: 'running', + previousStatus: 'idle', + currentPromptId: promptId, + }, + { type: 'messageCreated', message: cloneMessage(msg) }, + ]; + } + + function project( + rawType: string, + payload: unknown, + sessionId: string, + meta?: ProjectMeta, + ): AppEvent[] { try { - return _project(rawType, payload, sessionId); + return _project(rawType, payload, sessionId, meta); } catch (err) { // Defensive: log but never crash the caller console.error('[agentProjector] Error projecting event:', rawType, err instanceof Error ? err.message : err); @@ -256,7 +327,25 @@ export function createAgentProjector(): AgentProjector { } } - function _project(rawType: string, payload: unknown, sessionId: string): AppEvent[] { + /** + * Align a live text-delta against the per-turn accumulated length using the + * wire `offset`. Returns 'skip' for duplicates (offset behind local state), + * 'gap' when deltas were missed (offset ahead — trigger a re-snapshot), and + * 'append' otherwise. + */ + function alignDelta(localLen: number, offset: number | undefined): 'append' | 'skip' | 'gap' { + if (offset === undefined) return 'append'; + if (offset < localLen) return 'skip'; + if (offset > localLen) return 'gap'; + return 'append'; + } + + function _project( + rawType: string, + payload: unknown, + sessionId: string, + meta?: ProjectMeta, + ): AppEvent[] { const s = getOrCreate(sessionId); // eslint-disable-next-line @typescript-eslint/no-explicit-any const p = payload as any; @@ -285,6 +374,9 @@ export function createAgentProjector(): AgentProjector { if (turnId !== undefined) { s.turnPromptId.set(turnId, existingPromptId); } + // Fresh turn → fresh per-turn stream offsets. + s.turnTextLen = 0; + s.turnThinkLen = 0; out.push({ type: 'sessionStatusChanged', @@ -324,8 +416,16 @@ export function createAgentProjector(): AgentProjector { const delta: string = p?.delta ?? ''; if (!delta) break; + const align = alignDelta(s.turnThinkLen, meta?.offset); + if (align === 'skip') break; + if (align === 'gap') { + out.push({ type: 'historyCompacted', sessionId, beforeSeq: 0, reason: 'delta_gap' }); + break; + } + const thinkIdx = appendAssistantDelta(s, msgId, 'thinking', delta); if (thinkIdx < 0) break; + s.turnThinkLen += delta.length; out.push({ type: 'assistantDelta', sessionId, @@ -343,8 +443,19 @@ export function createAgentProjector(): AgentProjector { const delta: string = p?.delta ?? ''; if (!delta) break; + const align = alignDelta(s.turnTextLen, meta?.offset); + if (align === 'skip') break; + if (align === 'gap') { + // Deltas were missed in the snapshot↔subscribe window — the only + // exact recovery is a fresh snapshot. historyCompacted is routed to + // onResync by the client wrapper, which reloads via snapshot. + out.push({ type: 'historyCompacted', sessionId, beforeSeq: 0, reason: 'delta_gap' }); + break; + } + const textIdx = appendAssistantDelta(s, msgId, 'text', delta); if (textIdx < 0) break; + s.turnTextLen += delta.length; out.push({ type: 'assistantDelta', sessionId, @@ -675,7 +786,7 @@ export function createAgentProjector(): AgentProjector { return out; } - return { project, bindNextPromptId, reset }; + return { project, bindNextPromptId, seedInFlight, reset }; } // --------------------------------------------------------------------------- diff --git a/apps/kimi-web/src/api/daemon/client.ts b/apps/kimi-web/src/api/daemon/client.ts index a37619919..afd9a967f 100644 --- a/apps/kimi-web/src/api/daemon/client.ts +++ b/apps/kimi-web/src/api/daemon/client.ts @@ -9,7 +9,9 @@ import type { AppModel, AppProvider, AppSession, + AppSessionCursor, AppSessionRuntimeStatus, + AppSessionSnapshot, AppSessionStatus, AppTask, AppTaskStatus, @@ -29,11 +31,13 @@ import type { import { createAgentProjector } from './agentEventProjector'; import { DaemonHttpClient } from './http'; import { + toAppApprovalRequest, toAppEvent, toAppFsEntry, toAppMessage, toAppModel, toAppProvider, + toAppQuestionRequest, toAppSession, toAppTask, toWireApprovalResponse, @@ -62,6 +66,7 @@ import type { WireProvider, WireSession, WireSessionRuntimeStatus, + WireSessionSnapshot, WireWorkspace, WireLogoutResult, } from './wire'; @@ -336,6 +341,41 @@ export class DaemonKimiWebApi implements KimiWebApi { }; } + /** + * v2 initial sync: atomic session state at an `as_of_seq` watermark. + * Rebuild flow: getSessionSnapshot() → seedSnapshot() → subscribe(cursor). + */ + async getSessionSnapshot(sessionId: string): Promise { + const data = await this.http.get( + `/sessions/${encodeURIComponent(sessionId)}/snapshot`, + ); + return { + asOfSeq: data.as_of_seq, + epoch: data.epoch, + session: toAppSession(data.session), + // Snapshot messages are already chronological ascending. + messages: data.messages.items.map(toAppMessage), + hasMoreMessages: data.messages.has_more, + inFlightTurn: + data.in_flight_turn === null + ? null + : { + turnId: data.in_flight_turn.turn_id, + assistantText: data.in_flight_turn.assistant_text, + thinkingText: data.in_flight_turn.thinking_text, + runningTools: data.in_flight_turn.running_tools.map((t) => ({ + toolCallId: t.tool_call_id, + name: t.name, + args: t.args, + description: t.description, + lastProgress: t.last_progress, + })), + }, + pendingApprovals: data.pending_approvals.map(toAppApprovalRequest), + pendingQuestions: data.pending_questions.map(toAppQuestionRequest), + }; + } + // ------------------------------------------------------------------------- // Prompt // ------------------------------------------------------------------------- @@ -881,8 +921,8 @@ export class DaemonKimiWebApi implements KimiWebApi { // Raw agent-core frames — client-side projection path (real daemon) // ----------------------------------------------------------------------- onRawAgentEvent: (frame) => { - const { type, seq, session_id: sessionId, payload } = frame; - const appEvents = projector.project(type, payload, sessionId); + const { type, seq, session_id: sessionId, payload, offset } = frame; + const appEvents = projector.project(type, payload, sessionId, { offset }); for (const appEvent of appEvents) { // Auto-compaction: the projector can't see the wire seq, so it emits // historyCompacted with beforeSeq:0. Route it to onResync using the @@ -894,10 +934,10 @@ export class DaemonKimiWebApi implements KimiWebApi { } }, - onResync: (sessionId: string, currentSeq: number) => { + onResync: (sessionId: string, currentSeq: number, epoch?: string) => { // Reset per-session projector state on resync projector.reset(sessionId); - handlers.onResync(sessionId, currentSeq); + handlers.onResync(sessionId, currentSeq, epoch); }, onConnectionState: (connected: boolean) => { @@ -912,17 +952,33 @@ export class DaemonKimiWebApi implements KimiWebApi { socket.connect(); return { - subscribe(sessionId: string, lastSeq?: number): void { + subscribe(sessionId: string, cursor?: AppSessionCursor): void { // Do NOT reset projector state here: every sidebar click re-subscribes // the (possibly running) session, and a reset wipes the turn/prompt // bindings — the remainder of an in-flight turn would be dropped on // the floor. The projector starts sessions fresh on first sight, and // onResync (below) resets explicitly before messages are reloaded. - socket.subscribe(sessionId, lastSeq ?? 0); + socket.subscribe(sessionId, cursor ?? { seq: 0 }); }, unsubscribe(sessionId: string): void { socket.unsubscribe(sessionId); }, + seedSnapshot(sessionId: string, snapshot: AppSessionSnapshot): void { + // Rebuild the projector's mid-turn state from the snapshot. The + // resulting AppEvents (running status + partially-streamed assistant + // message) flow through the SAME onEvent path as live events, so the + // rendering layer needs no special handling. When there is no + // in-flight turn we only reset, so stale turn state can't leak into + // the freshly-loaded message list. + if (snapshot.inFlightTurn === null) { + projector.reset(sessionId); + return; + } + const appEvents = projector.seedInFlight(sessionId, snapshot.inFlightTurn); + for (const appEvent of appEvents) { + handlers.onEvent(appEvent, { sessionId, seq: snapshot.asOfSeq }); + } + }, bindNextPromptId(sessionId: string, promptId: string): void { // Wire the real daemon prompt_id into the projector so turn.started // uses it instead of a synthetic ulid('pr_'). Without this, the diff --git a/apps/kimi-web/src/api/daemon/eventReducer.ts b/apps/kimi-web/src/api/daemon/eventReducer.ts index 0387a2985..8cf6939f8 100644 --- a/apps/kimi-web/src/api/daemon/eventReducer.ts +++ b/apps/kimi-web/src/api/daemon/eventReducer.ts @@ -218,8 +218,8 @@ export function reduceAppEvent( id: optimistic.id, promptId: event.message.promptId ?? optimistic.promptId, metadata: { - ...(event.message.metadata ?? {}), - ...(optimistic.metadata ?? {}), + ...event.message.metadata, + ...optimistic.metadata, }, }; next.messagesBySession[sid] = updated; diff --git a/apps/kimi-web/src/api/daemon/wire.ts b/apps/kimi-web/src/api/daemon/wire.ts index 18d85ee04..b6e86daf9 100644 --- a/apps/kimi-web/src/api/daemon/wire.ts +++ b/apps/kimi-web/src/api/daemon/wire.ts @@ -420,11 +420,54 @@ export interface WireResyncRequired { timestamp: string; payload: { session_id: string; - reason: 'buffer_overflow' | 'session_recreated'; + reason: 'buffer_overflow' | 'session_recreated' | 'epoch_changed'; current_seq: number; + /** Current journal epoch — adopt it after resyncing (v2 sync protocol). */ + epoch?: string; }; } +// --------------------------------------------------------------------------- +// v2 sync protocol: cursors + session snapshot +// --------------------------------------------------------------------------- + +/** Per-session sync cursor: durable seq + journal epoch. */ +export interface WireSessionCursor { + seq: number; + epoch?: string; +} + +export interface WireInFlightToolCall { + tool_call_id: string; + name: string; + args?: unknown; + description?: string; + display?: unknown; + last_progress?: { + kind: 'stdout' | 'stderr' | 'progress' | 'status' | 'custom'; + text?: string; + percent?: number; + }; +} + +export interface WireInFlightTurn { + turn_id: number; + assistant_text: string; + thinking_text: string; + running_tools: WireInFlightToolCall[]; +} + +/** `GET /sessions/{sid}/snapshot` — atomic rebuild state at a watermark. */ +export interface WireSessionSnapshot { + as_of_seq: number; + epoch: string; + session: WireSession; + messages: { items: WireMessage[]; has_more: boolean }; + in_flight_turn: WireInFlightTurn | null; + pending_approvals: WireApprovalRequest[]; + pending_questions: WireQuestionRequest[]; +} + export interface WireErrorFrame { type: 'error'; timestamp: string; @@ -454,7 +497,7 @@ export interface WireClientHello { payload: { client_id: string; subscriptions: string[]; - last_seq_by_session?: Record; + cursors?: Record; }; } @@ -463,7 +506,7 @@ export interface WireSubscribe { id: string; payload: { session_ids: string[]; - last_seq_by_session?: Record; + cursors?: Record; }; } diff --git a/apps/kimi-web/src/api/daemon/ws.ts b/apps/kimi-web/src/api/daemon/ws.ts index 3b334f010..fe8fc254a 100644 --- a/apps/kimi-web/src/api/daemon/ws.ts +++ b/apps/kimi-web/src/api/daemon/ws.ts @@ -16,11 +16,20 @@ export interface DaemonEventSocketHandlers { /** * Called for raw agent-core frames (type does NOT start with "event." and * is not a control frame). The full parsed frame object is passed so the - * caller can extract type / seq / session_id / timestamp / payload. + * caller can extract type / seq / session_id / timestamp / payload, plus + * the v2 envelope extras (volatile / offset). */ - onRawAgentEvent?(frame: { type: string; seq: number; session_id: string; timestamp: string; payload: unknown }): void; + onRawAgentEvent?(frame: { + type: string; + seq: number; + session_id: string; + timestamp: string; + payload: unknown; + volatile?: boolean; + offset?: number; + }): void; /** Called when server says client is out of sync for a session */ - onResync(sessionId: string, currentSeq: number): void; + onResync(sessionId: string, currentSeq: number, epoch?: string): void; /** Called when the WS connection opens or closes */ onConnectionState(connected: boolean): void; /** Called on error frames or JSON parse failures */ @@ -31,9 +40,15 @@ export interface DaemonEventSocketHandlers { // DaemonEventSocket // --------------------------------------------------------------------------- +/** v2 sync cursor: durable seq + journal epoch. */ +export interface SessionCursor { + seq: number; + epoch?: string; +} + interface PendingSubscription { sessionId: string; - lastSeq: number; + cursor: SessionCursor; } export class DaemonEventSocket { @@ -41,8 +56,8 @@ export class DaemonEventSocket { private connected = false; private closed = false; - /** subscriptions we manage: sessionId → last known seq */ - private readonly subscriptions = new Map(); + /** subscriptions we manage: sessionId → last known cursor {seq, epoch} */ + private readonly subscriptions = new Map(); /** subscriptions queued while not yet connected */ private readonly pendingSubscriptions: PendingSubscription[] = []; @@ -109,19 +124,19 @@ export class DaemonEventSocket { } /** - * Subscribe to events for a session. + * Subscribe to events for a session at a `{seq, epoch}` cursor. * If connected, sends immediately; otherwise queues until after server_hello. */ - subscribe(sessionId: string, lastSeq = 0): void { - this.subscriptions.set(sessionId, lastSeq); + subscribe(sessionId: string, cursor: SessionCursor = { seq: 0 }): void { + this.subscriptions.set(sessionId, { ...cursor }); if (this.connected) { - this.sendSubscribe([sessionId], { [sessionId]: lastSeq }); + this.sendSubscribe([sessionId], { [sessionId]: cursor }); } else { // Remove any earlier pending entry for this session, then enqueue const idx = this.pendingSubscriptions.findIndex((p) => p.sessionId === sessionId); if (idx !== -1) this.pendingSubscriptions.splice(idx, 1); - this.pendingSubscriptions.push({ sessionId, lastSeq }); + this.pendingSubscriptions.push({ sessionId, cursor: { ...cursor } }); } } @@ -182,9 +197,15 @@ export class DaemonEventSocket { this.send({ type: 'pong', payload: { nonce: frame.payload.nonce } }); break; - case 'resync_required': - this.handlers.onResync(frame.payload.session_id, frame.payload.current_seq); + case 'resync_required': { + const sid = frame.payload.session_id as string; + const epoch = frame.payload.epoch as string | undefined; + // Adopt the announced cursor so the next reconnect handshake doesn't + // re-trigger the same resync before the snapshot reload lands. + this.subscriptions.set(sid, { seq: frame.payload.current_seq, epoch }); + this.handlers.onResync(sid, frame.payload.current_seq, epoch); break; + } case 'error': { // A session-scoped error (has top-level session_id) is a real agent-core @@ -211,6 +232,12 @@ export class DaemonEventSocket { break; default: { + // Track the per-session cursor from durable event envelopes so the + // reconnect handshake resumes from the freshest watermark. Volatile + // frames carry the same watermark (never ahead), so skipping them is + // safe and avoids regressing the cursor. + this.trackCursor(frame as Record); + // Classify the frame into protocol vs agent-core. Robust to all three // shapes: raw agent-core, "event."-prefixed agent-core, and genuine // projected "event.*" protocol events. See classifyFrame() for rules. @@ -237,12 +264,15 @@ export class DaemonEventSocket { timestamp: string; payload: unknown; }; + const extras = frame as { volatile?: boolean; offset?: number }; this.handlers.onRawAgentEvent({ type: decision.agentType, seq: f.seq, session_id: f.session_id, timestamp: f.timestamp, payload: f.payload, + ...(extras.volatile !== undefined ? { volatile: extras.volatile } : {}), + ...(extras.offset !== undefined ? { offset: extras.offset } : {}), }); } break; @@ -263,15 +293,15 @@ export class DaemonEventSocket { const allSessionIds = Array.from(this.subscriptions.keys()); // Drain pending: merge into subscriptions map (pending overrides if seq differs) for (const p of this.pendingSubscriptions) { - this.subscriptions.set(p.sessionId, p.lastSeq); + this.subscriptions.set(p.sessionId, p.cursor); if (!allSessionIds.includes(p.sessionId)) allSessionIds.push(p.sessionId); } this.pendingSubscriptions.length = 0; - // Build last_seq_by_session from subscriptions - const lastSeqBySession: Record = {}; - for (const [sid, seq] of this.subscriptions.entries()) { - lastSeqBySession[sid] = seq; + // Build cursors from subscriptions + const cursors: Record = {}; + for (const [sid, cursor] of this.subscriptions.entries()) { + cursors[sid] = cursor; } this.send({ @@ -280,22 +310,39 @@ export class DaemonEventSocket { payload: { client_id: this.clientId, subscriptions: allSessionIds, - last_seq_by_session: lastSeqBySession, + cursors, }, }); } - private sendSubscribe(sessionIds: string[], lastSeqBySession: Record): void { + private sendSubscribe(sessionIds: string[], cursors: Record): void { this.send({ type: 'subscribe', id: this.nextId(), payload: { session_ids: sessionIds, - last_seq_by_session: lastSeqBySession, + cursors, }, }); } + /** + * Advance the tracked cursor from a durable event envelope (seq + epoch). + * Volatile frames are skipped (their seq is the same watermark, and a + * volatile frame can never carry a NEWER seq than the last durable one). + */ + private trackCursor(frame: Record): void { + if (frame['volatile'] === true) return; + const sid = frame['session_id']; + const seq = frame['seq']; + if (typeof sid !== 'string' || typeof seq !== 'number') return; + const existing = this.subscriptions.get(sid); + if (!existing) return; // not a session we manage + if (seq <= existing.seq && existing.epoch !== undefined) return; + const epoch = typeof frame['epoch'] === 'string' ? (frame['epoch'] as string) : existing.epoch; + this.subscriptions.set(sid, { seq: Math.max(seq, existing.seq), epoch }); + } + private send(msg: unknown): void { if (!this.ws || this.ws.readyState !== WebSocket.OPEN) return; try { diff --git a/apps/kimi-web/src/api/types.ts b/apps/kimi-web/src/api/types.ts index 034c375a9..afff0d79d 100644 --- a/apps/kimi-web/src/api/types.ts +++ b/apps/kimi-web/src/api/types.ts @@ -305,15 +305,54 @@ export type AppEvent = // WebSocket connection helpers // --------------------------------------------------------------------------- +/** Per-session sync cursor (v2): durable seq + journal epoch. */ +export interface AppSessionCursor { + seq: number; + epoch?: string; +} + +/** In-flight (mid-turn) state recovered from the session snapshot. */ +export interface AppInFlightToolCall { + toolCallId: string; + name: string; + args?: unknown; + description?: string; + lastProgress?: { kind: string; text?: string; percent?: number }; +} + +export interface AppInFlightTurn { + turnId: number; + assistantText: string; + thinkingText: string; + runningTools: AppInFlightToolCall[]; +} + +/** + * IM-style initial sync result: everything needed to rebuild a session's UI + * state, consistent at `asOfSeq`. The standard flow is + * `getSessionSnapshot()` → `subscribe(sessionId, {seq: asOfSeq, epoch})`. + */ +export interface AppSessionSnapshot { + asOfSeq: number; + epoch: string; + session: AppSession; + /** Most recent messages, chronological ascending. */ + messages: AppMessage[]; + hasMoreMessages: boolean; + inFlightTurn: AppInFlightTurn | null; + pendingApprovals: AppApprovalRequest[]; + pendingQuestions: AppQuestionRequest[]; +} + export interface KimiEventHandlers { onEvent(event: AppEvent, meta: { sessionId: string; seq: number }): void; - onResync(sessionId: string, currentSeq: number): void; + onResync(sessionId: string, currentSeq: number, epoch?: string): void; onError(code: number, msg: string, fatal: boolean): void; onConnectionChange(connected: boolean): void; } export interface KimiEventConnection { - subscribe(sessionId: string, lastSeq?: number): void; + subscribe(sessionId: string, cursor?: AppSessionCursor): void; unsubscribe(sessionId: string): void; /** * Bind the real daemon prompt_id to the next turn for a session, so the @@ -321,6 +360,13 @@ export interface KimiEventConnection { * Call right after submitPrompt() returns. */ bindNextPromptId(sessionId: string, promptId: string): void; + /** + * Seed the client-side projector with a snapshot's in-flight turn so a + * reconnecting client renders mid-turn state immediately; emits the + * corresponding AppEvents through `onEvent`. Resets per-session projector + * state first — call BEFORE subscribe(), with the snapshot's cursor. + */ + seedSnapshot(sessionId: string, snapshot: AppSessionSnapshot): void; abort(sessionId: string, promptId: string): void; close(): void; } @@ -375,6 +421,8 @@ export interface KimiWebApi { getSessionStatus(sessionId: string): Promise; deleteSession(sessionId: string): Promise<{ deleted: true }>; listMessages(sessionId: string, input?: PageRequest & { role?: AppMessageRole }): Promise>; + /** v2 initial sync: atomic session state + `asOfSeq` watermark + epoch. */ + getSessionSnapshot(sessionId: string): Promise; submitPrompt(sessionId: string, input: PromptSubmission): Promise; abortPrompt(sessionId: string, promptId: string): Promise<{ aborted: boolean; atSeq?: number }>; compactSession(sessionId: string, instruction?: string): Promise; diff --git a/apps/kimi-web/src/composables/useKimiWebClient.ts b/apps/kimi-web/src/composables/useKimiWebClient.ts index 22b81d9f5..731a9bbca 100644 --- a/apps/kimi-web/src/composables/useKimiWebClient.ts +++ b/apps/kimi-web/src/composables/useKimiWebClient.ts @@ -555,8 +555,12 @@ function connectEventsIfNeeded(): void { } }, - onResync(sessionId: string, currentSeq: number) { - void reloadAndResubscribe(sessionId, currentSeq); + onResync(sessionId: string, currentSeq: number, epoch?: string) { + // The server-announced cursor is only a hint; the snapshot fetch + // returns the authoritative {asOfSeq, epoch} and re-subscribes. + if (epoch !== undefined) epochBySession[sessionId] = epoch; + void currentSeq; + void syncSessionFromSnapshot(sessionId); }, onError(_code: number, msg: string, _fatal: boolean) { @@ -570,31 +574,49 @@ function connectEventsIfNeeded(): void { }); } -/** - * The daemon's GET /messages returns messages NEWEST-FIRST (confirmed: the - * assistant reply comes before its user prompt). For display we need - * chronological order (oldest first), otherwise a reloaded session renders - * upside down. Reverse (handles the common newest-first case + equal - * timestamps), then stable-sort by createdAt as a safety net. - */ -function orderMessages( - items: import('../api/types').AppMessage[], -): import('../api/types').AppMessage[] { - return [...items] - .reverse() - .sort((a, b) => new Date(a.createdAt).getTime() - new Date(b.createdAt).getTime()); -} +// Journal epoch per session, learned from snapshots / resync frames. Not +// reactive — only consulted when building the subscribe cursor. +const epochBySession: Record = {}; -async function loadMessagesForSession(sessionId: string): Promise { +/** + * v2 initial sync (IM-style rebuild): fetch the atomic session snapshot, + * install its state, seed the projector's in-flight turn, then subscribe the + * WS at the snapshot's `{seq: asOfSeq, epoch}` cursor. The watermark ties + * the REST snapshot to the event stream — no gap, no duplication. + */ +async function syncSessionFromSnapshot(sessionId: string): Promise { try { const api = getKimiWebApi(); - const page = await api.listMessages(sessionId, { pageSize: 100 }); + const snap = await api.getSessionSnapshot(sessionId); + + rawState.sessions = rawState.sessions.map((s) => (s.id === sessionId ? snap.session : s)); rawState.messagesBySession = { ...rawState.messagesBySession, - [sessionId]: orderMessages(page.items), + [sessionId]: snap.messages, }; + rawState.approvalsBySession = { + ...rawState.approvalsBySession, + [sessionId]: snap.pendingApprovals, + }; + rawState.questionsBySession = { + ...rawState.questionsBySession, + [sessionId]: snap.pendingQuestions, + }; + rawState.lastSeqBySession = { + ...rawState.lastSeqBySession, + [sessionId]: snap.asOfSeq, + }; + epochBySession[sessionId] = snap.epoch; + + connectEventsIfNeeded(); + if (eventConn) { + // Seed BEFORE subscribing: the in-flight assistant message must exist + // before live deltas (aligned by wire offset) start appending to it. + eventConn.seedSnapshot(sessionId, snap); + eventConn.subscribe(sessionId, { seq: snap.asOfSeq, epoch: snap.epoch }); + } } catch (err) { - rawState.warnings = [...rawState.warnings, `Failed to load messages: ${String(err)}`]; + rawState.warnings = [...rawState.warnings, `Failed to load session snapshot: ${String(err)}`]; } } @@ -611,14 +633,6 @@ async function loadTasksForSession(sessionId: string): Promise { } } -async function reloadAndResubscribe(sessionId: string, currentSeq: number): Promise { - await loadMessagesForSession(sessionId); - rawState.lastSeqBySession = { ...rawState.lastSeqBySession, [sessionId]: currentSeq }; - if (eventConn) { - eventConn.subscribe(sessionId, currentSeq); - } -} - function hasLoadedMessages(sessionId: string): boolean { return Object.prototype.hasOwnProperty.call(rawState.messagesBySession, sessionId); } @@ -626,8 +640,9 @@ function hasLoadedMessages(sessionId: string): boolean { function subscribeToSessionEvents(sessionId: string): void { connectEventsIfNeeded(); if (eventConn) { - const lastSeq = rawState.lastSeqBySession[sessionId] ?? 0; - eventConn.subscribe(sessionId, lastSeq); + const seq = rawState.lastSeqBySession[sessionId] ?? 0; + const epoch = epochBySession[sessionId]; + eventConn.subscribe(sessionId, { seq, epoch }); } } @@ -1522,9 +1537,13 @@ async function selectSession(sessionId: string): Promise { refreshSessionSidecars(sessionId); if (!messagesLoaded) { - await loadMessagesForSession(sessionId); + // First open: full snapshot → seed → subscribe(asOfSeq). + await syncSessionFromSnapshot(sessionId); + } else { + // Re-open: resume from the tracked cursor; the daemon replays any + // missed durable events (or answers resync_required → snapshot). + subscribeToSessionEvents(sessionId); } - subscribeToSessionEvents(sessionId); } catch (err) { rawState.warnings = [...rawState.warnings, `selectSession failed: ${String(err)}`]; } finally { diff --git a/apps/kimi-web/test/useKimiWebClient-session-cache.test.ts b/apps/kimi-web/test/useKimiWebClient-session-cache.test.ts index 8d9830eea..f7ef0aba6 100644 --- a/apps/kimi-web/test/useKimiWebClient-session-cache.test.ts +++ b/apps/kimi-web/test/useKimiWebClient-session-cache.test.ts @@ -51,6 +51,7 @@ async function setup(messages: AppMessage[] = []) { subscribe: vi.fn(), unsubscribe: vi.fn(), bindNextPromptId: vi.fn(), + seedSnapshot: vi.fn(), abort: vi.fn(), close: vi.fn(), }; @@ -58,6 +59,16 @@ async function setup(messages: AppMessage[] = []) { const api = { createSession: vi.fn(async () => created), listMessages: vi.fn(async () => ({ items: messages, hasMore: false })), + getSessionSnapshot: vi.fn(async () => ({ + asOfSeq: 0, + epoch: 'ep_test', + session: created, + messages, + hasMoreMessages: false, + inFlightTurn: null, + pendingApprovals: [], + pendingQuestions: [], + })), submitPrompt: vi.fn(async () => ({ promptId: 'pr_1', userMessageId: 'msg_real' })), listTasks: vi.fn(async () => []), getGitStatus: vi.fn(async () => ({ branch: 'main', ahead: 0, behind: 0, entries: {} })), @@ -102,15 +113,19 @@ describe('useKimiWebClient session memory cache', () => { const { api, client, eventConn } = await setup([]); await client.createSession('/repo'); - expect(api.listMessages).toHaveBeenCalledTimes(1); + expect(api.getSessionSnapshot).toHaveBeenCalledTimes(1); expect(client.sessionLoading.value).toBe(false); const secondSelect = client.selectSession('sess_1'); expect(client.sessionLoading.value).toBe(false); await secondSelect; - expect(api.listMessages).toHaveBeenCalledTimes(1); - expect(eventConn.subscribe).toHaveBeenLastCalledWith('sess_1', 0); + // L1 hit: no second snapshot fetch — re-subscribe at the tracked cursor. + expect(api.getSessionSnapshot).toHaveBeenCalledTimes(1); + expect(eventConn.subscribe).toHaveBeenLastCalledWith('sess_1', { + seq: 0, + epoch: 'ep_test', + }); }); it('re-subscribes an L1 hit with the reducer-maintained latest seq', async () => { @@ -118,8 +133,11 @@ describe('useKimiWebClient session memory cache', () => { const { api, client, eventConn, getHandlers } = await setup([initial]); await client.createSession('/repo'); - expect(api.listMessages).toHaveBeenCalledTimes(1); - expect(eventConn.subscribe).toHaveBeenLastCalledWith('sess_1', 0); + expect(api.getSessionSnapshot).toHaveBeenCalledTimes(1); + expect(eventConn.subscribe).toHaveBeenLastCalledWith('sess_1', { + seq: 0, + epoch: 'ep_test', + }); getHandlers().onEvent( { type: 'messageCreated', message: userMessage('sess_1', 'msg_2') }, @@ -128,8 +146,11 @@ describe('useKimiWebClient session memory cache', () => { await client.selectSession('sess_1'); - expect(api.listMessages).toHaveBeenCalledTimes(1); - expect(eventConn.subscribe).toHaveBeenLastCalledWith('sess_1', 7); + expect(api.getSessionSnapshot).toHaveBeenCalledTimes(1); + expect(eventConn.subscribe).toHaveBeenLastCalledWith('sess_1', { + seq: 7, + epoch: 'ep_test', + }); }); it('keeps the optimistic user turn key stable after submit resolves', async () => { diff --git a/packages/daemon-e2e/scenarios/03-refresh-replay.ts b/packages/daemon-e2e/scenarios/03-refresh-replay.ts index 24418dd85..03305b5e1 100644 --- a/packages/daemon-e2e/scenarios/03-refresh-replay.ts +++ b/packages/daemon-e2e/scenarios/03-refresh-replay.ts @@ -7,7 +7,7 @@ * daemon is already up. The phases (REST.md §3, WS.md §3) are: * * Phase 0 environment probes GET /healthz, /meta, /auth - * Phase 1 open WS BEFORE history server_hello → client_hello(last_seq_by_session) → ack + * Phase 1 open WS BEFORE history server_hello → client_hello(cursors) → ack * Phase 2 pull persisted snapshot GET /sessions/{sid}, /messages, /tasks * Phase 5 steady state POST /prompts → observe events on WS * @@ -21,12 +21,12 @@ * - A first WS session completes one prompt; we record the current ring-buffer * seq for the session. * - We close the WS, open a fresh one, and on `client_hello` pass - * `last_seq_by_session: { [sid]: currentSeq }` — the daemon should ack + * `cursors: { [sid]: { seq: currentSeq } }` — the daemon should ack * with `accepted_subscriptions: [sid]`, `resync_required: []`, and NOT * replay any events (we are caught up). * - We then open a THIRD connection, this time with - * `last_seq_by_session: { [sid]: 0 }` — the daemon should replay every - * buffered event (seq 1..N) BEFORE the ack lands. + * `cursors: { [sid]: { seq: 0 } }` — the daemon should replay every + * durable event (seq 1..N) BEFORE the ack lands. * - Phase 2 REST snapshot reflects the user + assistant messages persisted * during the first run. * - Phase 5: a new prompt over the third connection delivers events on WS. @@ -64,7 +64,7 @@ interface MetaResponse { interface ClientHelloPayload extends Record { client_id: string; subscriptions: string[]; - last_seq_by_session?: Record; + cursors?: Record; } interface AckPayload { @@ -118,7 +118,7 @@ async function openSocketWithHello({ client_id: `scenario-03-${process.pid}`, subscriptions: [sid], }; - if (lastSeq !== undefined) payload.last_seq_by_session = { [sid]: lastSeq }; + if (lastSeq !== undefined) payload.cursors = { [sid]: { seq: lastSeq } }; ws.send({ type: 'client_hello', id: helloId, payload }); // 3) wait for the matching ack diff --git a/packages/daemon-e2e/test/refresh-replay.test.ts b/packages/daemon-e2e/test/refresh-replay.test.ts index 8e31d7936..03eef3e22 100644 --- a/packages/daemon-e2e/test/refresh-replay.test.ts +++ b/packages/daemon-e2e/test/refresh-replay.test.ts @@ -3,20 +3,23 @@ * * Models the page-refresh path a web client takes when the daemon is already * up: hit `/healthz`, `/meta`, `/auth`, then open a fresh WebSocket and replay - * any missed events via `client_hello.last_seq_by_session` BEFORE pulling REST - * history (REST.md §3 + WS.md §3.2). + * any missed events via `client_hello.cursors` BEFORE pulling REST history + * (REST.md §3 + WS.md §3.2). * * What's asserted here (and NOT in `client.test.ts`): * 1. `/healthz` returns `{ok: true}`. - * 2. `/meta` exposes a non-empty `daemon_id` — the signal clients use to - * detect a daemon restart and flush their `last_seq_by_session` cache. + * 2. `/meta` exposes a non-empty `daemon_id`. (Since the v2 sync protocol, + * cursors carry a journal `epoch` and seq is durable across restarts — + * a stale cursor is detected server-side via `epoch_changed` instead of + * clients comparing `daemon_id`.) * 3. `/auth` returns the `AuthSummary` shape. - * 4. After running one prompt to populate the ring buffer, a fresh WS that - * passes `last_seq_by_session: { [sid]: currentSeq }` is acked with + * 4. After running one prompt to populate the journal, a fresh WS that + * passes `cursors: { [sid]: { seq: currentSeq } }` is acked with * `accepted_subscriptions: [sid]`, `resync_required: []`, and NO event * frames arrive between `server_hello` and the ack (caught-up replay). - * 5. A fresh WS that passes `last_seq_by_session: { [sid]: 0 }` triggers - * replay of every buffered event in order (seq 1..N) BEFORE the ack. + * 5. A fresh WS that passes `cursors: { [sid]: { seq: 0 } }` triggers + * replay of every durable event in order (seq 1..N) BEFORE the ack. + * Volatile frames (deltas/progress/status) are never replayed. * 6. After reconnect, `GET /messages` reflects the persisted state from * before the WS close. * @@ -101,7 +104,7 @@ async function openSocketWithHello(opts: { subscriptions: [opts.sid], }; if (opts.lastSeq !== undefined) { - payload['last_seq_by_session'] = { [opts.sid]: opts.lastSeq }; + payload['cursors'] = { [opts.sid]: { seq: opts.lastSeq } }; } opts.log?.('refresh ws client_hello', { id: helloId, payload }); ws.send({ type: 'client_hello', id: helloId, payload }); diff --git a/packages/daemon/src/routes/registerApiV1Routes.ts b/packages/daemon/src/routes/registerApiV1Routes.ts index 4d4ef50d4..0558a93a4 100644 --- a/packages/daemon/src/routes/registerApiV1Routes.ts +++ b/packages/daemon/src/routes/registerApiV1Routes.ts @@ -14,6 +14,7 @@ import { registerOAuthRoutes } from './oauth'; import { registerPromptsRoutes } from './prompts'; import { registerQuestionsRoutes } from './questions'; import { registerSessionsRoutes } from './sessions'; +import { registerSnapshotRoutes } from './snapshot'; import { registerTasksRoutes } from './tasks'; import { registerToolsRoutes } from './tools'; import { registerWorkspaceFsRoutes } from './workspaceFs'; @@ -65,6 +66,7 @@ export async function registerApiV1Routes( ix, ); registerSessionsRoutes(apiV1 as unknown as Parameters[0], ix); + registerSnapshotRoutes(apiV1 as unknown as Parameters[0], ix); registerMessagesRoutes(apiV1 as unknown as Parameters[0], ix); registerPromptsRoutes(apiV1 as unknown as Parameters[0], ix); registerApprovalsRoutes( diff --git a/packages/daemon/src/routes/snapshot.ts b/packages/daemon/src/routes/snapshot.ts new file mode 100644 index 000000000..3023ee0e7 --- /dev/null +++ b/packages/daemon/src/routes/snapshot.ts @@ -0,0 +1,133 @@ +/** + * `GET /sessions/{session_id}/snapshot` — IM-style initial sync. + * + * Assembles an atomic-at-a-watermark view for client rebuild: + * + * as_of_seq / epoch ← `IWSBroadcastService.getSnapshotState` + * session ← `ISessionService.get` + * messages (asc) ← `IMessageService.list` (most recent page) + * in_flight_turn ← broadcast's `InFlightTurnTracker` + * pending_approvals ← daemon `ApprovalService.listPending` + * pending_questions ← daemon `QuestionService.listPending` + * + * Watermark stability: the durable seq is read before and after assembly; + * if a durable event landed in between, assembly retries (bounded). Durable + * events are low-frequency (turn/tool boundaries — deltas are volatile and + * don't advance seq), so this converges almost immediately. After the + * retries are exhausted the latest watermark is returned — the client's + * seq-guard drops any overlap on replay. + * + * **Error mapping**: `SessionNotFoundError` → 40401; everything else falls + * through to the global error handler (→ 50001). + */ + +import { + ErrorCode, + sessionSnapshotResponseSchema, + type Message, + type Session, +} from '@moonshot-ai/protocol'; +import { + IApprovalService, + IMessageService, + IQuestionService, + ISessionService, + SessionNotFoundError, +} from '@moonshot-ai/services'; +import { z } from 'zod'; + +import type { IInstantiationService } from '@moonshot-ai/agent-core'; + +import { errEnvelope, okEnvelope } from '../envelope'; +import { defineRoute } from '../middleware/defineRoute'; +import type { ApprovalService } from '#/services/approval/approvalService'; +import type { QuestionService } from '#/services/question/questionService'; +import { IWSBroadcastService } from '#/services/gateway'; + +interface SnapshotRouteHost { + get( + path: string, + options: { preHandler: unknown[]; schema?: Record } | undefined, + handler: ( + req: { id: string; query: unknown; params: unknown }, + reply: { send(payload: unknown): unknown }, + ) => Promise | void, + ): unknown; +} + +const sessionIdParamSchema = z.object({ + session_id: z.string().min(1), +}); + +/** Messages included in the snapshot page (most recent, ascending order). */ +const SNAPSHOT_MESSAGE_PAGE_SIZE = 100; + +/** Bounded watermark-stability retries (see module header). */ +const MAX_ASSEMBLY_ATTEMPTS = 3; + +export function registerSnapshotRoutes( + app: SnapshotRouteHost, + ix: IInstantiationService, +): void { + const route = defineRoute( + { + method: 'GET', + path: '/sessions/{session_id}/snapshot', + params: sessionIdParamSchema, + success: { data: sessionSnapshotResponseSchema }, + description: + 'Atomic session snapshot for client rebuild: state + as_of_seq watermark + epoch', + tags: ['sessions'], + }, + async (req, reply) => { + try { + const { session_id } = req.params; + const data = await ix.invokeFunction(async (a) => { + const broadcast = a.get(IWSBroadcastService); + const sessionService = a.get(ISessionService); + const messageService = a.get(IMessageService); + const approvals = a.get(IApprovalService) as ApprovalService; + const questions = a.get(IQuestionService) as QuestionService; + + let snapState = await broadcast.getSnapshotState(session_id); + let session: Session | undefined; + let items: Message[] = []; + let hasMore = false; + + for (let attempt = 0; attempt < MAX_ASSEMBLY_ATTEMPTS; attempt++) { + session = await sessionService.get(session_id); + const page = await messageService.list(session_id, { + page_size: SNAPSHOT_MESSAGE_PAGE_SIZE, + }); + // IMessageService returns newest-first; snapshot serves ascending. + items = [...page.items].reverse(); + hasMore = page.has_more; + + const post = await broadcast.getSnapshotState(session_id); + const stable = post.seq === snapState.seq && post.epoch === snapState.epoch; + snapState = post; + if (stable) break; + } + + return { + as_of_seq: snapState.seq, + epoch: snapState.epoch, + session: session!, + messages: { items, has_more: hasMore }, + in_flight_turn: snapState.inFlightTurn, + pending_approvals: approvals.listPending(session_id), + pending_questions: questions.listPending(session_id), + }; + }); + reply.send(okEnvelope(data, req.id)); + } catch (err) { + if (err instanceof SessionNotFoundError) { + reply.send(errEnvelope(ErrorCode.SESSION_NOT_FOUND, err.message, req.id)); + return; + } + throw err; + } + }, + ); + app.get(route.path, route.options, route.handler as Parameters[2]); +} diff --git a/packages/daemon/src/services/gateway/inFlightTurnTracker.ts b/packages/daemon/src/services/gateway/inFlightTurnTracker.ts new file mode 100644 index 000000000..34f2e3a55 --- /dev/null +++ b/packages/daemon/src/services/gateway/inFlightTurnTracker.ts @@ -0,0 +1,138 @@ +/** + * `InFlightTurnTracker` — accumulates the current turn's volatile stream + * state per session so a reconnecting client can rebuild mid-turn UI from + * the session snapshot instead of replaying deltas (which are not journaled). + * + * Owned by `WSBroadcastService` and updated INSIDE its per-session dispatch + * queue — this keeps the accumulated text, the journal watermark, and the + * fan-out order mutually consistent without a second event subscription. + * + * `apply()` also returns the pre-append character offset for text-delta + * frames; the broadcast layer stamps it on the wire envelope so clients can + * align live deltas against snapshot text exactly (skip duplicates, detect + * gaps). + * + * Only main-agent activity is tracked: subagent deltas share the session id + * but describe a different stream and would corrupt the accumulation. + */ + +import type { Event, InFlightToolCall, InFlightTurn } from '@moonshot-ai/protocol'; + +const MAIN_AGENT_ID = 'main'; + +interface ToolAccum { + tool_call_id: string; + name: string; + args?: unknown; + description?: string; + display?: unknown; + last_progress?: { + kind: 'stdout' | 'stderr' | 'progress' | 'status' | 'custom'; + text?: string; + percent?: number; + }; +} + +interface TurnAccum { + turnId: number; + assistantText: string; + thinkingText: string; + tools: Map; +} + +export interface VolatileAnnotation { + /** Pre-append offset for text-delta frames. */ + offset?: number; +} + +export class InFlightTurnTracker { + private readonly bySession = new Map(); + + apply(sessionId: string, event: Event): VolatileAnnotation { + if (event.agentId !== MAIN_AGENT_ID) return {}; + + switch (event.type) { + case 'turn.started': { + this.bySession.set(sessionId, { + turnId: event.turnId, + assistantText: '', + thinkingText: '', + tools: new Map(), + }); + return {}; + } + case 'turn.ended': { + this.bySession.delete(sessionId); + return {}; + } + case 'assistant.delta': { + const turn = this.bySession.get(sessionId); + if (!turn || turn.turnId !== event.turnId) return {}; + const offset = turn.assistantText.length; + turn.assistantText += event.delta; + return { offset }; + } + case 'thinking.delta': { + const turn = this.bySession.get(sessionId); + if (!turn || turn.turnId !== event.turnId) return {}; + const offset = turn.thinkingText.length; + turn.thinkingText += event.delta; + return { offset }; + } + case 'tool.call.started': { + const turn = this.bySession.get(sessionId); + if (!turn || turn.turnId !== event.turnId) return {}; + turn.tools.set(event.toolCallId, { + tool_call_id: event.toolCallId, + name: event.name, + args: event.args, + ...(event.description !== undefined ? { description: event.description } : {}), + ...(event.display !== undefined ? { display: event.display } : {}), + }); + return {}; + } + case 'tool.progress': { + const turn = this.bySession.get(sessionId); + const tool = turn?.tools.get(event.toolCallId); + if (!tool) return {}; + const { kind, text, percent } = event.update; + if (kind === 'custom') return {}; + tool.last_progress = { + kind, + ...(text !== undefined ? { text } : {}), + ...(percent !== undefined ? { percent } : {}), + }; + return {}; + } + case 'tool.result': { + this.bySession.get(sessionId)?.tools.delete(event.toolCallId); + return {}; + } + default: + return {}; + } + } + + get(sessionId: string): InFlightTurn | null { + const turn = this.bySession.get(sessionId); + if (!turn) return null; + const running_tools: InFlightToolCall[] = Array.from(turn.tools.values()).map((t) => ({ + tool_call_id: t.tool_call_id, + name: t.name, + ...(t.args !== undefined ? { args: t.args } : {}), + ...(t.description !== undefined ? { description: t.description } : {}), + ...(t.display !== undefined ? { display: t.display } : {}), + ...(t.last_progress !== undefined ? { last_progress: t.last_progress } : {}), + })); + return { + turn_id: turn.turnId, + assistant_text: turn.assistantText, + thinking_text: turn.thinkingText, + running_tools, + }; + } + + clear(sessionId: string): void { + this.bySession.delete(sessionId); + } +} diff --git a/packages/daemon/src/services/gateway/sessionEventJournal.ts b/packages/daemon/src/services/gateway/sessionEventJournal.ts new file mode 100644 index 000000000..03c148188 --- /dev/null +++ b/packages/daemon/src/services/gateway/sessionEventJournal.ts @@ -0,0 +1,241 @@ +/** + * `SessionEventJournal` — per-session durable event log (the IM-style + * server-side message log that makes multi-device cursors meaningful). + * + * One JSONL file per session under `/daemon/events/.jsonl`: + * + * line 1 {"kind":"journal_header","version":1,"epoch":"ep_","created_at":...} + * line 2+ {"kind":"event","seq":N,"envelope":{...wire envelope...}} + * + * Invariants: + * - `seq` is assigned at append time, starts at 1, and is monotonic across + * daemon restarts (recovered by scanning the file on open). + * - `epoch` identifies this journal incarnation. It changes only when the + * file is unreadable/corrupt at open (we start a fresh journal) — clients + * holding cursors from the old epoch get `resync_required(epoch_changed)`. + * - Only durable events are written (volatile delta/progress/status frames + * never touch the journal; see `VOLATILE_EVENT_TYPES`). + * + * Durability model matches `FileSystemAgentRecordPersistence`: `append()` is + * synchronous (callers need the seq immediately for fan-out), bytes are + * flushed on a microtask-scheduled async batch. `readSince()` flushes first, + * so replay reads never miss queued lines. A torn trailing line from a crash + * is tolerated and ignored on open. + */ + +import { createReadStream } from 'node:fs'; +import { appendFile, mkdir } from 'node:fs/promises'; +import { dirname } from 'node:path'; +import { ulid } from 'ulid'; + +import type { ILogService } from '@moonshot-ai/services'; + +import type { EventEnvelope } from '#/ws/protocol'; + +const JOURNAL_VERSION = 1; + +interface JournalHeaderLine { + kind: 'journal_header'; + version: number; + epoch: string; + created_at: number; +} + +interface JournalEventLine { + kind: 'event'; + seq: number; + envelope: EventEnvelope; +} + +export interface JournalEntry { + seq: number; + envelope: EventEnvelope; +} + +export class SessionEventJournal { + private _seq: number; + private pendingLines: string[] = []; + private flushPromise: Promise | undefined; + private headerPending: boolean; + + private constructor( + private readonly filePath: string, + private readonly logger: ILogService, + public readonly epoch: string, + lastSeq: number, + isFresh: boolean, + ) { + this._seq = lastSeq; + this.headerPending = isFresh; + } + + /** Highest durable seq appended (0 if none). */ + get seq(): number { + return this._seq; + } + + /** + * Open (or create) the journal for `filePath`. Scans an existing file to + * recover `{epoch, lastSeq}`. A missing file or an unreadable header + * starts a fresh journal with a new epoch. + */ + static async open(filePath: string, logger: ILogService): Promise { + let epoch: string | undefined; + let lastSeq = 0; + let sawAnyLine = false; + + try { + for await (const raw of readLines(filePath)) { + sawAnyLine = true; + const parsed = parseJournalLine(raw); + if (parsed === undefined) continue; // torn/corrupt line — skip + if (parsed.kind === 'journal_header') { + if (epoch === undefined) epoch = parsed.epoch; + continue; + } + if (parsed.seq > lastSeq) lastSeq = parsed.seq; + } + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + if (code !== 'ENOENT') { + logger.warn( + { filePath, err: String(error) }, + 'event journal unreadable; starting a fresh epoch', + ); + } + } + + if (epoch === undefined) { + if (sawAnyLine) { + // File exists but has no parseable header — treat as corrupt and + // start a fresh incarnation. Old cursors will epoch-mismatch. + logger.warn({ filePath }, 'event journal missing header; rotating to a fresh epoch'); + } + return new SessionEventJournal(filePath, logger, `ep_${ulid()}`, 0, true); + } + return new SessionEventJournal(filePath, logger, epoch, lastSeq, false); + } + + /** Reserve the next durable seq. The caller must follow with `append()`. */ + nextSeq(): number { + this._seq += 1; + return this._seq; + } + + /** Queue a durable event line for write-behind flush. */ + append(seq: number, envelope: EventEnvelope): void { + const line: JournalEventLine = { kind: 'event', seq, envelope }; + this.pendingLines.push(JSON.stringify(line)); + this.scheduleFlush(); + } + + /** Read journal entries with `seq > fromSeqExclusive`, capped at `limit`. */ + async readSince(fromSeqExclusive: number, limit: number): Promise { + await this.flush(); + const out: JournalEntry[] = []; + try { + for await (const raw of readLines(this.filePath)) { + const parsed = parseJournalLine(raw); + if (parsed === undefined || parsed.kind !== 'event') continue; + if (parsed.seq <= fromSeqExclusive) continue; + out.push({ seq: parsed.seq, envelope: parsed.envelope }); + if (out.length >= limit) break; + } + } catch (error) { + const code = (error as NodeJS.ErrnoException).code; + if (code !== 'ENOENT') throw error; + } + return out; + } + + async flush(): Promise { + while (this.flushPromise !== undefined || this.pendingLines.length > 0) { + if (this.flushPromise === undefined) { + this.flushPromise = this.flushOnce().finally(() => { + this.flushPromise = undefined; + }); + } + await this.flushPromise; + } + } + + async close(): Promise { + await this.flush(); + } + + private scheduleFlush(): void { + if (this.flushPromise !== undefined) return; + this.flushPromise = this.flushOnce().finally(() => { + this.flushPromise = undefined; + }); + } + + private async flushOnce(): Promise { + // Take the queue snapshot first so appends during the await are picked + // up by the next flush round, never lost. + const lines: string[] = []; + if (this.headerPending) { + const header: JournalHeaderLine = { + kind: 'journal_header', + version: JOURNAL_VERSION, + epoch: this.epoch, + created_at: Date.now(), + }; + lines.push(JSON.stringify(header)); + this.headerPending = false; + } + lines.push(...this.pendingLines); + this.pendingLines = []; + if (lines.length === 0) return; + try { + await mkdir(dirname(this.filePath), { recursive: true }); + await appendFile(this.filePath, lines.join('\n') + '\n', 'utf8'); + } catch (error) { + this.logger.warn( + { filePath: this.filePath, err: String(error) }, + 'event journal write failed; events remain live-only this round', + ); + } + } +} + +function parseJournalLine(raw: string): JournalHeaderLine | JournalEventLine | undefined { + const trimmed = raw.endsWith('\r') ? raw.slice(0, -1) : raw; + if (trimmed.length === 0) return undefined; + let value: unknown; + try { + value = JSON.parse(trimmed); + } catch { + return undefined; + } + if (typeof value !== 'object' || value === null) return undefined; + const kind = (value as { kind?: unknown }).kind; + if (kind === 'journal_header') { + const epoch = (value as { epoch?: unknown }).epoch; + if (typeof epoch !== 'string' || epoch.length === 0) return undefined; + return value as JournalHeaderLine; + } + if (kind === 'event') { + const seq = (value as { seq?: unknown }).seq; + const envelope = (value as { envelope?: unknown }).envelope; + if (typeof seq !== 'number' || !Number.isInteger(seq) || seq <= 0) return undefined; + if (typeof envelope !== 'object' || envelope === null) return undefined; + return value as JournalEventLine; + } + return undefined; +} + +async function* readLines(filePath: string): AsyncIterable { + let buffered = ''; + const stream = createReadStream(filePath, { encoding: 'utf8' }); + for await (const chunk of stream) { + buffered += chunk; + let newlineIndex = buffered.indexOf('\n'); + while (newlineIndex !== -1) { + yield buffered.slice(0, newlineIndex); + buffered = buffered.slice(newlineIndex + 1); + newlineIndex = buffered.indexOf('\n'); + } + } + if (buffered.length > 0) yield buffered; +} diff --git a/packages/daemon/src/services/gateway/wsBroadcast.ts b/packages/daemon/src/services/gateway/wsBroadcast.ts index eefd1d2fd..274d287d5 100644 --- a/packages/daemon/src/services/gateway/wsBroadcast.ts +++ b/packages/daemon/src/services/gateway/wsBroadcast.ts @@ -1,72 +1,93 @@ /** * `IWSBroadcastService` — daemon-local transport layer that turns the * in-process `IEventService.onDidPublish` firehose into a WS broadcast + - * per-session ring buffer + replay surface. + * durable per-session event journal + replay surface. * - * Responsibilities (all daemon transport concerns; intentionally NOT on the - * `@moonshot-ai/services` cross-package `IEventService` contract): + * v2 (IM-style multi-device sync) responsibilities: * * 1. Extract `sessionId` from each published event (defensive: accepts * both camelCase `sessionId` and snake_case `session_id`). Events * without a session id are dropped with a warn log. - * 2. Maintain per-session monotonic `seq` (starts at 1). - * 3. Maintain per-session ring buffer (capped at `maxBufferSize`, - * tracking `oldestSeq` for replay/resync decisions). - * 4. Build the WS `EventEnvelope` (snake_case wire shape with - * `seq`, `session_id`, `timestamp`, `payload`). - * 5. Fan out the envelope to every `WsConnection` subscribed to the - * session (via `ISessionClientsService.getConnections`). - * 6. Expose replay queries (`getBufferedSince` / `currentSeq`) for the - * `client_hello.last_seq_by_session` and WS abort `at_seq` paths. + * 2. Classify events as durable vs volatile (`VOLATILE_EVENT_TYPES`). + * 3. Durable events: assign the next per-session `seq` (journal offset, + * monotonic ACROSS DAEMON RESTARTS), persist to the session's + * `SessionEventJournal`, cache in an in-memory tail buffer, fan out. + * 4. Volatile events: fan out live with the current durable watermark as + * `seq` and `volatile: true`. Never journaled, never replayed. + * 5. Expose replay (`getBufferedSince`) keyed by `{seq, epoch}` cursors: + * epoch mismatch or a cursor ahead of the journal → `epoch_changed` + * resync; a gap larger than the replay cap → `buffer_overflow` resync + * (the client should rebuild via `GET /sessions/{sid}/snapshot`); + * otherwise events come from the memory tail or the journal file. + * 6. Expose `getCursor` so the snapshot route / subscribe acks can hand + * clients an authoritative `{seq, epoch}` watermark. * * Wiring: the impl auto-subscribes to `IEventService.onDidPublish` in its * constructor — producers continue to call `eventService.publish(event)` * unchanged; the broadcast layer transparently lifts those events onto the * wire. * - * Decorator name `'wsBroadcastService'` surfaces in - * `CyclicDependencyError.path` and `'No service registered for identifier - * ...'` diagnostics. Replaces the prior `'eventReplayService'` (which was - * an alias for the same singleton under a narrower surface). - * * Dispose order: this service MUST dispose BEFORE `IEventService` so the * `onDidPublish` subscription is detached before the bus tears down its * emitter (reverse-construction order in `start.ts` is responsible). */ import { createDecorator } from '@moonshot-ai/agent-core'; +import type { InFlightTurn, SessionCursor } from '@moonshot-ai/protocol'; import type { EventEnvelope } from '#/ws/protocol'; +export type ResyncReason = 'buffer_overflow' | 'session_recreated' | 'epoch_changed'; + +export interface SessionSnapshotState { + seq: number; + epoch: string; + inFlightTurn: InFlightTurn | null; +} + export interface BufferedSinceResult { events: Array<{ seq: number; envelope: EventEnvelope }>; /** - * True iff `lastSeq + 1 < oldestSeq` (the client's gap is older than what - * the buffer retains). The connection should send a `resync_required` - * frame for this session and NOT replay events. + * Set when the cursor cannot be served incrementally — the client must + * rebuild from the session snapshot and re-subscribe at the returned + * `{currentSeq, epoch}`. */ - resyncRequired: boolean; - /** Highest dispatched `seq` for the session (0 if no events yet). */ + resyncRequired: ResyncReason | false; + /** Highest durable `seq` for the session (0 if no events yet). */ currentSeq: number; + /** Current journal epoch for the session. */ + epoch: string; } export interface IWSBroadcastService { readonly _serviceBrand: undefined; /** - * Fetch buffered events with `seq > lastSeq` for `sessionId`. + * Fetch durable events with `seq > cursor.seq` for `sessionId`. * - * Result interpretation (per WS.md §6): - * - `currentSeq == 0` → session has no events yet; empty replay. - * - `lastSeq >= currentSeq` → client is caught up. - * - `lastSeq + 1 < oldestSeq` → buffer evicted past client; resyncRequired. - * - otherwise → events with `seq > lastSeq`, in order. + * Result interpretation: + * - `cursor.epoch` set but ≠ journal epoch → resync `epoch_changed`. + * - `cursor.seq > currentSeq` (client ahead) → resync `epoch_changed` + * (stale/foreign cursor — e.g. a v1 cursor from before journaling). + * - `currentSeq - cursor.seq > replay cap` → resync `buffer_overflow`. + * - otherwise → durable events with `seq > cursor.seq`, in order, from + * the memory tail or the on-disk journal. */ - getBufferedSince(sessionId: string, lastSeq: number): BufferedSinceResult; + getBufferedSince(sessionId: string, cursor: SessionCursor): Promise; + + /** Authoritative `{seq, epoch}` watermark for the session. */ + getCursor(sessionId: string): Promise<{ seq: number; epoch: string }>; /** - * Highest dispatched `seq` for the session (0 if never published). - * Used by the WS abort ack to populate `at_seq` on idempotent calls. + * Watermark + accumulated in-flight turn state, read atomically with + * respect to the per-session dispatch queue. Backs + * `GET /sessions/{sid}/snapshot`. + */ + getSnapshotState(sessionId: string): Promise; + + /** + * Best-effort sync watermark (0 if the session's journal has not been + * touched this run). Used by the WS abort ack `at_seq` path. */ currentSeq(sessionId: string): number; } @@ -75,5 +96,9 @@ export interface IWSBroadcastService { export const IWSBroadcastService = createDecorator('wsBroadcastService'); -/** Default ring buffer cap (WS.md §3.1, §6). */ +/** + * Max durable events served by one incremental replay. Larger gaps get a + * `buffer_overflow` resync — at that point a snapshot rebuild is cheaper + * than streaming the backlog. Also sizes the in-memory tail cache. + */ export const DEFAULT_MAX_BUFFER_SIZE = 1000; diff --git a/packages/daemon/src/services/gateway/wsBroadcastService.ts b/packages/daemon/src/services/gateway/wsBroadcastService.ts index acbb4594d..37582f9fb 100644 --- a/packages/daemon/src/services/gateway/wsBroadcastService.ts +++ b/packages/daemon/src/services/gateway/wsBroadcastService.ts @@ -1,15 +1,20 @@ +import { join } from 'node:path'; + import { Disposable } from '@moonshot-ai/agent-core'; -import type { Event } from '@moonshot-ai/protocol'; +import { isVolatileEventType, type Event, type SessionCursor } from '@moonshot-ai/protocol'; import { IEventService } from '@moonshot-ai/services'; -import { ILogService } from '@moonshot-ai/services'; +import { IEnvironmentService, ILogService } from '@moonshot-ai/services'; +import { InFlightTurnTracker } from './inFlightTurnTracker'; import { ISessionClientsService } from './sessionClients'; +import { SessionEventJournal } from './sessionEventJournal'; import { DEFAULT_MAX_BUFFER_SIZE, IWSBroadcastService, type BufferedSinceResult, + type SessionSnapshotState, } from './wsBroadcast'; import { buildEventEnvelope, type EventEnvelope } from '#/ws/protocol'; @@ -20,12 +25,14 @@ interface BufferEntry { } interface SessionState { - - seq: number; - - buffer: BufferEntry[]; - - oldestSeq: number; + /** Resolves when the journal file has been opened/recovered. */ + ready: Promise; + /** Set once `ready` resolves — for sync best-effort reads. */ + journal: SessionEventJournal | undefined; + /** In-memory tail cache of the most recent durable envelopes. */ + tail: BufferEntry[]; + /** Per-session dispatch chain: keeps journal append + fan-out ordered. */ + queue: Promise; } export class WSBroadcastService extends Disposable implements IWSBroadcastService { @@ -33,14 +40,18 @@ export class WSBroadcastService extends Disposable implements IWSBroadcastServic private readonly _sessions = new Map(); private readonly _maxBufferSize: number; + private readonly _journalDir: string; + private readonly _turnTracker = new InFlightTurnTracker(); constructor( @IEventService eventService: IEventService, @ILogService private readonly logger: ILogService, @ISessionClientsService private readonly sessionClients: ISessionClientsService, + @IEnvironmentService env: IEnvironmentService, ) { super(); this._maxBufferSize = DEFAULT_MAX_BUFFER_SIZE; + this._journalDir = join(env.homeDir, 'daemon', 'events'); this._register( eventService.onDidPublish((event) => { @@ -56,72 +67,160 @@ export class WSBroadcastService extends Disposable implements IWSBroadcastServic if (!sid) { this.logger.warn( { eventType: evType, eventKeys: Object.keys(event as object) }, - '[DBG wsBroadcast.onEvent] event has no session_id; dropping', + 'wsBroadcast: event has no session_id; dropping', ); return; } const state = this._getOrCreateSession(sid); - state.seq += 1; - const envelope = buildEventEnvelope(state.seq, sid, event); - state.buffer.push({ seq: state.seq, envelope }); + state.queue = state.queue + .then(() => this._dispatch(sid, state, event)) + .catch((err: unknown) => { + this.logger.warn({ sid, eventType: evType, err: String(err) }, 'wsBroadcast dispatch failed'); + }); + } - while (state.buffer.length > this._maxBufferSize) { - const evicted = state.buffer.shift(); - if (evicted) state.oldestSeq = evicted.seq + 1; + private async _dispatch(sid: string, state: SessionState, event: Event): Promise { + if (this._store.isDisposed) return; + const journal = await state.ready; + const evType = (event as { type?: string }).type ?? 'event.unknown'; + + // Track in-flight turn state inside the dispatch queue so accumulated + // text, the journal watermark, and fan-out order stay consistent. For + // text deltas this also yields the pre-append offset for the envelope. + const annotation = this._turnTracker.apply(sid, event); + + let envelope: EventEnvelope; + if (isVolatileEventType(evType)) { + // Volatile frames ride the current durable watermark and are never + // journaled or replayed; reconnecting clients recover their state from + // the session snapshot instead. + envelope = buildEventEnvelope(journal.seq, sid, event, { + epoch: journal.epoch, + volatile: true, + ...(annotation.offset !== undefined ? { offset: annotation.offset } : {}), + }); + } else { + const seq = journal.nextSeq(); + envelope = buildEventEnvelope(seq, sid, event, { epoch: journal.epoch }); + journal.append(seq, envelope); + state.tail.push({ seq, envelope }); + while (state.tail.length > this._maxBufferSize) { + state.tail.shift(); + } } - const targets = Array.from(this.sessionClients.getConnections(sid)); - this.logger.info( - { eventType: evType, sessionId: sid, seq: state.seq, targetCount: targets.length }, - '[DBG wsBroadcast.onEvent] fan-out', - ); - for (const conn of targets) { + if (this._store.isDisposed) return; + for (const conn of this.sessionClients.getConnections(sid)) { conn.send(envelope); } } - getBufferedSince(sid: string, lastSeq: number): BufferedSinceResult { - const state = this._sessions.get(sid); - if (!state) { - return { events: [], resyncRequired: false, currentSeq: 0 }; + async getBufferedSince(sid: string, cursor: SessionCursor): Promise { + const state = this._getOrCreateSession(sid); + const journal = await state.ready; + // Drain in-flight dispatches so the watermark reflects everything + // published before this call. + await state.queue; + + const currentSeq = journal.seq; + const epoch = journal.epoch; + + if (cursor.epoch !== undefined && cursor.epoch !== epoch) { + return { events: [], resyncRequired: 'epoch_changed', currentSeq, epoch }; } - if (lastSeq >= state.seq) { - return { events: [], resyncRequired: false, currentSeq: state.seq }; + if (cursor.seq > currentSeq) { + // Client is ahead of the journal — a cursor from another incarnation + // (e.g. pre-journal v1 daemon). Without a matching epoch we cannot + // trust it; force a snapshot rebuild. + return { events: [], resyncRequired: 'epoch_changed', currentSeq, epoch }; } - if (lastSeq + 1 < state.oldestSeq) { - return { events: [], resyncRequired: true, currentSeq: state.seq }; + if (cursor.seq === currentSeq) { + return { events: [], resyncRequired: false, currentSeq, epoch }; } - const events = state.buffer.filter((e) => e.seq > lastSeq); - return { events, resyncRequired: false, currentSeq: state.seq }; + if (currentSeq - cursor.seq > this._maxBufferSize) { + return { events: [], resyncRequired: 'buffer_overflow', currentSeq, epoch }; + } + + const tail = state.tail; + if (tail.length > 0 && tail[0]!.seq <= cursor.seq + 1) { + const events = tail.filter((e) => e.seq > cursor.seq); + return { events, resyncRequired: false, currentSeq, epoch }; + } + + // Gap reaches behind the memory tail (e.g. first subscribe after a + // daemon restart) — serve from the on-disk journal. + const events = await journal.readSince(cursor.seq, this._maxBufferSize); + return { events, resyncRequired: false, currentSeq, epoch }; + } + + async getCursor(sid: string): Promise<{ seq: number; epoch: string }> { + const state = this._getOrCreateSession(sid); + const journal = await state.ready; + await state.queue; + return { seq: journal.seq, epoch: journal.epoch }; + } + + async getSnapshotState(sid: string): Promise { + const state = this._getOrCreateSession(sid); + const journal = await state.ready; + await state.queue; + // Sync reads after the drain — seq and in-flight state form a + // consistent pair (no dispatch can interleave a sync section). + return { + seq: journal.seq, + epoch: journal.epoch, + inFlightTurn: this._turnTracker.get(sid), + }; } currentSeq(sid: string): number { - return this._sessions.get(sid)?.seq ?? 0; + return this._sessions.get(sid)?.journal?.seq ?? 0; } _currentSeqForTest(sid: string): number { - return this._sessions.get(sid)?.seq ?? 0; + return this.currentSeq(sid); } _bufferLengthForTest(sid: string): number { - return this._sessions.get(sid)?.buffer.length ?? 0; + return this._sessions.get(sid)?.tail.length ?? 0; } - _oldestSeqForTest(sid: string): number { - return this._sessions.get(sid)?.oldestSeq ?? 0; + /** Settles when every queued dispatch for `sid` has completed. */ + async _drainForTest(sid: string): Promise { + const state = this._sessions.get(sid); + if (!state) return; + await state.ready; + await state.queue; } private _getOrCreateSession(sid: string): SessionState { let state = this._sessions.get(sid); if (!state) { - state = { seq: 0, buffer: [], oldestSeq: 1 }; - this._sessions.set(sid, state); + const filePath = join(this._journalDir, `${sanitizeFileName(sid)}.jsonl`); + const created: SessionState = { + ready: SessionEventJournal.open(filePath, this.logger), + journal: undefined, + tail: [], + queue: Promise.resolve(), + }; + created.ready = created.ready.then((journal) => { + created.journal = journal; + return journal; + }); + this._sessions.set(sid, created); + state = created; } return state; } override dispose(): void { if (this._store.isDisposed) return; + for (const state of this._sessions.values()) { + const journal = state.journal; + if (journal) { + void journal.close().catch(() => {}); + } + } this._sessions.clear(); super.dispose(); } @@ -134,3 +233,8 @@ function extractSessionId(event: Event): string | undefined { if (typeof snake === 'string' && snake.length > 0) return snake; return undefined; } + +/** Session ids are ULID-ish, but never trust an id used as a path segment. */ +function sanitizeFileName(sid: string): string { + return sid.replace(/[^A-Za-z0-9._-]/g, '_'); +} diff --git a/packages/daemon/src/ws/connection.ts b/packages/daemon/src/ws/connection.ts index f80288abc..d24e973f4 100644 --- a/packages/daemon/src/ws/connection.ts +++ b/packages/daemon/src/ws/connection.ts @@ -5,13 +5,17 @@ import { ulid } from 'ulid'; import { ErrorCode, - clientControlMessageSchema, + WS_PROTOCOL_VERSION, type AbortMessage, type ClientHelloMessage, + type ClientControlMessage, + type CursorsBySession, + type SessionCursor, type SubscribeMessage, type UnsubscribeMessage, type WatchFsAddMessage, type WatchFsRemoveMessage, + getClientControlOperation, } from '@moonshot-ai/protocol'; import type { ILogService } from '@moonshot-ai/services'; @@ -29,12 +33,15 @@ import { rawDataToString } from './rawData'; export interface BufferReplaySource { getBufferedSince( sessionId: string, - lastSeq: number, - ): { + cursor: SessionCursor, + ): Promise<{ events: Array<{ seq: number; envelope: EventEnvelope }>; - resyncRequired: boolean; + resyncRequired: 'buffer_overflow' | 'session_recreated' | 'epoch_changed' | false; currentSeq: number; - }; + epoch: string; + }>; + + getCursor(sessionId: string): Promise<{ seq: number; epoch: string }>; } export interface AbortHandler { @@ -94,7 +101,8 @@ export class WsConnection { public readonly subscriptions = new Set(); - public readonly lastSeqBySession = new Map(); + /** Last cursor each subscribed session was synced from (client-claimed). */ + public readonly cursorsBySession = new Map(); private readonly socket: WebSocket; private readonly logger: ILogService; @@ -126,6 +134,7 @@ export class WsConnection { this.send( buildServerHello({ ws_connection_id: this.id, + protocol_version: WS_PROTOCOL_VERSION, heartbeat_ms: this.pingIntervalMs, max_event_buffer_size: this.maxEventBufferSize, capabilities: { event_batching: false, compression: false }, @@ -148,21 +157,35 @@ export class WsConnection { this.logger.warn('non-json ws frame; ignoring'); return; } - const result = clientControlMessageSchema.safeParse(parsed); + const type = frameType(parsed); + if (type === undefined) { + this.logger.warn('invalid control message type'); + return; + } + const operation = getClientControlOperation(type); + if (operation === undefined) { + this.logger.warn({ type }, 'unknown control message type'); + return; + } + const result = operation.messageSchema.safeParse(parsed); if (!result.success) { this.logger.warn({ issues: result.error.issues.length }, 'invalid control message'); return; } - const msg = result.data; + const msg = result.data as ClientControlMessage; switch (msg.type) { case 'client_hello': - this.onClientHello(msg); + void this.onClientHello(msg).catch((err: unknown) => { + this.logger.warn({ err: String(err) }, 'client_hello handler failed'); + }); break; case 'pong': this.onPong(); break; case 'subscribe': - this.onSubscribe(msg); + void this.onSubscribe(msg).catch((err: unknown) => { + this.logger.warn({ err: String(err) }, 'subscribe handler failed'); + }); break; case 'unsubscribe': this.onUnsubscribe(msg); @@ -184,72 +207,69 @@ export class WsConnection { } } - private onClientHello(msg: ClientHelloMessage): void { + private async onClientHello(msg: ClientHelloMessage): Promise { this.gotClientHello = true; - const { subscriptions, last_seq_by_session } = msg.payload; - const accepted: string[] = []; - const resyncRequired: string[] = []; + const { subscriptions, cursors } = msg.payload; - for (const sid of subscriptions) { - this.subscribe(sid); - accepted.push(sid); - } - - if (last_seq_by_session) { - for (const [sid, lastSeq] of Object.entries(last_seq_by_session)) { - this.lastSeqBySession.set(sid, lastSeq); - - if (!this.subscriptions.has(sid)) { - this.subscribe(sid); - accepted.push(sid); - } - const result = this.wsBroadcast.getBufferedSince(sid, lastSeq); - if (result.resyncRequired) { - this.send(buildResyncRequired(sid, 'buffer_overflow', result.currentSeq)); - resyncRequired.push(sid); - } else { - for (const entry of result.events) { - this.send(entry.envelope); - } - } - } - } + const sync = await this.syncSessions(subscriptions, cursors); this.logger.info( { - acceptedCount: accepted.length, - resyncRequiredCount: resyncRequired.length, + acceptedCount: sync.accepted.length, + resyncRequiredCount: sync.resyncRequired.length, }, 'client hello', ); this.send( buildAck(msg.id, 0, 'success', { - accepted_subscriptions: accepted, - resync_required: resyncRequired, + accepted_subscriptions: sync.accepted, + resync_required: sync.resyncRequired, + cursors: sync.serverCursors, }), ); } - private onSubscribe(msg: SubscribeMessage): void { - const { session_ids, last_seq_by_session, watch_fs } = msg.payload; - this.logger.info( - { sessionIds: session_ids, lastSeqBySession: last_seq_by_session, hasWatchFs: !!watch_fs }, - '[DBG ws.onSubscribe] received subscribe', - ); + /** + * Shared client_hello/subscribe session sync: + * 1. register the subscription FIRST (live events flow immediately; + * the client dedups overlap by seq), + * 2. replay durable events past the client's cursor, or emit + * `resync_required` when the cursor cannot be served, + * 3. report the server-side `{seq, epoch}` cursor for every accepted + * session so the client can adopt the current epoch. + */ + private async syncSessions( + sessionIds: readonly string[], + cursors: CursorsBySession | undefined, + ): Promise<{ + accepted: string[]; + resyncRequired: string[]; + serverCursors: CursorsBySession; + }> { const accepted: string[] = []; const resyncRequired: string[] = []; + const serverCursors: CursorsBySession = {}; - for (const sid of session_ids) { - this.subscribe(sid); - accepted.push(sid); + for (const sid of sessionIds) { + if (!this.subscriptions.has(sid)) { + this.subscribe(sid); + } + if (!accepted.includes(sid)) accepted.push(sid); } - if (last_seq_by_session) { - for (const [sid, lastSeq] of Object.entries(last_seq_by_session)) { - this.lastSeqBySession.set(sid, lastSeq); - const result = this.wsBroadcast.getBufferedSince(sid, lastSeq); - if (result.resyncRequired) { - this.send(buildResyncRequired(sid, 'buffer_overflow', result.currentSeq)); + if (cursors) { + for (const [sid, cursor] of Object.entries(cursors)) { + this.cursorsBySession.set(sid, cursor); + if (!this.subscriptions.has(sid)) { + this.subscribe(sid); + } + if (!accepted.includes(sid)) accepted.push(sid); + + const result = await this.wsBroadcast.getBufferedSince(sid, cursor); + if (result.resyncRequired !== false) { + this.send( + buildResyncRequired(sid, result.resyncRequired, result.currentSeq, result.epoch), + ); resyncRequired.push(sid); } else { for (const entry of result.events) { @@ -259,6 +279,26 @@ export class WsConnection { } } + for (const sid of accepted) { + try { + serverCursors[sid] = await this.wsBroadcast.getCursor(sid); + } catch (err) { + this.logger.warn({ sid, err: String(err) }, 'getCursor failed for ack'); + } + } + + return { accepted, resyncRequired, serverCursors }; + } + + private async onSubscribe(msg: SubscribeMessage): Promise { + const { session_ids, cursors, watch_fs } = msg.payload; + this.logger.info( + { sessionIds: session_ids, cursors, hasWatchFs: !!watch_fs }, + 'ws subscribe', + ); + + const sync = await this.syncSessions(session_ids, cursors); + if (watch_fs && this.fsWatchHandler !== undefined) { for (const [sid, cfg] of Object.entries(watch_fs)) { if (cfg.paths.length === 0) continue; @@ -284,9 +324,10 @@ export class WsConnection { this.send( buildAck(msg.id, 0, 'success', { - accepted, + accepted: sync.accepted, not_found: [], - resync_required: resyncRequired, + resync_required: sync.resyncRequired, + cursors: sync.serverCursors, }), ); } @@ -295,6 +336,7 @@ export class WsConnection { const { session_ids } = msg.payload; for (const sid of session_ids) { this.unsubscribe(sid); + this.cursorsBySession.delete(sid); if (this.fsWatchHandler !== undefined) { @@ -525,3 +567,11 @@ export class WsConnection { return this.gotClientHello; } } + +function frameType(value: unknown): string | undefined { + if (typeof value !== 'object' || value === null || !('type' in value)) { + return undefined; + } + const type = (value as { type?: unknown }).type; + return typeof type === 'string' ? type : undefined; +} diff --git a/packages/daemon/src/ws/protocol.ts b/packages/daemon/src/ws/protocol.ts index 1af6ab205..223a8377a 100644 --- a/packages/daemon/src/ws/protocol.ts +++ b/packages/daemon/src/ws/protocol.ts @@ -25,6 +25,7 @@ import { ulid } from 'ulid'; /** WS.md §3.1: `server_hello.payload`. */ export interface ServerHelloPayload { ws_connection_id: string; + protocol_version: number; heartbeat_ms: number; max_event_buffer_size: number; capabilities: { @@ -93,6 +94,16 @@ export function buildAck

(id: string, code: number, msg: string, payload: P): export interface EventEnvelope

{ type: string; seq: number; + /** Journal epoch the seq belongs to (cursor invalidation across restarts). */ + epoch?: string; + /** + * True for ephemeral frames (deltas / progress / periodic status): they + * carry the current durable watermark as `seq`, do not advance it, and are + * never replayed after a reconnect. + */ + volatile?: boolean; + /** Pre-append stream offset for volatile text-delta frames. */ + offset?: number; session_id: string; timestamp: string; payload: P; @@ -102,11 +113,15 @@ export function buildEventEnvelope( seq: number, sessionId: string, event: Event, + opts: { epoch?: string; volatile?: boolean; offset?: number } = {}, ): EventEnvelope { const type = (event as { type?: string }).type ?? 'event.unknown'; return { type, seq, + ...(opts.epoch !== undefined ? { epoch: opts.epoch } : {}), + ...(opts.volatile === true ? { volatile: true } : {}), + ...(opts.offset !== undefined ? { offset: opts.offset } : {}), session_id: sessionId, timestamp: new Date().toISOString(), payload: event, @@ -115,32 +130,37 @@ export function buildEventEnvelope( /** * WS.md §3.6: `resync_required` system message (S→C). Sent when the - * client's claimed `last_seq` for a session is older than the ring buffer - * retains (`lastSeq + 1 < oldestSeq`). The client should drop its local - * cache for that session and `GET /sessions/{id}/messages` to rebuild, - * then re-`subscribe` with `last_seq_by_session[sid] = current_seq`. - * - * `reason` is `'buffer_overflow'` unless session deletion / re-creation with - * the same id requires `'session_recreated'`. + * client's cursor cannot be served incrementally: the gap exceeds the replay + * cap (`buffer_overflow`), the cursor's journal epoch does not match + * (`epoch_changed`), or the session was recreated (`session_recreated`). + * The client should rebuild via `GET /sessions/{id}/snapshot` and + * re-`subscribe` with `cursors[sid] = { seq: as_of_seq, epoch }`. */ export interface ResyncRequiredFrame { type: 'resync_required'; timestamp: string; payload: { session_id: string; - reason: 'buffer_overflow' | 'session_recreated'; + reason: 'buffer_overflow' | 'session_recreated' | 'epoch_changed'; current_seq: number; + epoch?: string; }; } export function buildResyncRequired( sessionId: string, - reason: 'buffer_overflow' | 'session_recreated', + reason: 'buffer_overflow' | 'session_recreated' | 'epoch_changed', currentSeq: number, + epoch?: string, ): ResyncRequiredFrame { return { type: 'resync_required', timestamp: new Date().toISOString(), - payload: { session_id: sessionId, reason, current_seq: currentSeq }, + payload: { + session_id: sessionId, + reason, + current_seq: currentSeq, + ...(epoch !== undefined ? { epoch } : {}), + }, }; } diff --git a/packages/daemon/test/services.test.ts b/packages/daemon/test/services.test.ts index 520d0e3c2..fccf81d2e 100644 --- a/packages/daemon/test/services.test.ts +++ b/packages/daemon/test/services.test.ts @@ -1,5 +1,10 @@ +import { mkdtempSync, rmSync } from 'node:fs'; +import { mkdtemp, rm } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { @@ -17,6 +22,7 @@ import { ILogService, IQuestionService, type FsWatcherServiceOptions, + type IEnvironmentService, type ILogService as ILoggerT, type ISessionService, } from '@moonshot-ai/services'; @@ -121,6 +127,15 @@ type TestFsWatcher = ReturnType { testLogger = new TestLogger(); const collection = new ServiceCollection([ILogService, testLogger]); @@ -129,10 +144,29 @@ beforeEach(() => { afterEach(() => { ix.dispose(); + for (const dir of tmpHomeDirs.splice(0)) { + rmSync(dir, { recursive: true, force: true }); + } }); describe('WSBroadcastService (WS transport pump)', () => { - it('publishes event with seq=1, broadcasts to subscribers, advances seq monotonically per session', () => { + let homeDir: string; + + const makeEnv = (): IEnvironmentService => ({ + _serviceBrand: undefined, + homeDir, + configPath: `${homeDir}/config.toml`, + }); + + beforeEach(async () => { + homeDir = await mkdtemp(join(tmpdir(), 'kimi-ws-broadcast-')); + }); + + afterEach(async () => { + await rm(homeDir, { recursive: true, force: true }); + }); + + it('publishes event with seq=1, broadcasts to subscribers, advances seq monotonically per session', async () => { const clients = new FakeSessionClients(); const c1 = fakeConn('conn_a'); const c2 = fakeConn('conn_b'); @@ -140,24 +174,26 @@ describe('WSBroadcastService (WS transport pump)', () => { clients.subscribe(c2, 'sid_test'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); - bus.publish({ type: 'fake.x', sessionId: 'sid_test' } as unknown as Event); - bus.publish({ type: 'fake.y', sessionId: 'sid_test' } as unknown as Event); + const broadcast = new WSBroadcastService(bus, testLogger, clients, makeEnv()); + bus.publish({ type: 'fake.x', sessionId: 'sid_test', agentId: 'main' } as unknown as Event); + bus.publish({ type: 'fake.y', sessionId: 'sid_test', agentId: 'main' } as unknown as Event); + await broadcast._drainForTest('sid_test'); expect(c1.sent.length).toBe(2); expect(c2.sent.length).toBe(2); - const env1 = c1.sent[0] as { seq: number; session_id: string; type: string }; + const env1 = c1.sent[0] as { seq: number; session_id: string; type: string; epoch?: string }; const env2 = c1.sent[1] as { seq: number; session_id: string; type: string }; expect(env1.seq).toBe(1); expect(env1.session_id).toBe('sid_test'); expect(env1.type).toBe('fake.x'); + expect(env1.epoch).toMatch(/^ep_/); expect(env2.seq).toBe(2); expect(env2.type).toBe('fake.y'); broadcast.dispose(); bus.dispose(); }); - it('per-session seq counters are independent', () => { + it('per-session seq counters are independent', async () => { const clients = new FakeSessionClients(); const cA = fakeConn('conn_a'); const cB = fakeConn('conn_b'); @@ -165,10 +201,12 @@ describe('WSBroadcastService (WS transport pump)', () => { clients.subscribe(cB, 'sid_b'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); - bus.publish({ type: 'e1', sessionId: 'sid_a' } as unknown as Event); - bus.publish({ type: 'e1', sessionId: 'sid_b' } as unknown as Event); - bus.publish({ type: 'e2', sessionId: 'sid_a' } as unknown as Event); + const broadcast = new WSBroadcastService(bus, testLogger, clients, makeEnv()); + bus.publish({ type: 'e1', sessionId: 'sid_a', agentId: 'main' } as unknown as Event); + bus.publish({ type: 'e1', sessionId: 'sid_b', agentId: 'main' } as unknown as Event); + bus.publish({ type: 'e2', sessionId: 'sid_a', agentId: 'main' } as unknown as Event); + await broadcast._drainForTest('sid_a'); + await broadcast._drainForTest('sid_b'); const aSeqs = cA.sent.map((m) => (m as { seq: number }).seq); const bSeqs = cB.sent.map((m) => (m as { seq: number }).seq); @@ -180,7 +218,7 @@ describe('WSBroadcastService (WS transport pump)', () => { bus.dispose(); }); - it('does not broadcast to connections subscribed to a different session', () => { + it('does not broadcast to connections subscribed to a different session', async () => { const clients = new FakeSessionClients(); const onA = fakeConn('conn_a'); const onOther = fakeConn('conn_other'); @@ -188,8 +226,9 @@ describe('WSBroadcastService (WS transport pump)', () => { clients.subscribe(onOther, 'sid_other'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); - bus.publish({ type: 'evt', sessionId: 'sid_a' } as unknown as Event); + const broadcast = new WSBroadcastService(bus, testLogger, clients, makeEnv()); + bus.publish({ type: 'evt', sessionId: 'sid_a', agentId: 'main' } as unknown as Event); + await broadcast._drainForTest('sid_a'); expect(onA.sent.length).toBe(1); expect(onOther.sent.length).toBe(0); broadcast.dispose(); @@ -203,7 +242,7 @@ describe('WSBroadcastService (WS transport pump)', () => { const warnSpy = vi.spyOn(testLogger, 'warn'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); + const broadcast = new WSBroadcastService(bus, testLogger, clients, makeEnv()); bus.publish({ type: 'no_sid' } as unknown as Event); expect(c.sent.length).toBe(0); @@ -212,45 +251,168 @@ describe('WSBroadcastService (WS transport pump)', () => { bus.dispose(); }); - it('post-dispose, publish reaches no subscribers (broadcast unsubscribed)', () => { + it('post-dispose, publish reaches no subscribers (broadcast unsubscribed)', async () => { const clients = new FakeSessionClients(); const c = fakeConn(); clients.subscribe(c, 'sid_x'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); + const broadcast = new WSBroadcastService(bus, testLogger, clients, makeEnv()); broadcast.dispose(); - bus.publish({ type: 'late', sessionId: 'sid_x' } as unknown as Event); + bus.publish({ type: 'late', sessionId: 'sid_x', agentId: 'main' } as unknown as Event); + await new Promise((r) => setTimeout(r, 10)); expect(c.sent.length).toBe(0); bus.dispose(); }); - it('getBufferedSince returns events with seq > lastSeq when buffer covers the gap', () => { + it('getBufferedSince returns events with seq > cursor.seq when the gap is serveable', async () => { const clients = new FakeSessionClients(); const c = fakeConn(); clients.subscribe(c, 'sid_test'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); + const broadcast = new WSBroadcastService(bus, testLogger, clients, makeEnv()); for (let i = 0; i < 5; i++) { - bus.publish({ type: `e${i}`, sessionId: 'sid_test' } as unknown as Event); + bus.publish({ type: `e${i}`, sessionId: 'sid_test', agentId: 'main' } as unknown as Event); } - const replay = broadcast.getBufferedSince('sid_test', 2); + const replay = await broadcast.getBufferedSince('sid_test', { seq: 2 }); expect(replay.resyncRequired).toBe(false); expect(replay.events.map((e) => e.seq)).toEqual([3, 4, 5]); expect(replay.currentSeq).toBe(5); + expect(replay.epoch).toMatch(/^ep_/); broadcast.dispose(); bus.dispose(); }); - it('getBufferedSince returns empty + currentSeq=0 for a never-seen session', () => { + it('getBufferedSince forces a resync for a cursor ahead of the journal (stale v1 cursor)', async () => { const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, new FakeSessionClients()); - const replay = broadcast.getBufferedSince('sid_new', 5); + const broadcast = new WSBroadcastService(bus, testLogger, new FakeSessionClients(), makeEnv()); + const replay = await broadcast.getBufferedSince('sid_new', { seq: 5 }); expect(replay.events).toEqual([]); - expect(replay.resyncRequired).toBe(false); + expect(replay.resyncRequired).toBe('epoch_changed'); expect(replay.currentSeq).toBe(0); broadcast.dispose(); bus.dispose(); }); + + it('getBufferedSince forces a resync on epoch mismatch', async () => { + const bus = new EventService(); + const broadcast = new WSBroadcastService(bus, testLogger, new FakeSessionClients(), makeEnv()); + bus.publish({ type: 'e', sessionId: 'sid_e', agentId: 'main' } as unknown as Event); + await broadcast._drainForTest('sid_e'); + const replay = await broadcast.getBufferedSince('sid_e', { seq: 0, epoch: 'ep_other' }); + expect(replay.resyncRequired).toBe('epoch_changed'); + broadcast.dispose(); + bus.dispose(); + }); + + it('seq and epoch survive a daemon restart (journal recovery) and serve replay from disk', async () => { + const bus1 = new EventService(); + const b1 = new WSBroadcastService(bus1, testLogger, new FakeSessionClients(), makeEnv()); + for (let i = 0; i < 3; i++) { + bus1.publish({ type: `e${i}`, sessionId: 'sid_p', agentId: 'main' } as unknown as Event); + } + const before = await b1.getCursor('sid_p'); + expect(before.seq).toBe(3); + b1.dispose(); + bus1.dispose(); + // Let the write-behind flush settle. + await new Promise((r) => setTimeout(r, 50)); + + const bus2 = new EventService(); + const b2 = new WSBroadcastService(bus2, testLogger, new FakeSessionClients(), makeEnv()); + const after = await b2.getCursor('sid_p'); + expect(after.seq).toBe(3); + expect(after.epoch).toBe(before.epoch); + + // Replay across the restart comes from the on-disk journal. + const replay = await b2.getBufferedSince('sid_p', { seq: 1, epoch: before.epoch }); + expect(replay.resyncRequired).toBe(false); + expect(replay.events.map((e) => e.seq)).toEqual([2, 3]); + + // New events continue the persisted seq. + bus2.publish({ type: 'e3', sessionId: 'sid_p', agentId: 'main' } as unknown as Event); + await b2._drainForTest('sid_p'); + expect(b2._currentSeqForTest('sid_p')).toBe(4); + b2.dispose(); + bus2.dispose(); + }); + + it('volatile events ride the watermark, are flagged, and are not journaled or replayed', async () => { + const clients = new FakeSessionClients(); + const c = fakeConn(); + clients.subscribe(c, 'sid_v'); + const bus = new EventService(); + const broadcast = new WSBroadcastService(bus, testLogger, clients, makeEnv()); + + bus.publish({ + type: 'turn.started', + sessionId: 'sid_v', + agentId: 'main', + turnId: 1, + origin: { kind: 'user' }, + } as unknown as Event); + bus.publish({ + type: 'assistant.delta', + sessionId: 'sid_v', + agentId: 'main', + turnId: 1, + delta: 'hel', + } as unknown as Event); + bus.publish({ + type: 'assistant.delta', + sessionId: 'sid_v', + agentId: 'main', + turnId: 1, + delta: 'lo', + } as unknown as Event); + await broadcast._drainForTest('sid_v'); + + expect(c.sent.length).toBe(3); + const turnStarted = c.sent[0] as { seq: number; volatile?: boolean }; + const delta1 = c.sent[1] as { seq: number; volatile?: boolean; offset?: number }; + const delta2 = c.sent[2] as { seq: number; volatile?: boolean; offset?: number }; + expect(turnStarted.seq).toBe(1); + expect(turnStarted.volatile).toBeUndefined(); + expect(delta1.volatile).toBe(true); + expect(delta1.seq).toBe(1); // watermark, not advanced + expect(delta1.offset).toBe(0); + expect(delta2.offset).toBe(3); + + // Replay from 0 returns only the durable event. + const replay = await broadcast.getBufferedSince('sid_v', { seq: 0 }); + expect(replay.events.map((e) => e.seq)).toEqual([1]); + expect(replay.currentSeq).toBe(1); + + // The in-flight turn snapshot has the accumulated text. + const snap = await broadcast.getSnapshotState('sid_v'); + expect(snap.seq).toBe(1); + expect(snap.inFlightTurn?.assistant_text).toBe('hello'); + broadcast.dispose(); + bus.dispose(); + }); + + it('getSnapshotState clears the in-flight turn after turn.ended', async () => { + const bus = new EventService(); + const broadcast = new WSBroadcastService(bus, testLogger, new FakeSessionClients(), makeEnv()); + bus.publish({ + type: 'turn.started', + sessionId: 'sid_t', + agentId: 'main', + turnId: 1, + origin: { kind: 'user' }, + } as unknown as Event); + bus.publish({ + type: 'turn.ended', + sessionId: 'sid_t', + agentId: 'main', + turnId: 1, + reason: 'completed', + } as unknown as Event); + const snap = await broadcast.getSnapshotState('sid_t'); + expect(snap.inFlightTurn).toBeNull(); + expect(snap.seq).toBe(2); + broadcast.dispose(); + bus.dispose(); + }); }); describe('FsWatcherService', () => { @@ -328,7 +490,7 @@ describe('ApprovalService (broadcasts + resolve-by-approval_id)', () => { const conn = fakeConn('conn_subscriber'); clients.subscribe(conn, 'sess_1'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); + const broadcast = new WSBroadcastService(bus, testLogger, clients, tmpEnv()); const broker = new ApprovalService(testLogger, bus); return { broker, bus, broadcast, clients, conn }; } @@ -353,6 +515,7 @@ describe('ApprovalService (broadcasts + resolve-by-approval_id)', () => { action: 'Run', display: { kind: 'generic', summary: 'test' }, } as Parameters[0]); + await broadcast._drainForTest('sess_1'); const approvalId = extractApprovalId(conn.sent); expect(approvalId).toBeDefined(); @@ -361,6 +524,7 @@ describe('ApprovalService (broadcasts + resolve-by-approval_id)', () => { const response: ApprovalResponse = { decision: 'approved' }; broker.resolve(approvalId!, response); await expect(pending).resolves.toEqual(response); + await broadcast._drainForTest('sess_1'); const resolvedFrame = conn.sent.find( (f) => (f as { type: string }).type === 'event.approval.resolved', @@ -389,6 +553,7 @@ describe('ApprovalService (broadcasts + resolve-by-approval_id)', () => { await expect(pending).rejects.toMatchObject({ name: 'ApprovalExpiredError', }); + await broadcast._drainForTest('sess_1'); const expiredFrame = conn.sent.find( (f) => (f as { type: string }).type === 'event.approval.expired', ); @@ -446,7 +611,7 @@ describe('QuestionService (broadcasts + dismiss)', () => { const conn = fakeConn('conn_q_subscriber'); clients.subscribe(conn, 's'); const bus = new EventService(); - const broadcast = new WSBroadcastService(bus, testLogger, clients); + const broadcast = new WSBroadcastService(bus, testLogger, clients, tmpEnv()); const broker = new QuestionService(testLogger, bus); return { broker, bus, broadcast, clients, conn }; } @@ -474,6 +639,7 @@ describe('QuestionService (broadcasts + dismiss)', () => { }, ], } as Parameters[0]); + await broadcast._drainForTest('s'); const questionId = extractQuestionId(conn.sent); expect(questionId).toBeDefined(); @@ -482,6 +648,7 @@ describe('QuestionService (broadcasts + dismiss)', () => { const response: QuestionResult = { answers: { q_0: 'opt_0_0' } }; broker.resolve(questionId!, response); await expect(pending).resolves.toEqual(response); + await broadcast._drainForTest('s'); const answeredFrame = conn.sent.find( (f) => (f as { type: string }).type === 'event.question.answered', @@ -505,12 +672,14 @@ describe('QuestionService (broadcasts + dismiss)', () => { }, ], } as Parameters[0]); + await broadcast._drainForTest('s'); const questionId = extractQuestionId(conn.sent); expect(questionId).toBeDefined(); broker.dismiss(questionId!); await expect(pending).resolves.toBeNull(); + await broadcast._drainForTest('s'); const dismissedFrame = conn.sent.find( (f) => (f as { type: string }).type === 'event.question.dismissed', @@ -536,6 +705,7 @@ describe('QuestionService (broadcasts + dismiss)', () => { } as Parameters[0]); await expect(pending).rejects.toMatchObject({ name: 'QuestionExpiredError' }); + await broadcast._drainForTest('s'); const expiredFrame = conn.sent.find( (f) => (f as { type: string }).type === 'event.question.expired', ); diff --git a/packages/daemon/test/snapshot.e2e.test.ts b/packages/daemon/test/snapshot.e2e.test.ts new file mode 100644 index 000000000..bbe288df0 --- /dev/null +++ b/packages/daemon/test/snapshot.e2e.test.ts @@ -0,0 +1,216 @@ +/** + * `GET /api/v1/sessions/{sid}/snapshot` end-to-end tests (v2 sync protocol). + * + * Bootstrap mirrors `messages.e2e.test.ts`: real daemon (port 0, tmp lock + + * bridge home), endpoints exercised via `app.inject(...)`. + * + * Coverage: + * - Fresh idle session → as_of_seq=0, ep_* epoch, empty messages, null + * in_flight_turn, empty pending lists. + * - Durable events published → as_of_seq advances and matches the WS + * broadcast cursor (the snapshot↔stream alignment invariant). + * - Mid-turn snapshot → in_flight_turn carries accumulated assistant text + * and running tools; turn.ended clears it. + * - Unknown session id → 40401. + */ + +import { mkdtempSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; + +import { pino } from 'pino'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import type { Event, SessionSnapshotResponse } from '@moonshot-ai/protocol'; +import { IEventService } from '@moonshot-ai/services'; + +import { IRestGateway, IWSBroadcastService, startDaemon, type RunningDaemon } from '../src'; +import { WSBroadcastService } from '#/services/gateway/wsBroadcastService'; + +let tmpDir: string; +let lockPath: string; +let bridgeHome: string; +let daemon: RunningDaemon | undefined; + +beforeEach(() => { + tmpDir = mkdtempSync(join(tmpdir(), 'kimi-daemon-snapshot-test-')); + lockPath = join(tmpDir, 'lock'); + bridgeHome = mkdtempSync(join(tmpdir(), 'kimi-daemon-snapshot-home-')); +}); + +afterEach(async () => { + try { + await daemon?.close(); + } catch { + // ignore + } + daemon = undefined; + rmSync(tmpDir, { recursive: true, force: true }); + rmSync(bridgeHome, { recursive: true, force: true }); +}); + +async function bootDaemon(): Promise { + daemon = await startDaemon({ + host: '127.0.0.1', + port: 0, + lockPath, + logger: pino({ level: 'silent' }), + coreProcessOptions: { homeDir: bridgeHome }, + }); + return daemon; +} + +function appOf(r: RunningDaemon): { + inject: (req: unknown) => Promise<{ statusCode: number; json: () => unknown }>; +} { + return r.services.invokeFunction((a) => { + const gw = a.get(IRestGateway); + return gw.app as unknown as { + inject: (req: unknown) => Promise<{ statusCode: number; json: () => unknown }>; + }; + }); +} + +function envelopeOf(body: unknown): { + code: number; + msg: string; + data: T | null; + request_id: string; +} { + return body as { code: number; msg: string; data: T | null; request_id: string }; +} + +async function createSession(r: RunningDaemon): Promise { + const res = await appOf(r).inject({ + method: 'POST', + url: '/api/v1/sessions', + payload: { metadata: { cwd: join(tmpDir, 'workspace') } }, + }); + const env = envelopeOf<{ id: string }>(res.json()); + if (env.code !== 0 || env.data === null) { + throw new Error(`failed to create session: ${JSON.stringify(env)}`); + } + return env.data.id; +} + +async function getSnapshot( + r: RunningDaemon, + sid: string, +): Promise<{ statusCode: number; env: ReturnType> }> { + const res = await appOf(r).inject({ + method: 'GET', + url: `/api/v1/sessions/${sid}/snapshot`, + }); + return { statusCode: res.statusCode, env: envelopeOf(res.json()) }; +} + +describe('GET /api/v1/sessions/{sid}/snapshot (v2 initial sync)', () => { + it('idle fresh session → empty snapshot at the session-creation watermark', async () => { + const r = await bootDaemon(); + const sid = await createSession(r); + + const { statusCode, env } = await getSnapshot(r, sid); + expect(statusCode).toBe(200); + expect(env.code).toBe(0); + const data = env.data!; + expect(data.epoch).toMatch(/^ep_/); + expect(data.session.id).toBe(sid); + expect(data.messages.items).toEqual([]); + expect(data.messages.has_more).toBe(false); + expect(data.in_flight_turn).toBeNull(); + expect(data.pending_approvals).toEqual([]); + expect(data.pending_questions).toEqual([]); + }); + + it('as_of_seq matches the WS broadcast cursor after durable events', async () => { + const r = await bootDaemon(); + const sid = await createSession(r); + const bus = r.services.invokeFunction((a) => a.get(IEventService)); + const broadcast = r.services.invokeFunction( + (a) => a.get(IWSBroadcastService), + ) as WSBroadcastService; + + const baseline = (await broadcast.getCursor(sid)).seq; + bus.publish({ type: 'evt.x', sessionId: sid, agentId: 'main' } as unknown as Event); + bus.publish({ type: 'evt.y', sessionId: sid, agentId: 'main' } as unknown as Event); + + const { env } = await getSnapshot(r, sid); + const cursor = await broadcast.getCursor(sid); + expect(env.data!.as_of_seq).toBe(baseline + 2); + expect(env.data!.as_of_seq).toBe(cursor.seq); + expect(env.data!.epoch).toBe(cursor.epoch); + }); + + it('mid-turn snapshot carries accumulated text + running tools; turn.ended clears it', async () => { + const r = await bootDaemon(); + const sid = await createSession(r); + const bus = r.services.invokeFunction((a) => a.get(IEventService)); + + bus.publish({ + type: 'turn.started', + sessionId: sid, + agentId: 'main', + turnId: 1, + origin: { kind: 'user' }, + } as unknown as Event); + bus.publish({ + type: 'assistant.delta', + sessionId: sid, + agentId: 'main', + turnId: 1, + delta: 'partial ans', + } as unknown as Event); + bus.publish({ + type: 'tool.call.started', + sessionId: sid, + agentId: 'main', + turnId: 1, + toolCallId: 'call_1', + name: 'Bash', + args: { command: 'ls' }, + } as unknown as Event); + bus.publish({ + type: 'tool.progress', + sessionId: sid, + agentId: 'main', + turnId: 1, + toolCallId: 'call_1', + update: { kind: 'stdout', text: 'src\n' }, + } as unknown as Event); + + const { env } = await getSnapshot(r, sid); + const turn = env.data!.in_flight_turn; + expect(turn).not.toBeNull(); + expect(turn!.turn_id).toBe(1); + expect(turn!.assistant_text).toBe('partial ans'); + expect(turn!.running_tools).toHaveLength(1); + expect(turn!.running_tools[0]!.tool_call_id).toBe('call_1'); + expect(turn!.running_tools[0]!.last_progress?.text).toBe('src\n'); + + bus.publish({ + type: 'tool.result', + sessionId: sid, + agentId: 'main', + turnId: 1, + toolCallId: 'call_1', + output: 'src', + } as unknown as Event); + bus.publish({ + type: 'turn.ended', + sessionId: sid, + agentId: 'main', + turnId: 1, + reason: 'completed', + } as unknown as Event); + + const after = await getSnapshot(r, sid); + expect(after.env.data!.in_flight_turn).toBeNull(); + }); + + it('returns 40401 for an unknown session id', async () => { + const r = await bootDaemon(); + const { env } = await getSnapshot(r, 'sess_missing'); + expect(env.code).toBe(40401); + expect(env.data).toBeNull(); + }); +}); diff --git a/packages/daemon/test/ws-resync.e2e.test.ts b/packages/daemon/test/ws-resync.e2e.test.ts index e4152524e..01622c6cd 100644 --- a/packages/daemon/test/ws-resync.e2e.test.ts +++ b/packages/daemon/test/ws-resync.e2e.test.ts @@ -1,27 +1,26 @@ /** - * WS ring buffer + resync_required e2e (W5.3 / P0.17). + * WS durable journal + resync_required e2e (v2 sync protocol). * - * Three flows per WS.md §6: + * Flows: * * 1. **Replay**: publish N events; client A disconnects; publish M more; - * A reconnects with `client_hello.last_seq_by_session[sid]=N`; assert - * A receives exactly events N+1..N+M in order. + * A reconnects with `client_hello.cursors[sid] = {seq: N}`; assert A + * receives exactly events N+1..N+M in order. * - * 2. **Resync**: force the buffer to overflow (publish >1000 events). - * Client B connects with a stale `last_seq` (older than `oldestSeq`). - * Assert B receives a `resync_required` frame for that session, NOT - * events. + * 2. **Resync**: publish more than the replay cap (1000). Client B + * connects with a cursor whose gap exceeds the cap. Assert B receives + * `resync_required(buffer_overflow)`, NOT events. * - * 3. **No-op**: client C connects with `last_seq == current_seq`. Assert - * no replay events arrive on the first frames (only the normal ack + - * empty `resync_required`). + * 3. **No-op**: client C connects with `cursor.seq == current_seq`. Assert + * no replay events arrive (only the ack with empty `resync_required`). * - * `maxBufferSize` is the default 1000 baked into `WSBroadcastService`. The - * resync flow publishes 1005 events to trigger eviction; the buffer keeps - * the last 1000. Publishes go through `IEventService.publish(...)`; ring - * buffer state and replay queries are read off the daemon-local - * `IWSBroadcastService` (cast to the concrete class for `_*ForTest` - * helpers). + * 4. **Replay-cap boundaries**: a gap of exactly 1000 is served; 1001 is + * a resync. Memory-tail eviction no longer forces resyncs — the gap is + * served from the on-disk journal when it reaches behind the tail. + * + * Publishes go through `IEventService.publish(...)`; dispatch is async + * (per-session queue), so tests drain via `_drainForTest` before asserting + * watermark state. */ import { mkdtempSync, rmSync } from 'node:fs'; @@ -91,6 +90,8 @@ interface WsFrame { id?: string; code?: number; seq?: number; + epoch?: string; + volatile?: boolean; session_id?: string; [k: string]: unknown; } @@ -163,8 +164,8 @@ async function waitFor(cond: () => boolean, timeoutMs = 2000): Promise { throw new Error(`waitFor: condition not satisfied within ${timeoutMs}ms`); } -describe('WS ring buffer + resync_required (W5.3)', () => { - it('reconnect with last_seq replays buffered events in order', async () => { +describe('WS durable journal + resync_required (v2)', () => { + it('reconnect with a cursor replays buffered events in order', async () => { const r = await spawn(); // Client A: connect, subscribe to sid_test, capture seq up to 5. @@ -192,6 +193,7 @@ describe('WS ring buffer + resync_required (W5.3)', () => { for (let i = 1; i <= 5; i++) { const ev = await receiveType(a1, `evt.${i}`, 1000); expect(ev.seq).toBe(i); + expect(ev.epoch).toMatch(/^ep_/); } // Disconnect A1. @@ -207,7 +209,7 @@ describe('WS ring buffer + resync_required (W5.3)', () => { bus.publish({ type: `evt.${i}`, sessionId: 'sid_test' } as unknown as Event); } - // Reconnect with last_seq=5 — should replay 6, 7, 8 in order. + // Reconnect with cursor seq=5 — should replay 6, 7, 8 in order. const a2 = await openConn(wsUrl(r.address)); await receiveType(a2, 'server_hello', 1000); a2.ws.send( @@ -217,7 +219,7 @@ describe('WS ring buffer + resync_required (W5.3)', () => { payload: { client_id: 'A', subscriptions: ['sid_test'], - last_seq_by_session: { sid_test: 5 }, + cursors: { sid_test: { seq: 5 } }, }, }), ); @@ -231,17 +233,22 @@ describe('WS ring buffer + resync_required (W5.3)', () => { const ack = await receiveType(a2, 'ack', 1000); expect(ack.code).toBe(0); - const ackPayload = ack.payload as { resync_required: string[] }; + const ackPayload = ack.payload as { + resync_required: string[]; + cursors?: Record; + }; expect(ackPayload.resync_required).toEqual([]); + expect(ackPayload.cursors?.['sid_test']?.seq).toBe(8); + expect(ackPayload.cursors?.['sid_test']?.epoch).toMatch(/^ep_/); a2.ws.close(); }); - it('client connects with last_seq beyond ring-buffer retention → resync_required', async () => { + it('client connects with a gap beyond the replay cap → resync_required(buffer_overflow)', async () => { const r = await spawn(); - // Force the buffer to overflow. With the spec-faithful 1000-cap, we - // publish 1005 events. After that, oldestSeq is 6 (events 1..5 evicted). + // Publish past the replay cap. With the 1000-event cap, a client at + // seq=3 faces a 1002-event gap — snapshot rebuild is cheaper. const bus = r.services.invokeFunction((acc) => acc.get(IEventService)); const broadcast = r.services.invokeFunction( (acc) => acc.get(IWSBroadcastService), @@ -249,11 +256,10 @@ describe('WS ring buffer + resync_required (W5.3)', () => { for (let i = 1; i <= 1005; i++) { bus.publish({ type: 'evt', sessionId: 'sid_test' } as unknown as Event); } + await broadcast._drainForTest('sid_test'); expect(broadcast._currentSeqForTest('sid_test')).toBe(1005); expect(broadcast._bufferLengthForTest('sid_test')).toBe(1000); - expect(broadcast._oldestSeqForTest('sid_test')).toBe(6); - // Client connects with last_seq=3 — gap is too big (events 4, 5 are gone). const conn = await openConn(wsUrl(r.address)); await receiveType(conn, 'server_hello', 1000); conn.ws.send( @@ -263,7 +269,7 @@ describe('WS ring buffer + resync_required (W5.3)', () => { payload: { client_id: 'C', subscriptions: ['sid_test'], - last_seq_by_session: { sid_test: 3 }, + cursors: { sid_test: { seq: 3 } }, }, }), ); @@ -273,10 +279,12 @@ describe('WS ring buffer + resync_required (W5.3)', () => { session_id: string; reason: string; current_seq: number; + epoch?: string; }; expect(resyncPayload.session_id).toBe('sid_test'); expect(resyncPayload.reason).toBe('buffer_overflow'); expect(resyncPayload.current_seq).toBe(1005); + expect(resyncPayload.epoch).toMatch(/^ep_/); const ack = await receiveType(conn, 'ack', 1000); const ackPayload = ack.payload as { resync_required: string[] }; @@ -285,12 +293,47 @@ describe('WS ring buffer + resync_required (W5.3)', () => { conn.ws.close(); }); - it('caught-up client (last_seq == current_seq) gets no replay, just empty ack', async () => { + it('cursor from a different epoch → resync_required(epoch_changed)', async () => { const r = await spawn(); const bus = r.services.invokeFunction((acc) => acc.get(IEventService)); + const broadcast = r.services.invokeFunction( + (acc) => acc.get(IWSBroadcastService), + ) as WSBroadcastService; + bus.publish({ type: 'evt.a', sessionId: 'sid_epoch' } as unknown as Event); + await broadcast._drainForTest('sid_epoch'); + + const conn = await openConn(wsUrl(r.address)); + await receiveType(conn, 'server_hello', 1000); + conn.ws.send( + JSON.stringify({ + type: 'client_hello', + id: 'cli_epoch', + payload: { + client_id: 'E', + subscriptions: ['sid_epoch'], + cursors: { sid_epoch: { seq: 1, epoch: 'ep_FROM_ANOTHER_LIFE' } }, + }, + }), + ); + + const resync = await receiveType(conn, 'resync_required', 1000); + const resyncPayload = resync.payload as { reason: string; current_seq: number }; + expect(resyncPayload.reason).toBe('epoch_changed'); + expect(resyncPayload.current_seq).toBe(1); + + conn.ws.close(); + }); + + it('caught-up client (cursor.seq == current_seq) gets no replay, just empty ack', async () => { + const r = await spawn(); + const bus = r.services.invokeFunction((acc) => acc.get(IEventService)); + const broadcast = r.services.invokeFunction( + (acc) => acc.get(IWSBroadcastService), + ) as WSBroadcastService; bus.publish({ type: 'evt.a', sessionId: 'sid_test' } as unknown as Event); bus.publish({ type: 'evt.b', sessionId: 'sid_test' } as unknown as Event); bus.publish({ type: 'evt.c', sessionId: 'sid_test' } as unknown as Event); + await broadcast._drainForTest('sid_test'); const conn = await openConn(wsUrl(r.address)); await receiveType(conn, 'server_hello', 1000); @@ -301,7 +344,7 @@ describe('WS ring buffer + resync_required (W5.3)', () => { payload: { client_id: 'D', subscriptions: ['sid_test'], - last_seq_by_session: { sid_test: 3 }, // == current_seq + cursors: { sid_test: { seq: 3 } }, // == current_seq }, }), ); @@ -323,30 +366,29 @@ describe('WS ring buffer + resync_required (W5.3)', () => { conn.ws.close(); }); - it('ring buffer evicts oldest event when capacity is exceeded', async () => { + it('replay-cap boundaries: gap of exactly 1000 served, 1001 resyncs', async () => { const r = await spawn(); const bus = r.services.invokeFunction((acc) => acc.get(IEventService)); const broadcast = r.services.invokeFunction( (acc) => acc.get(IWSBroadcastService), ) as WSBroadcastService; - // Publish 1002 — buffer should retain seq 3..1002, oldestSeq=3. for (let i = 1; i <= 1002; i++) { bus.publish({ type: 'evt', sessionId: 'sid_evict' } as unknown as Event); } + await broadcast._drainForTest('sid_evict'); expect(broadcast._currentSeqForTest('sid_evict')).toBe(1002); expect(broadcast._bufferLengthForTest('sid_evict')).toBe(1000); - expect(broadcast._oldestSeqForTest('sid_evict')).toBe(3); - // getBufferedSince(sid, 2) → resyncRequired (lastSeq+1=3, oldestSeq=3 → NOT resync; - // lastSeq+1=3 == oldestSeq=3 → NOT resync). Verify boundary. - const replay = broadcast.getBufferedSince('sid_evict', 2); + // Gap of exactly 1000 (cursor.seq=2) → served; first event past the + // memory tail (seq 3..1002 retained) is still seq 3. + const replay = await broadcast.getBufferedSince('sid_evict', { seq: 2 }); expect(replay.resyncRequired).toBe(false); expect(replay.events[0]?.seq).toBe(3); expect(replay.events.length).toBe(1000); - // lastSeq=1 → lastSeq+1=2 < oldestSeq=3 → resync. - const replay2 = broadcast.getBufferedSince('sid_evict', 1); - expect(replay2.resyncRequired).toBe(true); + // Gap of 1001 (cursor.seq=1) → buffer_overflow resync. + const replay2 = await broadcast.getBufferedSince('sid_evict', { seq: 1 }); + expect(replay2.resyncRequired).toBe('buffer_overflow'); expect(replay2.events.length).toBe(0); }); }); diff --git a/packages/protocol/src/__tests__/events.test.ts b/packages/protocol/src/__tests__/events.test.ts index 79c70841d..ac5fdace8 100644 --- a/packages/protocol/src/__tests__/events.test.ts +++ b/packages/protocol/src/__tests__/events.test.ts @@ -4,6 +4,12 @@ import { fileURLToPath } from 'node:url'; import { describe, it, expect } from 'vitest'; +import { + agentEventSchema, + assistantDeltaEventSchema, + eventSchema, + toolCallStartedEventSchema, +} from '../events'; import type { Event } from '../events'; import type { ToolInputDisplay } from '../display'; @@ -49,4 +55,51 @@ describe('events / display re-exports', () => { it('ToolInputDisplay re-export is non-never (12-arm union preserved)', () => { expect(_assertDisplay).toBe(true); }); + + it('validates concrete agent event payloads with Zod schemas', () => { + expect( + assistantDeltaEventSchema.parse({ + type: 'assistant.delta', + turnId: 1, + delta: 'hello', + }), + ).toEqual({ + type: 'assistant.delta', + turnId: 1, + delta: 'hello', + }); + + expect( + toolCallStartedEventSchema.safeParse({ + type: 'tool.call.started', + turnId: 1, + toolCallId: 'call_1', + name: 'bash', + args: { command: 'pwd' }, + display: { kind: 'command', command: 'pwd', language: 'bash' }, + }).success, + ).toBe(true); + }); + + it('rejects unknown event types through the full agent event union', () => { + expect( + agentEventSchema.safeParse({ + type: 'unknown.event', + turnId: 1, + }).success, + ).toBe(false); + }); + + it('validates session-scoped daemon events with agentId and sessionId', () => { + const parsed = eventSchema.parse({ + type: 'turn.started', + agentId: 'agent_1', + sessionId: 'sess_1', + turnId: 1, + origin: { kind: 'user' }, + }); + + expect(parsed.agentId).toBe('agent_1'); + expect(parsed.sessionId).toBe('sess_1'); + }); }); diff --git a/packages/protocol/src/__tests__/snapshot.test.ts b/packages/protocol/src/__tests__/snapshot.test.ts new file mode 100644 index 000000000..ff909d6f9 --- /dev/null +++ b/packages/protocol/src/__tests__/snapshot.test.ts @@ -0,0 +1,130 @@ +import { describe, expect, it } from 'vitest'; + +import { isVolatileEventType, VOLATILE_EVENT_TYPES } from '../events'; +import { + inFlightTurnSchema, + sessionSnapshotResponseSchema, +} from '../rest/snapshot'; + +const TS = '2026-06-11T10:30:00.000Z'; + +const SESSION = { + id: 'sess_1', + workspace_id: 'wd_demo_0123456789ab', + title: 'demo', + created_at: TS, + updated_at: TS, + status: 'running', + metadata: { cwd: '/tmp/demo' }, + agent_config: { model: 'kimi' }, + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_tokens: 0, + cache_creation_tokens: 0, + total_cost_usd: 0, + context_tokens: 0, + context_limit: 0, + turn_count: 0, + }, + permission_rules: [], + message_count: 2, + last_seq: 12, +}; + +describe('rest/snapshot — session snapshot', () => { + it('parses a full snapshot with an in-flight turn', () => { + const result = sessionSnapshotResponseSchema.safeParse({ + as_of_seq: 12, + epoch: 'ep_01ABC', + session: SESSION, + messages: { + items: [ + { + id: 'msg_sess_1_000000', + session_id: 'sess_1', + role: 'user', + content: [{ type: 'text', text: 'hi' }], + created_at: TS, + }, + ], + has_more: false, + }, + in_flight_turn: { + turn_id: 3, + assistant_text: 'partial answer…', + thinking_text: '', + running_tools: [ + { + tool_call_id: 'call_1', + name: 'Bash', + args: { command: 'ls' }, + last_progress: { kind: 'stdout', text: 'src\n' }, + }, + ], + }, + pending_approvals: [], + pending_questions: [], + }); + expect(result.success).toBe(true); + }); + + it('parses an idle snapshot (no in-flight turn)', () => { + const result = sessionSnapshotResponseSchema.safeParse({ + as_of_seq: 0, + epoch: 'ep_01ABC', + session: SESSION, + messages: { items: [], has_more: false }, + in_flight_turn: null, + pending_approvals: [], + pending_questions: [], + }); + expect(result.success).toBe(true); + }); + + it('rejects a snapshot missing the watermark', () => { + const result = sessionSnapshotResponseSchema.safeParse({ + epoch: 'ep_01ABC', + session: SESSION, + messages: { items: [], has_more: false }, + in_flight_turn: null, + pending_approvals: [], + pending_questions: [], + }); + expect(result.success).toBe(false); + }); + + it('in_flight_turn requires accumulated text fields', () => { + expect( + inFlightTurnSchema.safeParse({ turn_id: 1, running_tools: [] }).success, + ).toBe(false); + }); +}); + +describe('events — volatile classification', () => { + it('classifies stream fragments and periodic status as volatile', () => { + for (const type of [ + 'assistant.delta', + 'thinking.delta', + 'tool.call.delta', + 'tool.progress', + 'agent.status.updated', + ]) { + expect(isVolatileEventType(type)).toBe(true); + } + expect(VOLATILE_EVENT_TYPES).toHaveLength(5); + }); + + it('keeps timeline-bearing events durable', () => { + for (const type of [ + 'turn.started', + 'turn.ended', + 'tool.call.started', + 'tool.result', + 'session.meta.updated', + 'error', + ]) { + expect(isVolatileEventType(type)).toBe(false); + } + }); +}); diff --git a/packages/protocol/src/__tests__/ws-control.test.ts b/packages/protocol/src/__tests__/ws-control.test.ts index 8cd7ee2b1..22b50e9c8 100644 --- a/packages/protocol/src/__tests__/ws-control.test.ts +++ b/packages/protocol/src/__tests__/ws-control.test.ts @@ -2,17 +2,27 @@ import { describe, expect, it } from 'vitest'; import { abortMessageSchema, + abortAckMessageSchema, clientControlMessageSchema, + clientControlOperations, clientHelloMessageSchema, + clientHelloAckMessageSchema, + getClientControlOperation, pingMessageSchema, pongMessageSchema, resyncRequiredMessageSchema, serverHelloMessageSchema, + serverSystemOperations, serverSystemMessageSchema, + sessionEventMessageSchema, + subscribeAckMessageSchema, subscribeMessageSchema, + unsubscribeAckMessageSchema, unsubscribeMessageSchema, + watchFsAckMessageSchema, watchFsAddMessageSchema, watchFsRemoveMessageSchema, + wsOperations, wsAckEnvelopeSchema, wsControlEnvelopeSchema, wsErrorMessageSchema, @@ -35,6 +45,21 @@ describe('ws-control — generic envelopes', () => { expect(parsed.seq).toBe(42); }); + it('wsEventEnvelopeSchema accepts a volatile frame carrying the watermark', () => { + const schema = wsEventEnvelopeSchema(z.object({ delta: z.string() })); + const parsed = schema.parse({ + type: 'assistant.delta', + seq: 42, + epoch: 'ep_01ABC', + volatile: true, + session_id: 'sess_1', + timestamp: TS, + payload: { delta: 'hi' }, + }); + expect(parsed.volatile).toBe(true); + expect(parsed.epoch).toBe('ep_01ABC'); + }); + it('wsControlEnvelopeSchema accepts an id-less message', () => { const schema = wsControlEnvelopeSchema(z.object({})); expect(schema.safeParse({ type: 'pong', payload: {} }).success).toBe(true); @@ -59,6 +84,7 @@ describe('ws-control — §3.1 server_hello', () => { timestamp: TS, payload: { ws_connection_id: 'conn_local', + protocol_version: 2, heartbeat_ms: 30000, max_event_buffer_size: 1000, capabilities: { event_batching: false, compression: false }, @@ -67,12 +93,27 @@ describe('ws-control — §3.1 server_hello', () => { expect(result.success).toBe(true); }); + it('rejects a server_hello missing protocol_version', () => { + const result = serverHelloMessageSchema.safeParse({ + type: 'server_hello', + timestamp: TS, + payload: { + ws_connection_id: 'conn_local', + heartbeat_ms: 30000, + max_event_buffer_size: 1000, + capabilities: { event_batching: false, compression: false }, + }, + }); + expect(result.success).toBe(false); + }); + it('rejects a server_hello missing capabilities', () => { const result = serverHelloMessageSchema.safeParse({ type: 'server_hello', timestamp: TS, payload: { ws_connection_id: 'conn_local', + protocol_version: 2, heartbeat_ms: 30000, max_event_buffer_size: 1000, }, @@ -89,12 +130,38 @@ describe('ws-control — §3.2 client_hello', () => { payload: { client_id: 'web_abc', subscriptions: ['sess_1', 'sess_2'], - last_seq_by_session: { sess_1: 99 }, + cursors: { sess_1: { seq: 99, epoch: 'ep_01ABC' } }, }, }); expect(result.success).toBe(true); }); + it('client_hello accepts an epoch-less fresh cursor', () => { + const result = clientHelloMessageSchema.safeParse({ + type: 'client_hello', + id: 'c1', + payload: { + client_id: 'web_abc', + subscriptions: [], + cursors: { sess_1: { seq: 0 } }, + }, + }); + expect(result.success).toBe(true); + }); + + it('client_hello rejects the v1 bare-seq cursor map', () => { + const result = clientHelloMessageSchema.safeParse({ + type: 'client_hello', + id: 'c1', + payload: { + client_id: 'web_abc', + subscriptions: [], + cursors: { sess_1: 99 }, + }, + }); + expect(result.success).toBe(false); + }); + it('rejects a client_hello missing payload.client_id', () => { const result = clientHelloMessageSchema.safeParse({ type: 'client_hello', @@ -249,6 +316,20 @@ describe('ws-control — §3.6 resync_required', () => { expect(result.success).toBe(true); }); + it('parses an epoch_changed resync with the new epoch', () => { + const result = resyncRequiredMessageSchema.safeParse({ + type: 'resync_required', + timestamp: TS, + payload: { + session_id: 'sess_1', + reason: 'epoch_changed', + current_seq: 12, + epoch: 'ep_01DEF', + }, + }); + expect(result.success).toBe(true); + }); + it('rejects an unknown reason', () => { const result = resyncRequiredMessageSchema.safeParse({ type: 'resync_required', @@ -315,3 +396,105 @@ describe('ws-control — discriminated unions', () => { ).toBe(true); }); }); + +describe('ws-control — operation registry', () => { + it('covers every client control frame with a message schema and ack schema', () => { + expect(clientControlOperations.map((op) => op.type)).toEqual([ + 'client_hello', + 'subscribe', + 'unsubscribe', + 'watch_fs_add', + 'watch_fs_remove', + 'abort', + 'pong', + ]); + + for (const op of clientControlOperations) { + expect(op.direction).toBe('client_to_server'); + expect(op.messageSchema).toBeDefined(); + if (op.type !== 'pong') { + expect(op.ackSchema).toBeDefined(); + } + } + }); + + it('looks up client control operations by frame type', () => { + expect(getClientControlOperation('subscribe')?.messageSchema).toBe(subscribeMessageSchema); + expect(getClientControlOperation('launch_missiles')).toBeUndefined(); + }); + + it('defines typed ack message schemas for control responses', () => { + expect( + clientHelloAckMessageSchema.safeParse({ + type: 'ack', + id: 'c1', + code: 0, + msg: 'success', + payload: { accepted_subscriptions: ['sess_1'], resync_required: [] }, + }).success, + ).toBe(true); + expect( + subscribeAckMessageSchema.safeParse({ + type: 'ack', + id: 'c2', + code: 0, + msg: 'success', + payload: { accepted: ['sess_1'], not_found: [], resync_required: [] }, + }).success, + ).toBe(true); + expect( + unsubscribeAckMessageSchema.safeParse({ + type: 'ack', + id: 'c3', + code: 0, + msg: 'success', + payload: { accepted: ['sess_1'], not_found: [], resync_required: [] }, + }).success, + ).toBe(true); + expect( + watchFsAckMessageSchema.safeParse({ + type: 'ack', + id: 'c4', + code: 0, + msg: 'success', + payload: { watched_paths: ['src'], current_count: 1 }, + }).success, + ).toBe(true); + expect( + abortAckMessageSchema.safeParse({ + type: 'ack', + id: 'c5', + code: 0, + msg: 'success', + payload: { aborted: true, at_seq: 10 }, + }).success, + ).toBe(true); + }); + + it('covers server system frames and the session event stream', () => { + expect(serverSystemOperations.map((op) => op.type)).toEqual([ + 'server_hello', + 'ping', + 'resync_required', + 'error', + ]); + + expect( + sessionEventMessageSchema.safeParse({ + type: 'assistant.delta', + seq: 1, + session_id: 'sess_1', + timestamp: TS, + payload: { + type: 'assistant.delta', + agentId: 'agent_1', + sessionId: 'sess_1', + turnId: 1, + delta: 'hello', + }, + }).success, + ).toBe(true); + + expect(wsOperations.some((op) => op.type === 'session_event')).toBe(true); + }); +}); diff --git a/packages/protocol/src/events.ts b/packages/protocol/src/events.ts index c2d703e36..918adea29 100644 --- a/packages/protocol/src/events.ts +++ b/packages/protocol/src/events.ts @@ -1,4 +1,6 @@ -import type { ToolInputDisplay } from './display'; +import { z } from 'zod'; + +import { ToolInputDisplaySchema, type ToolInputDisplay } from './display'; export interface TokenUsage { readonly inputOther: number; @@ -570,3 +572,616 @@ export type AgentEvent = | CronFiredEvent; export type Event = AgentEvent & { agentId: string; sessionId: string }; + +export const tokenUsageSchema = z.object({ + inputOther: z.number(), + output: z.number(), + inputCacheRead: z.number(), + inputCacheCreation: z.number(), +}) satisfies z.ZodType; + +export const finishReasonSchema = z.enum([ + 'completed', + 'tool_calls', + 'truncated', + 'filtered', + 'paused', + 'other', +]) satisfies z.ZodType; + +export const usageStatusSchema = z.object({ + byModel: z.record(z.string(), tokenUsageSchema).optional(), + currentTurn: tokenUsageSchema.optional(), + total: tokenUsageSchema.optional(), +}) satisfies z.ZodType; + +export const permissionModeSchema = z.enum(['manual', 'yolo', 'auto']) satisfies z.ZodType; + +export const skillSourceSchema = z.enum(['project', 'user', 'extra', 'builtin']) satisfies z.ZodType; + +export const userPromptOriginSchema = z.object({ + kind: z.literal('user'), +}) satisfies z.ZodType; + +export const skillActivationOriginSchema = z.object({ + kind: z.literal('skill_activation'), + activationId: z.string(), + skillName: z.string(), + skillArgs: z.string().optional(), + trigger: z.enum(['user-slash', 'model-tool', 'nested-skill']), + skillType: z.string().optional(), + skillPath: z.string().optional(), + skillSource: skillSourceSchema.optional(), +}) satisfies z.ZodType; + +export const injectionOriginSchema = z.object({ + kind: z.literal('injection'), + variant: z.string(), +}) satisfies z.ZodType; + +export const compactionSummaryOriginSchema = z.object({ + kind: z.literal('compaction_summary'), +}) satisfies z.ZodType; + +export const systemTriggerOriginSchema = z.object({ + kind: z.literal('system_trigger'), + name: z.string(), +}) satisfies z.ZodType; + +export const agentCoreBackgroundTaskStatusSchema = z.enum([ + 'running', + 'completed', + 'failed', + 'timed_out', + 'killed', + 'lost', +]) satisfies z.ZodType; + +export const backgroundTaskOriginSchema = z.object({ + kind: z.literal('background_task'), + taskId: z.string(), + status: agentCoreBackgroundTaskStatusSchema, + notificationId: z.string(), +}) satisfies z.ZodType; + +export const cronJobOriginSchema = z.object({ + kind: z.literal('cron_job'), + jobId: z.string(), + cron: z.string(), + recurring: z.boolean(), + coalescedCount: z.number(), + stale: z.boolean(), +}) satisfies z.ZodType; + +export const cronMissedOriginSchema = z.object({ + kind: z.literal('cron_missed'), + count: z.number(), +}) satisfies z.ZodType; + +export const hookResultOriginSchema = z.object({ + kind: z.literal('hook_result'), + event: z.string(), + blocked: z.boolean().optional(), +}) satisfies z.ZodType; + +export const retryOriginSchema = z.object({ + kind: z.literal('retry'), + trigger: z.string().optional(), +}) satisfies z.ZodType; + +export const promptOriginSchema = z.discriminatedUnion('kind', [ + userPromptOriginSchema, + skillActivationOriginSchema, + injectionOriginSchema, + compactionSummaryOriginSchema, + systemTriggerOriginSchema, + backgroundTaskOriginSchema, + cronJobOriginSchema, + cronMissedOriginSchema, + hookResultOriginSchema, + retryOriginSchema, +]) satisfies z.ZodType; + +export const goalStatusSchema = z.enum(['active', 'paused', 'blocked', 'complete']) satisfies z.ZodType; + +export const goalActorSchema = z.enum(['user', 'model', 'runtime', 'system']) satisfies z.ZodType; + +export const goalBudgetLimitsSchema = z.object({ + tokenBudget: z.number().optional(), + turnBudget: z.number().optional(), + wallClockBudgetMs: z.number().optional(), +}) satisfies z.ZodType; + +export const goalBudgetReportSchema = z.object({ + tokenBudget: z.number().nullable(), + turnBudget: z.number().nullable(), + wallClockBudgetMs: z.number().nullable(), + remainingTokens: z.number().nullable(), + remainingTurns: z.number().nullable(), + remainingWallClockMs: z.number().nullable(), + tokenBudgetReached: z.boolean(), + turnBudgetReached: z.boolean(), + wallClockBudgetReached: z.boolean(), + overBudget: z.boolean(), +}) satisfies z.ZodType; + +export const goalSnapshotSchema = z.object({ + goalId: z.string(), + objective: z.string(), + completionCriterion: z.string().optional(), + status: goalStatusSchema, + turnsUsed: z.number(), + tokensUsed: z.number(), + wallClockMs: z.number(), + budget: goalBudgetReportSchema, + terminalReason: z.string().optional(), +}) satisfies z.ZodType; + +export const goalToolResultSchema = z.object({ + goal: goalSnapshotSchema.nullable(), +}) satisfies z.ZodType; + +export const goalChangeStatsSchema = z.object({ + turnsUsed: z.number(), + tokensUsed: z.number(), + wallClockMs: z.number(), +}) satisfies z.ZodType; + +export const goalChangeKindSchema = z.enum(['lifecycle', 'completion']) satisfies z.ZodType; + +export const goalChangeSchema = z.object({ + kind: goalChangeKindSchema, + status: goalStatusSchema.optional(), + reason: z.string().optional(), + stats: goalChangeStatsSchema.optional(), + actor: goalActorSchema.optional(), +}) satisfies z.ZodType; + +export const kimiErrorCodeSchema = z.enum([ + 'config.invalid', + 'session.not_found', + 'session.already_exists', + 'session.id_invalid', + 'session.id_required', + 'session.id_empty', + 'session.title_empty', + 'session.state_not_found', + 'session.state_invalid', + 'session.fork_active_turn', + 'session.export_not_found', + 'session.export_missing_version', + 'session.closed', + 'session.permission_mode_invalid', + 'session.thinking_empty', + 'session.model_empty', + 'session.plan_mode_invalid', + 'session.approval_handler_error', + 'session.question_handler_error', + 'session.init_failed', + 'agent.not_found', + 'turn.agent_busy', + 'goal.already_exists', + 'goal.not_found', + 'goal.objective_empty', + 'goal.objective_too_long', + 'goal.status_invalid', + 'goal.metadata_reserved', + 'goal.not_resumable', + 'model.not_configured', + 'model.config_invalid', + 'auth.login_required', + 'context.overflow', + 'loop.max_steps_exceeded', + 'provider.api_error', + 'provider.rate_limit', + 'provider.auth_error', + 'provider.connection_error', + 'skill.not_found', + 'skill.type_unsupported', + 'skill.name_empty', + 'records.write_failed', + 'compaction.failed', + 'compaction.unable', + 'background.task_id_empty', + 'mcp.server_not_found', + 'mcp.server_disabled', + 'mcp.startup_failed', + 'mcp.tool_name_collision', + 'plugin.not_found', + 'plugin.load_failed', + 'request.invalid', + 'request.work_dir_required', + 'request.prompt_input_empty', + 'shell.git_bash_not_found', + 'not_implemented', + 'internal', +]) satisfies z.ZodType; + +export const kimiErrorPayloadSchema = z.object({ + code: kimiErrorCodeSchema, + message: z.string(), + name: z.string().optional(), + details: z.record(z.string(), z.unknown()).optional(), + retryable: z.boolean(), +}) satisfies z.ZodType; + +export const backgroundTaskInfoBaseSchema = z.object({ + taskId: z.string(), + description: z.string(), + status: agentCoreBackgroundTaskStatusSchema, + startedAt: z.number(), + endedAt: z.number().nullable(), + stopReason: z.string().optional(), + terminalNotificationSuppressed: z.boolean().optional(), + timeoutMs: z.number().optional(), +}) satisfies z.ZodType; + +export const processBackgroundTaskInfoSchema = backgroundTaskInfoBaseSchema.extend({ + kind: z.literal('process'), + command: z.string(), + pid: z.number(), + exitCode: z.number().nullable(), +}) satisfies z.ZodType; + +export const agentBackgroundTaskInfoSchema = backgroundTaskInfoBaseSchema.extend({ + kind: z.literal('agent'), + agentId: z.string().optional(), + subagentType: z.string().optional(), +}) satisfies z.ZodType; + +export const questionBackgroundTaskInfoSchema = backgroundTaskInfoBaseSchema.extend({ + kind: z.literal('question'), + questionCount: z.number(), + toolCallId: z.string().optional(), +}) satisfies z.ZodType; + +export const backgroundTaskInfoSchema = z.discriminatedUnion('kind', [ + processBackgroundTaskInfoSchema, + agentBackgroundTaskInfoSchema, + questionBackgroundTaskInfoSchema, +]) satisfies z.ZodType; + +export const compactionResultSchema = z.object({ + summary: z.string(), + compactedCount: z.number(), + tokensBefore: z.number(), + tokensAfter: z.number(), +}) satisfies z.ZodType; + +export const toolUpdateSchema = z.object({ + kind: z.enum(['stdout', 'stderr', 'progress', 'status', 'custom']), + text: z.string().optional(), + percent: z.number().optional(), + customKind: z.string().optional(), + customData: z.unknown().optional(), +}) satisfies z.ZodType; + +export const mcpOAuthAuthorizationUrlUpdateDataSchema = z.object({ + serverName: z.string(), + authorizationUrl: z.string(), +}) satisfies z.ZodType; + +export const turnEndReasonSchema = z.enum(['completed', 'cancelled', 'failed']) satisfies z.ZodType; + +export const agentStatusUpdatedEventSchema = z.object({ + type: z.literal('agent.status.updated'), + model: z.string().optional(), + contextTokens: z.number().optional(), + maxContextTokens: z.number().optional(), + contextUsage: z.number().optional(), + planMode: z.boolean().optional(), + swarmMode: z.boolean().optional(), + permission: permissionModeSchema.optional(), + usage: usageStatusSchema.optional(), +}) satisfies z.ZodType; + +export const sessionMetaUpdatedEventSchema = z.object({ + type: z.literal('session.meta.updated'), + title: z.string().optional(), + patch: z.record(z.string(), z.unknown()).optional(), +}) satisfies z.ZodType; + +export const goalUpdatedEventSchema = z.object({ + type: z.literal('goal.updated'), + snapshot: goalSnapshotSchema.nullable(), + change: goalChangeSchema.optional(), +}) satisfies z.ZodType; + +export const skillActivatedEventSchema = z.object({ + type: z.literal('skill.activated'), + activationId: z.string(), + skillName: z.string(), + skillArgs: z.string().optional(), + trigger: z.enum(['user-slash', 'model-tool', 'nested-skill']), + skillPath: z.string().optional(), + skillSource: skillSourceSchema.optional(), +}) satisfies z.ZodType; + +export const errorEventSchema = kimiErrorPayloadSchema.extend({ + type: z.literal('error'), +}) satisfies z.ZodType; + +export const warningEventSchema = z.object({ + type: z.literal('warning'), + message: z.string(), + code: z.string().optional(), +}) satisfies z.ZodType; + +export const turnStartedEventSchema = z.object({ + type: z.literal('turn.started'), + turnId: z.number(), + origin: promptOriginSchema, +}) satisfies z.ZodType; + +export const turnEndedEventSchema = z.object({ + type: z.literal('turn.ended'), + turnId: z.number(), + reason: turnEndReasonSchema, + error: kimiErrorPayloadSchema.optional(), +}) satisfies z.ZodType; + +export const turnStepStartedEventSchema = z.object({ + type: z.literal('turn.step.started'), + turnId: z.number(), + step: z.number(), + stepId: z.string().optional(), +}) satisfies z.ZodType; + +export const turnStepCompletedEventSchema = z.object({ + type: z.literal('turn.step.completed'), + turnId: z.number(), + step: z.number(), + stepId: z.string().optional(), + usage: tokenUsageSchema.optional(), + finishReason: z.string().optional(), + llmFirstTokenLatencyMs: z.number().optional(), + llmStreamDurationMs: z.number().optional(), + providerFinishReason: finishReasonSchema.optional(), + rawFinishReason: z.string().optional(), +}) satisfies z.ZodType; + +export const turnStepRetryingEventSchema = z.object({ + type: z.literal('turn.step.retrying'), + turnId: z.number(), + step: z.number(), + stepId: z.string().optional(), + failedAttempt: z.number(), + nextAttempt: z.number(), + maxAttempts: z.number(), + delayMs: z.number(), + errorName: z.string(), + errorMessage: z.string(), + statusCode: z.number().optional(), +}) satisfies z.ZodType; + +export const turnStepInterruptedEventSchema = z.object({ + type: z.literal('turn.step.interrupted'), + turnId: z.number(), + step: z.number(), + stepId: z.string().optional(), + reason: z.string(), + message: z.string().optional(), +}) satisfies z.ZodType; + +export const assistantDeltaEventSchema = z.object({ + type: z.literal('assistant.delta'), + turnId: z.number(), + delta: z.string(), +}) satisfies z.ZodType; + +export const hookResultEventSchema = z.object({ + type: z.literal('hook.result'), + turnId: z.number(), + hookEvent: z.string(), + content: z.string(), + blocked: z.boolean().optional(), +}) satisfies z.ZodType; + +export const thinkingDeltaEventSchema = z.object({ + type: z.literal('thinking.delta'), + turnId: z.number(), + delta: z.string(), +}) satisfies z.ZodType; + +export const toolCallDeltaEventSchema = z.object({ + type: z.literal('tool.call.delta'), + turnId: z.number(), + toolCallId: z.string(), + name: z.string().optional(), + argumentsPart: z.string().optional(), +}) satisfies z.ZodType; + +export const toolCallStartedEventSchema = z.object({ + type: z.literal('tool.call.started'), + turnId: z.number(), + toolCallId: z.string(), + name: z.string(), + args: z.unknown(), + description: z.string().optional(), + display: ToolInputDisplaySchema.optional(), +}) satisfies z.ZodType; + +export const toolProgressEventSchema = z.object({ + type: z.literal('tool.progress'), + turnId: z.number(), + toolCallId: z.string(), + update: toolUpdateSchema, +}) satisfies z.ZodType; + +export const toolResultEventSchema = z.object({ + type: z.literal('tool.result'), + turnId: z.number(), + toolCallId: z.string(), + output: z.unknown(), + isError: z.boolean().optional(), + synthetic: z.boolean().optional(), +}) satisfies z.ZodType; + +export const subagentSpawnedEventSchema = z.object({ + type: z.literal('subagent.spawned'), + subagentId: z.string(), + subagentName: z.string(), + parentToolCallId: z.string(), + parentToolCallUuid: z.string().optional(), + parentAgentId: z.string().optional(), + description: z.string().optional(), + swarmIndex: z.number().optional(), + runInBackground: z.boolean(), +}) satisfies z.ZodType; + +export const subagentStartedEventSchema = z.object({ + type: z.literal('subagent.started'), + subagentId: z.string(), +}) satisfies z.ZodType; + +export const subagentSuspendedEventSchema = z.object({ + type: z.literal('subagent.suspended'), + subagentId: z.string(), + reason: z.string(), +}) satisfies z.ZodType; + +export const subagentCompletedEventSchema = z.object({ + type: z.literal('subagent.completed'), + subagentId: z.string(), + resultSummary: z.string(), + usage: tokenUsageSchema.optional(), + contextTokens: z.number().optional(), +}) satisfies z.ZodType; + +export const subagentFailedEventSchema = z.object({ + type: z.literal('subagent.failed'), + subagentId: z.string(), + error: z.string(), +}) satisfies z.ZodType; + +export const compactionStartedEventSchema = z.object({ + type: z.literal('compaction.started'), + trigger: z.enum(['manual', 'auto']), + instruction: z.string().optional(), +}) satisfies z.ZodType; + +export const compactionBlockedEventSchema = z.object({ + type: z.literal('compaction.blocked'), + turnId: z.number().optional(), +}) satisfies z.ZodType; + +export const compactionCancelledEventSchema = z.object({ + type: z.literal('compaction.cancelled'), +}) satisfies z.ZodType; + +export const compactionCompletedEventSchema = z.object({ + type: z.literal('compaction.completed'), + result: compactionResultSchema, +}) satisfies z.ZodType; + +export const backgroundTaskStartedEventSchema = z.object({ + type: z.literal('background.task.started'), + info: backgroundTaskInfoSchema, +}) satisfies z.ZodType; + +export const backgroundTaskTerminatedEventSchema = z.object({ + type: z.literal('background.task.terminated'), + info: backgroundTaskInfoSchema, +}) satisfies z.ZodType; + +export const cronFiredEventSchema = z.object({ + type: z.literal('cron.fired'), + origin: cronJobOriginSchema, + prompt: z.string(), +}) satisfies z.ZodType; + +export const toolListUpdatedReasonSchema = z.enum([ + 'mcp.connected', + 'mcp.disconnected', + 'mcp.failed', +]) satisfies z.ZodType; + +export const toolListUpdatedEventSchema = z.object({ + type: z.literal('tool.list.updated'), + reason: toolListUpdatedReasonSchema, + serverName: z.string(), +}) satisfies z.ZodType; + +export const mcpServerStatusPayloadSchema = z.object({ + name: z.string(), + transport: z.enum(['stdio', 'http']), + status: z.enum(['pending', 'connected', 'failed', 'disabled', 'needs-auth']), + toolCount: z.number(), + error: z.string().optional(), +}) satisfies z.ZodType; + +export const mcpServerStatusEventSchema = z.object({ + type: z.literal('mcp.server.status'), + server: mcpServerStatusPayloadSchema, +}) satisfies z.ZodType; + +export const agentEventSchema = z.discriminatedUnion('type', [ + errorEventSchema, + warningEventSchema, + agentStatusUpdatedEventSchema, + sessionMetaUpdatedEventSchema, + goalUpdatedEventSchema, + skillActivatedEventSchema, + turnStartedEventSchema, + turnEndedEventSchema, + turnStepStartedEventSchema, + turnStepCompletedEventSchema, + turnStepRetryingEventSchema, + turnStepInterruptedEventSchema, + assistantDeltaEventSchema, + hookResultEventSchema, + thinkingDeltaEventSchema, + toolCallDeltaEventSchema, + toolCallStartedEventSchema, + toolProgressEventSchema, + toolResultEventSchema, + toolListUpdatedEventSchema, + mcpServerStatusEventSchema, + subagentSpawnedEventSchema, + subagentStartedEventSchema, + subagentSuspendedEventSchema, + subagentCompletedEventSchema, + subagentFailedEventSchema, + compactionStartedEventSchema, + compactionBlockedEventSchema, + compactionCancelledEventSchema, + compactionCompletedEventSchema, + backgroundTaskStartedEventSchema, + backgroundTaskTerminatedEventSchema, + cronFiredEventSchema, +]) satisfies z.ZodType; + +export const eventSchema = agentEventSchema.and( + z.object({ + agentId: z.string(), + sessionId: z.string(), + }), +) satisfies z.ZodType; + +/** + * Volatile (ephemeral) event types — the IM-style "typing indicator" class. + * + * Volatile events are NOT journaled and do NOT advance the per-session + * durable `seq`. They are fanned out live with the current durable watermark + * (`seq` = last durable seq, `volatile: true` on the envelope) and are never + * replayed after a reconnect. Clients recover any state they convey from the + * session snapshot (`GET /sessions/{sid}/snapshot` → `in_flight_turn`) or + * other REST surfaces instead of delta replay. + * + * Everything not listed here is durable: journaled, seq-bearing, replayable. + */ +export const VOLATILE_EVENT_TYPES = [ + 'assistant.delta', + 'thinking.delta', + 'tool.call.delta', + 'tool.progress', + 'agent.status.updated', +] as const satisfies readonly AgentEvent['type'][]; + +export type VolatileEventType = (typeof VOLATILE_EVENT_TYPES)[number]; + +const volatileEventTypeSet: ReadonlySet = new Set(VOLATILE_EVENT_TYPES); + +export function isVolatileEventType(type: string): type is VolatileEventType { + return volatileEventTypeSet.has(type); +} diff --git a/packages/protocol/src/index.ts b/packages/protocol/src/index.ts index 3a63d49e4..22b8efa86 100644 --- a/packages/protocol/src/index.ts +++ b/packages/protocol/src/index.ts @@ -22,6 +22,7 @@ export * from './rest/meta'; export * from './rest/auth'; export * from './rest/oauth'; export * from './rest/session'; +export * from './rest/snapshot'; export * from './rest/workspace'; export * from './rest/fsBrowse'; export * from './rest/message'; diff --git a/packages/protocol/src/rest/snapshot.ts b/packages/protocol/src/rest/snapshot.ts new file mode 100644 index 000000000..cf002e6e4 --- /dev/null +++ b/packages/protocol/src/rest/snapshot.ts @@ -0,0 +1,75 @@ +/** + * `GET /v1/sessions/{session_id}/snapshot` — IM-style "initial sync". + * + * Returns an atomic-at-a-watermark view of everything a client needs to + * rebuild a session's UI state, so the standard multi-device rebuild flow is: + * + * 1. `GET /sessions/{sid}/snapshot` → state + `as_of_seq` + `epoch` + * 2. WS `subscribe` with `cursors[sid] = { seq: as_of_seq, epoch }` + * 3. apply live durable events (`seq > as_of_seq`) on top + * + * No gap and no duplication by construction: the watermark ties the REST + * snapshot to the WS event stream. + * + * `in_flight_turn` carries the accumulated state of a currently-running turn + * (volatile deltas are not replayable; this is how a reconnecting client + * recovers mid-turn assistant/thinking text and running tool calls). + * + * The server reads the watermark, assembles the snapshot, then re-reads the + * watermark and retries assembly if a durable event landed in between + * (bounded retries). Durable events are low-frequency (turn/tool boundaries, + * not deltas), so this converges almost immediately. + */ + +import { z } from 'zod'; + +import { approvalRequestSchema } from '../approval'; +import { messageSchema } from '../message'; +import { questionRequestSchema } from '../question'; +import { sessionSchema } from '../session'; + +export const inFlightToolCallSchema = z.object({ + tool_call_id: z.string().min(1), + name: z.string().min(1), + args: z.unknown().optional(), + description: z.string().optional(), + /** Display payload from `tool.call.started` (ToolInputDisplay). */ + display: z.unknown().optional(), + /** Most recent `tool.progress` update, if any. */ + last_progress: z + .object({ + kind: z.enum(['stdout', 'stderr', 'progress', 'status', 'custom']), + text: z.string().optional(), + percent: z.number().optional(), + }) + .optional(), +}); +export type InFlightToolCall = z.infer; + +export const inFlightTurnSchema = z.object({ + turn_id: z.number().int().nonnegative(), + /** Assistant text accumulated from `assistant.delta` so far. */ + assistant_text: z.string(), + /** Thinking text accumulated from `thinking.delta` so far. */ + thinking_text: z.string(), + /** Tool calls started but without a `tool.result` yet. */ + running_tools: z.array(inFlightToolCallSchema), +}); +export type InFlightTurn = z.infer; + +export const sessionSnapshotResponseSchema = z.object({ + /** Durable event watermark this snapshot is consistent with. */ + as_of_seq: z.number().int().nonnegative(), + /** Journal epoch — pass back via the WS cursor for invalidation detection. */ + epoch: z.string().min(1), + session: sessionSchema, + /** Most recent messages (chronological ascending), bounded page. */ + messages: z.object({ + items: z.array(messageSchema), + has_more: z.boolean(), + }), + in_flight_turn: inFlightTurnSchema.nullable(), + pending_approvals: z.array(approvalRequestSchema), + pending_questions: z.array(questionRequestSchema), +}); +export type SessionSnapshotResponse = z.infer; diff --git a/packages/protocol/src/ws-control.ts b/packages/protocol/src/ws-control.ts index 5578a3325..77d2033d3 100644 --- a/packages/protocol/src/ws-control.ts +++ b/packages/protocol/src/ws-control.ts @@ -5,12 +5,51 @@ */ import { z } from 'zod'; +import { eventSchema } from './events'; import { isoDateTimeSchema } from './time'; +/** + * WS protocol version. v2 (breaking, IM-style multi-device sync): + * - per-session cursors are `{ seq, epoch }` instead of a bare seq + * - `seq` is durable (journal offset, survives daemon restarts) + * - volatile events carry `volatile: true` and do not advance `seq` + * - `resync_required` gains the `epoch_changed` reason + `epoch` field + */ +export const WS_PROTOCOL_VERSION = 2; + +/** + * Per-session sync cursor. `seq` is the last durable event seq the client + * has applied (journal offset). `epoch` identifies the journal incarnation + * (changes when a session's journal is recreated); a cursor whose epoch does + * not match the server's current epoch is invalid and triggers + * `resync_required(epoch_changed)`. `epoch` is absent on a fresh cursor. + */ +export const sessionCursorSchema = z.object({ + seq: z.number().int().nonnegative(), + epoch: z.string().min(1).optional(), +}); + +export type SessionCursor = z.infer; + +export const cursorsBySessionSchema = z.record(z.string(), sessionCursorSchema); + +export type CursorsBySession = z.infer; + export const wsEventEnvelopeSchema = (payload: T) => z.object({ type: z.string(), seq: z.number().int().nonnegative(), + epoch: z.string().optional(), + volatile: z.boolean().optional(), + /** + * For volatile text-delta frames (`assistant.delta` / `thinking.delta`): + * the cumulative character offset of this delta within the in-flight + * turn's accumulated stream. Clients align against + * `snapshot.in_flight_turn.*_text.length` — `offset < local length` is a + * duplicate (skip), `offset > local length` means deltas were missed + * (re-snapshot). + */ + offset: z.number().int().nonnegative().optional(), session_id: z.string().optional(), timestamp: isoDateTimeSchema, payload, @@ -34,6 +73,7 @@ export const wsAckEnvelopeSchema = (payload: T) => export const serverHelloPayloadSchema = z.object({ ws_connection_id: z.string(), + protocol_version: z.number().int().positive(), heartbeat_ms: z.number().int().positive(), max_event_buffer_size: z.number().int().positive(), capabilities: z.object({ @@ -53,7 +93,7 @@ export type ServerHelloMessage = z.infer; export const clientHelloPayloadSchema = z.object({ client_id: z.string(), subscriptions: z.array(z.string()), - last_seq_by_session: z.record(z.string(), z.number().int().nonnegative()).optional(), + cursors: cursorsBySessionSchema.optional(), }); export const clientHelloMessageSchema = z.object({ @@ -64,13 +104,17 @@ export const clientHelloMessageSchema = z.object({ export type ClientHelloMessage = z.infer; -export const helloAckPayloadSchema = z.object({ - accepted_subscriptions: z.array(z.string()).optional(), - accepted: z.array(z.string()).optional(), - not_found: z.array(z.string()).optional(), +export const clientHelloAckPayloadSchema = z.object({ + accepted_subscriptions: z.array(z.string()), resync_required: z.array(z.string()), + /** Server-side current cursor per accepted session ({seq, epoch}). */ + cursors: cursorsBySessionSchema.optional(), }); +export const helloAckPayloadSchema = clientHelloAckPayloadSchema; + +export const clientHelloAckMessageSchema = wsAckEnvelopeSchema(clientHelloAckPayloadSchema); + export const watchFsConfigSchema = z.object({ paths: z.array(z.string()), recursive: z.boolean().optional(), @@ -78,7 +122,7 @@ export const watchFsConfigSchema = z.object({ export const subscribePayloadSchema = z.object({ session_ids: z.array(z.string()), - last_seq_by_session: z.record(z.string(), z.number().int().nonnegative()).optional(), + cursors: cursorsBySessionSchema.optional(), watch_fs: z.record(z.string(), watchFsConfigSchema).optional(), }); @@ -90,6 +134,16 @@ export const subscribeMessageSchema = z.object({ export type SubscribeMessage = z.infer; +export const subscribeAckPayloadSchema = z.object({ + accepted: z.array(z.string()), + not_found: z.array(z.string()), + resync_required: z.array(z.string()), + /** Server-side current cursor per accepted session ({seq, epoch}). */ + cursors: cursorsBySessionSchema.optional(), +}); + +export const subscribeAckMessageSchema = wsAckEnvelopeSchema(subscribeAckPayloadSchema); + export const unsubscribePayloadSchema = z.object({ session_ids: z.array(z.string()), }); @@ -102,6 +156,10 @@ export const unsubscribeMessageSchema = z.object({ export type UnsubscribeMessage = z.infer; +export const unsubscribeAckPayloadSchema = subscribeAckPayloadSchema; + +export const unsubscribeAckMessageSchema = wsAckEnvelopeSchema(unsubscribeAckPayloadSchema); + export const watchFsAddPayloadSchema = z.object({ session_id: z.string(), paths: z.array(z.string()), @@ -134,6 +192,8 @@ export const watchFsAckPayloadSchema = z.object({ current_count: z.number().int().nonnegative().optional(), }); +export const watchFsAckMessageSchema = wsAckEnvelopeSchema(watchFsAckPayloadSchema); + export const abortPayloadSchema = z.object({ session_id: z.string(), prompt_id: z.string(), @@ -152,6 +212,8 @@ export const abortAckPayloadSchema = z.object({ at_seq: z.number().int().nonnegative().optional(), }); +export const abortAckMessageSchema = wsAckEnvelopeSchema(abortAckPayloadSchema); + export const pingPayloadSchema = z.object({ nonce: z.string(), }); @@ -177,8 +239,10 @@ export type PongMessage = z.infer; export const resyncRequiredPayloadSchema = z.object({ session_id: z.string(), - reason: z.enum(['buffer_overflow', 'session_recreated']), + reason: z.enum(['buffer_overflow', 'session_recreated', 'epoch_changed']), current_seq: z.number().int().nonnegative(), + /** Current journal epoch — the client should adopt it after resyncing. */ + epoch: z.string().min(1).optional(), }); export const resyncRequiredMessageSchema = z.object({ @@ -205,6 +269,8 @@ export const wsErrorMessageSchema = z.object({ export type WsErrorMessage = z.infer; +export const sessionEventMessageSchema = wsEventEnvelopeSchema(eventSchema); + export const clientControlMessageSchema = z.discriminatedUnion('type', [ clientHelloMessageSchema, subscribeMessageSchema, @@ -225,3 +291,125 @@ export const serverSystemMessageSchema = z.discriminatedUnion('type', [ ]); export type ServerSystemMessage = z.infer; + +export type WsOperationDirection = 'client_to_server' | 'server_to_client'; + +export type WsOperationKind = 'control' | 'system' | 'event'; + +export interface WsOperationDefinition { + readonly type: string; + readonly direction: WsOperationDirection; + readonly kind: WsOperationKind; + readonly messageSchema: z.ZodTypeAny; + readonly ackSchema?: z.ZodTypeAny; + readonly description: string; +} + +export const clientControlOperations = [ + { + type: 'client_hello', + direction: 'client_to_server', + kind: 'control', + messageSchema: clientHelloMessageSchema, + ackSchema: clientHelloAckMessageSchema, + description: 'Start a client session and optionally subscribe to existing daemon sessions.', + }, + { + type: 'subscribe', + direction: 'client_to_server', + kind: 'control', + messageSchema: subscribeMessageSchema, + ackSchema: subscribeAckMessageSchema, + description: 'Subscribe the connection to one or more session event streams.', + }, + { + type: 'unsubscribe', + direction: 'client_to_server', + kind: 'control', + messageSchema: unsubscribeMessageSchema, + ackSchema: unsubscribeAckMessageSchema, + description: 'Remove one or more session event stream subscriptions.', + }, + { + type: 'watch_fs_add', + direction: 'client_to_server', + kind: 'control', + messageSchema: watchFsAddMessageSchema, + ackSchema: watchFsAckMessageSchema, + description: 'Add filesystem watch paths for a subscribed session.', + }, + { + type: 'watch_fs_remove', + direction: 'client_to_server', + kind: 'control', + messageSchema: watchFsRemoveMessageSchema, + ackSchema: watchFsAckMessageSchema, + description: 'Remove filesystem watch paths for a subscribed session.', + }, + { + type: 'abort', + direction: 'client_to_server', + kind: 'control', + messageSchema: abortMessageSchema, + ackSchema: abortAckMessageSchema, + description: 'Abort a running prompt in a session.', + }, + { + type: 'pong', + direction: 'client_to_server', + kind: 'control', + messageSchema: pongMessageSchema, + description: 'Reply to a server ping with the same nonce.', + }, +] as const satisfies readonly WsOperationDefinition[]; + +export const serverSystemOperations = [ + { + type: 'server_hello', + direction: 'server_to_client', + kind: 'system', + messageSchema: serverHelloMessageSchema, + description: 'Initial server greeting sent immediately after the socket opens.', + }, + { + type: 'ping', + direction: 'server_to_client', + kind: 'system', + messageSchema: pingMessageSchema, + description: 'Heartbeat ping sent by the server; clients must answer with pong.', + }, + { + type: 'resync_required', + direction: 'server_to_client', + kind: 'system', + messageSchema: resyncRequiredMessageSchema, + description: 'Signals that a client must rebuild local session state from REST history.', + }, + { + type: 'error', + direction: 'server_to_client', + kind: 'system', + messageSchema: wsErrorMessageSchema, + description: 'Server-side WebSocket protocol or runtime error.', + }, +] as const satisfies readonly WsOperationDefinition[]; + +export const sessionEventOperation = { + type: 'session_event', + direction: 'server_to_client', + kind: 'event', + messageSchema: sessionEventMessageSchema, + description: 'Session-scoped agent event envelope; frame type is the payload event type.', +} as const satisfies WsOperationDefinition; + +export const wsOperations = [ + ...clientControlOperations, + ...serverSystemOperations, + sessionEventOperation, +] as const satisfies readonly WsOperationDefinition[]; + +export function getClientControlOperation( + type: string, +): (typeof clientControlOperations)[number] | undefined { + return clientControlOperations.find((operation) => operation.type === type); +}