From b16088810168d26c9bc6b09ef6b2ec34ad9bc2ae Mon Sep 17 00:00:00 2001 From: weishu Date: Tue, 30 Dec 2025 13:22:27 +0800 Subject: [PATCH] feat: implement message backfill to prevent message loss on reconnection --- cli/src/api/apiSession.ts | 141 ++++++++++++++++++++++++++++++---- cli/src/api/types.ts | 12 +++ server/src/store/index.ts | 11 +++ server/src/sync/syncEngine.ts | 11 +++ server/src/web/routes/cli.ts | 26 +++++++ 5 files changed, 187 insertions(+), 14 deletions(-) diff --git a/cli/src/api/apiSession.ts b/cli/src/api/apiSession.ts index 90d95183..137c0ab7 100644 --- a/cli/src/api/apiSession.ts +++ b/cli/src/api/apiSession.ts @@ -1,6 +1,7 @@ import { EventEmitter } from 'node:events' import { randomUUID } from 'node:crypto' import { io, type Socket } from 'socket.io-client' +import axios from 'axios' import type { ZodType } from 'zod' import { logger } from '@/ui/logger' import { backoff } from '@/utils/time' @@ -20,7 +21,7 @@ import type { Update, UserMessage } from './types' -import { AgentStateSchema, MetadataSchema, UserMessageSchema } from './types' +import { AgentStateSchema, CliMessagesResponseSchema, MetadataSchema, UserMessageSchema } from './types' import { RpcHandlerManager } from './rpc/RpcHandlerManager' import { registerCommonHandlers } from '../modules/common/registerCommonHandlers' import { TerminalManager } from '@/terminal/TerminalManager' @@ -41,6 +42,10 @@ export class ApiSessionClient extends EventEmitter { private readonly socket: Socket private pendingMessages: UserMessage[] = [] private pendingMessageCallback: ((message: UserMessage) => void) | null = null + private lastSeenMessageSeq: number | null = null + private backfillInFlight: Promise | null = null + private needsBackfill = false + private hasConnectedOnce = false readonly rpcHandlerManager: RpcHandlerManager private readonly terminalManager: TerminalManager private agentStateLock = new AsyncLock() @@ -91,6 +96,11 @@ export class ApiSessionClient extends EventEmitter { this.socket.on('connect', () => { logger.debug('Socket connected successfully') this.rpcHandlerManager.onSocketConnect(this.socket) + if (this.hasConnectedOnce) { + this.needsBackfill = true + } + void this.backfillIfNeeded() + this.hasConnectedOnce = true }) this.socket.on('rpc-request', async (data: { method: string; params: string }, callback: (response: string) => void) => { @@ -101,6 +111,9 @@ export class ApiSessionClient extends EventEmitter { logger.debug('[API] Socket disconnected:', reason) this.rpcHandlerManager.onSocketDisconnect() this.terminalManager.closeAll() + if (this.hasConnectedOnce) { + this.needsBackfill = true + } }) this.socket.on('connect_error', (error) => { @@ -143,19 +156,7 @@ export class ApiSessionClient extends EventEmitter { if (!data.body) return if (data.body.t === 'new-message') { - const content = data.body.message.content - - const userResult = UserMessageSchema.safeParse(content) - if (userResult.success) { - if (this.pendingMessageCallback) { - this.pendingMessageCallback(userResult.data) - } else { - this.pendingMessages.push(userResult.data) - } - return - } - - this.emit('message', content) + this.handleIncomingMessage(data.body.message) return } @@ -202,6 +203,118 @@ export class ApiSessionClient extends EventEmitter { } } + private enqueueUserMessage(message: UserMessage): void { + if (this.pendingMessageCallback) { + this.pendingMessageCallback(message) + } else { + this.pendingMessages.push(message) + } + } + + private handleIncomingMessage(message: { seq?: number; content: unknown }): void { + const seq = typeof message.seq === 'number' ? message.seq : null + if (seq !== null) { + if (this.lastSeenMessageSeq !== null && seq <= this.lastSeenMessageSeq) { + return + } + this.lastSeenMessageSeq = seq + } + + const userResult = UserMessageSchema.safeParse(message.content) + if (userResult.success) { + this.enqueueUserMessage(userResult.data) + return + } + + this.emit('message', message.content) + } + + private async backfillIfNeeded(): Promise { + if (!this.needsBackfill) { + return + } + try { + await this.backfillMessages() + this.needsBackfill = false + } catch (error) { + logger.debug('[API] Backfill failed', error) + this.needsBackfill = true + } + } + + private async backfillMessages(): Promise { + if (this.backfillInFlight) { + await this.backfillInFlight + return + } + + const startSeq = this.lastSeenMessageSeq + if (startSeq === null) { + logger.debug('[API] Skipping backfill because no last-seen message sequence is available') + return + } + + const limit = 200 + const run = async () => { + let cursor = startSeq + while (true) { + const response = await axios.get( + `${configuration.serverUrl}/cli/sessions/${encodeURIComponent(this.sessionId)}/messages`, + { + params: { afterSeq: cursor, limit }, + headers: { + Authorization: `Bearer ${this.token}`, + 'Content-Type': 'application/json' + }, + timeout: 15_000 + } + ) + + const parsed = CliMessagesResponseSchema.safeParse(response.data) + if (!parsed.success) { + throw new Error('Invalid /cli/sessions/:id/messages response') + } + + const messages = parsed.data.messages + if (messages.length === 0) { + break + } + + let maxSeq = cursor + for (const message of messages) { + if (typeof message.seq === 'number') { + if (message.seq > maxSeq) { + maxSeq = message.seq + } + } + this.handleIncomingMessage(message) + } + + const observedSeq = this.lastSeenMessageSeq ?? maxSeq + const nextCursor = Math.max(maxSeq, observedSeq) + if (nextCursor <= cursor) { + logger.debug('[API] Backfill stopped due to non-advancing cursor', { + cursor, + maxSeq, + observedSeq + }) + break + } + + cursor = nextCursor + if (messages.length < limit) { + break + } + } + } + + this.backfillInFlight = run().finally(() => { + this.backfillInFlight = null + }) + + await this.backfillInFlight + } + sendClaudeSessionMessage(body: RawJSONLines): void { let content: MessageContent diff --git a/cli/src/api/types.ts b/cli/src/api/types.ts index 9dbdf84e..cb6a39c6 100644 --- a/cli/src/api/types.ts +++ b/cli/src/api/types.ts @@ -241,6 +241,18 @@ export const UpdateSchema = z.object({ export type Update = z.infer +export const CliMessagesResponseSchema = z.object({ + messages: z.array(z.object({ + id: z.string(), + seq: z.number(), + createdAt: z.number(), + localId: z.string().nullable().optional(), + content: z.unknown() + })) +}) + +export type CliMessagesResponse = z.infer + export const CreateSessionResponseSchema = z.object({ session: z.object({ id: z.string(), diff --git a/server/src/store/index.ts b/server/src/store/index.ts index 88c5e949..87099d98 100644 --- a/server/src/store/index.ts +++ b/server/src/store/index.ts @@ -549,4 +549,15 @@ export class Store { return rows.reverse().map(toStoredMessage) } + + getMessagesAfter(sessionId: string, afterSeq: number, limit: number = 200): StoredMessage[] { + const safeLimit = Number.isFinite(limit) ? Math.max(1, Math.min(200, limit)) : 200 + const safeAfterSeq = Number.isFinite(afterSeq) ? afterSeq : 0 + + const rows = this.db.prepare( + 'SELECT * FROM messages WHERE session_id = ? AND seq > ? ORDER BY seq ASC LIMIT ?' + ).all(sessionId, safeAfterSeq, safeLimit) as DbMessageRow[] + + return rows.map(toStoredMessage) + } } diff --git a/server/src/sync/syncEngine.ts b/server/src/sync/syncEngine.ts index 12de5e0f..5f6173fa 100644 --- a/server/src/sync/syncEngine.ts +++ b/server/src/sync/syncEngine.ts @@ -298,6 +298,17 @@ export class SyncEngine { } } + getMessagesAfter(sessionId: string, options: { afterSeq: number; limit: number }): DecryptedMessage[] { + const stored = this.store.getMessagesAfter(sessionId, options.afterSeq, options.limit) + return stored.map((m) => ({ + id: m.id, + seq: m.seq, + localId: m.localId, + content: m.content, + createdAt: m.createdAt + })) + } + handleRealtimeEvent(event: SyncEvent): void { if (event.type === 'session-updated' && event.sessionId) { this.refreshSession(event.sessionId) diff --git a/server/src/web/routes/cli.ts b/server/src/web/routes/cli.ts index fa15be5d..83206a60 100644 --- a/server/src/web/routes/cli.ts +++ b/server/src/web/routes/cli.ts @@ -18,6 +18,11 @@ const createOrLoadMachineSchema = z.object({ daemonState: z.unknown().nullable().optional() }) +const getMessagesQuerySchema = z.object({ + afterSeq: z.coerce.number().int().min(0), + limit: z.coerce.number().int().min(1).max(200).optional() +}) + export function createCliRoutes(getSyncEngine: () => SyncEngine | null): Hono { const app = new Hono() @@ -68,6 +73,27 @@ export function createCliRoutes(getSyncEngine: () => SyncEngine | null): Hono { return c.json({ session }) }) + app.get('/sessions/:id/messages', (c) => { + const engine = getSyncEngine() + if (!engine) { + return c.json({ error: 'Not ready' }, 503) + } + const sessionId = c.req.param('id') + const session = engine.getSession(sessionId) + if (!session) { + return c.json({ error: 'Session not found' }, 404) + } + + const parsed = getMessagesQuerySchema.safeParse(c.req.query()) + if (!parsed.success) { + return c.json({ error: 'Invalid query' }, 400) + } + + const limit = parsed.data.limit ?? 200 + const messages = engine.getMessagesAfter(sessionId, { afterSeq: parsed.data.afterSeq, limit }) + return c.json({ messages }) + }) + app.post('/machines', async (c) => { const engine = getSyncEngine() if (!engine) {