From c4d4d35a926ab092b003078c1b48e4cfffe0732b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 02:50:41 -0700 Subject: [PATCH 01/12] chore(mothership): sync the worker protocol for the read-only run replay Adds StreamReplayRequest and StreamReplayEnd from the worker's contracts (bun run contracts:sync). --- apps/sim/lib/mothership/generated/protocol.ts | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) 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; From 748e7e7ae1c774ab0d9b12ff701bf29aa874cf74 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 02:50:43 -0700 Subject: [PATCH 02/12] fix(mothership): re-sync a reconnect the replay ring cannot serve from the worker log A reconnect whose cursor fell behind the ring, a fresh tab reading a ring that lost its head, and a cursor ahead of a ring whose numbering restarted now stream the run from the worker's read-only replay instead of ending the turn with replay_gap or replaying a partial response. - The reconnect route opens POST /api/streams/replay with no receipt and forwards it under its own cursors from 1, with x-mothership-stream-replay: log so the client rebuilds the turn from an empty response. A response on the replay is never handed back to the ring, which shares no position with the log. - A parked replay holds the response open until the run resumes; a capped, stalled, or cut replay ends it without a terminal so the client re-attaches. - A batch read the ring cannot serve returns no ring events. - A live tail whose ring restarts under it ends without a terminal, so the client re-attaches and is re-synced. - An unknown run keeps the replay_gap terminal; an unreachable worker answers 503 so the client retries. --- .../app/api/copilot/chat/stream/route.test.ts | 15 +- apps/sim/app/api/copilot/chat/stream/route.ts | 116 ++++++- .../home/hooks/use-chat.dom.test.tsx | 64 ++++ .../[workspaceId]/home/hooks/use-chat.ts | 14 +- apps/sim/lib/mothership/constants.ts | 7 + .../lib/mothership/request/session/index.ts | 14 +- .../mothership/request/session/recovery.ts | 141 +++++---- .../request/session/replay-gap.integration.ts | 296 ++++++++++++++++++ .../mothership/request/session/run-replay.ts | 102 ++++++ 9 files changed, 694 insertions(+), 75 deletions(-) create mode 100644 apps/sim/lib/mothership/request/session/replay-gap.integration.ts create mode 100644 apps/sim/lib/mothership/request/session/run-replay.ts 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..d3b2bb5954c 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,9 @@ vi.mock('@/lib/mothership/request/session', () => ({ status === 'complete' || status === 'error' || status === 'cancelled', readEvents, readFilePreviewSessions, - checkForReplayGap, + findReplayGap, + getLatestSeq: async () => null, + replayGapTerminal: async () => ({ gapDetected: true, envelopes: [] }), createEvent: (event: Record) => ({ stream: { streamId: event.streamId, @@ -82,7 +85,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 () => { diff --git a/apps/sim/app/api/copilot/chat/stream/route.ts b/apps/sim/app/api/copilot/chat/stream/route.ts index b1023a791b9..933a0740c37 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, @@ -24,15 +25,21 @@ import { import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1' import { readChatStream } from '@/lib/mothership/request/application/recover-stream' +import { FatalSseEventError } from '@/lib/mothership/request/go/parser' import { contextFromRequestHeaders } from '@/lib/mothership/request/go/propagation' import { getCopilotTracer, markSpanForError } from '@/lib/mothership/request/otel' import { - checkForReplayGap, createEvent, encodeSSEEnvelope, + findReplayGap, + getLatestSeq, isTerminalStreamStatus, + openRunReplay, + RunReplayUnavailableError, readEvents, readFilePreviewSessions, + readRunReplay, + replayGapTerminal, SSE_RESPONSE_HEADERS, } from '@/lib/mothership/request/session' import { toReplayEnvelope, toStreamBatchEvent } from '@/lib/mothership/request/session/types' @@ -229,7 +236,8 @@ async function handleResumeRequestBody({ if (batchMode) { const afterSeq = afterCursor || '0' - const [events, previewSessions] = await Promise.all([ + const [gap, events, previewSessions] = await Promise.all([ + findReplayGap(streamId, afterSeq, extractRunRequestId(run)), readEvents(streamId, afterSeq), readFilePreviewSessions(streamId).catch((error) => { logger.warn('Failed to read preview sessions for stream batch', { @@ -239,7 +247,8 @@ async function handleResumeRequestBody({ return [] }), ]) - const batchEvents = events.map(toStreamBatchEvent) + /* The ring cannot serve this reader; the live tail re-syncs it from the worker log. */ + const batchEvents = gap ? [] : events.map(toStreamBatchEvent) logger.info('[Resume] Batch response', { streamId, afterCursor: afterSeq, @@ -266,6 +275,37 @@ 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 gap = await findReplayGap(streamId, afterCursor || '0', extractRunRequestId(run)) + 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 (gap && 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 @@ -380,12 +420,62 @@ async function handleResumeRequestBody({ } } + /** + * Forwards the worker's replay under this response's own cursors, starting at 1. + * A parked run holds the response open until it resumes, so the client does not + * replay the whole log again for every poll; any other end short of the terminal + * (the worker's cap, a stalled run, a cut connection) ends it without one and the + * client re-attaches, reaching recovery through the reconnect route. + */ + const streamRunReplay = async (body: ReadableStream) => { + let seq = 0 + 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 readRunReplay(body, replaySignal, (event) => { + seq += 1 + const envelope = toReplayEnvelope({ + ...event, + seq, + stream: { ...event.stream, streamId, cursor: String(seq) }, + }) + if (!enqueueEvent(envelope)) return false + totalEventsFlushed += 1 + cursor = String(seq) + if (envelope.type === MothershipStreamV1EventType.complete) sawTerminalEvent = true + return true + }).catch((error: unknown) => { + if (error instanceof FatalSseEventError) throw error + logger.warn('Run replay connection ended early', { + streamId, + error: getErrorMessage(error), + }) + return 'closed' as const + }) + logger.info('[Resume] Run replay ended', { streamId, end, eventCount: seq }) + while (end === 'parked' && !controllerClosed && Date.now() - startTime < MAX_STREAM_MS) { + const current = await readRun().catch(() => null) + if (current?.status !== 'paused_waiting_for_tool') break + await sleep(POLL_INTERVAL_MAX_MS) + } + } finally { + clearInterval(keepalive) + replayAbort.abort() + } + } + try { enqueueComment('accepted') - const gap = await checkForReplayGap(streamId, afterCursor, currentRequestId) + if (replayBody) { + await streamRunReplay(replayBody) + return + } if (gap) { - for (const envelope of gap.envelopes) { + const terminal = await replayGapTerminal(streamId, gap, currentRequestId) + for (const envelope of terminal.envelopes) { if (!enqueueEvent(envelope)) { break } @@ -428,6 +518,16 @@ async function handleResumeRequestBody({ pollDelayMs = flushed > 0 ? POLL_INTERVAL_MS : Math.min(pollDelayMs * 2, POLL_INTERVAL_MAX_MS) + /* A ring whose numbering restarted after it expired never passes this cursor + again; end without a terminal so the client re-attaches and is re-synced. */ + if (flushed === 0 && !sawTerminalEvent) { + const latestSeq = await getLatestSeq(streamId) + if (latestSeq !== null && Number(cursor) > latestSeq) { + logger.warn('Replay ring restarted under a live tail', { streamId, cursor, latestSeq }) + break + } + } + if (controllerClosed) { break } @@ -488,5 +588,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/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.dom.test.tsx index cebc1911396..7fec1809866 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,69 @@ describe('useChat remount send recovery', () => { } }) + it('rebuilds the turn from an empty response when a reconnect is re-synced from the log', async () => { + vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }) + try { + let tails = 0 + 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) + 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.') + } 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..cc03044276d 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' @@ -2578,6 +2582,14 @@ export function useChat( return { error: false, aborted: true } } + /* The ring could not serve this cursor, so the turn is re-sent from the + worker's log with cursors restarting at 1: rebuild it from empty. */ + if (sseRes.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER) === 'log') { + const reset = applyReconnectReplaySelection(streamId, '0') + latestCursor = reset.afterCursor + preserveNextReplayState = reset.preserveExistingState + } + setTransportStreaming() const liveResult = await processSSEStreamRef.current( 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/request/session/index.ts b/apps/sim/lib/mothership/request/session/index.ts index aa9b22d1606..0f458f6f400 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 { + checkForReplayGap, + findReplayGap, + type ReplayGap, + type ReplayGapResult, + replayGapTerminal, +} from './recovery' +export { + openRunReplay, + type RunReplayEnd, + RunReplayUnavailableError, + readRunReplay, +} 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.ts b/apps/sim/lib/mothership/request/session/recovery.ts index bb0153c8a14..d0b34a4b4ee 100644 --- a/apps/sim/lib/mothership/request/session/recovery.ts +++ b/apps/sim/lib/mothership/request/session/recovery.ts @@ -18,18 +18,26 @@ export interface ReplayGapResult { envelopes: ReturnType[] } -export async function checkForReplayGap( +/** Where the replay ring stands relative to a reader it can no longer serve. */ +export interface ReplayGap { + requestedAfterSeq: number + oldestSeq: number + latestSeq: number +} + +/** + * Whether the ring can serve a reader from `afterCursor`. 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. + */ +export async function findReplayGap( streamId: string, afterCursor: string, requestId?: string -): Promise { +): 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 - } - return withCopilotSpan( TraceSpan.CopilotRecoveryCheckReplayGap, { @@ -38,70 +46,81 @@ export async function checkForReplayGap( ...(requestId ? { [TraceAttr.RequestId]: requestId } : {}), }, async (span) => { - const oldestSeq = await getOldestSeq(streamId) - const latestSeq = await getLatestSeq(streamId) + const [oldestSeq, latestSeq] = await Promise.all([ + getOldestSeq(streamId), + getLatestSeq(streamId), + ]) span.setAttributes({ [TraceAttr.CopilotRecoveryOldestSeq]: oldestSeq ?? -1, [TraceAttr.CopilotRecoveryLatestSeq]: latestSeq ?? -1, }) - - /* 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) + latestSeq === null || + latestSeq <= 0 || + oldestSeq === null || + (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], - } + 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, + requestedAfterSeq, + oldestAvailableSeq: oldestSeq, + latestSeq, + }) + span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.GapDetected) + return { requestedAfterSeq, oldestSeq, latestSeq } } ) } +/** Ends a reader's view with `replay_gap` when nothing can re-sync it. */ +export async function replayGapTerminal( + streamId: string, + gap: ReplayGap, + requestId?: string +): Promise { + const { latestSeq, oldestSeq, requestedAfterSeq } = gap + const resolvedRequestId = await resolveReplayGapRequestId(streamId, latestSeq, requestId) + 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] } +} + +export async function checkForReplayGap( + streamId: string, + afterCursor: string, + requestId?: string +): Promise { + const gap = await findReplayGap(streamId, afterCursor, requestId) + return gap ? replayGapTerminal(streamId, gap, requestId) : null +} + async function resolveReplayGapRequestId( streamId: string, latestSeq: number, 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..1dc7a265673 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts @@ -0,0 +1,296 @@ +/** + * 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) { + return streamGET( + new NextRequest( + `http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=${after}${batch ? '&batch=true' : ''}` + ), + { 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' }), + ] +} + +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() + 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 + } + }) + + 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) + + expect(frames.map((frame) => frame.type)).toEqual(['text']) + expect(Date.now() - startedAt).toBeGreaterThanOrEqual(1_500) + expect(worker.requests).toHaveLength(1) + }) + + 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('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('keeps the replay_gap terminal when the worker has no such run', async () => { + const { streamId } = await liveRunWithTrimmedRing() + worker.reply.status = 404 + + 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.ts b/apps/sim/lib/mothership/request/session/run-replay.ts new file mode 100644 index 00000000000..cbf7fdb6b12 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -0,0 +1,102 @@ +import { isRecordLike } from '@sim/utils/object' +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 { + type PersistedStreamEventEnvelope, + parsePersistedStreamEventEnvelope, +} from '@/lib/mothership/request/session/contract' +import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url' + +const REPLAY_PATH = '/api/streams/replay' + +/** 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 knows no such run for this chat and user. + */ +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 }) + let response: Response + try { + response = await fetchGo(`${baseUrl}${REPLAY_PATH}`, { + method: 'POST', + headers: mothershipRequestHeaders(), + body: JSON.stringify(StreamReplayRequest.parse({ streamId, chatId, userId })), + 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 }) + } + if (response.status === 404) { + 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 +} + +function replayEndReason(value: unknown): StreamReplayEnd['reason'] | null { + if (!isRecordLike(value) || value.type !== MothershipStreamV1EventType.run) return null + const payload = value.payload + if (!isRecordLike(payload) || payload.kind !== 'replay_end') return null + return payload.reason === 'parked' || payload.reason === 'cap' || payload.reason === 'stalled' + ? payload.reason + : null +} + +/** + * Reads a replay leg, handing each stream event to `onEvent` in order. The leg's + * `replay_end` frame is worker-to-Sim control and never reaches `onEvent`. Returning + * false from `onEvent` stops the read (the reader went away). + */ +export async function readRunReplay( + body: ReadableStream, + signal: AbortSignal, + onEvent: (event: PersistedStreamEventEnvelope) => boolean +): Promise { + let end: RunReplayEnd = 'closed' + await processSSEStream(body.getReader(), signal, (raw) => { + const reason = replayEndReason(raw) + if (reason) { + end = reason + 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 + }) + return end +} From 74185d19ae13f79eaa358ef406fcac2ae1638418 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 01:38:05 -0700 Subject: [PATCH 03/12] fix(mothership): trim the replay ring by bytes so a long run slides instead of refusing The append script pruned the oldest events by count only, so any stream averaging more than ~335 B per event reached the 32 MiB owner budget before 100k events and the refusal ended the turn. The ring now also trims its oldest events until the retained bytes fit three quarters of the owner ceiling, refunding exactly what it drops in the same script. A byte trim never drops a member the write adds, so an inflated counter still refuses rather than silently discarding the new frame. --- apps/sim/lib/core/redis/byte-budget.server.ts | 12 +-- .../mothership/request/session/buffer.test.ts | 8 +- .../lib/mothership/request/session/buffer.ts | 87 ++++++++++++++----- .../session/replay-budget.integration.ts | 57 ++++++++++++ 4 files changed, 133 insertions(+), 31 deletions(-) diff --git a/apps/sim/lib/core/redis/byte-budget.server.ts b/apps/sim/lib/core/redis/byte-budget.server.ts index 5728d69960a..bcc66e9d5ed 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. * @@ -79,8 +81,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/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..bfde6eef063 100644 --- a/apps/sim/lib/mothership/request/session/buffer.ts +++ b/apps/sim/lib/mothership/request/session/buffer.ts @@ -24,6 +24,13 @@ 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 type RedisOperationMetadata = { operation: string @@ -204,25 +211,39 @@ 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 never + * drops a member this write adds; only the count limit can, when a replay + * reintroduces a member below the retained ring. + * * 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 +258,46 @@ 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 +while prune_count < count_excess or pruned_bytes < byte_excess 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 + local existing = nil + if page_index <= #page then + existing = {member = page[page_index], score = tonumber(page[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 incoming = new_members[next_new] + if incoming and (not existing or ranks_before(incoming, existing)) then + if 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 +306,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 +358,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 +423,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/replay-budget.integration.ts b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts index a55706e2468..5b962d2df28 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -667,6 +667,63 @@ 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() + ) + const inRange = await reconnect(oldestSeq - 1) + expect(inRange.map((frame) => frame.seq)).toEqual(seqs) + 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) + }, 180_000) + 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 () => { From b59d4b611c48213fef28c7eec2ae32c8795d59d6 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 02:07:09 -0700 Subject: [PATCH 04/12] fix(mothership): never rebuild or paint a turn from a ring that lost its head A byte trim advances the ring past seq 1, and three callers read it from seq 0 assuming the head was there: stream recovery rebuilt a controller's context from the tail, and both chat snapshot routes painted a truncated turn. One predicate, startsAtReplayHead, now guards them and the reconnect route's gap check: - recovery refuses with StreamReplayHeadTrimmedError instead of persisting a truncated turn; - the snapshot routes skip the snapshot, so the client reconnects; - the reconnect route re-syncs the view from the worker's log (stream) or serves no tail events (batch). When recovery was refused, a parked or stalled replay ends the view with recovery_unavailable instead of re-attaching forever. The append script also trims a replayed member that lands below the ring for bytes, keeping the tail contiguous, and caps byte trimming at 4096 members per append so an oversized ring catches up over several appends. The budget docs now say the user counter bounds bytes held, not bytes written per hour. --- apps/sim/app/api/copilot/chat/queries.ts | 22 +- .../app/api/copilot/chat/stream/route.test.ts | 1 + apps/sim/app/api/copilot/chat/stream/route.ts | 45 ++- .../api/mothership/chats/[chatId]/route.ts | 22 +- apps/sim/lib/core/redis/byte-budget.server.ts | 7 + .../request/application/recover-stream.ts | 16 ++ .../lib/mothership/request/session/buffer.ts | 18 +- .../mothership/request/session/recovery.ts | 12 +- .../session/replay-budget.integration.ts | 262 +++++++++++++++++- 9 files changed, 370 insertions(+), 35 deletions(-) diff --git a/apps/sim/app/api/copilot/chat/queries.ts b/apps/sim/app/api/copilot/chat/queries.ts index 105bcac4011..6608d5d9e0f 100644 --- a/apps/sim/app/api/copilot/chat/queries.ts +++ b/apps/sim/app/api/copilot/chat/queries.ts @@ -18,6 +18,7 @@ import { } from '@/lib/mothership/request/http' import { readFilePreviewSessions } from '@/lib/mothership/request/session' import { readEvents } from '@/lib/mothership/request/session/buffer' +import { startsAtReplayHead } from '@/lib/mothership/request/session/recovery' import { toStreamBatchEvent } from '@/lib/mothership/request/session/types' import { assertActiveWorkspaceAccess, @@ -114,15 +115,18 @@ export async function GET(req: NextRequest) { }), ]) - streamSnapshot = { - events: events.map(toStreamBatchEvent), - previewSessions, - status: - typeof run?.status === 'string' - ? run.status - : events.length > 0 - ? 'active' - : 'unknown', + // A trimmed head would paint a truncated turn; the resume route reports the gap. + if (startsAtReplayHead(events[0]?.seq)) { + streamSnapshot = { + events: events.map(toStreamBatchEvent), + previewSessions, + status: + typeof run?.status === 'string' + ? run.status + : events.length > 0 + ? 'active' + : 'unknown', + } } } catch (error) { logger.warn('Failed to load copilot chat stream snapshot', { 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 d3b2bb5954c..c1ce582de94 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.test.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.test.ts @@ -27,6 +27,7 @@ const { getLatestRunForStream, readEvents, readFilePreviewSessions, findReplayGa vi.mock('@/lib/mothership/request/application/recover-stream', () => ({ readChatStream: { execute: getLatestRunForStream }, + StreamReplayHeadTrimmedError: class extends Error {}, })) vi.mock('@/lib/mothership/request/session', () => ({ diff --git a/apps/sim/app/api/copilot/chat/stream/route.ts b/apps/sim/app/api/copilot/chat/stream/route.ts index 933a0740c37..1e0e5d1097f 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.ts @@ -24,7 +24,10 @@ import { } from '@/lib/mothership/generated/trace-attribute-values-v1' import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1' -import { readChatStream } from '@/lib/mothership/request/application/recover-stream' +import { + readChatStream, + StreamReplayHeadTrimmedError, +} from '@/lib/mothership/request/application/recover-stream' import { FatalSseEventError } from '@/lib/mothership/request/go/parser' import { contextFromRequestHeaders } from '@/lib/mothership/request/go/propagation' import { getCopilotTracer, markSpanForError } from '@/lib/mothership/request/otel' @@ -214,11 +217,22 @@ async function handleResumeRequestBody({ rootSpan: Span rootContext: Context }) { + /** + * Set when the run needs a new controller but recovery refused because the ring lost + * its head. Its view is still re-synced from the worker log; see `streamRunReplay`. + */ + let recoveryRefused = false const readRun = () => - readChatStream.execute({ - principal, - input: { streamId }, - }) + readChatStream + .execute({ + principal, + input: { streamId }, + }) + .catch((error: unknown) => { + if (!(error instanceof StreamReplayHeadTrimmedError)) throw error + recoveryRefused = true + return error.run + }) const run = await readRun() logger.info('[Resume] Stream lookup', { streamId, @@ -425,7 +439,8 @@ async function handleResumeRequestBody({ * A parked run holds the response open until it resumes, so the client does not * replay the whole log again for every poll; any other end short of the terminal * (the worker's cap, a stalled run, a cut connection) ends it without one and the - * client re-attaches, reaching recovery through the reconnect route. + * client re-attaches, reaching recovery through the reconnect route. When that + * recovery was refused, a park or stall ends the view with an error instead. */ const streamRunReplay = async (body: ReadableStream) => { let seq = 0 @@ -455,11 +470,25 @@ async function handleResumeRequestBody({ return 'closed' as const }) logger.info('[Resume] Run replay ended', { streamId, end, eventCount: seq }) - while (end === 'parked' && !controllerClosed && Date.now() - startTime < MAX_STREAM_MS) { + while ( + end === 'parked' && + !recoveryRefused && + !controllerClosed && + Date.now() - startTime < MAX_STREAM_MS + ) { const current = await readRun().catch(() => null) if (current?.status !== 'paused_waiting_for_tool') break await sleep(POLL_INTERVAL_MAX_MS) } + /* No controller drives the run and none can be rebuilt from a trimmed ring, so a + park or stall is final for this view; re-attaching would only replay it again. */ + if (recoveryRefused && (end === 'parked' || end === 'stalled')) { + emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, { + message: 'This response can no longer continue. Send a message to pick up from here.', + code: 'recovery_unavailable', + reason: 'recovery_unavailable', + }) + } } finally { clearInterval(keepalive) replayAbort.abort() @@ -500,6 +529,8 @@ async function handleResumeRequestBody({ }) return null }) + /* The ring lost its head under this tail; the re-attach re-syncs from the log. */ + if (recoveryRefused) break if (!currentRun) { emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, { message: 'The stream could not be recovered because its run metadata is unavailable.', diff --git a/apps/sim/app/api/mothership/chats/[chatId]/route.ts b/apps/sim/app/api/mothership/chats/[chatId]/route.ts index f39b55f32af..3825c38aa67 100644 --- a/apps/sim/app/api/mothership/chats/[chatId]/route.ts +++ b/apps/sim/app/api/mothership/chats/[chatId]/route.ts @@ -28,6 +28,7 @@ import { 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' import { captureServerEvent } from '@/lib/posthog/server' @@ -89,15 +90,18 @@ export const GET = withRouteHandler( return null }) - liveTurnSnapshot = { - events: events.map(toStreamBatchEvent), - previewSessions, - status: - typeof run?.status === 'string' - ? run.status - : events.length > 0 - ? 'active' - : 'unknown', + // A trimmed head would paint a truncated turn; the resume route reports the gap. + if (startsAtReplayHead(events[0]?.seq)) { + liveTurnSnapshot = { + events: events.map(toStreamBatchEvent), + previewSessions, + status: + typeof run?.status === 'string' + ? run.status + : events.length > 0 + ? 'active' + : 'unknown', + } } } catch (error) { logger.warn('Failed to read stream snapshot for mothership chat', { diff --git a/apps/sim/lib/core/redis/byte-budget.server.ts b/apps/sim/lib/core/redis/byte-budget.server.ts index bcc66e9d5ed..e6124bbd00d 100644 --- a/apps/sim/lib/core/redis/byte-budget.server.ts +++ b/apps/sim/lib/core/redis/byte-budget.server.ts @@ -69,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 diff --git a/apps/sim/lib/mothership/request/application/recover-stream.ts b/apps/sim/lib/mothership/request/application/recover-stream.ts index 441750260e8..9d9fbde594c 100644 --- a/apps/sim/lib/mothership/request/application/recover-stream.ts +++ b/apps/sim/lib/mothership/request/application/recover-stream.ts @@ -27,10 +27,25 @@ import { import { 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') +type RecoverableRun = NonNullable>> + +/** + * The run's replay ring no longer holds its first events, so its context cannot be + * rebuilt from Redis. The reconnect reports `replay_gap` instead of recovering from the + * tail, which would persist a truncated turn. + */ +export class StreamReplayHeadTrimmedError extends Error { + constructor(readonly run: RecoverableRun) { + super(`Replay head trimmed for stream ${run.streamId}`) + this.name = 'StreamReplayHeadTrimmedError' + } +} + export const readChatStream = defineAuthorizedChatUseCase({ operation: defineWorkspaceOperation({ id: 'mothership.runs.reconnect', @@ -120,6 +135,7 @@ export const readChatStream = defineAuthorizedChatUseCase({ ]) if (workspaceId && !userPermission) throw new OrchestrationError('forbidden', 'Workspace access revoked') + if (!startsAtReplayHead(events[0]?.seq)) throw new StreamReplayHeadTrimmedError(run) const requestId = typeof saved?.requestId === 'string' ? saved.requestId : generateId() const completion = { chatId, diff --git a/apps/sim/lib/mothership/request/session/buffer.ts b/apps/sim/lib/mothership/request/session/buffer.ts index bfde6eef063..9e0f1559f80 100644 --- a/apps/sim/lib/mothership/request/session/buffer.ts +++ b/apps/sim/lib/mothership/request/session/buffer.ts @@ -31,6 +31,12 @@ const RETRY_DELAYS_MS = [0, 50, 150] as const 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 @@ -214,9 +220,10 @@ export async function scheduleBufferCleanup( * 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 never - * drops a member this write adds; only the count limit can, when a replay - * reintroduces a member below the retained ring. + * 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 @@ -276,7 +283,8 @@ local next_new = 1 local page = {} local page_index = 1 local fetched = 0 -while prune_count < count_excess or pruned_bytes < byte_excess do +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 @@ -288,7 +296,7 @@ while prune_count < count_excess or pruned_bytes < byte_excess do end local incoming = new_members[next_new] if incoming and (not existing or ranks_before(incoming, existing)) then - if prune_count >= count_excess then break end + 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 diff --git a/apps/sim/lib/mothership/request/session/recovery.ts b/apps/sim/lib/mothership/request/session/recovery.ts index d0b34a4b4ee..6d9adee79d8 100644 --- a/apps/sim/lib/mothership/request/session/recovery.ts +++ b/apps/sim/lib/mothership/request/session/recovery.ts @@ -25,6 +25,16 @@ export interface ReplayGap { 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 +} + /** * Whether the ring can serve a reader from `afterCursor`. It cannot once it has lost * its head: the events before its oldest are gone, and a cursor that was served from @@ -58,7 +68,7 @@ export async function findReplayGap( latestSeq === null || latestSeq <= 0 || oldestSeq === null || - (oldestSeq <= 1 && requestedAfterSeq <= latestSeq) + (startsAtReplayHead(oldestSeq) && requestedAfterSeq <= latestSeq) ) { span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.InRange) return null 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 5b962d2df28..203e0d1d409 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,6 +57,7 @@ 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[], @@ -105,13 +123,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 +142,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 +364,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() @@ -708,8 +794,16 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge ) ).text() ) + /* A ring without its head is re-synced from the worker log whatever the cursor; + this worker does not know the run, so the replay_gap terminal applies. */ + worker.replayRequests.length = 0 const inRange = await reconnect(oldestSeq - 1) - expect(inRange.map((frame) => frame.seq)).toEqual(seqs) + 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( [ @@ -722,8 +816,168 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge 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'])( + 'shows the full turn from the worker log without recovering from the tail (%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, frame.payload.code ?? frame.payload.status ?? null]) + ).toEqual([ + ['text', null], + ['error', 'recovery_unavailable'], + ['complete', 'error'], + ]) + expect(frames[0].payload.text).toBe('part 1 part 2 part 3 part 4 ') + expect(worker.runs).toEqual([]) + expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + } 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).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 () => { From 88700b87f0cb7bcbec3a11b0403f7400f9f95411 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 03:54:35 -0700 Subject: [PATCH 05/12] fix(mothership): recover a headless ring from an empty context and keep log readers on the log - Recovery no longer refuses a run whose ring lost its head, which orphaned long runs after a Sim deploy. It treats that ring like an expired one: the new controller starts from an empty context at the ring's latest seq, and re-attaches with an empty receipt. The worker then re-sends the whole response and re-hands its parked calls. Usage stays with the worker's per-run settlement; a re-attach under the same message identity is never a second run. - A client re-synced from the log sends source=log from then on, so the ring never serves its log cursors, even after it restarts and grows past them. - A live tail ends without a terminal as soon as its ring loses its head or restarts, so it re-attaches and re-syncs instead of reading re-sent text. - A replay that ends short of the terminal and cap holds its response at least 10 s (longer while parked), so a stalled run is not replayed every second. - replay_end is parsed with a schema tied to the protocol's reasons; an unknown reason ends the replay instead of passing as a run event. - The replay forwarder moves into session/run-replay.ts, the chat snapshot reader is shared by both chat routes, and checkForReplayGap is removed. --- apps/sim/app/api/copilot/chat/queries.ts | 49 +--- .../app/api/copilot/chat/stream/route.test.ts | 4 +- apps/sim/app/api/copilot/chat/stream/route.ts | 138 ++++------- .../api/mothership/chats/[chatId]/route.ts | 50 +--- .../home/hooks/use-chat.dom.test.tsx | 10 +- .../[workspaceId]/home/hooks/use-chat.ts | 19 +- apps/sim/lib/api/contracts/copilot.ts | 2 + .../sim/lib/api/contracts/mothership-chats.ts | 1 + .../lib/mothership/chat/live-turn-snapshot.ts | 51 ++++ .../request/application/recover-stream.ts | 32 ++- .../request/session/buffer-ttl.integration.ts | 4 +- .../lib/mothership/request/session/index.ts | 8 +- .../request/session/recovery.test.ts | 8 +- .../mothership/request/session/recovery.ts | 83 +++---- .../session/replay-budget.integration.ts | 36 +-- .../request/session/replay-gap.integration.ts | 91 +++++++- .../mothership/request/session/run-replay.ts | 107 +++++++-- .../session/stream-recovery.integration.ts | 220 ++++++++++++++++++ 18 files changed, 628 insertions(+), 285 deletions(-) create mode 100644 apps/sim/lib/mothership/chat/live-turn-snapshot.ts create mode 100644 apps/sim/lib/mothership/request/session/stream-recovery.integration.ts diff --git a/apps/sim/app/api/copilot/chat/queries.ts b/apps/sim/app/api/copilot/chat/queries.ts index 6608d5d9e0f..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,10 +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 { startsAtReplayHead } from '@/lib/mothership/request/session/recovery' -import { toStreamBatchEvent } from '@/lib/mothership/request/session/types' import { assertActiveWorkspaceAccess, isWorkspaceAccessDeniedError, @@ -88,46 +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 - }), - ]) - - // A trimmed head would paint a truncated turn; the resume route reports the gap. - if (startsAtReplayHead(events[0]?.seq)) { - 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 c1ce582de94..5726ce344d1 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.test.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.test.ts @@ -27,7 +27,6 @@ const { getLatestRunForStream, readEvents, readFilePreviewSessions, findReplayGa vi.mock('@/lib/mothership/request/application/recover-stream', () => ({ readChatStream: { execute: getLatestRunForStream }, - StreamReplayHeadTrimmedError: class extends Error {}, })) vi.mock('@/lib/mothership/request/session', () => ({ @@ -36,7 +35,8 @@ vi.mock('@/lib/mothership/request/session', () => ({ readEvents, readFilePreviewSessions, findReplayGap, - getLatestSeq: async () => null, + readRingPosition: async () => ({ requestedAfterSeq: 0, oldestSeq: 0, latestSeq: 0 }), + ringCanServe: () => true, replayGapTerminal: async () => ({ gapDetected: true, envelopes: [] }), createEvent: (event: Record) => ({ stream: { diff --git a/apps/sim/app/api/copilot/chat/stream/route.ts b/apps/sim/app/api/copilot/chat/stream/route.ts index 1e0e5d1097f..b57ccf1870a 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.ts @@ -24,25 +24,22 @@ import { } from '@/lib/mothership/generated/trace-attribute-values-v1' import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1' -import { - readChatStream, - StreamReplayHeadTrimmedError, -} from '@/lib/mothership/request/application/recover-stream' -import { FatalSseEventError } from '@/lib/mothership/request/go/parser' +import { readChatStream } from '@/lib/mothership/request/application/recover-stream' import { contextFromRequestHeaders } from '@/lib/mothership/request/go/propagation' import { getCopilotTracer, markSpanForError } from '@/lib/mothership/request/otel' import { createEvent, encodeSSEEnvelope, findReplayGap, - getLatestSeq, + forwardRunReplay, isTerminalStreamStatus, openRunReplay, RunReplayUnavailableError, readEvents, readFilePreviewSessions, - readRunReplay, + readRingPosition, replayGapTerminal, + ringCanServe, SSE_RESPONSE_HEADERS, } from '@/lib/mothership/request/session' import { toReplayEnvelope, toStreamBatchEvent } from '@/lib/mothership/request/session/types' @@ -139,7 +136,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 }) @@ -183,6 +180,7 @@ export const GET = withRouteHandler(async (request: NextRequest) => { streamId, afterCursor, batchMode, + fromLog: source === 'log', principal, rootSpan, rootContext, @@ -205,6 +203,7 @@ async function handleResumeRequestBody({ streamId, afterCursor, batchMode, + fromLog, principal, rootSpan, rootContext, @@ -213,26 +212,17 @@ 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 }) { - /** - * Set when the run needs a new controller but recovery refused because the ring lost - * its head. Its view is still re-synced from the worker log; see `streamRunReplay`. - */ - let recoveryRefused = false const readRun = () => - readChatStream - .execute({ - principal, - input: { streamId }, - }) - .catch((error: unknown) => { - if (!(error instanceof StreamReplayHeadTrimmedError)) throw error - recoveryRefused = true - return error.run - }) + readChatStream.execute({ + principal, + input: { streamId }, + }) const run = await readRun() logger.info('[Resume] Stream lookup', { streamId, @@ -251,7 +241,7 @@ async function handleResumeRequestBody({ if (batchMode) { const afterSeq = afterCursor || '0' const [gap, events, previewSessions] = await Promise.all([ - findReplayGap(streamId, afterSeq, extractRunRequestId(run)), + fromLog ? null : findReplayGap(streamId, afterSeq, extractRunRequestId(run)), readEvents(streamId, afterSeq), readFilePreviewSessions(streamId).catch((error) => { logger.warn('Failed to read preview sessions for stream batch', { @@ -261,8 +251,8 @@ async function handleResumeRequestBody({ return [] }), ]) - /* The ring cannot serve this reader; the live tail re-syncs it from the worker log. */ - const batchEvents = gap ? [] : events.map(toStreamBatchEvent) + // A reader the ring cannot serve is re-synced from the worker log by the live tail. + const batchEvents = fromLog || gap ? [] : events.map(toStreamBatchEvent) logger.info('[Resume] Batch response', { streamId, afterCursor: afterSeq, @@ -295,12 +285,15 @@ async function handleResumeRequestBody({ * 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 gap = await findReplayGap(streamId, afterCursor || '0', extractRunRequestId(run)) + const gap = fromLog + ? null + : await findReplayGap(streamId, afterCursor || '0', extractRunRequestId(run)) + 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 (gap && run.chatId) { + if (resyncFromLog && run.chatId) { try { replayBody = await openRunReplay({ streamId, @@ -323,7 +316,7 @@ async function handleResumeRequestBody({ 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)) }, @@ -434,61 +427,30 @@ async function handleResumeRequestBody({ } } - /** - * Forwards the worker's replay under this response's own cursors, starting at 1. - * A parked run holds the response open until it resumes, so the client does not - * replay the whole log again for every poll; any other end short of the terminal - * (the worker's cap, a stalled run, a cut connection) ends it without one and the - * client re-attaches, reaching recovery through the reconnect route. When that - * recovery was refused, a park or stall ends the view with an error instead. - */ + /** Forwards the worker's replay, keeping the response alive while it waits. */ const streamRunReplay = async (body: ReadableStream) => { - let seq = 0 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 readRunReplay(body, replaySignal, (event) => { - seq += 1 - const envelope = toReplayEnvelope({ - ...event, - seq, - stream: { ...event.stream, streamId, cursor: String(seq) }, - }) - if (!enqueueEvent(envelope)) return false - totalEventsFlushed += 1 - cursor = String(seq) - if (envelope.type === MothershipStreamV1EventType.complete) sawTerminalEvent = true - return true - }).catch((error: unknown) => { - if (error instanceof FatalSseEventError) throw error - logger.warn('Run replay connection ended early', { - streamId, - error: getErrorMessage(error), - }) - return 'closed' as const + 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 + }, + isParked: async () => + (await readRun().catch(() => null))?.status === 'paused_waiting_for_tool', + isClosed: () => controllerClosed, + deadlineAt: startTime + MAX_STREAM_MS, }) - logger.info('[Resume] Run replay ended', { streamId, end, eventCount: seq }) - while ( - end === 'parked' && - !recoveryRefused && - !controllerClosed && - Date.now() - startTime < MAX_STREAM_MS - ) { - const current = await readRun().catch(() => null) - if (current?.status !== 'paused_waiting_for_tool') break - await sleep(POLL_INTERVAL_MAX_MS) - } - /* No controller drives the run and none can be rebuilt from a trimmed ring, so a - park or stall is final for this view; re-attaching would only replay it again. */ - if (recoveryRefused && (end === 'parked' || end === 'stalled')) { - emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, { - message: 'This response can no longer continue. Send a message to pick up from here.', - code: 'recovery_unavailable', - reason: 'recovery_unavailable', - }) - } + logger.info('[Resume] Run replay ended', { streamId, end, eventCount: totalEventsFlushed }) } finally { clearInterval(keepalive) replayAbort.abort() @@ -502,8 +464,9 @@ async function handleResumeRequestBody({ await streamRunReplay(replayBody) return } - if (gap) { - const terminal = await replayGapTerminal(streamId, gap, currentRequestId) + 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 @@ -529,8 +492,11 @@ async function handleResumeRequestBody({ }) return null }) - /* The ring lost its head under this tail; the re-attach re-syncs from the log. */ - if (recoveryRefused) break + // The ring lost its head or restarted under this tail; the re-attach re-syncs. + if (!ringCanServe(await readRingPosition(streamId, cursor))) { + logger.warn('Replay ring can no longer serve a live tail', { streamId, cursor }) + break + } if (!currentRun) { emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, { message: 'The stream could not be recovered because its run metadata is unavailable.', @@ -549,16 +515,6 @@ async function handleResumeRequestBody({ pollDelayMs = flushed > 0 ? POLL_INTERVAL_MS : Math.min(pollDelayMs * 2, POLL_INTERVAL_MAX_MS) - /* A ring whose numbering restarted after it expired never passes this cursor - again; end without a terminal so the client re-attaches and is re-synced. */ - if (flushed === 0 && !sawTerminalEvent) { - const latestSeq = await getLatestSeq(streamId) - if (latestSeq !== null && Number(cursor) > latestSeq) { - logger.warn('Replay ring restarted under a live tail', { streamId, cursor, latestSeq }) - break - } - } - if (controllerClosed) { break } diff --git a/apps/sim/app/api/mothership/chats/[chatId]/route.ts b/apps/sim/app/api/mothership/chats/[chatId]/route.ts index 3825c38aa67..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,11 +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 { startsAtReplayHead } from '@/lib/mothership/request/session/recovery' -import { type StreamBatchEvent, toStreamBatchEvent } from '@/lib/mothership/request/session/types' import { captureServerEvent } from '@/lib/posthog/server' const logger = createLogger('MothershipChatAPI') @@ -56,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 }], @@ -70,39 +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 - }) - - // A trimmed head would paint a truncated turn; the resume route reports the gap. - if (startsAtReplayHead(events[0]?.seq)) { - 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 7fec1809866..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 @@ -1158,10 +1158,11 @@ describe('useChat remount send recovery', () => { } }) - it('rebuilds the turn from an empty response when a reconnect is re-synced from the log', async () => { + 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', @@ -1188,6 +1189,7 @@ describe('useChat remount send recovery', () => { 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' }) } @@ -1216,6 +1218,12 @@ describe('useChat remount send recovery', () => { 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() } 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 cc03044276d..a9c927ef54b 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -306,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. @@ -956,6 +965,8 @@ export function useChat( const streamRequestIdRef = useRef(undefined) const locallyTerminalStreamIdRef = useRef(undefined) const lastCursorRef = useRef('0') + const logResyncedStreamsRef = useRef | null>(null) + logResyncedStreamsRef.current ??= new Set() const activeStreamReturnRecoveryRef = useRef(null) const sendingRef = useRef(false) const streamGenRef = useRef(0) @@ -2371,7 +2382,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, logResyncedStreamsRef.current?.has(streamId) ?? false)}&batch=true`, { signal: fetchSignal, ...(streamTraceparentRef.current @@ -2563,7 +2574,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, logResyncedStreamsRef.current?.has(streamId) ?? false)}`, { signal: activeAbort.signal, ...(streamTraceparentRef.current @@ -2582,9 +2593,9 @@ export function useChat( return { error: false, aborted: true } } - /* The ring could not serve this cursor, so the turn is re-sent from the - worker's log with cursors restarting at 1: rebuild it from empty. */ + // Re-sent from the worker's log with cursors restarting at 1: rebuild from empty. if (sseRes.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER) === 'log') { + logResyncedStreamsRef.current?.add(streamId) const reset = applyReconnectReplaySelection(streamId, '0') latestCursor = reset.afterCursor preserveNextReplayState = reset.preserveExistingState 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/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/request/application/recover-stream.ts b/apps/sim/lib/mothership/request/application/recover-stream.ts index 9d9fbde594c..61eff64d315 100644 --- a/apps/sim/lib/mothership/request/application/recover-stream.ts +++ b/apps/sim/lib/mothership/request/application/recover-stream.ts @@ -24,7 +24,7 @@ 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' @@ -32,20 +32,6 @@ import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils' const logger = createLogger('MothershipStreamRecovery') -type RecoverableRun = NonNullable>> - -/** - * The run's replay ring no longer holds its first events, so its context cannot be - * rebuilt from Redis. The reconnect reports `replay_gap` instead of recovering from the - * tail, which would persist a truncated turn. - */ -export class StreamReplayHeadTrimmedError extends Error { - constructor(readonly run: RecoverableRun) { - super(`Replay head trimmed for stream ${run.streamId}`) - this.name = 'StreamReplayHeadTrimmedError' - } -} - export const readChatStream = defineAuthorizedChatUseCase({ operation: defineWorkspaceOperation({ id: 'mothership.runs.reconnect', @@ -135,7 +121,17 @@ export const readChatStream = defineAuthorizedChatUseCase({ ]) if (workspaceId && !userPermission) throw new OrchestrationError('forbidden', 'Workspace access revoked') - if (!startsAtReplayHead(events[0]?.seq)) throw new StreamReplayHeadTrimmedError(run) + /** + * 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, @@ -166,7 +162,7 @@ export const readChatStream = defineAuthorizedChatUseCase({ currentChat: null, message: '', titleModel: '', - resumeSeq: events.at(-1)?.seq ?? 0, + resumeSeq, orchestrateOptions: { userId, workspaceId, @@ -184,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/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/index.ts b/apps/sim/lib/mothership/request/session/index.ts index 0f458f6f400..d31150c625c 100644 --- a/apps/sim/lib/mothership/request/session/index.ts +++ b/apps/sim/lib/mothership/request/session/index.ts @@ -69,17 +69,17 @@ export { isFilePreviewSession, } from './file-preview-session-contract' export { - checkForReplayGap, findReplayGap, - type ReplayGap, - type ReplayGapResult, + type RingPosition, + readRingPosition, replayGapTerminal, + ringCanServe, } from './recovery' export { + forwardRunReplay, openRunReplay, type RunReplayEnd, RunReplayUnavailableError, - readRunReplay, } from './run-replay' export { encodeSSEEnvelope, SSE_RESPONSE_HEADERS } from './sse' export type { StreamBatchEvent } from './types' diff --git a/apps/sim/lib/mothership/request/session/recovery.test.ts b/apps/sim/lib/mothership/request/session/recovery.test.ts index b81ade42bc0..725de468a6f 100644 --- a/apps/sim/lib/mothership/request/session/recovery.test.ts +++ b/apps/sim/lib/mothership/request/session/recovery.test.ts @@ -12,9 +12,9 @@ vi.mock('./buffer', () => ({ readEvents, })) -import { checkForReplayGap } from './recovery' +import { findReplayGap, replayGapTerminal } 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,7 +24,9 @@ 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) diff --git a/apps/sim/lib/mothership/request/session/recovery.ts b/apps/sim/lib/mothership/request/session/recovery.ts index 6d9adee79d8..4c1e2b643c2 100644 --- a/apps/sim/lib/mothership/request/session/recovery.ts +++ b/apps/sim/lib/mothership/request/session/recovery.ts @@ -18,8 +18,8 @@ export interface ReplayGapResult { envelopes: ReturnType[] } -/** Where the replay ring stands relative to a reader it can no longer serve. */ -export interface ReplayGap { +/** 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 @@ -35,68 +35,78 @@ export function startsAtReplayHead(firstSeq: number | undefined): boolean { return firstSeq === undefined || firstSeq <= 1 } +export async function readRingPosition( + streamId: string, + 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 `afterCursor`. 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 + * 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. */ +export function ringCanServe({ requestedAfterSeq, oldestSeq, latestSeq }: RingPosition): boolean { + return latestSeq <= 0 || (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 { - const requestedAfterSeq = Number(afterCursor || '0') +): Promise { return withCopilotSpan( TraceSpan.CopilotRecoveryCheckReplayGap, { [TraceAttr.StreamId]: streamId, - [TraceAttr.CopilotRecoveryRequestedAfterSeq]: requestedAfterSeq, + [TraceAttr.CopilotRecoveryRequestedAfterSeq]: Number(afterCursor || '0'), ...(requestId ? { [TraceAttr.RequestId]: requestId } : {}), }, async (span) => { - const [oldestSeq, latestSeq] = await Promise.all([ - getOldestSeq(streamId), - 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, }) - if ( - latestSeq === null || - latestSeq <= 0 || - oldestSeq === null || - (startsAtReplayHead(oldestSeq) && requestedAfterSeq <= latestSeq) - ) { + if (ringCanServe(position)) { span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.InRange) return null } logger.warn('Replay gap detected: the ring cannot serve the requested cursor', { streamId, - requestedAfterSeq, - oldestAvailableSeq: oldestSeq, - latestSeq, + ...position, }) span.setAttribute(TraceAttr.CopilotRecoveryOutcome, CopilotRecoveryOutcome.GapDetected) - return { requestedAfterSeq, oldestSeq, latestSeq } + return position } ) } -/** Ends a reader's view with `replay_gap` when nothing can re-sync it. */ +/** + * 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, - gap: ReplayGap, + position: RingPosition, requestId?: string ): Promise { - const { latestSeq, oldestSeq, requestedAfterSeq } = gap + const { latestSeq, oldestSeq, requestedAfterSeq } = position + const baseSeq = Math.max(latestSeq, requestedAfterSeq) const resolvedRequestId = await resolveReplayGapRequestId(streamId, latestSeq, requestId) const gapEnvelope = createEvent({ streamId, - cursor: String(latestSeq + 1), - seq: latestSeq + 1, + cursor: String(baseSeq + 1), + seq: baseSeq + 1, requestId: resolvedRequestId, type: MothershipStreamV1EventType.error, payload: { @@ -110,8 +120,8 @@ export async function replayGapTerminal( }) const terminalEnvelope = createEvent({ streamId, - cursor: String(latestSeq + 2), - seq: latestSeq + 2, + cursor: String(baseSeq + 2), + seq: baseSeq + 2, requestId: resolvedRequestId, type: MothershipStreamV1EventType.complete, payload: { @@ -122,15 +132,6 @@ export async function replayGapTerminal( return { gapDetected: true, envelopes: [gapEnvelope, terminalEnvelope] } } -export async function checkForReplayGap( - streamId: string, - afterCursor: string, - requestId?: string -): Promise { - const gap = await findReplayGap(streamId, afterCursor, requestId) - return gap ? replayGapTerminal(streamId, gap, requestId) : null -} - 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 203e0d1d409..376286c0b57 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -62,7 +62,11 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { /** 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[] + }>, }, } }) @@ -75,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') { @@ -794,8 +806,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge ) ).text() ) - /* A ring without its head is re-synced from the worker log whatever the cursor; - this worker does not know the run, so the replay_gap terminal applies. */ + // 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( @@ -892,7 +903,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge }) it.each(['after=4', 'after=0'])( - 'shows the full turn from the worker log without recovering from the tail (%s)', + '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 = { @@ -923,16 +934,9 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge try { const frames = dataFrames(await (await resume(streamId, query)).text()) - expect( - frames.map((frame) => [frame.type, frame.payload.code ?? frame.payload.status ?? null]) - ).toEqual([ - ['text', null], - ['error', 'recovery_unavailable'], - ['complete', 'error'], - ]) + 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).toEqual([]) - expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + expect(worker.runs.map((run) => run.recoveredEvents)).toEqual([[]]) } finally { worker.hooks.replay = { status: 404, frames: [] } } @@ -945,7 +949,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge const batch = await (await resume(streamId, 'after=0&batch=true')).json() expect(batch.events).toEqual([]) - expect(worker.runs).toEqual([]) + expect(worker.runs.map((run) => run.recoveredEvents)).toEqual([[]]) }) it.each([ diff --git a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts index 1dc7a265673..7dd491cacda 100644 --- a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts @@ -113,10 +113,10 @@ async function liveRunWithTrimmedRing(): Promise<{ streamId: string; runId: stri return { streamId, runId } } -function reconnect(streamId: string, after: string, batch = false) { +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' : ''}` + `http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=${after}${batch ? '&batch=true' : ''}${extra}` ), { params: Promise.resolve({}) } ) @@ -260,7 +260,7 @@ describe.runIf(Boolean(redisUrl))('reconnects past the replay ring', () => { const response = await reconnect(streamId, '4') const body = response.text() - /* Let the tail reach its poll loop, past the reconnect-time gap check. */ + // Let the tail reach its poll loop, past the reconnect-time gap check. await sleep(500) await getRedisClient()!.del( `mothership_stream:${streamId}:events`, @@ -272,6 +272,91 @@ describe.runIf(Boolean(redisUrl))('reconnects past the replay ring', () => { 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('serves no ring events to a batch read the ring can no longer serve', async () => { const { streamId } = await liveRunWithTrimmedRing() diff --git a/apps/sim/lib/mothership/request/session/run-replay.ts b/apps/sim/lib/mothership/request/session/run-replay.ts index cbf7fdb6b12..8b757fb0d0b 100644 --- a/apps/sim/lib/mothership/request/session/run-replay.ts +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -1,4 +1,8 @@ -import { isRecordLike } from '@sim/utils/object' +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 { 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' @@ -9,9 +13,15 @@ import { 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' +/** 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 = 2_000 /** How a worker replay leg ended: at the run's terminal, short of it, or cut off. */ export type RunReplayEnd = 'complete' | StreamReplayEnd['reason'] | 'closed' @@ -63,30 +73,53 @@ export async function openRunReplay(params: { return response.body } -function replayEndReason(value: unknown): StreamReplayEnd['reason'] | null { +/** 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 = value.payload - if (!isRecordLike(payload) || payload.kind !== 'replay_end') return null - return payload.reason === 'parked' || payload.reason === 'cap' || payload.reason === 'stalled' - ? payload.reason - : 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 is worker-to-Sim control and never reaches `onEvent`. Returning - * false from `onEvent` stops the read (the reader went away). + * `replay_end` frame never reaches `onEvent`. Returning false from `onEvent` stops the + * read (the reader went away). */ -export async function readRunReplay( +async function readRunReplay( body: ReadableStream, signal: AbortSignal, onEvent: (event: PersistedStreamEventEnvelope) => boolean ): Promise { let end: RunReplayEnd = 'closed' await processSSEStream(body.getReader(), signal, (raw) => { - const reason = replayEndReason(raw) - if (reason) { - end = reason + const control = replayEnd(raw) + if (control) { + end = control return true } const parsed = parsePersistedStreamEventEnvelope(raw) @@ -100,3 +133,51 @@ export async function readRunReplay( }) 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 + /** Whether the run still waits on calls Sim executes. */ + isParked: () => 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. A park holds it while the + * run waits on Sim, and any other end holds it 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, isParked, 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 + while (!isClosed() && Date.now() < deadlineAt) { + const parked = end === 'parked' && (await isParked()) + 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..8d892acca63 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts @@ -0,0 +1,220 @@ +/** + * 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. + */ +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 = { + chatRequests: [] as Array>, + /** SSE frames for the re-attach, built once the stream id is known. */ + frames: [] as unknown[], + } + const server = createServer(async (request, response) => { + let body = '' + for await (const chunk of request) body += chunk + if (request.url === '/api/mothership') { + worker.chatRequests.push(JSON.parse(body)) + response.writeHead(200, { 'content-type': 'text/event-stream' }) + for (const frame of worker.frames) response.write(`data: ${JSON.stringify(frame)}\n\n`) + response.end('data: [DONE]\n\n') + return + } + response.writeHead(404, { 'content-type': 'application/json' }) + response.end(JSON.stringify({ error: 'Run not found' })) + }) + 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 { + copilotChats, + copilotMessages, + copilotRuns, + permissions, + user, + workspace, +} from '@sim/db/schema' +import { sleep } from '@sim/utils/helpers' +import { generateId } from '@sim/utils/id' +import { toRecord } from '@sim/utils/object' +import { eq } from 'drizzle-orm' +import { NextRequest } from 'next/server' +import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis' +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 chatId = generateId() +const FULL_TEXT = 'part 1 part 2 part 3 part 4 ' + +function workerFrame(streamId: string, seq: number, type: string, payload: unknown) { + return { v: 1, type, seq, ts: new Date().toISOString(), stream: { streamId, chatId }, payload } +} + +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 () => { + 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() + 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 + } + }) + + it('recovers from an empty context and persists the whole turn once', async () => { + 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, + 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) + worker.frames = [ + workerFrame(streamId, 1, 'session', { kind: 'start' }), + workerFrame(streamId, 2, 'text', { channel: 'assistant', text: FULL_TEXT, textOffset: 0 }), + workerFrame(streamId, 3, 'complete', { status: 'complete', textLength: FULL_TEXT.length }), + ] + + 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 run: { status: string } | undefined + for (let attempt = 0; attempt < 100 && run?.status !== 'complete'; attempt++) { + await sleep(100) + ;[run] = await db + .select({ status: copilotRuns.status }) + .from(copilotRuns) + .where(eq(copilotRuns.id, runId)) + } + expect(run?.status).toBe('complete') + + const [chat] = await db + .select({ conversationId: copilotChats.conversationId }) + .from(copilotChats) + .where(eq(copilotChats.id, chatId)) + expect(chat.conversationId).toBeNull() + const assistant = await db + .select({ content: copilotMessages.content }) + .from(copilotMessages) + .where(eq(copilotMessages.chatId, chatId)) + .then((rows) => + rows.map((row) => toRecord(row.content)).filter((m) => m.role === 'assistant') + ) + 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.chatRequests).toHaveLength(1) + expect(worker.chatRequests[0]).toMatchObject({ messageId: streamId, receivedTextChars: 0 }) + expect(await getRedisClient()!.get(chatStreamLockKey(chatId))).toBeNull() + }) +}) From 8b9493da909c687f98a66770a8fd932a9a13838b Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 07:50:10 -0700 Subject: [PATCH 06/12] fix(mothership): end a held replay as soon as its run resumes or finishes, and check a busy tail's ring less often - A replay held after a park now ends as soon as Sim sees the run leave its park, and any held replay ends as soon as the run reaches a terminal, so an approval no longer freezes the view for up to 10 s. A park Sim has not yet marked, a stall and a cut connection keep the 10 s floor. - A live tail checks that its ring can still serve it only after a quiet poll or every eighth busy one, instead of two Redis reads on every 250 ms poll. - The recovery integration test re-sends a replayed go tool and re-hands a Sim call the dead controller already ran: it is resumed with its stored result, never run again, and the turn keeps one block per tool. --- apps/sim/app/api/copilot/chat/stream/route.ts | 15 +- .../request/session/replay-gap.integration.ts | 20 +- .../mothership/request/session/run-replay.ts | 26 +- .../session/stream-recovery.integration.ts | 335 ++++++++++++------ 4 files changed, 277 insertions(+), 119 deletions(-) diff --git a/apps/sim/app/api/copilot/chat/stream/route.ts b/apps/sim/app/api/copilot/chat/stream/route.ts index b57ccf1870a..9a01b81c554 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.ts @@ -50,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 @@ -445,8 +447,7 @@ async function handleResumeRequestBody({ if (envelope.type === MothershipStreamV1EventType.complete) sawTerminalEvent = true return true }, - isParked: async () => - (await readRun().catch(() => null))?.status === 'paused_waiting_for_tool', + readRunStatus: async () => (await readRun().catch(() => null))?.status ?? null, isClosed: () => controllerClosed, deadlineAt: startTime + MAX_STREAM_MS, }) @@ -480,7 +481,7 @@ async function handleResumeRequestBody({ return } - await flushEvents() + let lastFlushed = await flushEvents() let pollDelayMs = POLL_INTERVAL_MS while (!controllerClosed && Date.now() - startTime < MAX_STREAM_MS) { @@ -492,8 +493,11 @@ async function handleResumeRequestBody({ }) return null }) - // The ring lost its head or restarted under this tail; the re-attach re-syncs. - if (!ringCanServe(await readRingPosition(streamId, cursor))) { + // The ring lost its head or restarted under this tail; the re-attach re-syncs. Only a + // quiet ring can have restarted or have a dead controller that recovery re-sends + // under, so a busy tail checks every few polls rather than on each one. + const checkRing = lastFlushed === 0 || pollIterations % RING_CHECK_EVERY_BUSY_POLLS === 0 + if (checkRing && !ringCanServe(await readRingPosition(streamId, cursor))) { logger.warn('Replay ring can no longer serve a live tail', { streamId, cursor }) break } @@ -509,6 +513,7 @@ async function handleResumeRequestBody({ currentRequestId = extractRunRequestId(currentRun) || currentRequestId const flushed = await flushEvents() + 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. */ diff --git a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts index 7dd491cacda..7b87ad6a16f 100644 --- a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts @@ -241,11 +241,29 @@ describe.runIf(Boolean(redisUrl))('reconnects past the replay ring', () => { 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(Date.now() - startedAt).toBeGreaterThanOrEqual(1_500) + 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({ diff --git a/apps/sim/lib/mothership/request/session/run-replay.ts b/apps/sim/lib/mothership/request/session/run-replay.ts index 8b757fb0d0b..642b0ec9a74 100644 --- a/apps/sim/lib/mothership/request/session/run-replay.ts +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -10,6 +10,7 @@ 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' @@ -21,7 +22,8 @@ const logger = createLogger('RunReplay') const REPLAY_PATH = '/api/streams/replay' /** 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 = 2_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' @@ -140,8 +142,8 @@ export interface ForwardRunReplayOptions { signal: AbortSignal /** Writes one event to the reader; false once the reader is gone. */ write: (event: PersistedStreamEventEnvelope) => boolean - /** Whether the run still waits on calls Sim executes. */ - isParked: () => Promise + /** 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 @@ -150,12 +152,13 @@ export interface ForwardRunReplayOptions { /** * 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. A park holds it while the - * run waits on Sim, and any other end holds it at least {@link REPLAY_MIN_RESPONSE_MS}, - * so a reader re-attaches, and replays the whole log again, at most that often. + * 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, isParked, isClosed, deadlineAt } = options + const { body, streamId, signal, write, readRunStatus, isClosed, deadlineAt } = options const startedAt = Date.now() let seq = 0 const end = await readRunReplay(body, signal, (event) => { @@ -173,8 +176,15 @@ export async function forwardRunReplay(options: ForwardRunReplayOptions): Promis 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 parked = end === 'parked' && (await isParked()) + 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)) diff --git a/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts b/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts index 8d892acca63..22af1e25f53 100644 --- a/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts +++ b/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts @@ -2,7 +2,8 @@ * 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. + * 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' @@ -11,22 +12,23 @@ 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 = { - chatRequests: [] as Array>, - /** SSE frames for the re-attach, built once the stream id is known. */ - frames: [] as unknown[], + 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 - if (request.url === '/api/mothership') { - worker.chatRequests.push(JSON.parse(body)) - response.writeHead(200, { 'content-type': 'text/event-stream' }) - for (const frame of worker.frames) response.write(`data: ${JSON.stringify(frame)}\n\n`) - response.end('data: [DONE]\n\n') + 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 } - response.writeHead(404, { 'content-type': 'application/json' }) - response.end(JSON.stringify({ error: 'Run not found' })) + 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 } @@ -44,6 +46,7 @@ vi.mock('@/lib/auth', () => authMock) import { db } from '@sim/db' import { + copilotAsyncToolCalls, copilotChats, copilotMessages, copilotRuns, @@ -53,10 +56,19 @@ import { } from '@sim/db/schema' import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' -import { toRecord } from '@sim/utils/object' +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' @@ -64,11 +76,103 @@ import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' const userId = generateId() const workspaceId = generateId() -const chatId = generateId() +const chatIds: string[] = [] const FULL_TEXT = 'part 1 part 2 part 3 part 4 ' -function workerFrame(streamId: string, seq: number, type: string, payload: unknown) { - return { v: 1, type, seq, ts: new Date().toISOString(), stream: { streamId, chatId }, payload } +/** 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') } describe.runIf(Boolean(redisUrl))('recovering a run whose ring lost its head', () => { @@ -102,9 +206,11 @@ describe.runIf(Boolean(redisUrl))('recovering a run whose ring lost its head', ( }) afterAll(async () => { - 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)) + 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)) @@ -117,104 +223,123 @@ describe.runIf(Boolean(redisUrl))('recovering a run whose ring lost its head', ( }) it('recovers from an empty context and persists the whole turn once', async () => { - 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, - 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) - worker.frames = [ - workerFrame(streamId, 1, 'session', { kind: 'start' }), - workerFrame(streamId, 2, 'text', { channel: 'assistant', text: FULL_TEXT, textOffset: 0 }), - workerFrame(streamId, 3, 'complete', { status: 'complete', textLength: FULL_TEXT.length }), - ] - - 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 run: { status: string } | undefined - for (let attempt = 0; attempt < 100 && run?.status !== 'complete'; attempt++) { - await sleep(100) - ;[run] = await db - .select({ status: copilotRuns.status }) - .from(copilotRuns) - .where(eq(copilotRuns.id, runId)) + 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(run?.status).toBe('complete') + + 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 db - .select({ content: copilotMessages.content }) - .from(copilotMessages) - .where(eq(copilotMessages.chatId, chatId)) - .then((rows) => - rows.map((row) => toRecord(row.content)).filter((m) => m.role === 'assistant') - ) + 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.chatRequests).toHaveLength(1) - expect(worker.chatRequests[0]).toMatchObject({ messageId: streamId, receivedTextChars: 0 }) + 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) + }) }) From f6c39f28f89b2db1f1c8a6308c7687f1fb178dc4 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 11:12:06 -0700 Subject: [PATCH 07/12] fix(mothership): fall back to replay_gap when the worker refuses a replay's key A deployment whose key may not call the worker's replay (401/403) now falls back to the replay_gap terminal as a missing run (404) does, instead of answering 503 until the client's reconnect budget runs out. A failed buffer TTL refresh during the chat-lock heartbeat is logged as such, not as a lock-extension failure. --- .../lib/mothership/request/session/abort.ts | 9 +++++++- .../request/session/replay-gap.integration.ts | 21 +++++++++++-------- .../mothership/request/session/run-replay.ts | 7 +++++-- 3 files changed, 25 insertions(+), 12 deletions(-) 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/replay-gap.integration.ts b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts index 7b87ad6a16f..28721e5f9e5 100644 --- a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts @@ -383,17 +383,20 @@ describe.runIf(Boolean(redisUrl))('reconnects past the replay ring', () => { expect(await response.json()).toMatchObject({ success: true, events: [], status: 'active' }) }) - it('keeps the replay_gap terminal when the worker has no such run', async () => { - const { streamId } = await liveRunWithTrimmedRing() - worker.reply.status = 404 + 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()) + const frames = dataFrames(await (await reconnect(streamId, '2')).text()) - expect(frames.map((frame) => [frame.type, frame.payload.code ?? frame.payload.status])).toEqual( - [ + 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.ts b/apps/sim/lib/mothership/request/session/run-replay.ts index 642b0ec9a74..f4b1379e2f4 100644 --- a/apps/sim/lib/mothership/request/session/run-replay.ts +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -20,6 +20,8 @@ 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 @@ -39,7 +41,8 @@ export class RunReplayUnavailableError extends Error { /** * 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 knows no such run for this chat and user. + * 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 @@ -64,7 +67,7 @@ export async function openRunReplay(params: { if (signal.aborted) throw error throw new RunReplayUnavailableError('The run replay could not be reached', { cause: error }) } - if (response.status === 404) { + if (REPLAY_REFUSED_STATUSES.has(response.status)) { await response.body?.cancel().catch(() => {}) return null } From 65d40881db6b074471749fcdeadcda9926d873b3 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 11:27:25 -0700 Subject: [PATCH 08/12] fix(mothership): bound a silent replay, never skip past a trimmed cursor, and re-sync an expired ring - The worker replay is bounded like a stream leg: no response headers, or no bytes including keepalives, for the idle timeout ends it so the reader re-attaches. - A ring read that starts after the reader's next event (the ring trimmed its head between the gap check and the read) is never delivered; the reader re-attaches and re-syncs from the log, in both the live tail and batch reads. - An empty ring serves only a reader starting from cursor 0; a live run whose buffer expired under a reader re-syncs from the log, and a finished one answers its terminal since its transcript is persisted. - A leg that delivers events after a failure starts a fresh 30 s reachable window; its three retries still refill only after five minutes of delivered events. - The two new integration suites close their worker server and restore env even when they are skipped. --- .../app/api/copilot/chat/stream/route.test.ts | 50 +++++++++++ apps/sim/app/api/copilot/chat/stream/route.ts | 49 ++++++++--- .../request/lifecycle/stream-retry.test.ts | 14 ++++ .../request/lifecycle/stream-retry.ts | 9 +- .../request/session/recovery.test.ts | 11 ++- .../mothership/request/session/recovery.ts | 5 +- .../request/session/replay-gap.integration.ts | 55 ++++++++++-- .../request/session/run-replay.test.ts | 83 +++++++++++++++++++ .../mothership/request/session/run-replay.ts | 42 +++++++++- .../session/stream-recovery.integration.ts | 14 ++-- 10 files changed, 303 insertions(+), 29 deletions(-) create mode 100644 apps/sim/lib/mothership/request/session/run-replay.test.ts 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 5726ce344d1..ea176606664 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.test.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.test.ts @@ -219,4 +219,54 @@ 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: [] }) + }) }) diff --git a/apps/sim/app/api/copilot/chat/stream/route.ts b/apps/sim/app/api/copilot/chat/stream/route.ts index 9a01b81c554..3bdaeecf49d 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.ts @@ -59,6 +59,15 @@ const RING_CHECK_EVERY_BUSY_POLLS = 8 */ 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 : '' } @@ -253,8 +262,10 @@ async function handleResumeRequestBody({ return [] }), ]) - // A reader the ring cannot serve is re-synced from the worker log by the live tail. - const batchEvents = fromLog || gap ? [] : 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, @@ -287,9 +298,12 @@ async function handleResumeRequestBody({ * 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 gap = fromLog + 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. */ @@ -380,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, @@ -482,6 +501,7 @@ async function handleResumeRequestBody({ } let lastFlushed = await flushEvents() + if (lastFlushed === null) return let pollDelayMs = POLL_INTERVAL_MS while (!controllerClosed && Date.now() - startTime < MAX_STREAM_MS) { @@ -493,14 +513,6 @@ async function handleResumeRequestBody({ }) return null }) - // The ring lost its head or restarted under this tail; the re-attach re-syncs. Only a - // quiet ring can have restarted or have a dead controller that recovery re-sends - // under, so a busy tail checks every few polls rather than on each one. - const checkRing = lastFlushed === 0 || pollIterations % RING_CHECK_EVERY_BUSY_POLLS === 0 - if (checkRing && !ringCanServe(await readRingPosition(streamId, cursor))) { - logger.warn('Replay ring can no longer serve a live tail', { streamId, cursor }) - break - } if (!currentRun) { emitTerminalIfMissing(MothershipStreamV1CompletionStatus.error, { message: 'The stream could not be recovered because its run metadata is unavailable.', @@ -509,10 +521,23 @@ 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 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/recovery.test.ts b/apps/sim/lib/mothership/request/session/recovery.test.ts index 725de468a6f..fffce906a72 100644 --- a/apps/sim/lib/mothership/request/session/recovery.test.ts +++ b/apps/sim/lib/mothership/request/session/recovery.test.ts @@ -12,7 +12,11 @@ vi.mock('./buffer', () => ({ readEvents, })) -import { findReplayGap, replayGapTerminal } from '@/lib/mothership/request/session/recovery' +import { + findReplayGap, + replayGapTerminal, + ringCanServe, +} from '@/lib/mothership/request/session/recovery' describe('replay gap', () => { it('uses the latest buffered request id when run metadata is missing it', async () => { @@ -33,4 +37,9 @@ describe('replay gap', () => { 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 4c1e2b643c2..83854f906cf 100644 --- a/apps/sim/lib/mothership/request/session/recovery.ts +++ b/apps/sim/lib/mothership/request/session/recovery.ts @@ -52,10 +52,11 @@ export async function readRingPosition( * 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. + * restarted after it expired produces, nor any cursor from a buffer that expired. */ export function ringCanServe({ requestedAfterSeq, oldestSeq, latestSeq }: RingPosition): boolean { - return latestSeq <= 0 || (startsAtReplayHead(oldestSeq) && requestedAfterSeq <= latestSeq) + if (latestSeq <= 0) return requestedAfterSeq <= 0 + return startsAtReplayHead(oldestSeq) && requestedAfterSeq <= latestSeq } /** The ring's position when it cannot serve `afterCursor` (see {@link ringCanServe}). */ diff --git a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts index 28721e5f9e5..f9e8be62fe5 100644 --- a/apps/sim/lib/mothership/request/session/replay-gap.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-gap.integration.ts @@ -134,6 +134,15 @@ function fullResponse(streamId: string) { ] } +/** 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() @@ -177,11 +186,6 @@ describe.runIf(Boolean(redisUrl))('reconnects past the replay ring', () => { await db.delete(workspace).where(eq(workspace.id, workspaceId)) await db.delete(user).where(eq(user.id, userId)) await closeRedisConnection() - 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 - } }) it.each([ @@ -375,6 +379,47 @@ describe.runIf(Boolean(redisUrl))('reconnects past the replay ring', () => { 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() 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..c73dea4f46d --- /dev/null +++ b/apps/sim/lib/mothership/request/session/run-replay.test.ts @@ -0,0 +1,83 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' + +vi.mock('@/lib/mothership/server/agent-url', async (importOriginal) => ({ + ...(await importOriginal()), + getMothershipBaseURL: async () => 'https://worker.test', +})) + +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 index f4b1379e2f4..11c90989b86 100644 --- a/apps/sim/lib/mothership/request/session/run-replay.ts +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -3,6 +3,7 @@ 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' @@ -52,13 +53,18 @@ export async function openRunReplay(params: { }): 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, + signal: AbortSignal.any([signal, unanswered.signal]), spanName: `sim → go ${REPLAY_PATH}`, operation: 'stream_replay', attributes: { [TraceAttr.StreamId]: streamId, [TraceAttr.ChatId]: chatId }, @@ -66,6 +72,8 @@ export async function openRunReplay(params: { } 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)) { await response.body?.cancel().catch(() => {}) @@ -78,6 +86,36 @@ export async function openRunReplay(params: { return response.body } +/** + * A replay reader that ends, as a cut connection does, once the worker sends nothing, + * not even a keepalive, for {@link WORKER_STREAM_IDLE_TIMEOUT_MS}. + */ +function idleBoundedReader( + reader: ReadableStreamDefaultReader +): ReadableStreamDefaultReader { + return { + async read() { + let timer: ReturnType | undefined + const silent = new Promise>((resolve) => { + timer = setTimeout(() => { + reader.cancel().catch(() => {}) + resolve({ done: true, value: undefined }) + }, WORKER_STREAM_IDLE_TIMEOUT_MS) + }) + try { + return await Promise.race([reader.read(), silent]) + } finally { + clearTimeout(timer) + } + }, + cancel: (reason) => reader.cancel(reason), + releaseLock: () => reader.releaseLock(), + get closed() { + return reader.closed + }, + } +} + /** Every reason the worker may end a replay with; a reason added to the contract fails here. */ const REPLAY_END_REASONS = { parked: true, @@ -121,7 +159,7 @@ async function readRunReplay( onEvent: (event: PersistedStreamEventEnvelope) => boolean ): Promise { let end: RunReplayEnd = 'closed' - await processSSEStream(body.getReader(), signal, (raw) => { + await processSSEStream(idleBoundedReader(body.getReader()), signal, (raw) => { const control = replayEnd(raw) if (control) { end = control diff --git a/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts b/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts index 22af1e25f53..6b714c19cb9 100644 --- a/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts +++ b/apps/sim/lib/mothership/request/session/stream-recovery.integration.ts @@ -175,6 +175,15 @@ async function assistantMessages(chatId: string) { 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() @@ -215,11 +224,6 @@ describe.runIf(Boolean(redisUrl))('recovering a run whose ring lost its head', ( await db.delete(workspace).where(eq(workspace.id, workspaceId)) await db.delete(user).where(eq(user.id, userId)) await closeRedisConnection() - 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 - } }) it('recovers from an empty context and persists the whole turn once', async () => { From 0b6935fa764199491d7cb38e2e389b3dd0270360 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 11:36:09 -0700 Subject: [PATCH 09/12] test(mothership): drive the run replay liveness tests through the central agent-url mock --- apps/sim/lib/mothership/request/session/run-replay.test.ts | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/apps/sim/lib/mothership/request/session/run-replay.test.ts b/apps/sim/lib/mothership/request/session/run-replay.test.ts index c73dea4f46d..36ff86dc1a6 100644 --- a/apps/sim/lib/mothership/request/session/run-replay.test.ts +++ b/apps/sim/lib/mothership/request/session/run-replay.test.ts @@ -1,9 +1,7 @@ +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', async (importOriginal) => ({ - ...(await importOriginal()), - getMothershipBaseURL: async () => 'https://worker.test', -})) +vi.mock('@/lib/mothership/server/agent-url', () => mothershipAgentUrlMock) import { forwardRunReplay, From 6a57d077efa4cf7536c3b87329c3586bba8659e1 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 12:31:25 -0700 Subject: [PATCH 10/12] fix(mothership): log a worker's refusal of the run replay, and test the mid-tail trim race - A 401 or 403 from the worker's replay endpoint is logged with its status, so a rotated or wrong worker key is visible instead of every reader silently falling back to replay_gap. - A live tail whose ring trims past its cursor between polls ends without a terminal and never delivers the events after the gap. --- .../app/api/copilot/chat/stream/route.test.ts | 27 +++++++++++++++++++ .../mothership/request/session/run-replay.ts | 7 +++++ 2 files changed, 34 insertions(+) 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 ea176606664..f2815f0e665 100644 --- a/apps/sim/app/api/copilot/chat/stream/route.test.ts +++ b/apps/sim/app/api/copilot/chat/stream/route.test.ts @@ -269,4 +269,31 @@ describe('copilot chat stream replay route', () => { 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/lib/mothership/request/session/run-replay.ts b/apps/sim/lib/mothership/request/session/run-replay.ts index 11c90989b86..22f3220f0f6 100644 --- a/apps/sim/lib/mothership/request/session/run-replay.ts +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -76,6 +76,13 @@ export async function openRunReplay(params: { 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 } From df43c8a6b7d91e3be47579e62b060efc63671510 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 12:57:56 -0700 Subject: [PATCH 11/12] fix(mothership): create the log re-sync set outside render --- .../workspace/[workspaceId]/home/hooks/use-chat.ts | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) 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 a9c927ef54b..a662fc071dc 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -1,5 +1,6 @@ import { type Dispatch, + type RefObject, type SetStateAction, useCallback, useEffect, @@ -311,6 +312,12 @@ const logger = createLogger('useChat') * 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. */ +/** The streams re-synced from the log, created on first use outside render. */ +function logResyncedStreams(ref: RefObject | null>): Set { + ref.current ??= new Set() + return ref.current +} + function streamReconnectQuery(streamId: string, afterCursor: string, fromLog: boolean): string { return `streamId=${encodeURIComponent(streamId)}&after=${encodeURIComponent(afterCursor)}${fromLog ? '&source=log' : ''}` } @@ -966,7 +973,6 @@ export function useChat( const locallyTerminalStreamIdRef = useRef(undefined) const lastCursorRef = useRef('0') const logResyncedStreamsRef = useRef | null>(null) - logResyncedStreamsRef.current ??= new Set() const activeStreamReturnRecoveryRef = useRef(null) const sendingRef = useRef(false) const streamGenRef = useRef(0) @@ -2382,7 +2388,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?${streamReconnectQuery(streamId, afterCursor, logResyncedStreamsRef.current?.has(streamId) ?? false)}&batch=true`, + `/api/mothership/chat/stream?${streamReconnectQuery(streamId, afterCursor, logResyncedStreams(logResyncedStreamsRef).has(streamId))}&batch=true`, { signal: fetchSignal, ...(streamTraceparentRef.current @@ -2574,7 +2580,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?${streamReconnectQuery(streamId, latestCursor, logResyncedStreamsRef.current?.has(streamId) ?? false)}`, + `/api/mothership/chat/stream?${streamReconnectQuery(streamId, latestCursor, logResyncedStreams(logResyncedStreamsRef).has(streamId))}`, { signal: activeAbort.signal, ...(streamTraceparentRef.current @@ -2595,7 +2601,7 @@ export function useChat( // Re-sent from the worker's log with cursors restarting at 1: rebuild from empty. if (sseRes.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER) === 'log') { - logResyncedStreamsRef.current?.add(streamId) + logResyncedStreams(logResyncedStreamsRef).add(streamId) const reset = applyReconnectReplaySelection(streamId, '0') latestCursor = reset.afterCursor preserveNextReplayState = reset.preserveExistingState From 26225d79a9016f9a6aac72aaeee6022ce668f215 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 13:12:35 -0700 Subject: [PATCH 12/12] refactor(mothership): reuse the SSE idle timeout for the run replay, and track the log re-sync as one stream id - processSSEStream takes an optional idle timeout and passes it to readSSELines, so the run replay uses the shared idle bound instead of its own reader wrapper. Callers that omit it are unchanged; the replay's header wait keeps its own timer. - A chat view re-syncs one stream at a time, so the log re-sync flag is the stream's id rather than a set. --- .../[workspaceId]/home/hooks/use-chat.ts | 15 ++--- apps/sim/lib/mothership/request/go/parser.ts | 6 +- .../mothership/request/session/run-replay.ts | 65 ++++++------------- 3 files changed, 29 insertions(+), 57 deletions(-) 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 a662fc071dc..0f506f151b8 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -1,6 +1,5 @@ import { type Dispatch, - type RefObject, type SetStateAction, useCallback, useEffect, @@ -312,12 +311,6 @@ const logger = createLogger('useChat') * 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. */ -/** The streams re-synced from the log, created on first use outside render. */ -function logResyncedStreams(ref: RefObject | null>): Set { - ref.current ??= new Set() - return ref.current -} - function streamReconnectQuery(streamId: string, afterCursor: string, fromLog: boolean): string { return `streamId=${encodeURIComponent(streamId)}&after=${encodeURIComponent(afterCursor)}${fromLog ? '&source=log' : ''}` } @@ -972,7 +965,7 @@ export function useChat( const streamRequestIdRef = useRef(undefined) const locallyTerminalStreamIdRef = useRef(undefined) const lastCursorRef = useRef('0') - const logResyncedStreamsRef = useRef | null>(null) + const logResyncedStreamIdRef = useRef(null) const activeStreamReturnRecoveryRef = useRef(null) const sendingRef = useRef(false) const streamGenRef = useRef(0) @@ -2388,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?${streamReconnectQuery(streamId, afterCursor, logResyncedStreams(logResyncedStreamsRef).has(streamId))}&batch=true`, + `/api/mothership/chat/stream?${streamReconnectQuery(streamId, afterCursor, logResyncedStreamIdRef.current === streamId)}&batch=true`, { signal: fetchSignal, ...(streamTraceparentRef.current @@ -2580,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?${streamReconnectQuery(streamId, latestCursor, logResyncedStreams(logResyncedStreamsRef).has(streamId))}`, + `/api/mothership/chat/stream?${streamReconnectQuery(streamId, latestCursor, logResyncedStreamIdRef.current === streamId)}`, { signal: activeAbort.signal, ...(streamTraceparentRef.current @@ -2601,7 +2594,7 @@ export function useChat( // Re-sent from the worker's log with cursors restarting at 1: rebuild from empty. if (sseRes.headers.get(MOTHERSHIP_STREAM_REPLAY_HEADER) === 'log') { - logResyncedStreams(logResyncedStreamsRef).add(streamId) + logResyncedStreamIdRef.current = streamId const reset = applyReconnectReplaySelection(streamId, '0') latestCursor = reset.afterCursor preserveNextReplayState = reset.preserveExistingState 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/session/run-replay.ts b/apps/sim/lib/mothership/request/session/run-replay.ts index 22f3220f0f6..992894b76d6 100644 --- a/apps/sim/lib/mothership/request/session/run-replay.ts +++ b/apps/sim/lib/mothership/request/session/run-replay.ts @@ -93,36 +93,6 @@ export async function openRunReplay(params: { return response.body } -/** - * A replay reader that ends, as a cut connection does, once the worker sends nothing, - * not even a keepalive, for {@link WORKER_STREAM_IDLE_TIMEOUT_MS}. - */ -function idleBoundedReader( - reader: ReadableStreamDefaultReader -): ReadableStreamDefaultReader { - return { - async read() { - let timer: ReturnType | undefined - const silent = new Promise>((resolve) => { - timer = setTimeout(() => { - reader.cancel().catch(() => {}) - resolve({ done: true, value: undefined }) - }, WORKER_STREAM_IDLE_TIMEOUT_MS) - }) - try { - return await Promise.race([reader.read(), silent]) - } finally { - clearTimeout(timer) - } - }, - cancel: (reason) => reader.cancel(reason), - releaseLock: () => reader.releaseLock(), - get closed() { - return reader.closed - }, - } -} - /** Every reason the worker may end a replay with; a reason added to the contract fails here. */ const REPLAY_END_REASONS = { parked: true, @@ -166,21 +136,26 @@ async function readRunReplay( onEvent: (event: PersistedStreamEventEnvelope) => boolean ): Promise { let end: RunReplayEnd = 'closed' - await processSSEStream(idleBoundedReader(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 - }) + 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 }