diff --git a/hub/src/store/messageStore.ts b/hub/src/store/messageStore.ts index 394a4972..38cd443e 100644 --- a/hub/src/store/messageStore.ts +++ b/hub/src/store/messageStore.ts @@ -10,6 +10,7 @@ import { getFirstMessages, getDeliverableMessagesAfter, getMessagesByPosition, + getLocalMessageStates, getUninvokedLocalMessages, getMatureScheduledMessages, getImmediateQueuedLocalMessages, @@ -23,6 +24,7 @@ import { getAllMessages, type CancelQueuedMessageResult, type LookupQueuedMessageResult, + type LocalMessageState, } from './messages' export class MessageStore { @@ -64,6 +66,10 @@ export class MessageStore { return getMessagesByPosition(this.db, sessionId, limit, before) } + getLocalMessageStates(sessionId: string, localIds: string[]): LocalMessageState[] { + return getLocalMessageStates(this.db, sessionId, localIds) + } + getUninvokedLocalMessages(sessionId: string): StoredMessage[] { return getUninvokedLocalMessages(this.db, sessionId) } diff --git a/hub/src/store/messages.ts b/hub/src/store/messages.ts index 4747c819..4994b3c4 100644 --- a/hub/src/store/messages.ts +++ b/hub/src/store/messages.ts @@ -258,6 +258,35 @@ export function getUninvokedLocalMessages( return rows.map(toStoredMessage) } +export type LocalMessageState = { + localId: string + invokedAt: number | null +} + +export function getLocalMessageStates( + db: Database, + sessionId: string, + localIds: string[] +): LocalMessageState[] { + if (localIds.length === 0) { + return [] + } + const placeholders = localIds.map(() => '?').join(', ') + const rows = db.prepare(` + SELECT local_id, invoked_at + FROM messages + WHERE session_id = ? AND local_id IN (${placeholders}) + ORDER BY seq ASC + `).all(sessionId, ...localIds) as Array<{ + local_id: string + invoked_at: number | null + }> + return rows.map((row) => ({ + localId: row.local_id, + invokedAt: row.invoked_at + })) +} + /** Returns scheduled messages across all sessions whose scheduled_at <= beforeTime * and have not yet been invoked. Used by the hub tick to emit mature messages to CLI. * Ordered by scheduled_at ASC (oldest first). */ diff --git a/hub/src/sync/messageService.test.ts b/hub/src/sync/messageService.test.ts index 00af8472..794ca0aa 100644 --- a/hub/src/sync/messageService.test.ts +++ b/hub/src/sync/messageService.test.ts @@ -307,6 +307,41 @@ describe('MessageService message pagination', () => { }) }) +describe('MessageService.getQueuedState', () => { + it('returns requested queued and invoked local IDs from the requested session', () => { + const store = makeStore() + const session = makeSession(store, 'queued-state') + const otherSession = makeSession(store, 'queued-state-other') + store.messages.addMessage(session.id, 'queued', 'local-queued') + store.messages.addMessage(session.id, 'invoked', 'local-invoked') + store.messages.addMessage( + session.id, + 'future scheduled', + 'local-future', + Date.now() + 60_000 + ) + store.messages.addMessage(otherSession.id, 'other session', 'local-other') + store.messages.markMessagesInvoked(session.id, ['local-invoked'], 1_000) + + const service = new MessageService(store, makeIo(() => {}), makePublisher() as any) + + expect(service.getQueuedState(session.id, [ + 'local-queued', + 'local-invoked', + 'local-absent', + 'local-future', + 'local-other' + ])).toEqual({ + queuedLocalIds: ['local-queued', 'local-future'], + invokedLocalMessages: [{ localId: 'local-invoked', invokedAt: 1_000 }] + }) + expect(service.getQueuedState(session.id, [])).toEqual({ + queuedLocalIds: [], + invokedLocalMessages: [] + }) + }) +}) + describe('MessageService.cancelQueuedMessage race scenarios', () => { describe('Race-A: CLI ack removed:true → DELETE + status=cancelled', () => { it('returns cancelled and emits message-cancelled SSE after CLI confirms removal', async () => { diff --git a/hub/src/sync/messageService.ts b/hub/src/sync/messageService.ts index a740f12d..fed7ab7d 100644 --- a/hub/src/sync/messageService.ts +++ b/hub/src/sync/messageService.ts @@ -10,6 +10,7 @@ import { unwrapRoleWrappedRecordEnvelope } from '@hapi/protocol/messages' import { isObject } from '@hapi/protocol' +import type { QueuedStateResponse } from '@hapi/protocol/apiTypes' import type { Server } from 'socket.io' import { randomUUID } from 'node:crypto' import type { Store, CancelQueuedMessageResult } from '../store' @@ -91,6 +92,18 @@ export class MessageService { return toVisibleDecryptedMessages(stored) } + getQueuedState(sessionId: string, localIds: string[]): QueuedStateResponse { + const states = this.store.messages.getLocalMessageStates(sessionId, localIds) + return { + queuedLocalIds: states + .filter((state) => state.invokedAt === null) + .map((state) => state.localId), + invokedLocalMessages: states.flatMap((state) => state.invokedAt === null + ? [] + : [{ localId: state.localId, invokedAt: state.invokedAt }]) + } + } + getSessionExport( sessionId: string, session: Session, diff --git a/hub/src/sync/syncEngine.ts b/hub/src/sync/syncEngine.ts index 7118e531..ffd5d766 100644 --- a/hub/src/sync/syncEngine.ts +++ b/hub/src/sync/syncEngine.ts @@ -8,7 +8,7 @@ */ import { isKnownFlavor, type LocalResumeTarget, type ResumableSession } from '@hapi/protocol' -import type { CursorChatStoreStatus, CursorMigrateOutcome, CursorMigrateToAcpRequest, SlashCommandsResponse } from '@hapi/protocol/apiTypes' +import type { CursorChatStoreStatus, CursorMigrateOutcome, CursorMigrateToAcpRequest, QueuedStateResponse, SlashCommandsResponse } from '@hapi/protocol/apiTypes' import type { AgentFlavor, CodexCollaborationMode, DecryptedMessage, PermissionMode, Session, SyncEvent } from '@hapi/protocol/types' import { unwrapRoleWrappedRecordEnvelope } from '@hapi/protocol/messages' import type { Server } from 'socket.io' @@ -333,6 +333,10 @@ export class SyncEngine { return this.messageService.getMessagesPage(sessionId, options) } + getQueuedState(sessionId: string, localIds: string[]): QueuedStateResponse { + return this.messageService.getQueuedState(sessionId, localIds) + } + getSessionExport(sessionId: string, session: Session): HapiSessionExportResult { return this.messageService.getSessionExport(sessionId, session) } diff --git a/hub/src/web/routes/messages.test.ts b/hub/src/web/routes/messages.test.ts index 64248810..346fc693 100644 --- a/hub/src/web/routes/messages.test.ts +++ b/hub/src/web/routes/messages.test.ts @@ -21,11 +21,25 @@ import { createMessagesRoutes } from './messages' function createApp(opts: { active?: boolean sendMessage?: (sessionId: string, payload: unknown) => Promise + getQueuedState?: (sessionId: string, localIds: string[]) => { + queuedLocalIds: string[] + invokedLocalMessages: Array<{ localId: string; invokedAt: number }> + } }) { const sentMessages: Array<{ sessionId: string; payload: unknown }> = [] + const queuedStateCalls: Array<{ sessionId: string; localIds: string[] }> = [] const sendMessage = opts.sendMessage ?? (async (sessionId: string, payload: unknown) => { sentMessages.push({ sessionId, payload }) }) + const getQueuedState = opts.getQueuedState ?? ((sessionId: string, localIds: string[]) => { + queuedStateCalls.push({ sessionId, localIds }) + return { + queuedLocalIds: localIds.filter((localId) => localId.startsWith('queued-')), + invokedLocalMessages: localIds + .filter((localId) => localId.startsWith('invoked-')) + .map((localId) => ({ localId, invokedAt: 1_000 })) + } + }) const engine = { resolveSessionAccess: () => ({ @@ -34,6 +48,7 @@ function createApp(opts: { session: { id: 'session-1', active: opts.active !== false } }), sendMessage, + getQueuedState, cancelQueuedMessage: async () => ({ status: 'cancelled' }), getMessagesPage: () => ({ messages: [], page: {} }), } as unknown as SyncEngine @@ -45,7 +60,7 @@ function createApp(opts: { }) app.route('/api', createMessagesRoutes(() => engine as SyncEngine)) - return { app, sentMessages } + return { app, sentMessages, queuedStateCalls } } // --------------------------------------------------------------------------- @@ -241,3 +256,59 @@ describe('POST /api/sessions/:id/messages — inactive session response shape', expect(sentMessages).toHaveLength(0) }) }) + +describe('POST /api/sessions/:id/messages/queued-state', () => { + it('deduplicates local IDs, forwards the session, and works for inactive sessions', async () => { + const { app, queuedStateCalls } = createApp({ active: false }) + + const response = await app.request('/api/sessions/session-1/messages/queued-state', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + localIds: ['queued-2', 'missing-1', 'queued-2', 'invoked-1', 'queued-1'] + }) + }) + + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ + queuedLocalIds: ['queued-2', 'queued-1'], + invokedLocalMessages: [{ localId: 'invoked-1', invokedAt: 1_000 }] + }) + expect(queuedStateCalls).toEqual([{ + sessionId: 'session-1', + localIds: ['queued-2', 'missing-1', 'invoked-1', 'queued-1'] + }]) + }) + + it('accepts an empty candidate list as a no-op', async () => { + const { app, queuedStateCalls } = createApp({}) + + const response = await app.request('/api/sessions/session-1/messages/queued-state', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ localIds: [] }) + }) + + expect(response.status).toBe(200) + expect(await response.json()).toEqual({ queuedLocalIds: [], invokedLocalMessages: [] }) + expect(queuedStateCalls).toHaveLength(0) + }) + + it.each([ + ['an empty local ID', { localIds: [''] }], + ['a non-array localIds value', { localIds: 'queued-1' }], + ['more than 1000 local IDs', { localIds: Array.from({ length: 1001 }, (_, i) => `local-${i}`) }] + ])('rejects %s', async (_label, body) => { + const { app, queuedStateCalls } = createApp({}) + + const response = await app.request('/api/sessions/session-1/messages/queued-state', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify(body) + }) + + expect(response.status).toBe(400) + expect(await response.json()).toMatchObject({ error: 'Invalid body' }) + expect(queuedStateCalls).toHaveLength(0) + }) +}) diff --git a/hub/src/web/routes/messages.ts b/hub/src/web/routes/messages.ts index 715950cc..f9bbc333 100644 --- a/hub/src/web/routes/messages.ts +++ b/hub/src/web/routes/messages.ts @@ -1,5 +1,5 @@ import { Hono } from 'hono' -import { MessagesQuerySchema, SendMessageRequestSchema } from '@hapi/protocol' +import { MessagesQuerySchema, QueuedStateRequestSchema, SendMessageRequestSchema } from '@hapi/protocol' import type { SyncEngine } from '../../sync/syncEngine' import type { WebAppEnv } from '../middleware/auth' import { requireSessionFromParam, requireSyncEngine } from './guards' @@ -48,6 +48,31 @@ export function createMessagesRoutes(getSyncEngine: () => SyncEngine | null): Ho return c.json(result) }) + app.post('/sessions/:id/messages/queued-state', async (c) => { + const engine = requireSyncEngine(c, getSyncEngine) + if (engine instanceof Response) { + return engine + } + + const sessionResult = requireSessionFromParam(c, engine) + if (sessionResult instanceof Response) { + return sessionResult + } + const sessionId = sessionResult.sessionId + + const body = await c.req.json().catch(() => null) + const parsed = QueuedStateRequestSchema.safeParse(body) + if (!parsed.success) { + return c.json({ error: 'Invalid body', issues: parsed.error.flatten() }, 400) + } + + const localIds = [...new Set(parsed.data.localIds)] + if (localIds.length === 0) { + return c.json({ queuedLocalIds: [], invokedLocalMessages: [] }) + } + return c.json(engine.getQueuedState(sessionId, localIds)) + }) + app.post('/sessions/:id/messages', async (c) => { const engine = requireSyncEngine(c, getSyncEngine) if (engine instanceof Response) { diff --git a/shared/src/apiTypes.ts b/shared/src/apiTypes.ts index fea310b2..f09a9708 100644 --- a/shared/src/apiTypes.ts +++ b/shared/src/apiTypes.ts @@ -251,6 +251,20 @@ export const SendMessageRequestSchema = z.object({ export type SendMessageRequest = z.infer +export const QueuedStateRequestSchema = z.object({ + localIds: z.array(z.string().min(1)).max(1000) +}) + +export type QueuedStateRequest = z.infer + +export type QueuedStateResponse = { + queuedLocalIds: string[] + invokedLocalMessages: Array<{ + localId: string + invokedAt: number + }> +} + export const SpawnSessionRequestSchema = z.object({ directory: z.string().min(1), agent: AgentFlavorSchema.optional(), diff --git a/web/src/App.tsx b/web/src/App.tsx index b8a8f010..1e519210 100644 --- a/web/src/App.tsx +++ b/web/src/App.tsx @@ -21,6 +21,7 @@ import { useTranslation } from '@/lib/use-translation' import { VoiceProvider } from '@/lib/voice-context' import { requireHubUrlForLogin } from '@/lib/runtime-config' import { getAppGlobalSseSubscription, getAppSessionSseSubscription } from '@/lib/appSseSubscriptions' +import { reconcileQueuedStateAfterConnect } from '@/lib/queued-state-reconciliation' import { LoginPrompt } from '@/components/LoginPrompt' import { InstallPrompt } from '@/components/InstallPrompt' import { OfflineBanner } from '@/components/OfflineBanner' @@ -260,6 +261,16 @@ function AppInner() { clearMessageWindow(event.sessionId) void fetchLatestMessages(api, event.sessionId) }, [api, selectedSessionId]) + + const handleSessionSseConnect = useCallback(() => { + if (!api || !selectedSessionId) { + return + } + void reconcileQueuedStateAfterConnect(api, selectedSessionId).catch((error) => { + console.error('Failed to reconcile queued state after SSE connect:', error) + }) + }, [api, selectedSessionId]) + const translateIncomingToast = useCallback((title: string, body: string): { title: string; body: string } => { const normalizedTitle = title.trim() const normalizedBody = body.trim() @@ -339,6 +350,7 @@ function AppInner() { baseUrl, subscription: sessionEventSubscription ?? undefined, scope: 'full', + onConnect: handleSessionSseConnect, onEvent: handleSseEvent }) diff --git a/web/src/api/client.test.ts b/web/src/api/client.test.ts index 3297c2c8..134bde6d 100644 --- a/web/src/api/client.test.ts +++ b/web/src/api/client.test.ts @@ -92,4 +92,27 @@ describe('ApiClient error mapping', () => { }) expect(fetchMock.mock.calls[0]?.[0]).toBe('/api/sessions/session%20cursor/cursor-chat-store') }) + + it('loads the authoritative queued state for encoded session IDs', async () => { + fetchMock.mockResolvedValueOnce( + new Response(JSON.stringify({ + queuedLocalIds: ['local-2'], + invokedLocalMessages: [{ localId: 'local-3', invokedAt: 1_000 }] + }), { status: 200 }) + ) + + const api = new ApiClient('test-token') + await expect(api.getQueuedState('session /?#', ['local-1', 'local-2'])).resolves.toEqual({ + queuedLocalIds: ['local-2'], + invokedLocalMessages: [{ localId: 'local-3', invokedAt: 1_000 }] + }) + + const [url, init] = fetchMock.mock.calls[0] ?? [] + expect(url).toBe('/api/sessions/session%20%2F%3F%23/messages/queued-state') + expect(init).toMatchObject({ + method: 'POST', + body: JSON.stringify({ localIds: ['local-1', 'local-2'] }) + }) + expect(new Headers(init?.headers).get('content-type')).toBe('application/json') + }) }) diff --git a/web/src/api/client.ts b/web/src/api/client.ts index a99dcfe7..7d3cd2a5 100644 --- a/web/src/api/client.ts +++ b/web/src/api/client.ts @@ -39,6 +39,7 @@ import type { MachinePathsExistsResponse, OpencodeModelsResponse, OpencodeReasoningEffortResponse, + QueuedStateResponse, ReopenSessionResponse, UploadFileResponse } from '@hapi/protocol/apiTypes' @@ -408,6 +409,16 @@ export class ApiClient { }) } + async getQueuedState(sessionId: string, localIds: string[]): Promise { + return await this.request( + `/api/sessions/${encodeURIComponent(sessionId)}/messages/queued-state`, + { + method: 'POST', + body: JSON.stringify({ localIds }) + } + ) + } + async cancelMessage(sessionId: string, messageId: string): Promise { const response = await this.request( `/api/sessions/${encodeURIComponent(sessionId)}/messages/${encodeURIComponent(messageId)}`, diff --git a/web/src/hooks/mutations/useSendMessage.test.tsx b/web/src/hooks/mutations/useSendMessage.test.tsx index 084f06a2..20659aeb 100644 --- a/web/src/hooks/mutations/useSendMessage.test.tsx +++ b/web/src/hooks/mutations/useSendMessage.test.tsx @@ -35,6 +35,14 @@ function createMockApi(sendMessage: (...args: unknown[]) => Promise = asyn return { sendMessage } as unknown as ApiClient } +function deferred() { + let resolve!: (value: T | PromiseLike) => void + const promise = new Promise((res) => { + resolve = res + }) + return { promise, resolve } +} + describe('useSendMessage', () => { beforeEach(() => { vi.clearAllMocks() @@ -58,6 +66,43 @@ describe('useSendMessage', () => { }) }) + it('keeps a thinking-session send in flight until the POST confirms it is queued', async () => { + const request = deferred() + const api = createMockApi(() => request.promise) + const { appendOptimisticMessage, updateMessageStatus } = await import('@/lib/message-window-store') + const appendMock = vi.mocked(appendOptimisticMessage) + const updateMock = vi.mocked(updateMessageStatus) + + const { result } = renderHook( + () => useSendMessage(api, 'session-A', { isSessionThinking: true }), + { wrapper: createWrapper() }, + ) + + act(() => { + void result.current.sendMessage('queue this') + }) + + await waitFor(() => { + expect(appendMock).toHaveBeenCalledWith( + 'session-A', + expect.objectContaining({ + localId: 'local-id-1', + status: 'sending', + }), + ) + }) + expect(updateMock).not.toHaveBeenCalledWith('session-A', 'local-id-1', 'queued') + + await act(async () => { + request.resolve() + await request.promise + }) + + await waitFor(() => { + expect(updateMock).toHaveBeenCalledWith('session-A', 'local-id-1', 'queued') + }) + }) + it('calls onSuccess with resolved session ID, not the original', async () => { const onSuccess = vi.fn() const api = createMockApi() diff --git a/web/src/hooks/mutations/useSendMessage.ts b/web/src/hooks/mutations/useSendMessage.ts index c37d3874..04325e47 100644 --- a/web/src/hooks/mutations/useSendMessage.ts +++ b/web/src/hooks/mutations/useSendMessage.ts @@ -161,15 +161,15 @@ export function useSendMessage( await api.sendMessage(input.sessionId, input.text, input.localId, input.attachments, input.scheduledAt) }, onMutate: async (input) => { - const status = isSessionThinkingRef.current ? 'queued' as const : 'sending' as const - appendOptimisticMessage(input.sessionId, createOptimisticMessage(input, status)) - return { status } + const successStatus = isSessionThinkingRef.current ? 'queued' as const : 'sent' as const + appendOptimisticMessage(input.sessionId, createOptimisticMessage(input, 'sending')) + return { successStatus } }, onSuccess: (_, input, context) => { updateMessageStatus( input.sessionId, input.localId, - context?.status === 'queued' ? 'queued' : 'sent' + context?.successStatus ?? 'sent' ) haptic.notification('success') options?.onSuccess?.(input.sessionId) diff --git a/web/src/lib/message-window-store.test.ts b/web/src/lib/message-window-store.test.ts index 24369531..859a0c34 100644 --- a/web/src/lib/message-window-store.test.ts +++ b/web/src/lib/message-window-store.test.ts @@ -6,10 +6,12 @@ import { clearMessageWindow, fetchLatestMessages, fetchOlderMessages, + getQueuedReconcileCandidateLocalIds, getMessageWindowState, ingestIncomingMessages, markMessagesConsumed, reconcileQueuedAgainstLatest, + reconcileQueuedLocalIds, removeOptimisticMessage, setAtBottom, VISIBLE_WINDOW_SIZE, @@ -215,6 +217,39 @@ describe('message-window-store async generations', () => { sessionStorage.clear() }) + it('waits for an in-flight latest refresh instead of resolving early', async () => { + const request = deferred>>() + const api = { + getMessages: vi.fn(() => request.promise) + } as Pick & { + getMessages: ReturnType + } + + const firstLoad = fetchLatestMessages(api as unknown as ApiClient, SESSION_ID) + let secondResolved = false + const secondLoad = fetchLatestMessages(api as unknown as ApiClient, SESSION_ID) + .then(() => { + secondResolved = true + }) + await Promise.resolve() + + expect(api.getMessages).toHaveBeenCalledTimes(1) + expect(secondResolved).toBe(false) + + request.resolve({ + messages: [], + page: { + limit: 50, + nextBeforeSeq: null, + nextBeforeAt: null, + hasMore: false + } + }) + await Promise.all([firstLoad, secondLoad]) + + expect(secondResolved).toBe(true) + }) + it('does not let a stale failed retry overwrite a newer reset-and-reload state', async () => { const firstRequest = deferred>>() const api = { @@ -452,6 +487,107 @@ describe('message-window-store status updates', () => { }) }) +describe('queued-state reconciliation', () => { + const CANDIDATE_SESSION_ID = 'session-queued-state-candidates-test' + const PERSISTED_SENDING_SESSION_ID = 'session-queued-state-persisted-sending-test' + const RECONCILE_SESSION_ID = 'session-queued-state-reconcile-test' + + function makeQueuedUserMessage(props: Parameters[0]): DecryptedMessage { + return { + ...makeUserMessage(props), + invokedAt: null, + } + } + + function hydrate(sessionId: string, messages: DecryptedMessage[], pending: DecryptedMessage[] = []): void { + sessionStorage.setItem(`hapi:message-window:v1:${sessionId}`, JSON.stringify({ + messages, + pending, + atBottom: true, + })) + } + + afterEach(() => { + clearMessageWindow(CANDIDATE_SESSION_ID) + clearMessageWindow(PERSISTED_SENDING_SESSION_ID) + clearMessageWindow(RECONCILE_SESSION_ID) + }) + + it('deduplicates queued candidates and excludes unsafe optimistic rows', () => { + hydrate(CANDIDATE_SESSION_ID, [ + makeQueuedUserMessage({ id: 'server-echo', localId: 'local-server' }), + makeQueuedUserMessage({ id: 'local-queued', localId: 'local-queued', status: 'queued' }), + makeQueuedUserMessage({ id: 'local-sent', localId: 'local-sent', status: 'sent' }), + makeQueuedUserMessage({ id: 'local-sending', localId: 'local-sending', status: 'queued' }), + makeQueuedUserMessage({ id: 'local-failed', localId: 'local-failed', status: 'failed' }), + { + ...makeQueuedUserMessage({ id: 'local-invoked', localId: 'local-invoked', status: 'sent' }), + invokedAt: 1_700_000_000_000, + }, + ], [ + makeQueuedUserMessage({ id: 'server-echo-duplicate', localId: 'local-server' }), + ]) + updateMessageStatus(CANDIDATE_SESSION_ID, 'local-sending', 'sending') + + expect(getQueuedReconcileCandidateLocalIds(CANDIDATE_SESSION_ID)).toEqual([ + 'local-server', + 'local-queued', + 'local-sent', + ]) + }) + + it('treats persisted sending rows as queued candidates after reload', () => { + hydrate(PERSISTED_SENDING_SESSION_ID, [ + makeQueuedUserMessage({ id: 'local-sending', localId: 'local-sending', status: 'sending' }), + ]) + + expect(getQueuedReconcileCandidateLocalIds(PERSISTED_SENDING_SESSION_ID)).toEqual(['local-sending']) + }) + + it('removes only snapshotted rows that are no longer authoritatively queued', () => { + hydrate(RECONCILE_SESSION_ID, [ + makeQueuedUserMessage({ id: 'stale-message', localId: 'local-stale-message' }), + makeQueuedUserMessage({ id: 'queued-message', localId: 'local-queued-message' }), + makeQueuedUserMessage({ id: 'new-message', localId: 'local-new-message' }), + makeQueuedUserMessage({ id: 'local-retry', localId: 'local-retry', status: 'sending' }), + { + ...makeQueuedUserMessage({ id: 'invoked-message', localId: 'local-invoked-message' }), + invokedAt: 1_700_000_000_000, + }, + ], [ + makeQueuedUserMessage({ id: 'stale-pending', localId: 'local-stale-pending' }), + makeQueuedUserMessage({ id: 'queued-pending', localId: 'local-queued-pending' }), + makeQueuedUserMessage({ id: 'new-pending', localId: 'local-new-pending' }), + ]) + updateMessageStatus(RECONCILE_SESSION_ID, 'local-retry', 'sending') + + reconcileQueuedLocalIds( + RECONCILE_SESSION_ID, + [ + 'local-stale-message', + 'local-queued-message', + 'local-retry', + 'local-invoked-message', + 'local-stale-pending', + 'local-queued-pending', + ], + ['local-queued-message', 'local-queued-pending'], + ) + + const state = getMessageWindowState(RECONCILE_SESSION_ID) + expect(state.messages.map((message) => message.id)).toEqual([ + 'queued-message', + 'new-message', + 'local-retry', + 'invoked-message', + ]) + expect(state.pending.map((message) => message.id)).toEqual([ + 'queued-pending', + 'new-pending', + ]) + }) +}) + describe('message-window-store visible trimming', () => { const SESSION_ID = 'session-message-window-trim-test' diff --git a/web/src/lib/message-window-store.ts b/web/src/lib/message-window-store.ts index 35d4c449..16c2ddb8 100644 --- a/web/src/lib/message-window-store.ts +++ b/web/src/lib/message-window-store.ts @@ -64,6 +64,7 @@ type PersistedMessageWindowState = { const states = new Map() const listeners = new Map void>>() const pendingVisibilityCacheBySession = new Map>() +const latestLoads = new Map>() // Throttled notification: coalesce rapid state updates into at most one // notification per NOTIFY_THROTTLE_MS during streaming. This prevents @@ -284,9 +285,20 @@ function hydrateState(sessionId: string): InternalState | null { return null } const base = createState(sessionId) + const restorePersistedMessage = (message: DecryptedMessage): DecryptedMessage => { + if (message.status !== 'sending') { + return message + } + // A page reload ends the in-flight POST, so persisted sending rows + // must re-enter authoritative queued-state reconciliation. + return { + ...message, + status: message.invokedAt === null ? 'queued' as const : 'sent' as const + } + } return buildState(base, { - messages: parsed.messages, - pending: parsed.pending, + messages: parsed.messages.map(restorePersistedMessage), + pending: parsed.pending.map(restorePersistedMessage), pendingOverflowCount: typeof parsed.pendingOverflowCount === 'number' ? parsed.pendingOverflowCount : 0, pendingOverflowVisibleCount: typeof parsed.pendingOverflowVisibleCount === 'number' ? parsed.pendingOverflowVisibleCount : 0, hasMore: parsed.hasMore === true, @@ -669,6 +681,16 @@ function isOptimisticMessage(message: DecryptedMessage): boolean { return Boolean(message.localId && message.id === message.localId) } +function isQueuedReconcileCandidate(message: DecryptedMessage): boolean { + if (!message.localId || !isQueuedForInvocation(message)) { + return false + } + if (!isOptimisticMessage(message)) { + return true + } + return message.status === 'queued' || message.status === 'sent' +} + /** * Drops phantom queued messages during an at-bottom full refresh. * @@ -764,6 +786,48 @@ export function getMessageWindowState(sessionId: string): MessageWindowState { return getState(sessionId) } +export function getQueuedReconcileCandidateLocalIds(sessionId: string): string[] { + const state = getState(sessionId) + const localIds = new Set() + for (const message of [...state.messages, ...state.pending]) { + if (isQueuedReconcileCandidate(message)) { + localIds.add(message.localId!) + } + } + return [...localIds] +} + +export function reconcileQueuedLocalIds( + sessionId: string, + candidateLocalIds: string[], + queuedLocalIds: string[] +): void { + if (candidateLocalIds.length === 0) { + return + } + const candidates = new Set(candidateLocalIds) + const queued = new Set(queuedLocalIds) + updateState(sessionId, (prev) => { + let changed = false + const reconcile = (messages: DecryptedMessage[]) => messages.filter((message) => { + if (!message.localId || !candidates.has(message.localId)) { + return true + } + if (queued.has(message.localId) || !isQueuedReconcileCandidate(message)) { + return true + } + changed = true + return false + }) + const messages = reconcile(prev.messages) + const pending = reconcile(prev.pending) + if (!changed) { + return prev + } + return buildState(prev, { messages, pending }) + }, true) +} + export function subscribeMessageWindow(sessionId: string, listener: () => void): () => void { const subs = listeners.get(sessionId) ?? new Set() subs.add(listener) @@ -780,6 +844,7 @@ export function subscribeMessageWindow(sessionId: string, listener: () => void): } export function clearMessageWindow(sessionId: string): void { + latestLoads.delete(sessionId) clearPendingVisibilityCache(sessionId) clearPersistedState(sessionId) const previous = states.get(sessionId) @@ -819,7 +884,23 @@ export function seedMessageWindowFromSession(fromSessionId: string, toSessionId: }) } -export async function fetchLatestMessages(api: ApiClient, sessionId: string): Promise { +export function fetchLatestMessages(api: ApiClient, sessionId: string): Promise { + const existing = latestLoads.get(sessionId) + if (existing) { + return existing + } + const load = fetchLatestMessagesOnce(api, sessionId) + latestLoads.set(sessionId, load) + const cleanup = () => { + if (latestLoads.get(sessionId) === load) { + latestLoads.delete(sessionId) + } + } + void load.then(cleanup, cleanup) + return load +} + +async function fetchLatestMessagesOnce(api: ApiClient, sessionId: string): Promise { const initial = getState(sessionId) if (initial.isLoading) { return diff --git a/web/src/lib/queued-state-reconciliation.test.ts b/web/src/lib/queued-state-reconciliation.test.ts new file mode 100644 index 00000000..96a342c0 --- /dev/null +++ b/web/src/lib/queued-state-reconciliation.test.ts @@ -0,0 +1,151 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' +import type { ApiClient } from '@/api/client' + +vi.mock('./message-window-store', () => ({ + fetchLatestMessages: vi.fn(), + getQueuedReconcileCandidateLocalIds: vi.fn(), + markMessagesConsumed: vi.fn(), + reconcileQueuedLocalIds: vi.fn(), +})) + +import { + fetchLatestMessages, + getQueuedReconcileCandidateLocalIds, + markMessagesConsumed, + reconcileQueuedLocalIds, +} from './message-window-store' +import { reconcileQueuedStateAfterConnect } from './queued-state-reconciliation' + +const mockFetchLatestMessages = vi.mocked(fetchLatestMessages) +const mockGetCandidates = vi.mocked(getQueuedReconcileCandidateLocalIds) +const mockMarkMessagesConsumed = vi.mocked(markMessagesConsumed) +const mockReconcileQueuedLocalIds = vi.mocked(reconcileQueuedLocalIds) + +function createMockApi( + getQueuedState: ApiClient['getQueuedState'] = async () => ({ + queuedLocalIds: [], + invokedLocalMessages: [] + }) +): ApiClient { + return { getQueuedState } as ApiClient +} + +describe('reconcileQueuedStateAfterConnect', () => { + beforeEach(() => { + vi.clearAllMocks() + mockFetchLatestMessages.mockResolvedValue(undefined) + mockGetCandidates.mockReturnValue([]) + }) + + it('waits for the latest messages before snapshotting and querying queued state', async () => { + let resolveRefresh: (() => void) | undefined + mockFetchLatestMessages.mockImplementationOnce( + () => new Promise((resolve) => { + resolveRefresh = resolve + }) + ) + mockGetCandidates.mockReturnValueOnce(['local-1']) + const getQueuedState = vi.fn(async () => ({ + queuedLocalIds: ['local-1'], + invokedLocalMessages: [] + })) + + const reconciliation = reconcileQueuedStateAfterConnect( + createMockApi(getQueuedState), + 'session-A' + ) + await Promise.resolve() + + expect(mockGetCandidates).not.toHaveBeenCalled() + expect(getQueuedState).not.toHaveBeenCalled() + + resolveRefresh?.() + await reconciliation + + expect(mockGetCandidates).toHaveBeenCalledWith('session-A') + expect(getQueuedState).toHaveBeenCalledWith('session-A', ['local-1']) + }) + + it('passes the exact candidate snapshot to the endpoint and reconciliation', async () => { + const candidateLocalIds = ['local-1', 'local-2'] + mockGetCandidates.mockReturnValueOnce(candidateLocalIds) + const getQueuedState = vi.fn(async () => ({ + queuedLocalIds: ['local-2'], + invokedLocalMessages: [] + })) + + await reconcileQueuedStateAfterConnect(createMockApi(getQueuedState), 'session-B') + + expect(getQueuedState).toHaveBeenCalledWith('session-B', candidateLocalIds) + expect(mockReconcileQueuedLocalIds).toHaveBeenCalledWith( + 'session-B', + candidateLocalIds, + ['local-2'], + ) + }) + + it('applies authoritative invoked timestamps before reconciling absent rows', async () => { + mockGetCandidates.mockReturnValueOnce(['local-1', 'local-2']) + const getQueuedState = vi.fn(async () => ({ + queuedLocalIds: ['local-2'], + invokedLocalMessages: [{ localId: 'local-1', invokedAt: 1_000 }] + })) + + await reconcileQueuedStateAfterConnect(createMockApi(getQueuedState), 'session-B') + + expect(mockMarkMessagesConsumed).toHaveBeenCalledWith('session-B', ['local-1'], 1_000) + expect(mockMarkMessagesConsumed.mock.invocationCallOrder[0]).toBeLessThan( + mockReconcileQueuedLocalIds.mock.invocationCallOrder[0] ?? Number.MAX_SAFE_INTEGER + ) + expect(mockReconcileQueuedLocalIds).toHaveBeenCalledWith( + 'session-B', + ['local-1', 'local-2'], + ['local-2'], + ) + }) + + it('chunks candidate IDs to stay within the endpoint limit', async () => { + const candidateLocalIds = Array.from({ length: 1001 }, (_, index) => `local-${index}`) + mockGetCandidates.mockReturnValueOnce(candidateLocalIds) + const getQueuedState = vi.fn(async (_sessionId: string, localIds: string[]) => ({ + queuedLocalIds: localIds, + invokedLocalMessages: [] + })) + + await reconcileQueuedStateAfterConnect(createMockApi(getQueuedState), 'session-E') + + expect(getQueuedState).toHaveBeenCalledTimes(2) + expect(getQueuedState.mock.calls[0]?.[1]).toHaveLength(1000) + expect(getQueuedState.mock.calls[1]?.[1]).toEqual(['local-1000']) + expect(mockReconcileQueuedLocalIds).toHaveBeenCalledWith( + 'session-E', + candidateLocalIds, + candidateLocalIds, + ) + }) + + it('skips the endpoint and reconciliation when there are no candidates', async () => { + const getQueuedState = vi.fn(async () => ({ + queuedLocalIds: [], + invokedLocalMessages: [] + })) + + await reconcileQueuedStateAfterConnect(createMockApi(getQueuedState), 'session-C') + + expect(getQueuedState).not.toHaveBeenCalled() + expect(mockReconcileQueuedLocalIds).not.toHaveBeenCalled() + }) + + it('propagates endpoint failures without reconciling', async () => { + const endpointError = new Error('queued state unavailable') + mockGetCandidates.mockReturnValueOnce(['local-1']) + const getQueuedState = vi.fn(async () => { + throw endpointError + }) + + await expect( + reconcileQueuedStateAfterConnect(createMockApi(getQueuedState), 'session-D') + ).rejects.toBe(endpointError) + expect(mockReconcileQueuedLocalIds).not.toHaveBeenCalled() + }) +}) diff --git a/web/src/lib/queued-state-reconciliation.ts b/web/src/lib/queued-state-reconciliation.ts new file mode 100644 index 00000000..7d67fe36 --- /dev/null +++ b/web/src/lib/queued-state-reconciliation.ts @@ -0,0 +1,38 @@ +import type { ApiClient } from '@/api/client' +import { + fetchLatestMessages, + getQueuedReconcileCandidateLocalIds, + markMessagesConsumed, + reconcileQueuedLocalIds, +} from './message-window-store' + +const QUEUED_STATE_BATCH_SIZE = 1000 + +export async function reconcileQueuedStateAfterConnect( + api: ApiClient, + sessionId: string +): Promise { + await fetchLatestMessages(api, sessionId) + const candidateLocalIds = getQueuedReconcileCandidateLocalIds(sessionId) + if (candidateLocalIds.length === 0) { + return + } + const queuedLocalIds: string[] = [] + const invokedLocalMessages: Array<{ localId: string; invokedAt: number }> = [] + for (let index = 0; index < candidateLocalIds.length; index += QUEUED_STATE_BATCH_SIZE) { + const batch = candidateLocalIds.slice(index, index + QUEUED_STATE_BATCH_SIZE) + const state = await api.getQueuedState(sessionId, batch) + queuedLocalIds.push(...state.queuedLocalIds) + invokedLocalMessages.push(...state.invokedLocalMessages) + } + const invokedByTimestamp = new Map() + for (const message of invokedLocalMessages) { + const localIds = invokedByTimestamp.get(message.invokedAt) ?? [] + localIds.push(message.localId) + invokedByTimestamp.set(message.invokedAt, localIds) + } + for (const [invokedAt, localIds] of invokedByTimestamp) { + markMessagesConsumed(sessionId, localIds, invokedAt) + } + reconcileQueuedLocalIds(sessionId, candidateLocalIds, queuedLocalIds) +}