diff --git a/apps/sim/app/api/copilot/chat/queries.ts b/apps/sim/app/api/copilot/chat/queries.ts index 105bcac4011..dde58162c57 100644 --- a/apps/sim/app/api/copilot/chat/queries.ts +++ b/apps/sim/app/api/copilot/chat/queries.ts @@ -5,9 +5,12 @@ import { authorizeWorkflowByWorkspacePermission } from '@sim/platform-authz/work import { toError } from '@sim/utils/errors' import { and, desc, eq, isNull } from 'drizzle-orm' import { type NextRequest, NextResponse } from 'next/server' -import { getLatestRunForStream } from '@/lib/mothership/async-runs/repository' import { buildEffectiveChatTranscript } from '@/lib/mothership/chat/effective-transcript' import { getAccessibleCopilotChat } from '@/lib/mothership/chat/lifecycle' +import { + type LiveTurnSnapshot, + readLiveTurnSnapshot, +} from '@/lib/mothership/chat/live-turn-snapshot' import { normalizeMessage } from '@/lib/mothership/chat/persisted-message' import { authenticateCopilotRequestSessionOnly, @@ -16,9 +19,6 @@ import { createInternalServerErrorResponse, createUnauthorizedResponse, } from '@/lib/mothership/request/http' -import { readFilePreviewSessions } from '@/lib/mothership/request/session' -import { readEvents } from '@/lib/mothership/request/session/buffer' -import { toStreamBatchEvent } from '@/lib/mothership/request/session/types' import { assertActiveWorkspaceAccess, isWorkspaceAccessDeniedError, @@ -87,43 +87,10 @@ export async function GET(req: NextRequest) { return NextResponse.json({ success: false, error: 'Chat not found' }, { status: 404 }) } - let streamSnapshot: { - events: ReturnType[] - previewSessions: Awaited> - status: string - } | null = null + let streamSnapshot: LiveTurnSnapshot | null = null if (chat.conversationId) { try { - const [events, previewSessions, run] = await Promise.all([ - readEvents(chat.conversationId, '0'), - readFilePreviewSessions(chat.conversationId).catch((error) => { - logger.warn('Failed to read preview sessions for copilot chat', { - chatId, - conversationId: chat.conversationId, - error: toError(error).message, - }) - return [] - }), - getLatestRunForStream(chat.conversationId, authenticatedUserId).catch((error) => { - logger.warn('Failed to fetch latest run for copilot chat snapshot', { - chatId, - conversationId: chat.conversationId, - error: toError(error).message, - }) - return null - }), - ]) - - streamSnapshot = { - events: events.map(toStreamBatchEvent), - previewSessions, - status: - typeof run?.status === 'string' - ? run.status - : events.length > 0 - ? 'active' - : 'unknown', - } + streamSnapshot = await readLiveTurnSnapshot(chat.conversationId, authenticatedUserId) } catch (error) { logger.warn('Failed to load copilot chat stream snapshot', { chatId, diff --git a/apps/sim/app/api/copilot/chat/stream/route.test.ts b/apps/sim/app/api/copilot/chat/stream/route.test.ts index 9a020efeda3..f2815f0e665 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.test.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.test.ts @@ -16,13 +16,14 @@ import { CopilotResumeOutcome } from '@/lib/mothership/generated/trace-attribute import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1' -const { getLatestRunForStream, readEvents, readFilePreviewSessions, checkForReplayGap } = - vi.hoisted(() => ({ +const { getLatestRunForStream, readEvents, readFilePreviewSessions, findReplayGap } = vi.hoisted( + () => ({ getLatestRunForStream: vi.fn(), readEvents: vi.fn(), readFilePreviewSessions: vi.fn(), - checkForReplayGap: vi.fn(), - })) + findReplayGap: vi.fn(), + }) +) vi.mock('@/lib/mothership/request/application/recover-stream', () => ({ readChatStream: { execute: getLatestRunForStream }, @@ -33,7 +34,10 @@ vi.mock('@/lib/mothership/request/session', () => ({ status === 'complete' || status === 'error' || status === 'cancelled', readEvents, readFilePreviewSessions, - checkForReplayGap, + findReplayGap, + readRingPosition: async () => ({ requestedAfterSeq: 0, oldestSeq: 0, latestSeq: 0 }), + ringCanServe: () => true, + replayGapTerminal: async () => ({ gapDetected: true, envelopes: [] }), createEvent: (event: Record) => ({ stream: { streamId: event.streamId, @@ -82,7 +86,7 @@ describe('copilot chat stream replay route', () => { }) readEvents.mockResolvedValue([]) readFilePreviewSessions.mockResolvedValue([]) - checkForReplayGap.mockResolvedValue(null) + findReplayGap.mockResolvedValue(null) }) it('refuses replay after organization membership is removed', async () => { @@ -215,4 +219,81 @@ describe('copilot chat stream replay route', () => { ) trace.disable() }) + + it('never delivers a ring read that starts past the reader cursor, and ends without a terminal', async () => { + getLatestRunForStream.mockResolvedValue({ + status: 'active', + executionId: 'exec-1', + id: 'run-1', + }) + readEvents.mockResolvedValue([ + { + stream: { streamId: 'stream-1', cursor: '5' }, + seq: 5, + trace: { requestId: 'req-1' }, + type: MothershipStreamV1EventType.text, + payload: { channel: 'assistant', text: 'the middle of the turn' }, + }, + ]) + + const response = await GET( + new NextRequest('http://localhost:3000/api/copilot/chat/stream?streamId=stream-1&after=0') + ) + const text = (await readAllChunks(response)).join('') + + expect(text).not.toContain('the middle of the turn') + expect(text).not.toContain(`"type":"${MothershipStreamV1EventType.complete}"`) + }) + + it('serves a batch read that starts past the reader cursor no events', async () => { + getLatestRunForStream.mockResolvedValue({ + status: 'active', + executionId: 'exec-1', + id: 'run-1', + }) + readEvents.mockResolvedValue([ + { + stream: { streamId: 'stream-1', cursor: '5' }, + seq: 5, + trace: { requestId: 'req-1' }, + type: MothershipStreamV1EventType.text, + payload: { channel: 'assistant', text: 'the middle of the turn' }, + }, + ]) + + const response = await GET( + new NextRequest( + 'http://localhost:3000/api/copilot/chat/stream?streamId=stream-1&after=0&batch=true' + ) + ) + + await expect(response.json()).resolves.toMatchObject({ success: true, events: [] }) + }) + + it('ends a live tail without a terminal when the ring trims past its cursor mid-tail', async () => { + getLatestRunForStream.mockResolvedValue({ + status: 'active', + executionId: 'exec-1', + id: 'run-1', + }) + const event = (seq: number, text: string) => ({ + stream: { streamId: 'stream-1', cursor: String(seq) }, + seq, + trace: { requestId: 'req-1' }, + type: MothershipStreamV1EventType.text, + payload: { channel: 'assistant', text }, + }) + readEvents + .mockResolvedValueOnce([event(1, 'the start of the turn')]) + .mockResolvedValue([event(5, 'past a trimmed gap')]) + + const response = await GET( + new NextRequest('http://localhost:3000/api/copilot/chat/stream?streamId=stream-1&after=0') + ) + const text = (await readAllChunks(response)).join('') + + expect(text).toContain('the start of the turn') + expect(text).not.toContain('past a trimmed gap') + expect(text).not.toContain(`"type":"${MothershipStreamV1EventType.complete}"`) + }) }) diff --git a/apps/sim/app/api/copilot/chat/stream/route.ts b/apps/sim/app/api/copilot/chat/stream/route.ts index b1023a791b9..3bdaeecf49d 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.ts @@ -13,6 +13,7 @@ import { } from '@/lib/api/server/routes' import { encodeSSEComment } from '@/lib/core/utils/sse' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' +import { MOTHERSHIP_STREAM_REPLAY_HEADER } from '@/lib/mothership/constants' import { MothershipStreamV1CompletionStatus, MothershipStreamV1EventType, @@ -27,12 +28,18 @@ import { readChatStream } from '@/lib/mothership/request/application/recover-str import { contextFromRequestHeaders } from '@/lib/mothership/request/go/propagation' import { getCopilotTracer, markSpanForError } from '@/lib/mothership/request/otel' import { - checkForReplayGap, createEvent, encodeSSEEnvelope, + findReplayGap, + forwardRunReplay, isTerminalStreamStatus, + openRunReplay, + RunReplayUnavailableError, readEvents, readFilePreviewSessions, + readRingPosition, + replayGapTerminal, + ringCanServe, SSE_RESPONSE_HEADERS, } from '@/lib/mothership/request/session' import { toReplayEnvelope, toStreamBatchEvent } from '@/lib/mothership/request/session/types' @@ -43,6 +50,8 @@ const logger = createLogger('CopilotChatStreamAPI') const POLL_INTERVAL_MS = 250 const POLL_INTERVAL_MAX_MS = 2_000 const REPLAY_KEEPALIVE_INTERVAL_MS = 15_000 +/** How often a tail that is still flushing events checks that its ring can serve it. */ +const RING_CHECK_EVERY_BUSY_POLLS = 8 /** * One replay response stays open at most this long, inside the route's `maxDuration`. * A run still going at the cap is not over: the response ends without a terminal @@ -50,6 +59,15 @@ const REPLAY_KEEPALIVE_INTERVAL_MS = 15_000 */ const MAX_STREAM_MS = 60 * 60 * 1000 - 60_000 +/** + * Whether ring events read after `cursor` start right after it. The ring can trim its + * head between a gap check and the read, and a read that starts later would silently + * skip part of the turn. + */ +function startsAfterCursor(events: readonly { seq: number }[], cursor: string): boolean { + return events.length === 0 || events[0].seq <= Number(cursor || '0') + 1 +} + function extractCanonicalRequestId(value: unknown): string { return typeof value === 'string' && value.length > 0 ? value : '' } @@ -129,7 +147,7 @@ export const GET = withRouteHandler(async (request: NextRequest) => { const parsed = await parseRequest(copilotChatStreamContract, request, {}) if (!parsed.success) return parsed.response - const { streamId, after: afterCursor, batch: batchMode } = parsed.data.query + const { streamId, after: afterCursor, batch: batchMode, source } = parsed.data.query if (!streamId) { return NextResponse.json({ error: 'streamId is required' }, { status: 400 }) @@ -173,6 +191,7 @@ export const GET = withRouteHandler(async (request: NextRequest) => { streamId, afterCursor, batchMode, + fromLog: source === 'log', principal, rootSpan, rootContext, @@ -195,6 +214,7 @@ async function handleResumeRequestBody({ streamId, afterCursor, batchMode, + fromLog, principal, rootSpan, rootContext, @@ -203,6 +223,8 @@ async function handleResumeRequestBody({ streamId: string afterCursor: string batchMode: boolean + /** The reader's cursor came from a log re-sync, so the ring never serves it. */ + fromLog: boolean principal: SessionPrincipal rootSpan: Span rootContext: Context @@ -229,7 +251,8 @@ async function handleResumeRequestBody({ if (batchMode) { const afterSeq = afterCursor || '0' - const [events, previewSessions] = await Promise.all([ + const [gap, events, previewSessions] = await Promise.all([ + fromLog ? null : findReplayGap(streamId, afterSeq, extractRunRequestId(run)), readEvents(streamId, afterSeq), readFilePreviewSessions(streamId).catch((error) => { logger.warn('Failed to read preview sessions for stream batch', { @@ -239,7 +262,10 @@ async function handleResumeRequestBody({ return [] }), ]) - const batchEvents = events.map(toStreamBatchEvent) + // A reader the ring cannot serve, or whose next event it trimmed after the gap check, + // is re-synced from the worker log by the live tail. + const batchEvents = + fromLog || gap || !startsAfterCursor(events, afterSeq) ? [] : events.map(toStreamBatchEvent) logger.info('[Resume] Batch response', { streamId, afterCursor: afterSeq, @@ -266,10 +292,47 @@ async function handleResumeRequestBody({ let totalEventsFlushed = 0 let pollIterations = 0 + /** + * A reader the ring cannot serve is re-synced from the worker's durable log for the + * rest of this response, never handed back to the ring: the log and the ring have + * no shared position to join on. The header tells the client to rebuild the turn + * from an empty response, since the replay's cursors restart at 1. + */ + const ringGap = fromLog + ? null + : await findReplayGap(streamId, afterCursor || '0', extractRunRequestId(run)) + // A finished run whose buffer expired answers its terminal; its transcript is persisted. + const gap = + ringGap && !(ringGap.latestSeq <= 0 && isTerminalStreamStatus(run.status)) ? ringGap : null + const resyncFromLog = fromLog || gap !== null + let replayBody: ReadableStream | null = null + /** Releases the worker's replay once this response ends; the request signal may never fire. */ + const replayAbort = new AbortController() + const replaySignal = AbortSignal.any([request.signal, replayAbort.signal]) + if (resyncFromLog && run.chatId) { + try { + replayBody = await openRunReplay({ + streamId, + chatId: run.chatId, + userId: principal.userId, + signal: replaySignal, + }) + } catch (error) { + if (!(error instanceof RunReplayUnavailableError)) throw error + logger.warn('Run replay unavailable; the client will retry', { + streamId, + error: getErrorMessage(error), + }) + markSpanForError(rootSpan, error) + rootSpan.end() + return NextResponse.json({ error: 'Stream replay is unavailable' }, { status: 503 }) + } + } + const stream = new ReadableStream({ async start(controller) { // Re-enter the root OTel context so any `withCopilotSpan` call below - // (inside flushEvents/checkForReplayGap/etc.) parents under + // (inside flushEvents/replayGapTerminal/etc.) parents under // copilot.resume.request instead of becoming an orphan. return otelContext.with(rootContext, () => startInner(controller)) }, @@ -331,8 +394,13 @@ async function handleResumeRequestBody({ } request.signal.addEventListener('abort', abortListener, { once: true }) - const flushEvents = async (): Promise => { + /** Delivers the ring's events after the cursor, or returns null if it trimmed the next one. */ + const flushEvents = async (): Promise => { const events = await readEvents(streamId, cursor) + if (!startsAfterCursor(events, cursor)) { + logger.warn('Replay ring trimmed past a reader cursor', { streamId, cursor }) + return null + } if (events.length > 0) { logger.debug('[Resume] Flushing events', { streamId, @@ -380,12 +448,46 @@ async function handleResumeRequestBody({ } } + /** Forwards the worker's replay, keeping the response alive while it waits. */ + const streamRunReplay = async (body: ReadableStream) => { + const keepalive = setInterval(() => { + if (Date.now() - lastWriteTime < REPLAY_KEEPALIVE_INTERVAL_MS) return + if (!enqueueComment('keepalive')) replayAbort.abort() + }, REPLAY_KEEPALIVE_INTERVAL_MS) + try { + const end = await forwardRunReplay({ + body, + streamId, + signal: replaySignal, + write: (envelope) => { + if (!enqueueEvent(envelope)) return false + totalEventsFlushed += 1 + cursor = envelope.stream.cursor ?? cursor + if (envelope.type === MothershipStreamV1EventType.complete) sawTerminalEvent = true + return true + }, + readRunStatus: async () => (await readRun().catch(() => null))?.status ?? null, + isClosed: () => controllerClosed, + deadlineAt: startTime + MAX_STREAM_MS, + }) + logger.info('[Resume] Run replay ended', { streamId, end, eventCount: totalEventsFlushed }) + } finally { + clearInterval(keepalive) + replayAbort.abort() + } + } + try { enqueueComment('accepted') - const gap = await checkForReplayGap(streamId, afterCursor, currentRequestId) - if (gap) { - for (const envelope of gap.envelopes) { + if (replayBody) { + await streamRunReplay(replayBody) + return + } + if (resyncFromLog) { + const position = gap ?? (await readRingPosition(streamId, cursor)) + const terminal = await replayGapTerminal(streamId, position, currentRequestId) + for (const envelope of terminal.envelopes) { if (!enqueueEvent(envelope)) { break } @@ -398,7 +500,8 @@ async function handleResumeRequestBody({ return } - await flushEvents() + let lastFlushed = await flushEvents() + if (lastFlushed === null) return let pollDelayMs = POLL_INTERVAL_MS while (!controllerClosed && Date.now() - startTime < MAX_STREAM_MS) { @@ -418,10 +521,24 @@ async function handleResumeRequestBody({ }) break } + // The ring lost its head, restarted or expired under this live tail; the re-attach + // re-syncs, and a finished run answers its terminal instead. Only a quiet ring can + // restart or be re-sent into by a recovery, so a busy tail checks every few polls. + const checkRing = lastFlushed === 0 || pollIterations % RING_CHECK_EVERY_BUSY_POLLS === 0 + if ( + checkRing && + !isTerminalStreamStatus(currentRun.status) && + !ringCanServe(await readRingPosition(streamId, cursor)) + ) { + logger.warn('Replay ring can no longer serve a live tail', { streamId, cursor }) + break + } currentRequestId = extractRunRequestId(currentRun) || currentRequestId const flushed = await flushEvents() + if (flushed === null) break + lastFlushed = flushed /* Adaptive tail: 4 Hz only while events are actually flowing; a quiet stream decays toward the cap so an attached client doesn't hammer Postgres + Redis at 4 Hz for up to an hour. Any flushed event snaps back to full rate. */ @@ -488,5 +605,9 @@ async function handleResumeRequestBody({ } } - return new Response(stream, { headers: SSE_RESPONSE_HEADERS }) + return new Response(stream, { + headers: replayBody + ? { ...SSE_RESPONSE_HEADERS, [MOTHERSHIP_STREAM_REPLAY_HEADER]: 'log' } + : SSE_RESPONSE_HEADERS, + }) } diff --git a/apps/sim/app/api/mothership/chats/[chatId]/route.ts b/apps/sim/app/api/mothership/chats/[chatId]/route.ts index f39b55f32af..bd975a2dd99 100644 --- a/apps/sim/app/api/mothership/chats/[chatId]/route.ts +++ b/apps/sim/app/api/mothership/chats/[chatId]/route.ts @@ -11,12 +11,15 @@ import { } from '@/lib/api/contracts/mothership-chats' import { parseRequest } from '@/lib/api/server' import { withRouteHandler } from '@/lib/core/utils/with-route-handler' -import { getLatestRunForStream } from '@/lib/mothership/async-runs/repository' import { buildEffectiveChatTranscript } from '@/lib/mothership/chat/effective-transcript' import { getAccessibleCopilotChatAuth, getAccessibleCopilotChatWithMessages, } from '@/lib/mothership/chat/lifecycle' +import { + type LiveTurnSnapshot, + readLiveTurnSnapshot, +} from '@/lib/mothership/chat/live-turn-snapshot' import { normalizeMessage } from '@/lib/mothership/chat/persisted-message' import { reconcileChatStreamMarkers } from '@/lib/mothership/chat/stream-liveness' import { publishChatStatusChanged } from '@/lib/mothership/chat-status' @@ -25,10 +28,6 @@ import { createInternalServerErrorResponse, createUnauthorizedResponse, } from '@/lib/mothership/request/http' -import type { FilePreviewSession } from '@/lib/mothership/request/session' -import { readEvents } from '@/lib/mothership/request/session/buffer' -import { readFilePreviewSessions } from '@/lib/mothership/request/session/file-preview-session' -import { type StreamBatchEvent, toStreamBatchEvent } from '@/lib/mothership/request/session/types' import { captureServerEvent } from '@/lib/posthog/server' const logger = createLogger('MothershipChatAPI') @@ -55,11 +54,7 @@ export const GET = withRouteHandler( // to the client: when `activeStreamId` is set, the client reconnects to // the replay buffer (from seq 0) via the stream resume endpoint, which // is the source of truth for streaming state. - let liveTurnSnapshot: { - events: StreamBatchEvent[] - previewSessions: FilePreviewSession[] - status: string - } | null = null + let liveTurnSnapshot: LiveTurnSnapshot | null = null const reconciledMarkers = await reconcileChatStreamMarkers( [{ chatId: chat.id, streamId: chat.conversationId }], @@ -69,36 +64,7 @@ export const GET = withRouteHandler( if (liveStreamId) { try { - const [events, previewSessions] = await Promise.all([ - readEvents(liveStreamId, '0'), - readFilePreviewSessions(liveStreamId).catch((error) => { - logger.warn('Failed to read preview sessions for mothership chat', { - chatId, - streamId: liveStreamId, - error: toError(error).message, - }) - return [] - }), - ]) - const run = await getLatestRunForStream(liveStreamId, userId).catch((error) => { - logger.warn('Failed to fetch latest run for mothership chat snapshot', { - chatId, - streamId: liveStreamId, - error: toError(error).message, - }) - return null - }) - - liveTurnSnapshot = { - events: events.map(toStreamBatchEvent), - previewSessions, - status: - typeof run?.status === 'string' - ? run.status - : events.length > 0 - ? 'active' - : 'unknown', - } + liveTurnSnapshot = await readLiveTurnSnapshot(liveStreamId, userId) } catch (error) { logger.warn('Failed to read stream snapshot for mothership chat', { chatId, diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx index cebc1911396..acd0efa6d15 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx @@ -73,6 +73,7 @@ import { seedDeploymentShape, } from '@/lib/core/config/deployment-shape' import { MothershipHandoffStorage } from '@/lib/core/utils/browser-storage' +import { MOTHERSHIP_STREAM_REPLAY_HEADER } from '@/lib/mothership/constants' import type { MothershipStreamV1EventEnvelope } from '@/lib/mothership/generated/mothership-stream-v1' import { getChatResourceSelectionId } from '@/lib/mothership/resources/types' import { collectCitedMessageSources } from '@/app/workspace/[workspaceId]/home/components/message-content/message-sources' @@ -1157,6 +1158,77 @@ describe('useChat remount send recovery', () => { } }) + it('rebuilds the turn from an empty response when a reconnect is re-synced from the log, and stays on the log', async () => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + try { + let tails = 0 + const streamUrls: string[] = [] + const history: MothershipChatHistory = { + id: 'chat-log-resync', + mode: 'agent', + title: 'Log re-sync', + messages: [], + activeStreamId: null, + resources: [], + } + mockRequestJson.mockImplementation(() => + Promise.resolve({ + chat: { ...history, activeStreamId: state.postBodies[0]?.userMessageId ?? null }, + }) + ) + state.postBehavior = 'accept' + const frame = (streamId: string, seq: number, text: string) => + `data: ${JSON.stringify({ + v: 1, + seq, + ts: new Date().toISOString(), + type: 'text', + stream: { streamId, cursor: String(seq) }, + payload: { channel: 'assistant', text }, + } satisfies MothershipStreamV1EventEnvelope)}\n\n` + vi.stubGlobal('fetch', async (input: RequestInfo | URL, init?: RequestInit) => { + const url = String(input) + if (!url.includes('/api/mothership/chat/stream')) return fetchStub(input, init) + streamUrls.push(url) + if (url.includes('batch=true')) { + return Response.json({ success: true, events: [], status: 'streaming' }) + } + tails++ + const streamId = state.postBodies[0]?.userMessageId ?? '' + if (tails === 1) { + return new Response([1, 2, 3].map((seq) => frame(streamId, seq, 'stale ')).join(''), { + headers: { 'Content-Type': 'text/event-stream' }, + }) + } + return new Response(frame(streamId, 1, 'Full response.'), { + headers: { + 'Content-Type': 'text/event-stream', + [MOTHERSHIP_STREAM_REPLAY_HEADER]: 'log', + }, + }) + }) + const { getResult } = renderUseChatInChat(history.id, history) + await act(async () => { + void getResult().sendMessage('Pick up where it left off') + }) + for (let second = 0; second < 10 && tails < 3; second++) { + await act(async () => vi.advanceTimersByTimeAsync(1_000)) + } + + const answer = getResult().messages.find((message) => message.role === 'assistant') + expect(tails).toBeGreaterThanOrEqual(3) + expect(answer?.content).toBe('Full response.') + const logResyncTail = streamUrls.findIndex( + (url) => url.includes('after=3') && !url.includes('batch=true') + ) + const afterLogResync = streamUrls.slice(logResyncTail + 1) + expect(afterLogResync.length).toBeGreaterThan(0) + expect(afterLogResync.every((url) => url.includes('source=log'))).toBe(true) + } finally { + vi.useRealTimers() + } + }) + it('sends a queued correction after stopping with more than 10 MiB of tool input', async () => { state.postBehavior = 'tool' state.toolInputPadding = 'x'.repeat(11 * 1024 * 1024) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index 69707e2bf95..0f506f151b8 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -54,7 +54,11 @@ import { type RevealedSimKeysByMessage, restoreRevealedSimKeysForMessage, } from '@/lib/mothership/chat/sim-key-redaction' -import { MOTHERSHIP_CHAT_API_PATH, MOTHERSHIP_CHAT_ID_HEADER } from '@/lib/mothership/constants' +import { + MOTHERSHIP_CHAT_API_PATH, + MOTHERSHIP_CHAT_ID_HEADER, + MOTHERSHIP_STREAM_REPLAY_HEADER, +} from '@/lib/mothership/constants' import { sendMothershipMessage } from '@/lib/mothership/events' import type { AssistantSearchLevel } from '@/lib/mothership/generated/assistant' import { resolveMothershipModelSettings } from '@/lib/mothership/model-options' @@ -302,6 +306,15 @@ const EMPTY_MESSAGE_QUEUE: QueuedMothershipMessage[] = [] const logger = createLogger('useChat') +/** + * The reconnect query for a stream. Once a stream was re-synced from the worker's log, + * its cursors are log positions, so every later read names the log as its source and + * is never served from the replay ring, even one that restarted and grew past them. + */ +function streamReconnectQuery(streamId: string, afterCursor: string, fromLog: boolean): string { + return `streamId=${encodeURIComponent(streamId)}&after=${encodeURIComponent(afterCursor)}${fromLog ? '&source=log' : ''}` +} + /** * Fire-and-forget desktop-surface handoff between chat scopes: drops an * abandoned pending scope (never a durable one) before activating the next. @@ -952,6 +965,7 @@ export function useChat( const streamRequestIdRef = useRef(undefined) const locallyTerminalStreamIdRef = useRef(undefined) const lastCursorRef = useRef('0') + const logResyncedStreamIdRef = useRef(null) const activeStreamReturnRecoveryRef = useRef(null) const sendingRef = useRef(false) const streamGenRef = useRef(0) @@ -2367,7 +2381,7 @@ export function useChat( ) // boundary-raw-fetch: stream-resume batch endpoint requires dynamic per-request traceparent header propagation that the contract layer does not model, and the response is consumed alongside live SSE tail fetches const response = await fetch( - `/api/mothership/chat/stream?streamId=${encodeURIComponent(streamId)}&after=${encodeURIComponent(afterCursor)}&batch=true`, + `/api/mothership/chat/stream?${streamReconnectQuery(streamId, afterCursor, logResyncedStreamIdRef.current === streamId)}&batch=true`, { signal: fetchSignal, ...(streamTraceparentRef.current @@ -2559,7 +2573,7 @@ export function useChat( // boundary-raw-fetch: live SSE tail endpoint streams events consumed via response.body.getReader() and processSSEStream const sseRes = await fetch( - `/api/mothership/chat/stream?streamId=${encodeURIComponent(streamId)}&after=${encodeURIComponent(latestCursor)}`, + `/api/mothership/chat/stream?${streamReconnectQuery(streamId, latestCursor, logResyncedStreamIdRef.current === streamId)}`, { signal: activeAbort.signal, ...(streamTraceparentRef.current @@ -2578,6 +2592,14 @@ export function useChat( return { error: false, aborted: true } } + // Re-sent from the worker's log with cursors restarting at 1: rebuild from empty. + if (sseRes.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER) === 'log') { + logResyncedStreamIdRef.current = streamId + const reset = applyReconnectReplaySelection(streamId, '0') + latestCursor = reset.afterCursor + preserveNextReplayState = reset.preserveExistingState + } + setTransportStreaming() const liveResult = await processSSEStreamRef.current( diff --git a/apps/sim/lib/api/contracts/copilot.ts b/apps/sim/lib/api/contracts/copilot.ts index d9fcd2ff2a5..7ab95c22b6d 100644 --- a/apps/sim/lib/api/contracts/copilot.ts +++ b/apps/sim/lib/api/contracts/copilot.ts @@ -200,6 +200,8 @@ export const copilotChatStreamQuerySchema = z.object({ .string() .optional() .transform((value) => value === 'true'), + /** `log` once the reader was re-synced from the worker log: its cursor is a log position. */ + source: z.enum(['ring', 'log']).optional(), }) export const copilotChatStopBodySchema = z.object({ diff --git a/apps/sim/lib/api/contracts/mothership-chats.ts b/apps/sim/lib/api/contracts/mothership-chats.ts index 16b0b2a6a7b..0370e50d74a 100644 --- a/apps/sim/lib/api/contracts/mothership-chats.ts +++ b/apps/sim/lib/api/contracts/mothership-chats.ts @@ -203,6 +203,7 @@ export const mothershipChatStreamQuerySchema = z streamId: z.string().optional(), after: z.string().optional(), batch: z.string().optional(), + source: z.enum(['ring', 'log']).optional(), }) .passthrough() diff --git a/apps/sim/lib/core/redis/byte-budget.server.ts b/apps/sim/lib/core/redis/byte-budget.server.ts index 5728d69960a..e6124bbd00d 100644 --- a/apps/sim/lib/core/redis/byte-budget.server.ts +++ b/apps/sim/lib/core/redis/byte-budget.server.ts @@ -8,9 +8,11 @@ import type { Logger } from '@sim/logger' * budget the execution event buffer has enforced since it was written, which the * copilot stream buffer now shares rather than inventing a bound of its own. * - * A quota is the right bound for a buffer whose contents must stay contiguous: the - * copilot replay chain and an execution's event history are read from a cursor, so - * the write that would breach the ceiling is refused and the buffer stops growing. + * A quota is the right bound for a buffer whose contents must stay contiguous: an + * execution's event history is read from a cursor, so the write that would breach + * the ceiling is refused and the buffer stops growing. The copilot replay ring trims + * its oldest events by bytes below its ceiling instead, refunding what it drops, so a + * long run slides rather than refuses; a reader behind the trim gets a replay gap. * A live-update feed is bounded differently — see `lib/realtime/event-log.ts`, whose * readers already handle a prune by refetching, so it drops oldest-first instead. * @@ -67,6 +69,13 @@ export interface RedisBudgetLimits { * already dropped and eventually pin the user at their ceiling until they went a full * window without writing. User counters therefore get a fixed window: set on * creation, never extended. + * + * Because a trim refunds both counters, the user counter bounds bytes HELD across a + * user's owners, not bytes written per hour: a single long copilot stream holds at most + * its ring's byte target however much it writes. Bytes of owners that ended stay counted + * until the window lapses. The reset is not reconciled with what is still held, so + * right after it a user can hold up to about twice the cap: the bytes the lapsed + * window counted plus a fresh cap. */ const REDIS_BUDGET_TTL_SECONDS = 60 * 60 @@ -79,8 +88,8 @@ const LIMITS: Record }, /** * A copilot turn streams text and tool frames, not payloads — a single frame past - * 1 MB is already pathological. The owner ceiling is what a long agentic session - * may retain for replay across its whole hour. + * 1 MB is already pathological. The owner ceiling bounds what one stream retains + * for replay; the ring trims its oldest events to stay below it. */ copilot_stream: { maxSingleWriteBytes: 1 * 1024 * 1024, diff --git a/apps/sim/lib/mothership/chat/live-turn-snapshot.ts b/apps/sim/lib/mothership/chat/live-turn-snapshot.ts new file mode 100644 index 00000000000..eead7d9bef8 --- /dev/null +++ b/apps/sim/lib/mothership/chat/live-turn-snapshot.ts @@ -0,0 +1,51 @@ +import { createLogger } from '@sim/logger' +import { toError } from '@sim/utils/errors' +import { getLatestRunForStream } from '@/lib/mothership/async-runs/repository' +import type { FilePreviewSession } from '@/lib/mothership/request/session' +import { readEvents } from '@/lib/mothership/request/session/buffer' +import { readFilePreviewSessions } from '@/lib/mothership/request/session/file-preview-session' +import { startsAtReplayHead } from '@/lib/mothership/request/session/recovery' +import { type StreamBatchEvent, toStreamBatchEvent } from '@/lib/mothership/request/session/types' + +const logger = createLogger('LiveTurnSnapshot') + +/** An in-flight turn's replay, for a chat load's first paint. */ +export interface LiveTurnSnapshot { + events: StreamBatchEvent[] + previewSessions: FilePreviewSession[] + status: string +} + +/** + * The in-flight turn of `streamId` as its replay ring holds it, or `null` once the ring + * lost its head: a truncated turn is never painted, and the client re-syncs it through + * the reconnect route instead. + */ +export async function readLiveTurnSnapshot( + streamId: string, + userId: string +): Promise { + const [events, previewSessions, run] = await Promise.all([ + readEvents(streamId, '0'), + readFilePreviewSessions(streamId).catch((error) => { + logger.warn('Failed to read preview sessions for a live turn', { + streamId, + error: toError(error).message, + }) + return [] + }), + getLatestRunForStream(streamId, userId).catch((error) => { + logger.warn('Failed to read the latest run for a live turn', { + streamId, + error: toError(error).message, + }) + return null + }), + ]) + if (!startsAtReplayHead(events[0]?.seq)) return null + return { + events: events.map(toStreamBatchEvent), + previewSessions, + status: typeof run?.status === 'string' ? run.status : events.length > 0 ? 'active' : 'unknown', + } +} diff --git a/apps/sim/lib/mothership/constants.ts b/apps/sim/lib/mothership/constants.ts index 2fda1e8ef9a..7f81ab2c245 100644 --- a/apps/sim/lib/mothership/constants.ts +++ b/apps/sim/lib/mothership/constants.ts @@ -67,6 +67,13 @@ export const STREAM_STORAGE_KEY = 'copilot_active_stream' /** POST — send a chat message through the unified mothership chat surface. */ export const MOTHERSHIP_CHAT_API_PATH = '/api/mothership/chat' +/** + * Set to `log` on a reconnect response the replay ring could not serve: the turn is + * re-sent from the worker's durable log with cursors restarting at 1, so the client + * rebuilds it from an empty response. + */ +export const MOTHERSHIP_STREAM_REPLAY_HEADER = 'x-mothership-stream-replay' + /** Durable chat identity returned after the send transaction commits, before SSE delivery. */ export const MOTHERSHIP_CHAT_ID_HEADER = 'x-mothership-chat-id' diff --git a/apps/sim/lib/mothership/generated/protocol.ts b/apps/sim/lib/mothership/generated/protocol.ts index c140abe0c68..0d6668c49ee 100644 --- a/apps/sim/lib/mothership/generated/protocol.ts +++ b/apps/sim/lib/mothership/generated/protocol.ts @@ -336,6 +336,39 @@ export interface ResumeResult { export const AbortRequest = z.strictObject({ messageId: z.uuid() }); export type AbortRequest = z.infer; +/** + * POST /api/streams/replay (sim's inbound key only): the run's response rebuilt from the + * durable log for a reader whose cursor fell off sim's replay ring. Read-only: the owner, + * its emitter, its lease and its parked calls are untouched. The SSE leg restores what + * the receipt lacks, follows the log, and ends with `complete` at a terminal or with one + * `run` frame of kind `replay_end` otherwise. + */ +export const StreamReplayRequest = z.strictObject({ + streamId: z.uuid(), + chatId: z.uuid(), + /** The chat's user; a mismatch answers 404 like an unknown run. */ + userId: z.string().min(1), + ...ResponseReceiptSchema.shape, +}); +export type StreamReplayRequest = z.infer; + +/** + * A replay leg that ended before the run's terminal. `parked`: the run waits on tool + * results the owner's leg handed sim. `cap`: the connection reached its length cap; + * replay again. `stalled`: no instance drives the run (stale heartbeat, not parked); + * replay again later, once a takeover can have resumed it. + * `textLength` is the main text this leg's log reached — diagnostic only. A later + * replay sends sim's OWN received length as `receivedTextChars`, never this value: the + * log trails live text, and resending from the durable end is what exposes divergence. + * Replayed tool activity is presentation only (`replay: true`) and never authorizes + * execution. + */ +export interface StreamReplayEnd extends StreamTextCompletion { + kind: "replay_end"; + reason: "parked" | "cap" | "stalled"; + textLength: number; +} + /** Accepted Stop intent is distinct from an observed terminal worker run. */ export interface AbortResponse { stopped: boolean; diff --git a/apps/sim/lib/mothership/request/application/recover-stream.ts b/apps/sim/lib/mothership/request/application/recover-stream.ts index 441750260e8..61eff64d315 100644 --- a/apps/sim/lib/mothership/request/application/recover-stream.ts +++ b/apps/sim/lib/mothership/request/application/recover-stream.ts @@ -24,9 +24,10 @@ import { getLocalChatStreamLease, releasePendingChatStream, } from '@/lib/mothership/request/session/abort' -import { readEvents } from '@/lib/mothership/request/session/buffer' +import { getLatestSeq, readEvents } from '@/lib/mothership/request/session/buffer' import { assertChatStreamLease } from '@/lib/mothership/request/session/controller-lease' import { eventToStreamEvent } from '@/lib/mothership/request/session/event' +import { startsAtReplayHead } from '@/lib/mothership/request/session/recovery' import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils' const logger = createLogger('MothershipStreamRecovery') @@ -120,6 +121,17 @@ export const readChatStream = defineAuthorizedChatUseCase({ ]) if (workspaceId && !userPermission) throw new OrchestrationError('forbidden', 'Workspace access revoked') + /** + * A ring that lost its head is treated like an expired one: the controller starts + * from an empty context and re-attaches with an empty receipt, so the worker re-sends + * the whole response and re-hands its parked calls. Rebuilding from the tail would + * persist a truncated turn. + */ + const ringIntact = startsAtReplayHead(events[0]?.seq) + const recoveredEvents = ringIntact ? events : [] + const resumeSeq = ringIntact + ? (events.at(-1)?.seq ?? 0) + : ((await getLatestSeq(run.streamId)) ?? 0) const requestId = typeof saved?.requestId === 'string' ? saved.requestId : generateId() const completion = { chatId, @@ -150,7 +162,7 @@ export const readChatStream = defineAuthorizedChatUseCase({ currentChat: null, message: '', titleModel: '', - resumeSeq: events.at(-1)?.seq ?? 0, + resumeSeq, orchestrateOptions: { userId, workspaceId, @@ -168,7 +180,7 @@ export const readChatStream = defineAuthorizedChatUseCase({ recovery: { ...config.data, streamId: run.streamId, - events: events.map(eventToStreamEvent), + events: recoveredEvents.map(eventToStreamEvent), }, onComplete: buildOnComplete(completion), onError: buildOnError(completion), diff --git a/apps/sim/lib/mothership/request/go/parser.ts b/apps/sim/lib/mothership/request/go/parser.ts index d37a71685e0..a5ef054e2cf 100644 --- a/apps/sim/lib/mothership/request/go/parser.ts +++ b/apps/sim/lib/mothership/request/go/parser.ts @@ -25,15 +25,19 @@ function createParseFailure(message: string, preview: string): FatalSseEventErro * all come from the shared engine. * * @param onEvent Called per parsed event. Return true to stop processing. + * @param idleTimeoutMs Fails the read once the stream sends nothing, comments included, + * for this long. */ export async function processSSEStream( reader: ReadableStreamDefaultReader, abortSignal: AbortSignal | undefined, - onEvent: (event: unknown) => boolean | undefined | Promise + onEvent: (event: unknown) => boolean | undefined | Promise, + idleTimeoutMs?: number ): Promise { try { await readSSELines(reader, { signal: abortSignal, + idleTimeoutMs, onData: async (jsonStr) => { let parsed: unknown try { diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index 558566dbe31..a1823759788 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -216,4 +216,18 @@ describe('stream recovery budget', () => { vi.advanceTimersByTime(30 * 60_000) expect(retry.nextDelay(error)).toBeNull() }) + + it('retries a leg that fails again minutes after it re-attached and delivered events', () => { + vi.useFakeTimers() + const error = new WorkerStreamInterruptedError(new Error('socket closed')) + const retry = new StreamRetryWindow() + const first = retry.nextDelay(error) + expect(first).not.toBeNull() + vi.advanceTimersByTime(first ?? 0) + for (let second = 0; second < 120; second += 10) { + retry.recovered() + vi.advanceTimersByTime(10_000) + } + expect(retry.nextDelay(error)).not.toBeNull() + }) }) diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts index 2650ac6f195..0bc4cfa1248 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -32,7 +32,8 @@ const HEALTHY_STREAM_REPLENISH_MS = 5 * 60_000 * Recovery is bounded independently of the healthy leg's lifetime, by two * budgets that never share state: an unreachable worker gets a two-minute * window from the moment it stopped answering, and any failure of a worker that - * did answer gets three retries within 30 s, replenished only after + * did answer gets three retries, each burst of them within 30 s of its first + * failure, replenished only after * {@link HEALTHY_STREAM_REPLENISH_MS} of healthy streaming. A leg has no deadline * unless the caller sets one. */ @@ -63,9 +64,13 @@ export class StreamRetryWindow { return remaining } - /** The worker delivered an event, so a later loss of it starts a fresh unreachable window. */ + /** + * The worker delivered an event: a later loss starts a fresh unreachable window, and + * a fresh 30 s reachable window. Only the three reachable retries carry over. + */ recovered(): void { this.resetUnreachable() + this.firstFailureAt = undefined this.lastEventAt = Date.now() this.streamingSince ??= this.lastEventAt } diff --git a/apps/sim/lib/mothership/request/session/abort.ts b/apps/sim/lib/mothership/request/session/abort.ts index 7ff29fa5000..b5c36fe3b22 100644 --- a/apps/sim/lib/mothership/request/session/abort.ts +++ b/apps/sim/lib/mothership/request/session/abort.ts @@ -397,7 +397,14 @@ export function startAbortPoller( }) return } - await refreshBufferTtl(streamId) + await refreshBufferTtl(streamId).catch((error) => { + logger.warn('Failed to refresh stream buffer TTL', { + chatId, + streamId, + ...(requestId ? { requestId } : {}), + error: toError(error).message, + }) + }) } catch (error) { logger.warn('Failed to extend chat stream lock TTL', { chatId, diff --git a/apps/sim/lib/mothership/request/session/buffer-ttl.integration.ts b/apps/sim/lib/mothership/request/session/buffer-ttl.integration.ts index 9fc83b57c48..59d9b49d6a5 100644 --- a/apps/sim/lib/mothership/request/session/buffer-ttl.integration.ts +++ b/apps/sim/lib/mothership/request/session/buffer-ttl.integration.ts @@ -33,7 +33,7 @@ import { scheduleBufferCleanup, } from '@/lib/mothership/request/session/buffer' import { createEvent } from '@/lib/mothership/request/session/event' -import { checkForReplayGap } from '@/lib/mothership/request/session/recovery' +import { findReplayGap } from '@/lib/mothership/request/session/recovery' async function appendText(streamId: string, text: string): Promise { const { seq, cursor } = await allocateCursor(streamId) @@ -104,6 +104,6 @@ describe.runIf(Boolean(redisUrl))('replay buffer lifetime', () => { await redis.del(`mothership_stream:${streamId}:events`, `mothership_stream:${streamId}:seq`) await appendText(streamId, 'after expiry') - expect(await checkForReplayGap(streamId, '5')).not.toBeNull() + expect(await findReplayGap(streamId, '5')).not.toBeNull() }) }) diff --git a/apps/sim/lib/mothership/request/session/buffer.test.ts b/apps/sim/lib/mothership/request/session/buffer.test.ts index b3e137dab51..1c614b27879 100644 --- a/apps/sim/lib/mothership/request/session/buffer.test.ts +++ b/apps/sim/lib/mothership/request/session/buffer.test.ts @@ -72,8 +72,8 @@ const createRedisStub = () => { const numKeys = Number(args[1]) const keys = args.slice(2, 2 + numKeys) as string[] const argv = args.slice(2 + numKeys) as Array - const leased = String(args[0]).includes("if redis.call('GET', KEYS[3]) ~= ARGV[7]") - if (leased && values.get(keys[2]) !== argv[6]) return Promise.resolve([-1]) + const leased = String(args[0]).includes("if redis.call('GET', KEYS[3]) ~= ARGV[8]") + if (leased && values.get(keys[2]) !== argv[7]) return Promise.resolve([-1]) if (api.budgetRefusal) return Promise.resolve(api.budgetRefusal) @@ -81,7 +81,7 @@ const createRedisStub = () => { const eventLimit = Number(argv[1]) const lastSeq = String(argv[5]) const entries = sortedSets.get(eventsKey) ?? [] - for (let i = leased ? 7 : 6; i < argv.length; i += 2) { + for (let i = leased ? 8 : 7; i < argv.length; i += 2) { const score = Number(argv[i]) const value = String(argv[i + 1]) if (!entries.some((entry) => entry.value === value)) entries.push({ score, value }) @@ -239,7 +239,7 @@ describe('mothership-stream-outbox', () => { expect(eventsKey).toBe('mothership_stream:stream-1:events') expect(seqKey).toBe('mothership_stream:stream-1:seq') expect(ownerKey).toBe('execution:redis-budget:copilot_stream:stream-1') - // ARGV: [ttl, eventLimit, ownerLimit, userLimit, budgetTtl, lastSeq, ...zaddArgs] + // ARGV: [ttl, eventLimit, ownerLimit, userLimit, budgetTtl, lastSeq, retainedBytes, ...zaddArgs] expect(argv[1]).toBe(100_000) }) diff --git a/apps/sim/lib/mothership/request/session/buffer.ts b/apps/sim/lib/mothership/request/session/buffer.ts index 01923ea3619..9e0f1559f80 100644 --- a/apps/sim/lib/mothership/request/session/buffer.ts +++ b/apps/sim/lib/mothership/request/session/buffer.ts @@ -24,6 +24,19 @@ const DEFAULT_TTL_SECONDS = 60 * 60 const DEFAULT_COMPLETED_TTL_SECONDS = 5 * 60 const DEFAULT_EVENT_LIMIT = 100_000 const RETRY_DELAYS_MS = [0, 50, 150] as const +/** + * Share of the owner ceiling the replay ring retains before trimming its oldest events. + * The rest is headroom, so a long run trims instead of being refused. + */ +const RETAINED_BYTES_FRACTION = 0.75 +/** Existing ring members read per page while choosing which to trim. */ +const TRIM_PAGE_SIZE = 256 +/** + * Most members one append trims for bytes beyond what the count limit requires. A ring + * already past its byte target (written before byte trimming existed) catches up over + * several appends instead of in one long script. + */ +const MAX_BYTE_TRIM_MEMBERS = 16 * TRIM_PAGE_SIZE type RedisOperationMetadata = { operation: string @@ -204,25 +217,40 @@ export async function scheduleBufferCleanup( * the stream's budget — in one script, so the reservation and the write it pays for * commit together. * + * The ring is a sliding window bounded by count and by bytes: the lowest-ranked + * members are trimmed until both fit, and exactly the trimmed bytes are refunded. + * The owner counter is the ring's byte total, so a stream of any length stays under + * its retained-bytes target and never reaches the owner ceiling. A byte trim drops an + * incoming member only when a replay reintroduces it below a retained one, which keeps + * the ring contiguous; it never drops the newest, so a counter already past the + * ceiling still refuses rather than silently discarding the write. + * * Entries already present are skipped when counting, which makes the script * idempotent: `withRedisRetry` may run it up to three times, and a retry after a * partial failure must not charge the same bytes twice. * - * KEYS: [events, seq, budgetOwner, budgetUser?] + * KEYS: [events, seq, lease?, budgetOwner, budgetUser?] * ARGV: [ttlSeconds, eventLimit, ownerLimit, userLimit, budgetTtlSeconds, lastSeq, - * score, member, ...] + * retainedBytesLimit, leaseValue?, score, member, ...] * Returns {1} on success, or {0, resource, currentBytes} when the budget refuses. */ function appendEventsScript(leased: boolean): string { - const firstMember = leased ? 8 : 7 + const firstMember = leased ? 9 : 8 + const ownerKey = `KEYS[${leased ? 4 : 3}]` return ` -${leased ? "if redis.call('GET', KEYS[3]) ~= ARGV[7] then return {-1} end" : ''} +${leased ? "if redis.call('GET', KEYS[3]) ~= ARGV[8] then return {-1} end" : ''} local ttl_seconds = tonumber(ARGV[1]) local event_limit = tonumber(ARGV[2]) local owner_limit = tonumber(ARGV[3]) local user_limit = tonumber(ARGV[4]) local budget_ttl_seconds = tonumber(ARGV[5]) local last_seq = ARGV[6] +local retained_bytes_limit = tonumber(ARGV[7]) + +local function ranks_before(a, b) + if a.score == b.score then return a.member < b.member end + return a.score < b.score +end local new_count = 0 local new_bytes = 0 @@ -237,28 +265,47 @@ for i = ${firstMember}, #ARGV, 2 do table.insert(new_members, {member = member, score = tonumber(ARGV[i])}) end end +table.sort(new_members, ranks_before) local current_count = redis.call('ZCARD', KEYS[1]) -local prune_count = math.max(current_count + new_count - event_limit, 0) +local count_excess = math.max(current_count + new_count - event_limit, 0) +local byte_excess = 0 +if retained_bytes_limit > 0 then + local retained_bytes = tonumber(redis.call('GET', ${ownerKey}) or '0') + byte_excess = math.max(retained_bytes + new_bytes - retained_bytes_limit, 0) +end + +-- Walk the union of the ring and this batch in rank order, paging the ring so a +-- trim reads only as many existing members as it removes. +local prune_count = 0 local pruned_bytes = 0 -if prune_count > 0 then - -- A replay can reintroduce an already-trimmed member before the retained ring. - -- Price the actual lowest-ranked union, not all existing members before new ones. - -- Only this many existing members can be pruned, so never scan the whole ring. - local existing_prune_count = math.min(prune_count, current_count) - if existing_prune_count > 0 then - local existing = redis.call('ZRANGE', KEYS[1], 0, existing_prune_count - 1, 'WITHSCORES') - for i = 1, #existing, 2 do - table.insert(new_members, {member = existing[i], score = tonumber(existing[i + 1])}) - end +local next_new = 1 +local page = {} +local page_index = 1 +local fetched = 0 +local max_prune_count = count_excess + ${MAX_BYTE_TRIM_MEMBERS} +while prune_count < count_excess or (pruned_bytes < byte_excess and prune_count < max_prune_count) do + if page_index > #page and fetched < current_count then + page = redis.call('ZRANGE', KEYS[1], fetched, fetched + ${TRIM_PAGE_SIZE} - 1, 'WITHSCORES') + fetched = fetched + #page / 2 + page_index = 1 end - table.sort(new_members, function(a, b) - if a.score == b.score then return a.member < b.member end - return a.score < b.score - end) - for i = 1, prune_count do - pruned_bytes = pruned_bytes + string.len(new_members[i].member) + local existing = nil + if page_index <= #page then + existing = {member = page[page_index], score = tonumber(page[page_index + 1])} end + local incoming = new_members[next_new] + if incoming and (not existing or ranks_before(incoming, existing)) then + if not existing and prune_count >= count_excess then break end + pruned_bytes = pruned_bytes + string.len(incoming.member) + next_new = next_new + 1 + elseif existing then + pruned_bytes = pruned_bytes + string.len(existing.member) + page_index = page_index + 2 + else + break + end + prune_count = prune_count + 1 end local net_bytes = new_bytes - pruned_bytes @@ -267,7 +314,9 @@ ${renderRedisBudgetLua(leased ? 3 : 2)} for i = ${firstMember}, #ARGV, 2 do redis.call('ZADD', KEYS[1], ARGV[i], ARGV[i + 1]) end -redis.call('ZREMRANGEBYRANK', KEYS[1], 0, -event_limit - 1) +if prune_count > 0 then + redis.call('ZREMRANGEBYRANK', KEYS[1], 0, prune_count - 1) +end redis.call('EXPIRE', KEYS[1], ttl_seconds) redis.call('SET', KEYS[2], last_seq, 'EX', ttl_seconds) return {1} @@ -317,6 +366,7 @@ export async function appendEvents( break that invariant silently. */ const budgetTtlSeconds = Math.max(limits.ttlSeconds, config.ttlSeconds) + const retainedBytesLimit = Math.floor(limits.maxOwnerBytes * RETAINED_BYTES_FRACTION) /* Redis measures a member in UTF-8 bytes, so the ceiling has to be measured the same @@ -381,6 +431,7 @@ export async function appendEvents( limits.maxUserBytes, budgetTtlSeconds, String(chunk.members[chunk.members.length - 1].seq), + retainedBytesLimit, ...(lease ? [lease.value] : []), ...zaddArgs ) diff --git a/apps/sim/lib/mothership/request/session/index.ts b/apps/sim/lib/mothership/request/session/index.ts index aa9b22d1606..d31150c625c 100644 --- a/apps/sim/lib/mothership/request/session/index.ts +++ b/apps/sim/lib/mothership/request/session/index.ts @@ -68,7 +68,19 @@ export { FILE_PREVIEW_SESSION_SCHEMA_VERSION, isFilePreviewSession, } from './file-preview-session-contract' -export { checkForReplayGap, type ReplayGapResult } from './recovery' +export { + findReplayGap, + type RingPosition, + readRingPosition, + replayGapTerminal, + ringCanServe, +} from './recovery' +export { + forwardRunReplay, + openRunReplay, + type RunReplayEnd, + RunReplayUnavailableError, +} from './run-replay' export { encodeSSEEnvelope, SSE_RESPONSE_HEADERS } from './sse' export type { StreamBatchEvent } from './types' export { StreamWriter, type StreamWriterOptions } from './writer' diff --git a/apps/sim/lib/mothership/request/session/recovery.test.ts b/apps/sim/lib/mothership/request/session/recovery.test.ts index b81ade42bc0..fffce906a72 100644 --- a/apps/sim/lib/mothership/request/session/recovery.test.ts +++ b/apps/sim/lib/mothership/request/session/recovery.test.ts @@ -12,9 +12,13 @@ vi.mock('./buffer', () => ({ readEvents, })) -import { checkForReplayGap } from './recovery' +import { + findReplayGap, + replayGapTerminal, + ringCanServe, +} from '@/lib/mothership/request/session/recovery' -describe('checkForReplayGap', () => { +describe('replay gap', () => { it('uses the latest buffered request id when run metadata is missing it', async () => { getOldestSeq.mockResolvedValue(10) getLatestSeq.mockResolvedValue(12) @@ -24,11 +28,18 @@ describe('checkForReplayGap', () => { }, ]) - const result = await checkForReplayGap('stream-1', '1') + const gap = await findReplayGap('stream-1', '1') + expect(gap).not.toBeNull() + const result = await replayGapTerminal('stream-1', gap!) expect(readEvents).toHaveBeenCalledWith('stream-1', '11') expect(result?.gapDetected).toBe(true) expect(result?.envelopes[0].trace.requestId).toBe('req-live-123') expect(result?.envelopes[1].trace.requestId).toBe('req-live-123') }) + + it('cannot serve a reader that holds a cursor from an empty ring', () => { + expect(ringCanServe({ requestedAfterSeq: 12, oldestSeq: 0, latestSeq: 0 })).toBe(false) + expect(ringCanServe({ requestedAfterSeq: 0, oldestSeq: 0, latestSeq: 0 })).toBe(true) + }) }) diff --git a/apps/sim/lib/mothership/request/session/recovery.ts b/apps/sim/lib/mothership/request/session/recovery.ts index bb0153c8a14..83854f906cf 100644 --- a/apps/sim/lib/mothership/request/session/recovery.ts +++ b/apps/sim/lib/mothership/request/session/recovery.ts @@ -18,90 +18,121 @@ export interface ReplayGapResult { envelopes: ReturnType[] } -export async function checkForReplayGap( +/** Where the replay ring stands relative to a reader's cursor; 0 marks an empty ring. */ +export interface RingPosition { + requestedAfterSeq: number + oldestSeq: number + latestSeq: number +} + +/** + * Whether a ring whose first retained event has `firstSeq` still holds the stream's + * first event. The ring trims its oldest events, so a read from cursor 0 can start + * mid-stream; anything that rebuilds a turn from such a read must not, and a reader + * of it is re-synced from the worker's log instead ({@link findReplayGap}). + */ +export function startsAtReplayHead(firstSeq: number | undefined): boolean { + return firstSeq === undefined || firstSeq <= 1 +} + +export async function readRingPosition( streamId: string, - afterCursor: string, - requestId?: string -): Promise { - const requestedAfterSeq = Number(afterCursor || '0') - if (requestedAfterSeq <= 0) { - // Fast path: no cursor → nothing to check. Skip the span to avoid - // emitting zero-work spans on every stream connect. - return null + afterCursor: string +): Promise { + const [oldestSeq, latestSeq] = await Promise.all([getOldestSeq(streamId), getLatestSeq(streamId)]) + return { + requestedAfterSeq: Number(afterCursor || '0'), + oldestSeq: oldestSeq ?? 0, + latestSeq: latestSeq ?? 0, } +} + +/** + * Whether the ring can serve a reader from its cursor. It cannot once it has lost its + * head: the events before its oldest are gone, and a cursor that was served from the + * worker's log instead is not a position in the ring, so no cursor is trusted. Nor can + * it serve a cursor ahead of its latest event, which only a buffer whose numbering + * restarted after it expired produces, nor any cursor from a buffer that expired. + */ +export function ringCanServe({ requestedAfterSeq, oldestSeq, latestSeq }: RingPosition): boolean { + if (latestSeq <= 0) return requestedAfterSeq <= 0 + return startsAtReplayHead(oldestSeq) && requestedAfterSeq <= latestSeq +} +/** The ring's position when it cannot serve `afterCursor` (see {@link ringCanServe}). */ +export async function findReplayGap( + streamId: string, + afterCursor: string, + requestId?: string +): Promise { return withCopilotSpan( TraceSpan.CopilotRecoveryCheckReplayGap, { [TraceAttr.StreamId]: streamId, - [TraceAttr.CopilotRecoveryRequestedAfterSeq]: requestedAfterSeq, + [TraceAttr.CopilotRecoveryRequestedAfterSeq]: Number(afterCursor || '0'), ...(requestId ? { [TraceAttr.RequestId]: requestId } : {}), }, async (span) => { - const oldestSeq = await getOldestSeq(streamId) - const latestSeq = await getLatestSeq(streamId) + const position = await readRingPosition(streamId, afterCursor) span.setAttributes({ - [TraceAttr.CopilotRecoveryOldestSeq]: oldestSeq ?? -1, - [TraceAttr.CopilotRecoveryLatestSeq]: latestSeq ?? -1, + [TraceAttr.CopilotRecoveryOldestSeq]: position.oldestSeq, + [TraceAttr.CopilotRecoveryLatestSeq]: position.latestSeq, }) - - /* Trimmed below the ring, or ahead of a buffer whose numbering restarted after it - expired: either way the events after the cursor are not the ones it names. */ - if ( - latestSeq !== null && - latestSeq > 0 && - oldestSeq !== null && - (requestedAfterSeq < oldestSeq - 1 || requestedAfterSeq > latestSeq) - ) { - const resolvedRequestId = await resolveReplayGapRequestId(streamId, latestSeq, requestId) - logger.warn('Replay gap detected: requested cursor is outside the retained events', { - streamId, - requestedAfterSeq, - oldestAvailableSeq: oldestSeq, - latestSeq, - }) - span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.GapDetected) - - const gapEnvelope = createEvent({ - streamId, - cursor: String(latestSeq + 1), - seq: latestSeq + 1, - requestId: resolvedRequestId, - type: MothershipStreamV1EventType.error, - payload: { - message: 'Replay history is no longer available. Some events may have been lost.', - code: 'replay_gap', - data: { - oldestAvailableSeq: oldestSeq, - requestedAfterSeq, - }, - }, - }) - - const terminalEnvelope = createEvent({ - streamId, - cursor: String(latestSeq + 2), - seq: latestSeq + 2, - requestId: resolvedRequestId, - type: MothershipStreamV1EventType.complete, - payload: { - status: MothershipStreamV1CompletionStatus.error, - reason: 'replay_gap', - }, - }) - - return { - gapDetected: true, - envelopes: [gapEnvelope, terminalEnvelope], - } + if (ringCanServe(position)) { + span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.InRange) + return null } - - span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.InRange) - return null + logger.warn('Replay gap detected: the ring cannot serve the requested cursor', { + streamId, + ...position, + }) + span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.GapDetected) + return position } ) } +/** + * Ends a reader's view with `replay_gap` when nothing can re-sync it, numbered past + * both the ring and the reader's cursor so the reader cannot drop it as already seen. + */ +export async function replayGapTerminal( + streamId: string, + position: RingPosition, + requestId?: string +): Promise { + const { latestSeq, oldestSeq, requestedAfterSeq } = position + const baseSeq = Math.max(latestSeq, requestedAfterSeq) + const resolvedRequestId = await resolveReplayGapRequestId(streamId, latestSeq, requestId) + const gapEnvelope = createEvent({ + streamId, + cursor: String(baseSeq + 1), + seq: baseSeq + 1, + requestId: resolvedRequestId, + type: MothershipStreamV1EventType.error, + payload: { + message: 'Replay history is no longer available. Some events may have been lost.', + code: 'replay_gap', + data: { + oldestAvailableSeq: oldestSeq, + requestedAfterSeq, + }, + }, + }) + const terminalEnvelope = createEvent({ + streamId, + cursor: String(baseSeq + 2), + seq: baseSeq + 2, + requestId: resolvedRequestId, + type: MothershipStreamV1EventType.complete, + payload: { + status: MothershipStreamV1CompletionStatus.error, + reason: 'replay_gap', + }, + }) + return { gapDetected: true, envelopes: [gapEnvelope, terminalEnvelope] } +} + async function resolveReplayGapRequestId( streamId: string, latestSeq: number, diff --git a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts index a55706e2468..376286c0b57 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -13,10 +13,27 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') const { createServer: createHttpServer } = await import('node:http') const abortRequests: Array> = [] - const hooks = { onAbort: undefined as (() => Promise) | undefined } + const hooks = { + onAbort: undefined as (() => Promise) | undefined, + /** The read-only replay's answer; a worker that does not know the run by default. */ + replay: { status: 404, frames: [] as unknown[] }, + } + const replayRequests: Array> = [] const server = createHttpServer(async (request, response) => { let body = '' for await (const chunk of request) body += chunk + if (request.url === '/api/streams/replay') { + replayRequests.push(JSON.parse(body)) + if (hooks.replay.status !== 200) { + response.writeHead(hooks.replay.status, { 'content-type': 'application/json' }) + response.end(JSON.stringify({ error: 'Run not found' })) + return + } + response.writeHead(200, { 'content-type': 'text/event-stream' }) + for (const frame of hooks.replay.frames) response.write(`data: ${JSON.stringify(frame)}\n\n`) + response.end('data: [DONE]\n\n') + return + } if (request.url === '/api/streams/explicit-abort') { abortRequests.push(JSON.parse(body)) await hooks.onAbort?.() @@ -40,11 +57,16 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { worker: { server, abortRequests, + replayRequests, hooks, /** Events, or steps to run between them, that the scripted worker streams in order. */ script: [] as unknown[], /** Controller lifecycles started, and what each sink call threw. */ - runs: [] as Array<{ dispatched: unknown[]; sinkErrors: unknown[] }>, + runs: [] as Array<{ + dispatched: unknown[] + sinkErrors: unknown[] + recoveredEvents?: unknown[] + }>, }, } }) @@ -57,9 +79,17 @@ vi.mock('@/lib/mothership/request/lifecycle/run', () => ({ */ runCopilotLifecycle: async ( _payload: unknown, - options: { onEvent?: (event: unknown) => Promise; abortSignal?: AbortSignal } + options: { + onEvent?: (event: unknown) => Promise + abortSignal?: AbortSignal + recovery?: { events: unknown[] } + } ) => { - const run = { dispatched: [] as unknown[], sinkErrors: [] as unknown[] } + const run = { + dispatched: [] as unknown[], + sinkErrors: [] as unknown[], + recoveredEvents: options.recovery?.events, + } worker.runs.push(run) for (const event of worker.script) { if (typeof event === 'function') { @@ -105,13 +135,13 @@ import { import { finalizeStream } from '@/lib/mothership/request/lifecycle/finalize' import { createSSEStream } from '@/lib/mothership/request/lifecycle/start' import { acquirePendingChatStream } from '@/lib/mothership/request/session/abort' -import { readEvents } from '@/lib/mothership/request/session/buffer' +import { appendEvents, readEvents } from '@/lib/mothership/request/session/buffer' import { type ChatStreamLease, chatStreamLockKey, StreamControllerSupersededError, } from '@/lib/mothership/request/session/controller-lease' -import { eventToStreamEvent } from '@/lib/mothership/request/session/event' +import { createEvent, eventToStreamEvent } from '@/lib/mothership/request/session/event' import { REPLAY_BUDGET_EXHAUSTED_CODE, StreamReplayBudgetExhaustedError, @@ -124,7 +154,9 @@ import { type PendingFileIntent, storeFileIntent, } from '@/lib/mothership/tools/server/files/file-intent-store' +import { GET as copilotChatGET } from '@/app/api/copilot/chat/queries' import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' +import { GET as mothershipChatGET } from '@/app/api/mothership/chats/[chatId]/route' const MB = 1024 * 1024 @@ -344,6 +376,72 @@ describe.runIf(Boolean(redisUrl))('leased Chat stream writer with Redis', () => }) }) +describe.runIf(Boolean(redisUrl))('the replay ring trimmed past its byte target', () => { + /** A ring whose counter is already past the byte target, as after a long run. */ + async function overTargetRing(seqs: number[]) { + const streamId = generateId() + const [ownerBudgetKey] = getRedisBudgetKeys({ kind: 'copilot_stream', id: streamId }) + const persisted = await appendEvents( + seqs.map((seq) => textEvent(streamId, seq)), + { streamId } + ) + expect(persisted).toEqual({ persisted: true }) + const { maxOwnerBytes } = getRedisBudgetLimits('copilot_stream') + await redis().set(ownerBudgetKey, maxOwnerBytes - MB, 'EX', 3600) + return { streamId, ownerBudgetKey } + } + + function textEvent(streamId: string, seq: number) { + return createEvent({ + streamId, + cursor: String(seq), + seq, + requestId: 'replay-trim', + type: 'text', + payload: { channel: 'assistant', text: `part ${String(seq).padStart(6, '0')}` }, + }) + } + + const seqsOf = async (streamId: string) => + (await storedMembers(streamId)).map((member) => JSON.parse(member).seq as number) + const bytesOf = (members: string[]) => + members.reduce((sum, member) => sum + Buffer.byteLength(member), 0) + + it('trims a replayed member below the ring with it, leaving a contiguous tail', async () => { + const { streamId, ownerBudgetKey } = await overTargetRing([2, 3]) + const before = Number(await redis().get(ownerBudgetKey)) + const dropped = await storedMembers(streamId) + const newest = textEvent(streamId, 4) + + expect(await appendEvents([textEvent(streamId, 1), newest], { streamId })).toEqual({ + persisted: true, + }) + + expect(await seqsOf(streamId)).toEqual([4]) + expect(Number(await redis().get(ownerBudgetKey))).toBe( + before + Buffer.byteLength(JSON.stringify(newest)) - bytesOf(dropped) + ) + }) + + it('catches an oversized ring up over several appends', async () => { + const seqs = Array.from({ length: 6000 }, (_, index) => index + 1) + const { streamId, ownerBudgetKey } = await overTargetRing(seqs) + + await appendEvents([textEvent(streamId, 6001)], { streamId }) + const afterFirst = await seqsOf(streamId) + expect(afterFirst.length).toBeGreaterThan(1) + expect(afterFirst).toEqual( + Array.from({ length: afterFirst.length }, (_, index) => 6001 - afterFirst.length + 1 + index) + ) + + await appendEvents([textEvent(streamId, 6002)], { streamId }) + expect(await seqsOf(streamId)).toEqual([6002]) + expect(Number(await redis().get(ownerBudgetKey))).toBeLessThan( + getRedisBudgetLimits('copilot_stream').maxOwnerBytes - MB + ) + }) +}) + describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budget', () => { const userId = generateId() const workspaceId = generateId() @@ -667,6 +765,223 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge expect(stored.status).toBe('complete') }) + it('streams far past the owner budget without refusing, retaining a bounded contiguous tail', async () => { + const { maxOwnerBytes, maxUserBytes } = getRedisBudgetLimits('copilot_stream') + const chunk = 'x'.repeat(4 * 1024) + const eventCount = Math.ceil((maxOwnerBytes * 1.3) / chunk.length) + const { streamId, runId, frames } = await runTurn( + Array.from({ length: eventCount }, (_, index) => text(`${index}:${chunk}`)) + ) + + expect(frames.map((frame) => frame.type)).not.toContain('error') + expect(frames.at(-1)).toMatchObject({ type: 'complete', payload: { status: 'complete' } }) + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('complete') + + const members = await storedMembers(streamId) + const retainedBytes = members.reduce((sum, member) => sum + Buffer.byteLength(member), 0) + const [ownerBudgetKey, userBudgetKey] = getRedisBudgetKeys({ + kind: 'copilot_stream', + id: streamId, + userId, + }) + expect(Number(await redis().get(ownerBudgetKey))).toBe(retainedBytes) + expect(retainedBytes).toBeLessThan(maxOwnerBytes) + expect(Number(await redis().get(userBudgetKey))).toBeLessThan(maxUserBytes) + + const seqs = members.map((member) => JSON.parse(member).seq as number) + const oldestSeq = seqs[0] + const latestSeq = seqs.at(-1)! + expect(oldestSeq).toBeGreaterThan(1) + expect(seqs).toEqual(Array.from({ length: seqs.length }, (_, index) => oldestSeq + index)) + + const reconnect = async (after: number) => + dataFrames( + await ( + await streamGET( + new NextRequest( + `http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=${after}` + ), + { params: Promise.resolve({}) } + ) + ).text() + ) + // A headless ring is re-synced from the log; this worker lacks the run, so replay_gap. + worker.replayRequests.length = 0 + const inRange = await reconnect(oldestSeq - 1) + expect( + inRange.map((frame) => [frame.type, frame.payload.code ?? frame.payload.reason]) + ).toEqual([ + ['error', 'replay_gap'], + ['complete', 'replay_gap'], + ]) + const behind = await reconnect(oldestSeq - 2) + expect(behind.map((frame) => [frame.type, frame.payload.code ?? frame.payload.reason])).toEqual( + [ + ['error', 'replay_gap'], + ['complete', 'replay_gap'], + ] + ) + expect(behind[0].payload.data).toEqual({ + oldestAvailableSeq: oldestSeq, + requestedAfterSeq: oldestSeq - 2, + }) + expect(behind[0].seq).toBe(latestSeq + 1) + expect(worker.replayRequests).toEqual([ + { streamId, chatId, userId }, + { streamId, chatId, userId }, + ]) + }, 180_000) + + /** + * An unfinished run with no live controller whose ring a byte trim has advanced past + * its head: seqs 1–2 are gone and 3–4 remain. + */ + async function trimmedRecoverableStream() { + const streamId = generateId() + const request = { + message: 'Summarize the logs', + userId, + messageId: streamId, + chatId, + workspaceId, + } + await db.insert(copilotRuns).values({ + id: generateId(), + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + requestContext: { + requestId: generateId(), + controllerToken: `dead\n${generateId()}`, + recovery: { + kind: 'interactive_stream', + request, + goRoute: '/api/mothership', + clientToolPickupExpected: false, + }, + }, + }) + await db + .update(copilotChats) + .set({ conversationId: streamId }) + .where(eq(copilotChats.id, chatId)) + const persisted = await appendEvents( + [1, 2, 3, 4].map((seq) => + createEvent({ + streamId, + cursor: String(seq), + seq, + requestId: generateId(), + type: 'text', + payload: { channel: 'assistant', text: `part ${seq} ` }, + }) + ), + { streamId } + ) + expect(persisted).toEqual({ persisted: true }) + await redis().zremrangebyrank(`mothership_stream:${streamId}:events`, 0, 1) + worker.runs.length = 0 + worker.script = [] + return streamId + } + + const resume = async (streamId: string, query: string) => + streamGET( + new NextRequest( + `http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&${query}` + ), + { params: Promise.resolve({}) } + ) + + describe('when a byte trim has removed the head of an unfinished turn', () => { + afterAll(async () => { + await db.update(copilotChats).set({ conversationId: null }).where(eq(copilotChats.id, chatId)) + await redis().del(chatStreamLockKey(chatId)) + }) + + it.each(['after=4', 'after=0'])( + 'recovers from an empty context, not the tail, and shows the full turn from the log (%s)', + async (query) => { + const streamId = await trimmedRecoverableStream() + worker.hooks.replay = { + status: 200, + frames: [ + { + v: 1, + type: 'text', + seq: 1, + ts: new Date().toISOString(), + stream: { streamId, chatId }, + payload: { + channel: 'assistant', + text: 'part 1 part 2 part 3 part 4 ', + textOffset: 0, + }, + }, + { + v: 1, + type: 'run', + seq: 2, + ts: new Date().toISOString(), + stream: { streamId, chatId }, + payload: { kind: 'replay_end', reason: 'stalled', textLength: 28 }, + }, + ], + } + try { + const frames = dataFrames(await (await resume(streamId, query)).text()) + + expect(frames.map((frame) => frame.type)).toEqual(['text']) + expect(frames[0].payload.text).toBe('part 1 part 2 part 3 part 4 ') + expect(worker.runs.map((run) => run.recoveredEvents)).toEqual([[]]) + } finally { + worker.hooks.replay = { status: 404, frames: [] } + } + } + ) + + it('serves a batch reconnect no tail events', async () => { + const streamId = await trimmedRecoverableStream() + + const batch = await (await resume(streamId, 'after=0&batch=true')).json() + + expect(batch.events).toEqual([]) + expect(worker.runs.map((run) => run.recoveredEvents)).toEqual([[]]) + }) + + it.each([ + [ + 'mothership chat', + () => + mothershipChatGET( + new NextRequest(`http://localhost:3000/api/mothership/chats/${chatId}`), + { + params: Promise.resolve({ chatId }), + } + ), + ], + [ + 'copilot chat', + () => + copilotChatGET( + new NextRequest(`http://localhost:3000/api/copilot/chat?chatId=${chatId}`) + ), + ], + ])('leaves the %s snapshot to the resume route', async (_label, load) => { + const streamId = await trimmedRecoverableStream() + + const body = await (await load()).json() + + expect(body.success).toBe(true) + expect(body.chat.streamSnapshot).toBeUndefined() + expect(JSON.stringify(body.chat.messages)).not.toContain('part 3') + expect(await storedMembers(streamId)).toHaveLength(2) + }) + }) + it("leaves its successor's stream untouched when the lease is lost while ending a refused turn", async () => { const successorToken = `successor\n${generateId()}` worker.hooks.onAbort = async () => { diff --git a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts new file mode 100644 index 00000000000..f9e8be62fe5 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts @@ -0,0 +1,447 @@ +/** + * Reconnects that Sim's replay ring can no longer serve, against real Redis and + * PostgreSQL through the production reconnect route. A local HTTP server stands in for + * the worker's read-only replay endpoint; everything on Sim's side is production code. + */ +import { authMock, authMockFns } from '@sim/testing/mocks/auth.mock' +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest' + +const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { + const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') + const { createServer } = await import('node:http') + const worker = { + requests: [] as Array>, + /** What the replay endpoint answers: an HTTP status, or SSE frames. */ + reply: { status: 200, frames: [] as unknown[] }, + } + const server = createServer(async (request, response) => { + let body = '' + for await (const chunk of request) body += chunk + if (request.url !== '/api/streams/replay') { + response.writeHead(404).end() + return + } + worker.requests.push(JSON.parse(body)) + if (worker.reply.status !== 200) { + response.writeHead(worker.reply.status, { 'content-type': 'application/json' }) + response.end(JSON.stringify({ error: 'Run not found' })) + return + } + response.writeHead(200, { 'content-type': 'text/event-stream' }) + for (const frame of worker.reply.frames) response.write(`data: ${JSON.stringify(frame)}\n\n`) + response.end('data: [DONE]\n\n') + }) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const { port } = server.address() as { port: number } + const url = readTestRedisUrl() + const inheritedEnv = { + REDIS_URL: process.env.REDIS_URL, + SIM_AGENT_API_URL: process.env.SIM_AGENT_API_URL, + COPILOT_STREAM_EVENT_LIMIT: process.env.COPILOT_STREAM_EVENT_LIMIT, + } + process.env.REDIS_URL = url + process.env.SIM_AGENT_API_URL = `http://127.0.0.1:${port}` + /** A ring this small loses the head of every stream below. */ + process.env.COPILOT_STREAM_EVENT_LIMIT = '5' + return { redisUrl: url, inheritedEnv, worker: Object.assign(worker, { server }) } +}) + +vi.mock('@/lib/auth', () => authMock) + +import { db } from '@sim/db' +import { copilotChats, copilotRuns, permissions, user, workspace } from '@sim/db/schema' +import { sleep } from '@sim/utils/helpers' +import { generateId } from '@sim/utils/id' +import { eq } from 'drizzle-orm' +import { NextRequest } from 'next/server' +import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis' +import { MOTHERSHIP_STREAM_REPLAY_HEADER } from '@/lib/mothership/constants' +import { allocateCursor, appendEvents } from '@/lib/mothership/request/session/buffer' +import { createEvent } from '@/lib/mothership/request/session/event' +import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' + +const userId = generateId() +const workspaceId = generateId() +const chatId = generateId() + +function dataFrames(body: string) { + return body + .split('\n\n') + .filter((frame) => frame.startsWith('data: ')) + .map((frame) => JSON.parse(frame.slice('data: '.length))) +} + +/** A worker frame as the replay endpoint writes it, with the worker's own sequence. */ +function workerFrame(streamId: string, seq: number, type: string, payload: unknown) { + return { + v: 1, + type, + seq, + ts: new Date().toISOString(), + stream: { streamId, chatId }, + payload, + } +} + +async function appendText(streamId: string, value: string): Promise { + const { seq, cursor } = await allocateCursor(streamId) + await appendEvents([ + createEvent({ + streamId, + cursor, + seq, + requestId: 'req-ring', + type: 'text', + payload: { channel: 'assistant', text: value }, + }), + ]) +} + +/** A live run whose ring holds only the last five of its ten events. */ +async function liveRunWithTrimmedRing(): Promise<{ streamId: string; runId: string }> { + const streamId = generateId() + const runId = generateId() + await db.insert(copilotRuns).values({ + id: runId, + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + }) + for (let index = 1; index <= 10; index++) await appendText(streamId, `part ${index} `) + return { streamId, runId } +} + +function reconnect(streamId: string, after: string, batch = false, extra = '') { + return streamGET( + new NextRequest( + `http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=${after}${batch ? '&batch=true' : ''}${extra}` + ), + { params: Promise.resolve({}) } + ) +} + +function fullResponse(streamId: string) { + return [ + workerFrame(streamId, 1, 'session', { kind: 'start' }), + workerFrame(streamId, 2, 'text', { + channel: 'assistant', + text: 'part 1 part 2 part 3 part 4 part 5 part 6 part 7 part 8 part 9 part 10 ', + textOffset: 0, + }), + workerFrame(streamId, 3, 'complete', { status: 'complete' }), + ] +} + +/** Runs whether or not the suite does, so a skipped suite never leaks the worker or env. */ +afterAll(async () => { + await new Promise((resolve) => worker.server.close(() => resolve())) + for (const [key, value] of Object.entries(inheritedEnv)) { + if (value === undefined) delete process.env[key] + else process.env[key] = value + } +}) + +describe.runIf(Boolean(redisUrl))('reconnects past the replay ring', () => { + beforeAll(async () => { + const now = new Date() + await db.insert(user).values({ + id: userId, + name: 'Replay gap fixture', + email: `${userId}@replay-gap.test`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(workspace).values({ + id: workspaceId, + name: 'Replay gap fixture', + ownerId: userId, + billedAccountUserId: userId, + }) + await db.insert(permissions).values({ + id: generateId(), + userId, + entityType: 'workspace', + entityId: workspaceId, + permissionType: 'admin', + }) + await db.insert(copilotChats).values({ id: chatId, userId, workspaceId, type: 'mothership' }) + authMockFns.mockGetSession.mockResolvedValue({ + user: { id: userId }, + session: { id: generateId() }, + }) + }) + + beforeEach(() => { + worker.requests.length = 0 + worker.reply = { status: 200, frames: [] } + }) + + afterAll(async () => { + await db.delete(copilotRuns).where(eq(copilotRuns.chatId, chatId)) + await db.delete(copilotChats).where(eq(copilotChats.id, chatId)) + await db.delete(permissions).where(eq(permissions.userId, userId)) + await db.delete(workspace).where(eq(workspace.id, workspaceId)) + await db.delete(user).where(eq(user.id, userId)) + await closeRedisConnection() + }) + + it.each([ + ['a fresh tab reconnecting from 0', '0'], + ['a cursor behind the retained ring', '2'], + ])( + 're-syncs %s from the worker log instead of a partial replay or an error', + async (_name, after) => { + const { streamId } = await liveRunWithTrimmedRing() + worker.reply.frames = fullResponse(streamId) + + const response = await reconnect(streamId, after) + const frames = dataFrames(await response.text()) + + expect(response.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER)).toBe('log') + expect(worker.requests).toEqual([{ streamId, chatId, userId }]) + expect(frames.map((frame) => [frame.seq, frame.stream.cursor, frame.type])).toEqual([ + [1, '1', 'session'], + [2, '2', 'text'], + [3, '3', 'complete'], + ]) + expect(frames[1].payload.text).toMatch(/^part 1 part 2 /) + } + ) + + it('re-syncs a cursor ahead of a ring whose numbering restarted', async () => { + const { streamId } = await liveRunWithTrimmedRing() + await getRedisClient()!.del( + `mothership_stream:${streamId}:events`, + `mothership_stream:${streamId}:seq` + ) + await appendText(streamId, 'after expiry ') + worker.reply.frames = fullResponse(streamId) + + const response = await reconnect(streamId, '4') + const frames = dataFrames(await response.text()) + + expect(response.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER)).toBe('log') + expect(frames.map((frame) => frame.type)).toEqual(['session', 'text', 'complete']) + }) + + it('ends a parked replay without a terminal once the run resumes', async () => { + const { streamId, runId } = await liveRunWithTrimmedRing() + await db + .update(copilotRuns) + .set({ status: 'paused_waiting_for_tool' }) + .where(eq(copilotRuns.id, runId)) + worker.reply.frames = [ + workerFrame(streamId, 1, 'text', { channel: 'assistant', text: 'so far', textOffset: 0 }), + workerFrame(streamId, 2, 'run', { kind: 'replay_end', reason: 'parked', textLength: 6 }), + ] + + const startedAt = Date.now() + const response = await reconnect(streamId, '0') + const body = response.text() + await sleep(1_500) + await db.update(copilotRuns).set({ status: 'active' }).where(eq(copilotRuns.id, runId)) + const frames = dataFrames(await body) + + const elapsed = Date.now() - startedAt + expect(frames.map((frame) => frame.type)).toEqual(['text']) + expect(elapsed).toBeGreaterThanOrEqual(1_500) + expect(elapsed).toBeLessThan(5_000) + expect(worker.requests).toHaveLength(1) + }) + + it('ends a stalled replay promptly once the run finishes', async () => { + const { streamId, runId } = await liveRunWithTrimmedRing() + worker.reply.frames = [ + workerFrame(streamId, 1, 'text', { channel: 'assistant', text: 'so far', textOffset: 0 }), + workerFrame(streamId, 2, 'run', { kind: 'replay_end', reason: 'stalled', textLength: 6 }), + ] + + const startedAt = Date.now() + const body = (await reconnect(streamId, '0')).text() + await sleep(1_000) + await db.update(copilotRuns).set({ status: 'complete' }).where(eq(copilotRuns.id, runId)) + await body + + expect(Date.now() - startedAt).toBeLessThan(5_000) + }) + + it('ends a live tail without a terminal when its ring restarts under it', async () => { + const streamId = generateId() + await db.insert(copilotRuns).values({ + id: generateId(), + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + }) + for (let index = 1; index <= 4; index++) await appendText(streamId, `part ${index} `) + + const response = await reconnect(streamId, '4') + const body = response.text() + // Let the tail reach its poll loop, past the reconnect-time gap check. + await sleep(500) + await getRedisClient()!.del( + `mothership_stream:${streamId}:events`, + `mothership_stream:${streamId}:seq` + ) + await appendText(streamId, 'after expiry ') + const frames = dataFrames(await body) + + expect(frames).toEqual([]) + }) + + it('keeps a reader re-synced from the log on the log once a restarted ring grows past its cursor', async () => { + const streamId = generateId() + await db.insert(copilotRuns).values({ + id: generateId(), + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + }) + for (let index = 1; index <= 4; index++) await appendText(streamId, `part ${index} `) + worker.reply.frames = fullResponse(streamId) + + const response = await reconnect(streamId, '2', false, '&source=log') + const frames = dataFrames(await response.text()) + + expect(response.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER)).toBe('log') + expect(frames.map((frame) => frame.type)).toEqual(['session', 'text', 'complete']) + }) + + it('serves no ring events to a batch read from a reader re-synced from the log', async () => { + const streamId = generateId() + await db.insert(copilotRuns).values({ + id: generateId(), + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + }) + for (let index = 1; index <= 4; index++) await appendText(streamId, `part ${index} `) + + const response = await reconnect(streamId, '2', true, '&source=log') + + expect(await response.json()).toMatchObject({ success: true, events: [], status: 'active' }) + }) + + it('holds a stalled replay open before the client re-attaches', async () => { + const { streamId } = await liveRunWithTrimmedRing() + worker.reply.frames = [ + workerFrame(streamId, 1, 'text', { channel: 'assistant', text: 'so far', textOffset: 0 }), + workerFrame(streamId, 2, 'run', { kind: 'replay_end', reason: 'stalled', textLength: 6 }), + ] + + const startedAt = Date.now() + const frames = dataFrames(await (await reconnect(streamId, '0')).text()) + + expect(frames.map((frame) => frame.type)).toEqual(['text']) + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(9_000) + }) + + it('ends a replay whose end reason it does not know without a run event or an error', async () => { + const { streamId } = await liveRunWithTrimmedRing() + worker.reply.frames = [ + workerFrame(streamId, 1, 'text', { channel: 'assistant', text: 'so far', textOffset: 0 }), + workerFrame(streamId, 2, 'run', { kind: 'replay_end', reason: 'drained', textLength: 6 }), + ] + + const frames = dataFrames(await (await reconnect(streamId, '0')).text()) + + expect(frames.map((frame) => frame.type)).toEqual(['text']) + }) + + it('ends a live tail without a terminal when its ring loses its head under it', async () => { + const streamId = generateId() + await db.insert(copilotRuns).values({ + id: generateId(), + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + }) + for (let index = 1; index <= 4; index++) await appendText(streamId, `part ${index} `) + + const response = await reconnect(streamId, '4') + const body = response.text() + await sleep(500) + for (let index = 5; index <= 7; index++) await appendText(streamId, `part ${index} `) + const frames = dataFrames(await body) + + expect(frames.map((frame) => frame.type)).not.toContain('complete') + expect(response.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER)).toBeNull() + }) + + it('re-syncs a live run whose buffer expired under a reader cursor', async () => { + const streamId = generateId() + await db.insert(copilotRuns).values({ + id: generateId(), + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + }) + worker.reply.frames = fullResponse(streamId) + + const response = await reconnect(streamId, '6') + const frames = dataFrames(await response.text()) + + expect(response.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER)).toBe('log') + expect(frames.map((frame) => frame.type)).toEqual(['session', 'text', 'complete']) + }) + + it('answers a finished run whose buffer expired with its terminal, not a replay', async () => { + const streamId = generateId() + await db.insert(copilotRuns).values({ + id: generateId(), + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + status: 'complete', + }) + + const response = await reconnect(streamId, '6') + const frames = dataFrames(await response.text()) + + expect(response.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER)).toBeNull() + expect(frames.map((frame) => [frame.type, frame.payload.status])).toEqual([ + ['complete', 'complete'], + ]) + expect(worker.requests).toEqual([]) + }) + + it('serves no ring events to a batch read the ring can no longer serve', async () => { + const { streamId } = await liveRunWithTrimmedRing() + + const response = await reconnect(streamId, '0', true) + + expect(await response.json()).toMatchObject({ success: true, events: [], status: 'active' }) + }) + + it.each([404, 401, 403])( + 'keeps the replay_gap terminal when the worker will not replay the run (%i)', + async (status) => { + const { streamId } = await liveRunWithTrimmedRing() + worker.reply.status = status + + const frames = dataFrames(await (await reconnect(streamId, '2')).text()) + + expect( + frames.map((frame) => [frame.type, frame.payload.code ?? frame.payload.status]) + ).toEqual([ + ['error', 'replay_gap'], + ['complete', 'error'], + ]) + } + ) +}) diff --git a/apps/sim/lib/mothership/request/session/run-replay.test.ts b/apps/sim/lib/mothership/request/session/run-replay.test.ts new file mode 100644 index 00000000000..36ff86dc1a6 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/run-replay.test.ts @@ -0,0 +1,81 @@ +import { mothershipAgentUrlMock } from '@sim/testing/mocks/mothership-agent-url.mock' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('@/lib/mothership/server/agent-url', () => mothershipAgentUrlMock) + +import { + forwardRunReplay, + openRunReplay, + RunReplayUnavailableError, +} from '@/lib/mothership/request/session/run-replay' + +/** Well past the worker idle bound, and under common intermediary idle cuts. */ +const INTERMEDIARY_IDLE_MS = 300_000 + +function settle(promise: Promise) { + const state: { done: boolean; value?: T; error?: unknown } = { done: false } + promise.then( + (value) => { + state.done = true + state.value = value + }, + (error: unknown) => { + state.done = true + state.error = error + } + ) + return state +} + +describe('run replay liveness', () => { + beforeEach(() => { + vi.useFakeTimers() + vi.stubGlobal('fetch', vi.fn()) + }) + + afterEach(() => { + vi.useRealTimers() + }) + + it('gives up on a worker that never answers the replay request', async () => { + vi.mocked(fetch).mockImplementationOnce( + (_url, init) => + new Promise((_resolve, reject) => { + init?.signal?.addEventListener('abort', () => reject(init.signal?.reason), { + once: true, + }) + }) + ) + const state = settle( + openRunReplay({ + streamId: '00000000-0000-4000-8000-000000000001', + chatId: '00000000-0000-4000-8000-000000000002', + userId: 'user-1', + signal: new AbortController().signal, + }) + ) + + await vi.advanceTimersByTimeAsync(INTERMEDIARY_IDLE_MS) + + expect(state.done).toBe(true) + expect(state.error).toBeInstanceOf(RunReplayUnavailableError) + }) + + it('ends a replay whose worker goes silent so the reader can re-attach', async () => { + const state = settle( + forwardRunReplay({ + body: new ReadableStream(), + streamId: 'stream-1', + signal: new AbortController().signal, + write: () => true, + readRunStatus: async () => 'active', + isClosed: () => false, + deadlineAt: Date.now() + 60 * 60_000, + }) + ) + + await vi.advanceTimersByTimeAsync(INTERMEDIARY_IDLE_MS) + + expect(state).toMatchObject({ done: true, value: 'closed' }) + }) +}) diff --git a/apps/sim/lib/mothership/request/session/run-replay.ts b/apps/sim/lib/mothership/request/session/run-replay.ts new file mode 100644 index 00000000000..992894b76d6 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -0,0 +1,216 @@ +import { createLogger } from '@sim/logger' +import { getErrorMessage } from '@sim/utils/errors' +import { sleep } from '@sim/utils/helpers' +import { isRecordLike, toRecord } from '@sim/utils/object' +import { z } from 'zod' +import { WORKER_STREAM_IDLE_TIMEOUT_MS } from '@/lib/mothership/constants' +import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' +import { type StreamReplayEnd, StreamReplayRequest } from '@/lib/mothership/generated/protocol' +import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' +import { fetchGo } from '@/lib/mothership/request/go/fetch' +import { FatalSseEventError, processSSEStream } from '@/lib/mothership/request/go/parser' +import { mothershipRequestHeaders } from '@/lib/mothership/request/headers' +import { + isTerminalStreamStatus, + type PersistedStreamEventEnvelope, + parsePersistedStreamEventEnvelope, +} from '@/lib/mothership/request/session/contract' +import { toReplayEnvelope } from '@/lib/mothership/request/session/types' +import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url' + +const logger = createLogger('RunReplay') + +const REPLAY_PATH = '/api/streams/replay' +/** Answers that no retry changes, so the reader falls back to `replay_gap`. */ +const REPLAY_REFUSED_STATUSES: ReadonlySet = new Set([401, 403, 404]) +/** A reader's replay response stays open at least this long unless the run ended. */ +const REPLAY_MIN_RESPONSE_MS = 10_000 +const REPLAY_HOLD_POLL_MS = 1_000 +const PARKED_RUN_STATUS = 'paused_waiting_for_tool' + +/** How a worker replay leg ended: at the run's terminal, short of it, or cut off. */ +export type RunReplayEnd = 'complete' | StreamReplayEnd['reason'] | 'closed' + +/** The worker could not serve the replay; the reader should retry later. */ +export class RunReplayUnavailableError extends Error { + constructor(message: string, options?: { cause?: unknown }) { + super(message, options) + this.name = 'RunReplayUnavailableError' + } +} + +/** + * Opens the worker's read-only replay of a run from its durable log, for a reader the + * replay ring can no longer serve. No receipt is sent: the reader starts from an empty + * response. Returns `null` when the worker will not replay it: it knows no such run for + * this chat and user, or this deployment's key may not call the replay at all. + */ +export async function openRunReplay(params: { + streamId: string + chatId: string + userId: string + signal: AbortSignal +}): Promise | null> { + const { streamId, chatId, userId, signal } = params + const baseUrl = await getMothershipBaseURL({ userId }) + const unanswered = new AbortController() + const headersTimer = setTimeout( + () => unanswered.abort(new Error('The worker did not answer the replay request')), + WORKER_STREAM_IDLE_TIMEOUT_MS + ) + let response: Response + try { + response = await fetchGo(`${baseUrl}${REPLAY_PATH}`, { + method: 'POST', + headers: mothershipRequestHeaders(), + body: JSON.stringify(StreamReplayRequest.parse({ streamId, chatId, userId })), + signal: AbortSignal.any([signal, unanswered.signal]), + spanName: `sim → go ${REPLAY_PATH}`, + operation: 'stream_replay', + attributes: { [TraceAttr.StreamId]: streamId, [TraceAttr.ChatId]: chatId }, + }) + } catch (error) { + if (signal.aborted) throw error + throw new RunReplayUnavailableError('The run replay could not be reached', { cause: error }) + } finally { + clearTimeout(headersTimer) + } + if (REPLAY_REFUSED_STATUSES.has(response.status)) { + // A key refusal is otherwise silent: every reader just falls back to replay_gap. + if (response.status !== 404) { + logger.warn('The worker refused this deployment the run replay', { + streamId, + status: response.status, + }) + } + await response.body?.cancel().catch(() => {}) + return null + } + if (!response.ok || !response.body) { + await response.body?.cancel().catch(() => {}) + throw new RunReplayUnavailableError(`The run replay failed with status ${response.status}`) + } + return response.body +} + +/** Every reason the worker may end a replay with; a reason added to the contract fails here. */ +const REPLAY_END_REASONS = { + parked: true, + cap: true, + stalled: true, +} as const satisfies Record + +const StreamReplayEndSchema = z.object({ + kind: z.literal('replay_end'), + reason: z + .string() + .refine((reason): reason is StreamReplayEnd['reason'] => + Object.hasOwn(REPLAY_END_REASONS, reason) + ), + textLength: z.number().int().nonnegative(), +}) satisfies z.ZodType + +/** + * The end a `replay_end` frame reports: its reason, `closed` for a reason this build + * does not know, or `null` for any other frame. It is worker-to-Sim control, never a + * stream event, whatever reason it carries. + */ +function replayEnd(value: unknown): RunReplayEnd | null { + if (!isRecordLike(value) || value.type !== MothershipStreamV1EventType.run) return null + const payload = toRecord(value.payload) + if (payload.kind !== 'replay_end') return null + const parsed = StreamReplayEndSchema.safeParse(payload) + if (parsed.success) return parsed.data.reason + logger.warn('Run replay ended with an unknown reason', { reason: payload.reason }) + return 'closed' +} + +/** + * Reads a replay leg, handing each stream event to `onEvent` in order. The leg's + * `replay_end` frame never reaches `onEvent`. Returning false from `onEvent` stops the + * read (the reader went away). + */ +async function readRunReplay( + body: ReadableStream, + signal: AbortSignal, + onEvent: (event: PersistedStreamEventEnvelope) => boolean +): Promise { + let end: RunReplayEnd = 'closed' + await processSSEStream( + body.getReader(), + signal, + (raw) => { + const control = replayEnd(raw) + if (control) { + end = control + return true + } + const parsed = parsePersistedStreamEventEnvelope(raw) + if (!parsed.ok) throw new FatalSseEventError(`Invalid run replay event: ${parsed.message}`) + if (!onEvent(parsed.event)) return true + if (parsed.event.type === MothershipStreamV1EventType.complete) { + end = 'complete' + return true + } + return undefined + }, + WORKER_STREAM_IDLE_TIMEOUT_MS + ) + return end +} + +export interface ForwardRunReplayOptions { + body: ReadableStream + streamId: string + signal: AbortSignal + /** Writes one event to the reader; false once the reader is gone. */ + write: (event: PersistedStreamEventEnvelope) => boolean + /** The run's current status, or null when it cannot be read. */ + readRunStatus: () => Promise + isClosed: () => boolean + /** When the reader's response must end regardless. */ + deadlineAt: number +} + +/** + * Forwards a replay leg to one reader under that response's own cursors, starting at + * 1, then decides how long the response stays open. A terminal or the worker's cap + * ends it at once: the cap came after minutes of progress. Otherwise it holds, ending + * as soon as the run reaches a terminal or, after a park, resumes; a stall or a cut + * connection holds at least {@link REPLAY_MIN_RESPONSE_MS}, so a reader re-attaches, + * and replays the whole log again, at most that often. + */ +export async function forwardRunReplay(options: ForwardRunReplayOptions): Promise { + const { body, streamId, signal, write, readRunStatus, isClosed, deadlineAt } = options + const startedAt = Date.now() + let seq = 0 + const end = await readRunReplay(body, signal, (event) => { + seq += 1 + return write( + toReplayEnvelope({ + ...event, + seq, + stream: { ...event.stream, streamId, cursor: String(seq) }, + }) + ) + }).catch((error: unknown) => { + if (error instanceof FatalSseEventError) throw error + logger.warn('Run replay connection ended early', { streamId, error: getErrorMessage(error) }) + return 'closed' as const + }) + if (end === 'complete' || end === 'cap') return end + // Sim may mark the park a moment after the worker ends on it, so a park only counts as + // resumed once Sim was seen parked; until then it holds like any other end. + let sawParked = false + while (!isClosed() && Date.now() < deadlineAt) { + const status = await readRunStatus() + if (isTerminalStreamStatus(status)) break + const parked = end === 'parked' && status === PARKED_RUN_STATUS + if (end === 'parked' && sawParked && !parked) break + sawParked ||= parked + const remaining = REPLAY_MIN_RESPONSE_MS - (Date.now() - startedAt) + if (!parked && remaining <= 0) break + await sleep(parked ? REPLAY_HOLD_POLL_MS : Math.min(REPLAY_HOLD_POLL_MS, remaining)) + } + return end +} diff --git a/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts b/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts new file mode 100644 index 00000000000..6b714c19cb9 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts @@ -0,0 +1,349 @@ +/** + * Recovery of a Chat run whose Sim controller died after its replay ring lost its + * head, against real Redis and PostgreSQL through the production reconnect route and + * chat lifecycle. A local HTTP server stands in for the worker: it answers the new + * controller's re-attach with the run's whole response, as its duplicate-send path does, + * and answers the tool resume that follows a re-handed call. + */ +import { authMock, authMockFns } from '@sim/testing/mocks/auth.mock' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' + +const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { + const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') + const { createServer } = await import('node:http') + const worker = { + requests: [] as Array<{ path: string; body: Record }>, + /** SSE frames per worker path, set by each test once its ids are known. */ + replies: {} as Record, + } + const server = createServer(async (request, response) => { + let body = '' + for await (const chunk of request) body += chunk + const frames = request.url ? worker.replies[request.url] : undefined + if (!request.url || !frames) { + response.writeHead(404, { 'content-type': 'application/json' }) + response.end(JSON.stringify({ error: 'Run not found' })) + return + } + worker.requests.push({ path: request.url, body: JSON.parse(body) }) + response.writeHead(200, { 'content-type': 'text/event-stream' }) + for (const frame of frames) response.write(`data: ${JSON.stringify(frame)}\n\n`) + response.end('data: [DONE]\n\n') + }) + await new Promise((resolve) => server.listen(0, '127.0.0.1', resolve)) + const { port } = server.address() as { port: number } + const url = readTestRedisUrl() + const inheritedEnv = { + REDIS_URL: process.env.REDIS_URL, + SIM_AGENT_API_URL: process.env.SIM_AGENT_API_URL, + } + process.env.REDIS_URL = url + process.env.SIM_AGENT_API_URL = `http://127.0.0.1:${port}` + return { redisUrl: url, inheritedEnv, worker: Object.assign(worker, { server }) } +}) + +vi.mock('@/lib/auth', () => authMock) + +import { db } from '@sim/db' +import { + copilotAsyncToolCalls, + copilotChats, + copilotMessages, + copilotRuns, + permissions, + user, + workspace, +} from '@sim/db/schema' +import { sleep } from '@sim/utils/helpers' +import { generateId } from '@sim/utils/id' +import { toArray, toRecord } from '@sim/utils/object' +import { eq } from 'drizzle-orm' +import { NextRequest } from 'next/server' +import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis' +import { + claimSimToolExecution, + completeOwnedSimToolCall, + upsertAsyncToolCall, +} from '@/lib/mothership/async-runs/repository' +import { + createProviderToolCallIdentity, + scopeProviderToolCallId, +} from '@/lib/mothership/request/go/tool-call-identity' +import { appendEvents } from '@/lib/mothership/request/session/buffer' +import { chatStreamLockKey } from '@/lib/mothership/request/session/controller-lease' +import { createEvent } from '@/lib/mothership/request/session/event' +import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' + +const userId = generateId() +const workspaceId = generateId() +const chatIds: string[] = [] +const FULL_TEXT = 'part 1 part 2 part 3 part 4 ' + +/** A live run with a dead controller whose ring kept only seqs 3–4 of its four events. */ +async function orphanedRunWithTrimmedRing() { + const chatId = generateId() + chatIds.push(chatId) + const streamId = generateId() + const runId = generateId() + const request = { + message: 'Summarize the logs', + userId, + messageId: streamId, + chatId, + workspaceId, + } + await db + .insert(copilotChats) + .values({ id: chatId, userId, workspaceId, type: 'mothership', conversationId: streamId }) + await db.insert(copilotMessages).values({ + chatId, + messageId: streamId, + role: 'user', + streamId, + seq: 0, + content: { id: streamId, role: 'user', content: request.message }, + }) + await db.insert(copilotRuns).values({ + id: runId, + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + toolExecutionVersion: 2, + requestContext: { + requestId: generateId(), + controllerToken: `dead\n${generateId()}`, + recovery: { + kind: 'interactive_stream', + request, + goRoute: '/api/mothership', + clientToolPickupExpected: false, + }, + }, + }) + await appendEvents( + [1, 2, 3, 4].map((seq) => + createEvent({ + streamId, + cursor: String(seq), + seq, + requestId: generateId(), + type: 'text', + payload: { channel: 'assistant', text: `part ${seq} ` }, + }) + ), + { streamId } + ) + await getRedisClient()!.set(`mothership_stream:${streamId}:seq`, '4') + await getRedisClient()!.zremrangebyrank(`mothership_stream:${streamId}:events`, 0, 1) + const frame = (seq: number, type: string, payload: unknown) => ({ + v: 1, + type, + seq, + ts: new Date().toISOString(), + stream: { streamId, chatId }, + payload, + }) + return { chatId, streamId, runId, frame } +} + +/** Opens a reconnect, which recovers the run, and waits for the recovered turn to finish. */ +async function recoverAndFinish(streamId: string, runId: string) { + const reconnect = await streamGET( + new NextRequest(`http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=4`), + { params: Promise.resolve({}) } + ) + await reconnect.body?.cancel() + let status: string | undefined + for (let attempt = 0; attempt < 150 && status !== 'complete'; attempt++) { + await sleep(100) + const [run] = await db + .select({ status: copilotRuns.status }) + .from(copilotRuns) + .where(eq(copilotRuns.id, runId)) + status = run?.status + } + return status +} + +async function assistantMessages(chatId: string) { + const rows = await db + .select({ content: copilotMessages.content }) + .from(copilotMessages) + .where(eq(copilotMessages.chatId, chatId)) + return rows.map((row) => toRecord(row.content)).filter((message) => message.role === 'assistant') +} + +/** Runs whether or not the suite does, so a skipped suite never leaks the worker or env. */ +afterAll(async () => { + await new Promise((resolve) => worker.server.close(() => resolve())) + for (const [key, value] of Object.entries(inheritedEnv)) { + if (value === undefined) delete process.env[key] + else process.env[key] = value + } +}) + +describe.runIf(Boolean(redisUrl))('recovering a run whose ring lost its head', () => { + beforeAll(async () => { + const now = new Date() + await db.insert(user).values({ + id: userId, + name: 'Stream recovery fixture', + email: `${userId}@stream-recovery.test`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(workspace).values({ + id: workspaceId, + name: 'Stream recovery fixture', + ownerId: userId, + billedAccountUserId: userId, + }) + await db.insert(permissions).values({ + id: generateId(), + userId, + entityType: 'workspace', + entityId: workspaceId, + permissionType: 'admin', + }) + authMockFns.mockGetSession.mockResolvedValue({ + user: { id: userId }, + session: { id: generateId() }, + }) + }) + + afterAll(async () => { + for (const chatId of chatIds) { + await db.delete(copilotMessages).where(eq(copilotMessages.chatId, chatId)) + await db.delete(copilotRuns).where(eq(copilotRuns.chatId, chatId)) + await db.delete(copilotChats).where(eq(copilotChats.id, chatId)) + } + await db.delete(permissions).where(eq(permissions.userId, userId)) + await db.delete(workspace).where(eq(workspace.id, workspaceId)) + await db.delete(user).where(eq(user.id, userId)) + await closeRedisConnection() + }) + + it('recovers from an empty context and persists the whole turn once', async () => { + const { chatId, streamId, runId, frame } = await orphanedRunWithTrimmedRing() + worker.requests.length = 0 + worker.replies = { + '/api/mothership': [ + frame(1, 'session', { kind: 'start' }), + frame(2, 'text', { channel: 'assistant', text: FULL_TEXT, textOffset: 0 }), + frame(3, 'complete', { status: 'complete', textLength: FULL_TEXT.length }), + ], + } + + expect(await recoverAndFinish(streamId, runId)).toBe('complete') + + const [chat] = await db + .select({ conversationId: copilotChats.conversationId }) + .from(copilotChats) + .where(eq(copilotChats.id, chatId)) + expect(chat.conversationId).toBeNull() + const assistant = await assistantMessages(chatId) + expect(assistant).toHaveLength(1) + expect(String(assistant[0].content).trim()).toBe(FULL_TEXT.trim()) + // One re-attach under the original identity and an empty receipt: the worker re-sends + // the response rather than running, or billing, the turn again. + expect(worker.requests.map((request) => request.path)).toEqual(['/api/mothership']) + expect(worker.requests[0].body).toMatchObject({ messageId: streamId, receivedTextChars: 0 }) + expect(await getRedisClient()!.get(chatStreamLockKey(chatId))).toBeNull() + }) + + it('shows replayed tools once and never re-runs a re-handed call the dead controller ran', async () => { + const { chatId, streamId, runId, frame } = await orphanedRunWithTrimmedRing() + const simCallId = scopeProviderToolCallId('sim-call', createProviderToolCallIdentity(runId)) + const storedResult = { servers: [] } + await upsertAsyncToolCall({ + runId, + toolCallId: simCallId, + toolName: 'list_workspace_mcp_servers', + args: {}, + }) + expect( + await claimSimToolExecution({ toolCallId: simCallId, runId, userId, ownerToken: 'dead' }) + ).toEqual({ outcome: 'claimed' }) + await completeOwnedSimToolCall( + { toolCallId: simCallId, status: 'completed', result: storedResult }, + 'dead' + ) + const [ran] = await db + .select({ startedAt: copilotAsyncToolCalls.executionStartedAt }) + .from(copilotAsyncToolCalls) + .where(eq(copilotAsyncToolCalls.toolCallId, simCallId)) + worker.requests.length = 0 + worker.replies = { + '/api/mothership': [ + frame(1, 'session', { kind: 'start' }), + frame(2, 'text', { channel: 'assistant', text: FULL_TEXT, textOffset: 0 }), + frame(3, 'tool', { + phase: 'call', + toolCallId: 'go-call', + toolName: 'search_online', + executor: 'go', + mode: 'sync', + arguments: { query: 'logs' }, + replay: true, + }), + frame(4, 'tool', { + phase: 'result', + toolCallId: 'go-call', + toolName: 'search_online', + executor: 'go', + mode: 'sync', + success: true, + output: { results: [] }, + replay: true, + }), + frame(5, 'tool', { + phase: 'call', + toolCallId: 'sim-call', + toolName: 'list_workspace_mcp_servers', + executor: 'sim', + mode: 'async', + arguments: {}, + }), + frame(6, 'run', { + kind: 'checkpoint_pause', + checkpointId: generateId(), + executionId: generateId(), + runId: generateId(), + pendingToolCallIds: ['sim-call'], + }), + ], + '/api/tools/resume': [ + frame(7, 'text', { channel: 'assistant', text: 'done', textOffset: FULL_TEXT.length }), + frame(8, 'complete', { status: 'complete', textLength: FULL_TEXT.length + 4 }), + ], + } + + expect(await recoverAndFinish(streamId, runId)).toBe('complete') + + const [settled] = await db + .select({ + startedAt: copilotAsyncToolCalls.executionStartedAt, + result: copilotAsyncToolCalls.result, + }) + .from(copilotAsyncToolCalls) + .where(eq(copilotAsyncToolCalls.toolCallId, simCallId)) + expect(settled.startedAt).toEqual(ran.startedAt) + expect(settled.result).toEqual(storedResult) + const resumes = worker.requests.filter((request) => request.path === '/api/tools/resume') + expect(resumes).toHaveLength(1) + expect(toArray(resumes[0].body.results)).toEqual([ + expect.objectContaining({ callId: 'sim-call', success: true }), + ]) + + const assistant = await assistantMessages(chatId) + expect(assistant).toHaveLength(1) + const toolIds = toArray(assistant[0].contentBlocks) + .map((block) => toRecord(toRecord(block).toolCall).id) + .filter((id): id is string => typeof id === 'string') + expect(toolIds).toHaveLength(2) + expect(new Set(toolIds).size).toBe(2) + }) +})