vrr/kimi-code
2
0
Fork 0
mirror of https://github.com/MoonshotAI/kimi-code.git synced 2026-08-26 17:14:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

feat(sync): add snapshot-based session resync

Browse source 
- 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
This commit is contained in:
haozhe.yang 2026-06-11 12:31:16 +08:00
parent 9e9777382b
commit 517433454c
28 changed files with 3042 additions and 308 deletions
Download patch file Download diff file Expand all files Collapse all files

123
apps/kimi-web/src/api/daemon/agentEventProjector.ts
View file

@ -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<string, number>;
@ -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 };
}
// ---------------------------------------------------------------------------

68
apps/kimi-web/src/api/daemon/client.ts
View file

@ -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<AppSessionSnapshot> {
const data = await this.http.get<WireSessionSnapshot>(
`/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

4
apps/kimi-web/src/api/daemon/eventReducer.ts
View file

@ -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;

49
apps/kimi-web/src/api/daemon/wire.ts
View file

@ -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<string, number>;
cursors?: Record<string, WireSessionCursor>;
};
}
@ -463,7 +506,7 @@ export interface WireSubscribe {
id: string;
payload: {
session_ids: string[];
last_seq_by_session?: Record<string, number>;
cursors?: Record<string, WireSessionCursor>;
};
}

89
apps/kimi-web/src/api/daemon/ws.ts
View file

@ -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<string, number>();
/** subscriptions we manage: sessionId → last known cursor {seq, epoch} */
private readonly subscriptions = new Map<string, SessionCursor>();
/** 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<string, unknown>);
// 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<string, number> = {};
for (const [sid, seq] of this.subscriptions.entries()) {
lastSeqBySession[sid] = seq;
// Build cursors from subscriptions
const cursors: Record<string, SessionCursor> = {};
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<string, number>): void {
private sendSubscribe(sessionIds: string[], cursors: Record<string, SessionCursor>): 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<string, unknown>): 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 {

52
apps/kimi-web/src/api/types.ts
View file

@ -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<AppSessionRuntimeStatus>;
deleteSession(sessionId: string): Promise<{ deleted: true }>;
listMessages(sessionId: string, input?: PageRequest & { role?: AppMessageRole }): Promise<Page<AppMessage>>;
/** v2 initial sync: atomic session state + `asOfSeq` watermark + epoch. */
getSessionSnapshot(sessionId: string): Promise<AppSessionSnapshot>;
submitPrompt(sessionId: string, input: PromptSubmission): Promise<PromptSubmitResult>;
abortPrompt(sessionId: string, promptId: string): Promise<{ aborted: boolean; atSeq?: number }>;
compactSession(sessionId: string, instruction?: string): Promise<void>;

83
apps/kimi-web/src/composables/useKimiWebClient.ts
View file

@ -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<string, string> = {};
async function loadMessagesForSession(sessionId: string): Promise<void> {
/**
* 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<void> {
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<void> {
}
}
async function reloadAndResubscribe(sessionId: string, currentSeq: number): Promise<void> {
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<void> {
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 {

35
apps/kimi-web/test/useKimiWebClient-session-cache.test.ts
View file

@ -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 () => {

12
packages/daemon-e2e/scenarios/03-refresh-replay.ts
View file

@ -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<string, unknown> {
client_id: string;
subscriptions: string[];
last_seq_by_session?: Record<string, number>;
cursors?: Record<string, { seq: number; epoch?: string }>;
}
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

21
packages/daemon-e2e/test/refresh-replay.test.ts
View file

@ -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 });

2
packages/daemon/src/routes/registerApiV1Routes.ts
View file

@ -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<typeof registerSessionsRoutes>[0], ix);
registerSnapshotRoutes(apiV1 as unknown as Parameters<typeof registerSnapshotRoutes>[0], ix);
registerMessagesRoutes(apiV1 as unknown as Parameters<typeof registerMessagesRoutes>[0], ix);
registerPromptsRoutes(apiV1 as unknown as Parameters<typeof registerPromptsRoutes>[0], ix);
registerApprovalsRoutes(

133
packages/daemon/src/routes/snapshot.ts Normal file
View file

@ -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<string, unknown> } | undefined,
handler: (
req: { id: string; query: unknown; params: unknown },
reply: { send(payload: unknown): unknown },
) => Promise<void> | 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<SnapshotRouteHost['get']>[2]);
}

138
packages/daemon/src/services/gateway/inFlightTurnTracker.ts Normal file
View file

@ -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<string, ToolAccum>;
}
export interface VolatileAnnotation {
/** Pre-append offset for text-delta frames. */
offset?: number;
}
export class InFlightTurnTracker {
private readonly bySession = new Map<string, TurnAccum>();
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);
}
}

241
packages/daemon/src/services/gateway/sessionEventJournal.ts Normal file
View file

@ -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 `<kimiHome>/daemon/events/<sessionId>.jsonl`:
*
* line 1 {"kind":"journal_header","version":1,"epoch":"ep_<ulid>","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<void> | 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<SessionEventJournal> {
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<JournalEntry[]> {
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<void> {
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<void> {
await this.flush();
}
private scheduleFlush(): void {
if (this.flushPromise !== undefined) return;
this.flushPromise = this.flushOnce().finally(() => {
this.flushPromise = undefined;
});
}
private async flushOnce(): Promise<void> {
// 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<string> {
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;
}

89
packages/daemon/src/services/gateway/wsBroadcast.ts
View file

@ -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<Event>` (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<BufferedSinceResult>;
/** 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<SessionSnapshotState>;
/**
* 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<IWSBroadcastService>('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;

180
packages/daemon/src/services/gateway/wsBroadcastService.ts
View file

@ -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<SessionEventJournal>;
/** 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<void>;
}
export class WSBroadcastService extends Disposable implements IWSBroadcastService {
@ -33,14 +40,18 @@ export class WSBroadcastService extends Disposable implements IWSBroadcastServic
private readonly _sessions = new Map<string, SessionState>();
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<void> {
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<BufferedSinceResult> {
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<SessionSnapshotState> {
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<void> {
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, '_');
}

168
packages/daemon/src/ws/connection.ts
View file

@ -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<string>();
public readonly lastSeqBySession = new Map<string, number>();
/** Last cursor each subscribed session was synced from (client-claimed). */
public readonly cursorsBySession = new Map<string, SessionCursor>();
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<void> {
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<void> {
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;
}

40
packages/daemon/src/ws/protocol.ts
View file

@ -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<P>(id: string, code: number, msg: string, payload: P):
export interface EventEnvelope<P = Event> {
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<Event> {
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 } : {}),
},
};
}

224
packages/daemon/test/services.test.ts
View file

@ -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<NonNullable<FsWatcherServiceOptions['watcherFact
let ix: InstantiationService;
let testLogger: TestLogger;
const tmpHomeDirs: string[] = [];
/** Throwaway `IEnvironmentService` whose homeDir is a fresh temp dir. */
function tmpEnv(): IEnvironmentService {
const dir = mkdtempSync(join(tmpdir(), 'kimi-daemon-test-'));
tmpHomeDirs.push(dir);
return { _serviceBrand: undefined, homeDir: dir, configPath: join(dir, 'config.toml') };
}
beforeEach(() => {
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<typeof broker.request>[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<typeof broker.request>[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<typeof broker.request>[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<typeof broker.request>[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',
);

216
packages/daemon/test/snapshot.e2e.test.ts Normal file
View file

@ -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<RunningDaemon> {
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<T>(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<string> {
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<typeof envelopeOf<SessionSnapshotResponse>> }> {
const res = await appOf(r).inject({
method: 'GET',
url: `/api/v1/sessions/${sid}/snapshot`,
});
return { statusCode: res.statusCode, env: envelopeOf<SessionSnapshotResponse>(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();
});
});

120
packages/daemon/test/ws-resync.e2e.test.ts
View file

@ -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<void> {
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<string, { seq: number; epoch?: string }>;
};
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);
});
});

53
packages/protocol/src/__tests__/events.test.ts
View file

@ -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');
});
});

130
packages/protocol/src/__tests__/snapshot.test.ts Normal file
View file

@ -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);
}
});
});

185
packages/protocol/src/__tests__/ws-control.test.ts
View file

@ -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);
});
});

617
packages/protocol/src/events.ts
View file

@ -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<TokenUsage>;
export const finishReasonSchema = z.enum([
'completed',
'tool_calls',
'truncated',
'filtered',
'paused',
'other',
]) satisfies z.ZodType<FinishReason>;
export const usageStatusSchema = z.object({
byModel: z.record(z.string(), tokenUsageSchema).optional(),
currentTurn: tokenUsageSchema.optional(),
total: tokenUsageSchema.optional(),
}) satisfies z.ZodType<UsageStatus>;
export const permissionModeSchema = z.enum(['manual', 'yolo', 'auto']) satisfies z.ZodType<PermissionMode>;
export const skillSourceSchema = z.enum(['project', 'user', 'extra', 'builtin']) satisfies z.ZodType<SkillSource>;
export const userPromptOriginSchema = z.object({
kind: z.literal('user'),
}) satisfies z.ZodType<UserPromptOrigin>;
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<SkillActivationOrigin>;
export const injectionOriginSchema = z.object({
kind: z.literal('injection'),
variant: z.string(),
}) satisfies z.ZodType<InjectionOrigin>;
export const compactionSummaryOriginSchema = z.object({
kind: z.literal('compaction_summary'),
}) satisfies z.ZodType<CompactionSummaryOrigin>;
export const systemTriggerOriginSchema = z.object({
kind: z.literal('system_trigger'),
name: z.string(),
}) satisfies z.ZodType<SystemTriggerOrigin>;
export const agentCoreBackgroundTaskStatusSchema = z.enum([
'running',
'completed',
'failed',
'timed_out',
'killed',
'lost',
]) satisfies z.ZodType<AgentCoreBackgroundTaskStatus>;
export const backgroundTaskOriginSchema = z.object({
kind: z.literal('background_task'),
taskId: z.string(),
status: agentCoreBackgroundTaskStatusSchema,
notificationId: z.string(),
}) satisfies z.ZodType<BackgroundTaskOrigin>;
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<CronJobOrigin>;
export const cronMissedOriginSchema = z.object({
kind: z.literal('cron_missed'),
count: z.number(),
}) satisfies z.ZodType<CronMissedOrigin>;
export const hookResultOriginSchema = z.object({
kind: z.literal('hook_result'),
event: z.string(),
blocked: z.boolean().optional(),
}) satisfies z.ZodType<HookResultOrigin>;
export const retryOriginSchema = z.object({
kind: z.literal('retry'),
trigger: z.string().optional(),
}) satisfies z.ZodType<RetryOrigin>;
export const promptOriginSchema = z.discriminatedUnion('kind', [
userPromptOriginSchema,
skillActivationOriginSchema,
injectionOriginSchema,
compactionSummaryOriginSchema,
systemTriggerOriginSchema,
backgroundTaskOriginSchema,
cronJobOriginSchema,
cronMissedOriginSchema,
hookResultOriginSchema,
retryOriginSchema,
]) satisfies z.ZodType<PromptOrigin>;
export const goalStatusSchema = z.enum(['active', 'paused', 'blocked', 'complete']) satisfies z.ZodType<GoalStatus>;
export const goalActorSchema = z.enum(['user', 'model', 'runtime', 'system']) satisfies z.ZodType<GoalActor>;
export const goalBudgetLimitsSchema = z.object({
tokenBudget: z.number().optional(),
turnBudget: z.number().optional(),
wallClockBudgetMs: z.number().optional(),
}) satisfies z.ZodType<GoalBudgetLimits>;
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<GoalBudgetReport>;
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<GoalSnapshot>;
export const goalToolResultSchema = z.object({
goal: goalSnapshotSchema.nullable(),
}) satisfies z.ZodType<GoalToolResult>;
export const goalChangeStatsSchema = z.object({
turnsUsed: z.number(),
tokensUsed: z.number(),
wallClockMs: z.number(),
}) satisfies z.ZodType<GoalChangeStats>;
export const goalChangeKindSchema = z.enum(['lifecycle', 'completion']) satisfies z.ZodType<GoalChangeKind>;
export const goalChangeSchema = z.object({
kind: goalChangeKindSchema,
status: goalStatusSchema.optional(),
reason: z.string().optional(),
stats: goalChangeStatsSchema.optional(),
actor: goalActorSchema.optional(),
}) satisfies z.ZodType<GoalChange>;
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<KimiErrorCode>;
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<KimiErrorPayload>;
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<BackgroundTaskInfoBase>;
export const processBackgroundTaskInfoSchema = backgroundTaskInfoBaseSchema.extend({
kind: z.literal('process'),
command: z.string(),
pid: z.number(),
exitCode: z.number().nullable(),
}) satisfies z.ZodType<ProcessBackgroundTaskInfo>;
export const agentBackgroundTaskInfoSchema = backgroundTaskInfoBaseSchema.extend({
kind: z.literal('agent'),
agentId: z.string().optional(),
subagentType: z.string().optional(),
}) satisfies z.ZodType<AgentBackgroundTaskInfo>;
export const questionBackgroundTaskInfoSchema = backgroundTaskInfoBaseSchema.extend({
kind: z.literal('question'),
questionCount: z.number(),
toolCallId: z.string().optional(),
}) satisfies z.ZodType<QuestionBackgroundTaskInfo>;
export const backgroundTaskInfoSchema = z.discriminatedUnion('kind', [
processBackgroundTaskInfoSchema,
agentBackgroundTaskInfoSchema,
questionBackgroundTaskInfoSchema,
]) satisfies z.ZodType<BackgroundTaskInfo>;
export const compactionResultSchema = z.object({
summary: z.string(),
compactedCount: z.number(),
tokensBefore: z.number(),
tokensAfter: z.number(),
}) satisfies z.ZodType<CompactionResult>;
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<ToolUpdate>;
export const mcpOAuthAuthorizationUrlUpdateDataSchema = z.object({
serverName: z.string(),
authorizationUrl: z.string(),
}) satisfies z.ZodType<McpOAuthAuthorizationUrlUpdateData>;
export const turnEndReasonSchema = z.enum(['completed', 'cancelled', 'failed']) satisfies z.ZodType<TurnEndReason>;
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<AgentStatusUpdatedEvent>;
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<SessionMetaUpdatedEvent>;
export const goalUpdatedEventSchema = z.object({
type: z.literal('goal.updated'),
snapshot: goalSnapshotSchema.nullable(),
change: goalChangeSchema.optional(),
}) satisfies z.ZodType<GoalUpdatedEvent>;
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<SkillActivatedEvent>;
export const errorEventSchema = kimiErrorPayloadSchema.extend({
type: z.literal('error'),
}) satisfies z.ZodType<ErrorEvent>;
export const warningEventSchema = z.object({
type: z.literal('warning'),
message: z.string(),
code: z.string().optional(),
}) satisfies z.ZodType<WarningEvent>;
export const turnStartedEventSchema = z.object({
type: z.literal('turn.started'),
turnId: z.number(),
origin: promptOriginSchema,
}) satisfies z.ZodType<TurnStartedEvent>;
export const turnEndedEventSchema = z.object({
type: z.literal('turn.ended'),
turnId: z.number(),
reason: turnEndReasonSchema,
error: kimiErrorPayloadSchema.optional(),
}) satisfies z.ZodType<TurnEndedEvent>;
export const turnStepStartedEventSchema = z.object({
type: z.literal('turn.step.started'),
turnId: z.number(),
step: z.number(),
stepId: z.string().optional(),
}) satisfies z.ZodType<TurnStepStartedEvent>;
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<TurnStepCompletedEvent>;
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<TurnStepRetryingEvent>;
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<TurnStepInterruptedEvent>;
export const assistantDeltaEventSchema = z.object({
type: z.literal('assistant.delta'),
turnId: z.number(),
delta: z.string(),
}) satisfies z.ZodType<AssistantDeltaEvent>;
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<HookResultEvent>;
export const thinkingDeltaEventSchema = z.object({
type: z.literal('thinking.delta'),
turnId: z.number(),
delta: z.string(),
}) satisfies z.ZodType<ThinkingDeltaEvent>;
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<ToolCallDeltaEvent>;
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<ToolCallStartedEvent>;
export const toolProgressEventSchema = z.object({
type: z.literal('tool.progress'),
turnId: z.number(),
toolCallId: z.string(),
update: toolUpdateSchema,
}) satisfies z.ZodType<ToolProgressEvent>;
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<ToolResultEvent>;
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<SubagentSpawnedEvent>;
export const subagentStartedEventSchema = z.object({
type: z.literal('subagent.started'),
subagentId: z.string(),
}) satisfies z.ZodType<SubagentStartedEvent>;
export const subagentSuspendedEventSchema = z.object({
type: z.literal('subagent.suspended'),
subagentId: z.string(),
reason: z.string(),
}) satisfies z.ZodType<SubagentSuspendedEvent>;
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<SubagentCompletedEvent>;
export const subagentFailedEventSchema = z.object({
type: z.literal('subagent.failed'),
subagentId: z.string(),
error: z.string(),
}) satisfies z.ZodType<SubagentFailedEvent>;
export const compactionStartedEventSchema = z.object({
type: z.literal('compaction.started'),
trigger: z.enum(['manual', 'auto']),
instruction: z.string().optional(),
}) satisfies z.ZodType<CompactionStartedEvent>;
export const compactionBlockedEventSchema = z.object({
type: z.literal('compaction.blocked'),
turnId: z.number().optional(),
}) satisfies z.ZodType<CompactionBlockedEvent>;
export const compactionCancelledEventSchema = z.object({
type: z.literal('compaction.cancelled'),
}) satisfies z.ZodType<CompactionCancelledEvent>;
export const compactionCompletedEventSchema = z.object({
type: z.literal('compaction.completed'),
result: compactionResultSchema,
}) satisfies z.ZodType<CompactionCompletedEvent>;
export const backgroundTaskStartedEventSchema = z.object({
type: z.literal('background.task.started'),
info: backgroundTaskInfoSchema,
}) satisfies z.ZodType<BackgroundTaskStartedEvent>;
export const backgroundTaskTerminatedEventSchema = z.object({
type: z.literal('background.task.terminated'),
info: backgroundTaskInfoSchema,
}) satisfies z.ZodType<BackgroundTaskTerminatedEvent>;
export const cronFiredEventSchema = z.object({
type: z.literal('cron.fired'),
origin: cronJobOriginSchema,
prompt: z.string(),
}) satisfies z.ZodType<CronFiredEvent>;
export const toolListUpdatedReasonSchema = z.enum([
'mcp.connected',
'mcp.disconnected',
'mcp.failed',
]) satisfies z.ZodType<ToolListUpdatedReason>;
export const toolListUpdatedEventSchema = z.object({
type: z.literal('tool.list.updated'),
reason: toolListUpdatedReasonSchema,
serverName: z.string(),
}) satisfies z.ZodType<ToolListUpdatedEvent>;
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<McpServerStatusPayload>;
export const mcpServerStatusEventSchema = z.object({
type: z.literal('mcp.server.status'),
server: mcpServerStatusPayloadSchema,
}) satisfies z.ZodType<McpServerStatusEvent>;
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<AgentEvent>;
export const eventSchema = agentEventSchema.and(
z.object({
agentId: z.string(),
sessionId: z.string(),
}),
) satisfies z.ZodType<Event>;
/**
* 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<string> = new Set(VOLATILE_EVENT_TYPES);
export function isVolatileEventType(type: string): type is VolatileEventType {
return volatileEventTypeSet.has(type);
}

1
packages/protocol/src/index.ts
View file

@ -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';

75
packages/protocol/src/rest/snapshot.ts Normal file
View file

@ -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<typeof inFlightToolCallSchema>;
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<typeof inFlightTurnSchema>;
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<typeof sessionSnapshotResponseSchema>;

202
packages/protocol/src/ws-control.ts
View file

@ -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<typeof sessionCursorSchema>;
export const cursorsBySessionSchema = z.record(z.string(), sessionCursorSchema);
export type CursorsBySession = z.infer<typeof cursorsBySessionSchema>;
export const wsEventEnvelopeSchema = <T extends z.ZodTypeAny>(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 = <T extends z.ZodTypeAny>(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<typeof serverHelloMessageSchema>;
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<typeof clientHelloMessageSchema>;
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<typeof subscribeMessageSchema>;
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<typeof unsubscribeMessageSchema>;
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<typeof pongMessageSchema>;
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<typeof wsErrorMessageSchema>;
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<typeof serverSystemMessageSchema>;
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);
}