From 3214612cef9073fcfd74892196b32e4dc00fec7a Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 1 Oct 2026 12:54:51 -0700 Subject: [PATCH 1/3] fix(mothership): bound stream recovery and hand off only on a lost lease A recovered Chat controller that failed to persist an event for any reason other than a budget refusal was treated as superseded, so it left its run recoverable without finalizing. Every reconnect poll then claimed the run again and re-POSTed it to the worker, for as long as a tab stayed open. - Only a lease another controller holds is a hand-off. A failed leased append latches StreamPersistenceFailedError in the writer, like a budget refusal: the turn ends as an error through the existing terminal path, which settles the run and stops the worker run. An unreadable lease proves nothing, so the fenced writes that follow decide ownership. - Refusals, persistence failures and recovery exhaustion share one StreamTurnFailure base that the lifecycle classifies as an error, never a cancellation or a hand-off. - Recovery takeovers carry a per-run budget in the run's request context, written in the same token-compared claim UPDATE, so it covers every pod, tab and caller. The first takeover is immediate; later ones wait an exponential backoff (1 s base, 60 s cap). After 5 takeovers without a controller staying alive for 5 minutes, the next claim ends the run as an error and stops the worker run. Each claim logs one line with its attempt number. --- .../async-runs/orphaned-runs.integration.ts | 7 + .../application/recover-stream.test.ts | 4 +- .../request/application/recover-stream.ts | 16 +- apps/sim/lib/mothership/request/go/parser.ts | 4 +- .../request/lifecycle/controller-ownership.ts | 60 ++- .../lib/mothership/request/lifecycle/run.ts | 24 +- .../lib/mothership/request/lifecycle/start.ts | 114 +++-- .../session/recovery-storm.integration.ts | 450 ++++++++++++++++++ .../session/replay-budget.integration.ts | 43 +- .../request/session/replay-budget.ts | 11 +- .../request/session/turn-failure.ts | 53 +++ .../mothership/request/session/writer.test.ts | 114 +++-- .../lib/mothership/request/session/writer.ts | 63 ++- 13 files changed, 810 insertions(+), 153 deletions(-) create mode 100644 apps/sim/lib/mothership/request/session/recovery-storm.integration.ts create mode 100644 apps/sim/lib/mothership/request/session/turn-failure.ts diff --git a/apps/sim/lib/mothership/async-runs/orphaned-runs.integration.ts b/apps/sim/lib/mothership/async-runs/orphaned-runs.integration.ts index 05ef2d58341..1a51dd8d062 100644 --- a/apps/sim/lib/mothership/async-runs/orphaned-runs.integration.ts +++ b/apps/sim/lib/mothership/async-runs/orphaned-runs.integration.ts @@ -69,6 +69,9 @@ import { chatStreamLockKey, } from '@/lib/mothership/request/session/controller-lease' +/** A recovering controller's first takeover of a run. */ +const FIRST_RECOVERY = { attempts: 1, claimedAt: 0, notBefore: 0 } + function redis() { const client = getRedisClient() if (!client) throw new Error('The integration suite requires TEST_REDIS_URL') @@ -466,6 +469,7 @@ describe.runIf(Boolean(redisUrl))('Chat runs no controller owns', () => { chatId: orphan.chatId, previousToken: orphan.controllerToken!, token: `${orphan.streamId}\n${generateId()}`, + recoveryBackoff: FIRST_RECOVERY, }), sweepOrphanedRuns(), ]) @@ -533,6 +537,7 @@ describe.runIf(Boolean(redisUrl))('Chat runs no controller owns', () => { chatId: orphan.chatId, previousToken: orphan.controllerToken!, token: `${orphan.streamId}\n${generateId()}`, + recoveryBackoff: FIRST_RECOVERY, }) ) ), @@ -570,6 +575,7 @@ describe.runIf(Boolean(redisUrl))('Chat runs no controller owns', () => { chatId: orphan.chatId, previousToken: orphan.controllerToken!, token: `${orphan.streamId}\n${generateId()}`, + recoveryBackoff: FIRST_RECOVERY, }), settleStoppedRunWithoutController(orphan.runId), ]) @@ -604,6 +610,7 @@ describe.runIf(Boolean(redisUrl))('Chat runs no controller owns', () => { chatId: orphan.chatId, previousToken: orphan.controllerToken!, token: lease.value, + recoveryBackoff: FIRST_RECOVERY, }) return { owned, claimed } } finally { diff --git a/apps/sim/lib/mothership/request/application/recover-stream.test.ts b/apps/sim/lib/mothership/request/application/recover-stream.test.ts index 67009e610e1..7b0d466bd35 100644 --- a/apps/sim/lib/mothership/request/application/recover-stream.test.ts +++ b/apps/sim/lib/mothership/request/application/recover-stream.test.ts @@ -43,7 +43,8 @@ vi.mock('@/lib/mothership/request/session/controller-lease', async (original) => ...(await original()), assertChatStreamLease: hoisted.assertLease, })) -vi.mock('@/lib/mothership/request/lifecycle/controller-ownership', () => ({ +vi.mock('@/lib/mothership/request/lifecycle/controller-ownership', async (original) => ({ + ...(await original()), claimRunController: hoisted.claim, })) vi.mock('@/lib/mothership/request/session/buffer', () => ({ @@ -187,6 +188,7 @@ describe('authorized chat stream recovery', () => { chatId: '22222222-2222-4222-8222-222222222222', previousToken: 'old-controller', token: 'stream\nnew-controller', + recoveryBackoff: expect.objectContaining({ attempts: 1 }), }) expect(mocks.start).toHaveBeenCalledOnce() const params = mocks.start.mock.calls[0][0] diff --git a/apps/sim/lib/mothership/request/application/recover-stream.ts b/apps/sim/lib/mothership/request/application/recover-stream.ts index 10b1c9ad390..95c3e8673bd 100644 --- a/apps/sim/lib/mothership/request/application/recover-stream.ts +++ b/apps/sim/lib/mothership/request/application/recover-stream.ts @@ -15,7 +15,10 @@ import { defineAuthorizedChatUseCase } from '@/lib/mothership/chat/application/a import { resolveOwnedChatContext } from '@/lib/mothership/chat/application/context' import { buildOnComplete, buildOnError } from '@/lib/mothership/chat/completion' import { restoreBillingAdmission } from '@/lib/mothership/request/lifecycle/admission' -import { claimRunController } from '@/lib/mothership/request/lifecycle/controller-ownership' +import { + claimRunController, + planRecovery, +} from '@/lib/mothership/request/lifecycle/controller-ownership' import { StreamRecoveryConfigSchema } from '@/lib/mothership/request/lifecycle/recovery-config' import { createSSEStream } from '@/lib/mothership/request/lifecycle/start' import { isTerminalStreamStatus } from '@/lib/mothership/request/session' @@ -28,6 +31,7 @@ import { getLatestSeq, readEvents } from '@/lib/mothership/request/session/buffe 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 { StreamRecoveryExhaustedError } from '@/lib/mothership/request/session/turn-failure' import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils' const logger = createLogger('MothershipStreamRecovery') @@ -81,6 +85,8 @@ export const readChatStream = defineAuthorizedChatUseCase({ ) { throw new OrchestrationError('validation', 'Saved stream identity does not match its chat') } + const plan = planRecovery(saved.recoveryBackoff, Date.now()) + if (plan.kind === 'wait') return run if (!(await acquirePendingChatStream(chatId, run.streamId, 0))) return run const lease = getLocalChatStreamLease(chatId, run.streamId)! try { @@ -91,11 +97,18 @@ export const readChatStream = defineAuthorizedChatUseCase({ chatId, previousToken: saved.controllerToken, token: lease.value, + recoveryBackoff: plan.backoff, })) ) { await releasePendingChatStream(chatId, run.streamId, lease) return (await getLatestRunForStream(run.streamId, userId)) ?? run } + logger.info('Claimed stream run for recovery', { + runId: run.id, + streamId: run.streamId, + attempt: plan.backoff.attempts, + exhausted: plan.kind === 'exhausted', + }) if (isHosted && !config.data.billingAdmission) throw new OrchestrationError( 'forbidden', @@ -163,6 +176,7 @@ export const readChatStream = defineAuthorizedChatUseCase({ message: '', titleModel: '', resumeSeq, + ...(plan.kind === 'exhausted' ? { failure: new StreamRecoveryExhaustedError() } : {}), orchestrateOptions: { userId, workspaceId, diff --git a/apps/sim/lib/mothership/request/go/parser.ts b/apps/sim/lib/mothership/request/go/parser.ts index a5ef054e2cf..f5a8af01d81 100644 --- a/apps/sim/lib/mothership/request/go/parser.ts +++ b/apps/sim/lib/mothership/request/go/parser.ts @@ -2,7 +2,7 @@ import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { readSSELines } from '@/lib/core/utils/sse' import { StreamControllerSupersededError } from '@/lib/mothership/request/session/controller-lease' -import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' +import { StreamTurnFailure } from '@/lib/mothership/request/session/turn-failure' const logger = createLogger('CopilotSseParser') @@ -55,7 +55,7 @@ export async function processSSEStream( if ( error instanceof FatalSseEventError || error instanceof StreamControllerSupersededError || - error instanceof StreamReplayBudgetExhaustedError + error instanceof StreamTurnFailure ) throw error logger.warn('Failed to handle SSE event', { diff --git a/apps/sim/lib/mothership/request/lifecycle/controller-ownership.ts b/apps/sim/lib/mothership/request/lifecycle/controller-ownership.ts index 585b61e350b..a00816b74ed 100644 --- a/apps/sim/lib/mothership/request/lifecycle/controller-ownership.ts +++ b/apps/sim/lib/mothership/request/lifecycle/controller-ownership.ts @@ -1,13 +1,66 @@ import { db } from '@sim/db' import { copilotChats, copilotRuns } from '@sim/db/schema' +import { backoffWithJitter } from '@sim/utils/retry' import { and, eq, notInArray, sql } from 'drizzle-orm' +import { z } from 'zod' -/** Serializes takeover with assistant persistence; Redis alone cannot fence a delayed DB write. */ +/** + * Takeovers of one run allowed in a row before it ends as an error. A controller keeps + * its chat lock while it streams and while it waits on parked, permission-gated or + * client-executed tools, so only a lost controller (a pod deploy or crash, or a lost + * lease) needs one. Five in a row is a crash loop, not a deploy. + */ +export const MAX_RECOVERY_ATTEMPTS = 5 + +/** A takeover this long after the previous one starts a fresh budget: that controller lived. */ +export const RECOVERY_BUDGET_RESET_MS = 5 * 60_000 + +const RECOVERY_BACKOFF = { baseMs: 1_000, maxMs: 60_000 } as const + +/** Stored on the run beside its controller token, in epoch milliseconds. */ +const RecoveryBackoffSchema = z.object({ + attempts: z.number().int().positive(), + claimedAt: z.number(), + notBefore: z.number(), +}) + +export type RecoveryBackoff = z.infer + +export type RecoveryPlan = + | { kind: 'wait' } + | { kind: 'claim' | 'exhausted'; backoff: RecoveryBackoff } + +/** + * Exponential backoff with a restart limit for taking over a run, as a supervisor + * restarts a crashing child. The first takeover is immediate; each later one waits + * until the previous one's `notBefore`. Seq progress does not reset the budget: a + * recovered leg re-persists the frames the worker replays since its last checkpoint, + * so a crash loop advances the replay ring on every attempt. + */ +export function planRecovery(saved: unknown, now: number): RecoveryPlan { + const previous = RecoveryBackoffSchema.safeParse(saved) + const fresh = !previous.success || now - previous.data.claimedAt >= RECOVERY_BUDGET_RESET_MS + if (previous.success && !fresh && now < previous.data.notBefore) return { kind: 'wait' } + const attempts = fresh ? 1 : previous.data.attempts + 1 + const backoff = { + attempts, + claimedAt: now, + notBefore: now + backoffWithJitter(attempts, null, RECOVERY_BACKOFF), + } + return { kind: attempts > MAX_RECOVERY_ATTEMPTS ? 'exhausted' : 'claim', backoff } +} + +/** + * Serializes takeover with assistant persistence; Redis alone cannot fence a delayed DB write. + * Every claim replaces the controller token it compares, so the recovery budget written + * beside it is as atomic as the claim, across pods, tabs and callers. + */ export async function claimRunController(input: { runId: string chatId: string previousToken: string token: string + recoveryBackoff: RecoveryBackoff }): Promise { return db.transaction(async (tx) => { await tx @@ -18,7 +71,10 @@ export async function claimRunController(input: { const [run] = await tx .update(copilotRuns) .set({ - requestContext: sql`jsonb_set(${copilotRuns.requestContext}, '{controllerToken}', ${JSON.stringify(input.token)}::jsonb)`, + requestContext: sql`${copilotRuns.requestContext} || ${JSON.stringify({ + controllerToken: input.token, + recoveryBackoff: input.recoveryBackoff, + })}::jsonb`, updatedAt: new Date(), }) .where( diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index e2a7913f8e4..5f8e24c18cb 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -53,7 +53,7 @@ import { StreamRetryWindow } from '@/lib/mothership/request/lifecycle/stream-ret import { recordDegraded } from '@/lib/mothership/request/metrics' import { AbortReason } from '@/lib/mothership/request/session/abort-reason' import { StreamControllerSupersededError } from '@/lib/mothership/request/session/controller-lease' -import { replayRefusal } from '@/lib/mothership/request/session/replay-budget' +import { turnFailure } from '@/lib/mothership/request/session/turn-failure' import { getToolCallTerminalData, requireToolCallStateResult, @@ -562,27 +562,27 @@ export async function runCopilotLifecycle( // the work the user watched succeed. const backendFinishedTurn = context.completionStatus === MothershipStreamV1CompletionStatus.complete - // A refused replay write aborts the turn to stop it, but the turn failed; it - // was not stopped by the user. - const refusal = replayRefusal(lifecycleOptions.abortSignal?.reason) + // A turn failure (such as a refused replay write) aborts the turn to stop it, but + // the turn failed; it was not stopped by the user. + const failure = turnFailure(lifecycleOptions.abortSignal?.reason) // Consult the lifecycle signal as well as the flag. `context.wasAborted` is // only reached from a fanout leg through the (deliberately asymmetric) merge // in `mergeResumeLegOutputs`, so a Stop landing mid-fanout could otherwise // classify the turn as a success. Mirrors the check already used below on // the throw path. const turnWasAborted = - !refusal && + !failure && (context.completionStatus === MothershipStreamV1CompletionStatus.cancelled || context.wasAborted || (lifecycleOptions.abortSignal?.aborted ?? false)) const succeeded = - !refusal && + !failure && !turnWasAborted && (backendFinishedTurn || (!context.completionStatus && context.errors.length === 0)) // The worker sends an error terminal with no `error` event only when it replays a run // that already ended (for example at its deadline) to a resume or reattach, because // that replay does not carry the run's stored reason. Say so rather than leave the turn - // to a generic failure; a reported reason or a replay refusal always wins. + // to a generic failure; a reported reason or a turn failure always wins. const endedWithoutReason = !turnWasAborted && context.completionStatus === MothershipStreamV1CompletionStatus.error && @@ -606,7 +606,7 @@ export async function runCopilotLifecycle( chatId: context.chatId, requestId: context.requestId, ...(endedWithoutReason ? { error: ENDED_RUN_MESSAGE } : {}), - ...(refusal ? { error: refusal.userMessage, errorCode: refusal.code } : {}), + ...(failure ? { error: failure.userMessage, errorCode: failure.code } : {}), errors: !succeeded && context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, @@ -645,8 +645,8 @@ export async function runCopilotLifecycle( // partial content can be appended. // Return `cancelled: true` so upstream classification stays // consistent with the success-path cancel result. - const refusal = replayRefusal(lifecycleOptions.abortSignal?.reason) - const wasCancelled = !refusal && (lifecycleOptions.abortSignal?.aborted ?? false) + const failure = turnFailure(lifecycleOptions.abortSignal?.reason) + const wasCancelled = !failure && (lifecycleOptions.abortSignal?.aborted ?? false) // Preserve whatever streamed before the throw for both terminals. A thrown // backend error (as opposed to an `error` SSE event that lets the loop finish // normally) must still carry the partial assistant turn so onError can @@ -661,8 +661,8 @@ export async function runCopilotLifecycle( toolCalls: buildToolCallSummaries(context), chatId: context.chatId, requestId: context.requestId, - error: refusal?.userMessage ?? err.message, - ...(refusal ? { errorCode: refusal.code } : {}), + error: failure?.userMessage ?? err.message, + ...(failure ? { errorCode: failure.code } : {}), errors: context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 4be603e2be8..a817cbf7ce2 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -60,11 +60,13 @@ import { holdsChatStreamLease, StreamControllerSupersededError, } from '@/lib/mothership/request/session/controller-lease' -import { - replayRefusal, - type StreamReplayBudgetExhaustedError, -} from '@/lib/mothership/request/session/replay-budget' +import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' import { SSE_RESPONSE_HEADERS } from '@/lib/mothership/request/session/sse' +import { + StreamPersistenceFailedError, + type StreamTurnFailure, + turnFailure, +} from '@/lib/mothership/request/session/turn-failure' import { TraceCollector } from '@/lib/mothership/request/trace' import type { OrchestratorResult } from '@/lib/mothership/request/types' import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url' @@ -95,6 +97,8 @@ export interface StreamingOrchestrationParams { /** Interactive admission commits before the HTTP stream is exposed. */ admittedRun?: Awaited> resumeSeq?: number + /** Ends the turn with this failure instead of contacting the worker. */ + failure?: StreamTurnFailure /** Pre-started root; omit to let the stream start its own root. */ otelRoot?: ReturnType } @@ -133,17 +137,32 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS }) const abortController = new AbortController() + if (params.failure) abortController.abort(params.failure) registerActiveStream(streamId, abortController) const lease = chatId ? getLocalChatStreamLease(chatId, streamId) : undefined + /** + * Only a lease another controller holds proves a takeover. A lease that cannot be read + * proves nothing: the fenced writes that follow (the leased replay append and the + * token-checked run updates) decide, instead of leaving the run to a successor. + */ const assertControllerOwnership = async () => { if (!chatId) return - try { - if (!lease || abortController.signal.reason instanceof StreamControllerSupersededError) { - throw new StreamControllerSupersededError() - } - await assertChatStreamLease(lease) - } catch { + const superseded = + !lease || + abortController.signal.reason instanceof StreamControllerSupersededError || + (await assertChatStreamLease(lease).then( + () => false, + (error: unknown) => { + if (error instanceof StreamControllerSupersededError) return true + logger.warn(`[${requestId}] Chat stream lease is unreadable`, { + streamId, + error: getErrorMessage(error), + }) + return false + } + )) + if (superseded) { const error = new StreamControllerSupersededError() abortController.abort(error) throw error @@ -211,11 +230,11 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS const collector = new TraceCollector() /** - * The replay refusal that ends this turn. Once the controller is aborted its - * reason is authoritative, so a refusal that follows a Stop stays a Stop. + * The failure that ends this turn. Once the controller is aborted its reason is + * authoritative, so a failure that follows a Stop stays a Stop. */ - const refusalOf = (thrown?: unknown) => - replayRefusal(abortController.signal.aborted ? abortController.signal.reason : thrown) + const failureOf = (thrown?: unknown) => + turnFailure(abortController.signal.aborted ? abortController.signal.reason : thrown) /** * Set once this controller ended the turn, even if publishing its terminal @@ -234,21 +253,26 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS } /** - * A refused replay write ends the turn as an error: every replacement would be - * refused the same event. The run is marked terminal before the lock is - * released, and the worker is told to stop so it does not wait out a tool call - * that was never dispatched. + * A turn failure ends the turn as an error: every replacement would fail the same + * way. The run is marked terminal before the lock is released, and the worker is told + * to stop so it does not wait out a tool call that was never dispatched. */ - const finalizeAfterReplayRefusal = async ( - refusal: StreamReplayBudgetExhaustedError, + const finalizeAfterTurnFailure = async ( + failure: StreamTurnFailure, result?: OrchestratorResult ) => { - logger.warn(`[${requestId}] Stream replay budget exhausted; ending the turn`, { + logger.warn(`[${requestId}] ${failure.message}; ending the turn`, { streamId, - resource: refusal.refusal.resource, - attemptedBytes: refusal.refusal.attemptedBytes, - currentBytes: refusal.refusal.currentBytes, - limitBytes: refusal.refusal.limitBytes, + code: failure.code, + ...(failure instanceof StreamReplayBudgetExhaustedError + ? { + resource: failure.refusal.resource, + attemptedBytes: failure.refusal.attemptedBytes, + currentBytes: failure.refusal.currentBytes, + limitBytes: failure.refusal.limitBytes, + } + : {}), + ...(failure.cause === undefined ? {} : { cause: getErrorMessage(failure.cause) }), }) await endTurn( { @@ -258,8 +282,8 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS ...result, success: false, cancelled: false, - error: refusal.userMessage, - errorCode: refusal.code, + error: failure.userMessage, + errorCode: failure.code, }, publisher, runId, @@ -272,7 +296,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS ) await requestExplicitStreamAbort({ streamId, userId, chatId }) } catch (error) { - logger.warn(`[${requestId}] Worker stop after replay refusal was not delivered`, { + logger.warn(`[${requestId}] Worker stop after a turn failure was not delivered`, { streamId, error: getErrorMessage(error), }) @@ -424,17 +448,17 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS await publisher.publish(event) } catch (error) { /* - A refused write is a terminal failure, not a handoff: leaving it - recoverable made each replacement re-receive and re-refuse the same - event. Any other failure to persist means this controller can no - longer prove ownership of the replay, so a successor takes over. + Only a lost lease hands the turn to a successor. A refusal, or a + Redis error that outlasted its retries, is terminal: leaving it + recoverable made each replacement re-receive and fail the same + event, re-POSTing the run on every reconnect poll. */ - if (!abortController.signal.aborted) { - abortController.abort( - replayRefusal(error) ?? new StreamControllerSupersededError() - ) - } - throw error + const failure = + error instanceof StreamControllerSupersededError + ? error + : (turnFailure(error) ?? new StreamPersistenceFailedError(error)) + if (!abortController.signal.aborted) abortController.abort(failure) + throw failure } }, onAbortObserved: (reason) => { @@ -445,11 +469,11 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS }) lifecycleResult = result - const refusal = refusalOf() + const failure = failureOf() // A completed result wins a late Stop; passive disconnection never cancels. outcome = result.success ? RequestTraceV1Outcome.success - : refusal + : failure ? RequestTraceV1Outcome.error : result.cancelled || abortController.signal.aborted ? RequestTraceV1Outcome.cancelled @@ -458,8 +482,8 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS cancelReason = recordCancelled() } await assertControllerOwnership() - if (refusal && !result.success) { - await finalizeAfterReplayRefusal(refusal, result) + if (failure && !result.success) { + await finalizeAfterTurnFailure(failure, result) } else { await endTurn(result, publisher, runId, outcome, requestId) } @@ -472,10 +496,10 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS return } await assertControllerOwnership() - const refusal = refusalOf(error) - if (refusal) { + const failure = failureOf(error) + if (failure) { outcome = RequestTraceV1Outcome.error - await finalizeAfterReplayRefusal(refusal) + await finalizeAfterTurnFailure(failure) return } const wasCancelled = abortController.signal.aborted diff --git a/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts b/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts new file mode 100644 index 00000000000..8b4deb26b12 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts @@ -0,0 +1,450 @@ +/** + * How often reconnecting to an orphaned Chat run re-POSTs it to the worker, against real + * Redis and PostgreSQL. The production reconnect route, stream recovery, chat lifecycle + * and finalization run unmodified; a local HTTP server stands in for the worker. Faults + * are injected at one seam, the leased replay append: an append that fails for a reason + * other than a lost lease, or a lease lost with no successor to take over. + */ +import { authMock, authMockFns } from '@sim/testing/mocks/auth.mock' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' + +const { redisUrl, inheritedEnv, worker, faults } = await vi.hoisted(async () => { + const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') + const { createServer } = await import('node:http') + const faults = { + /** Which leased appends fail, and how; `undefined` lets every append through. */ + append: undefined as + | undefined + | { frames: 'any_tool' | 'tool_result'; effect: 'throw' | 'lose_lease' }, + } + const worker = { + posts: [] as Array<{ path: string; at: number }>, + mode: 'frames' as 'frames' | 'json500' | 'drop', + frames: [] as unknown[], + } + const server = createServer(async (request, response) => { + await new Promise((resolve) => request.on('end', resolve).resume()) + worker.posts.push({ path: request.url ?? '', at: Date.now() }) + if (request.url === '/api/streams/explicit-abort') { + response.writeHead(200, { 'content-type': 'application/json' }) + response.end(JSON.stringify({ settled: true })) + return + } + if (worker.mode === 'json500') { + response.writeHead(500, { 'content-type': 'application/json' }) + response.end(JSON.stringify({ error: 'internal' })) + return + } + response.writeHead(200, { 'content-type': 'text/event-stream' }) + if (worker.mode === 'drop') { + response.write(`data: ${JSON.stringify(worker.frames[0])}\n\n`) + setTimeout(() => response.socket?.destroy(), 20) + return + } + for (const frame of worker.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, + } + /** The real Redis module and worker URL resolution read these at import. */ + 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 }), faults } +}) + +vi.mock('@/lib/auth', () => authMock) +vi.mock('@/lib/mothership/request/session/buffer', async (importOriginal) => { + const actual = await importOriginal() + const { getRedisClient } = await import('@/lib/core/config/redis') + return { + ...actual, + appendEvents: async (...args: Parameters) => { + const [envelopes, , lease] = args + const fault = faults.append + const hit = + lease && + fault && + envelopes.some( + (envelope) => + envelope.type === 'tool' && + (fault.frames === 'any_tool' || + (envelope.payload as { phase?: string }).phase === 'result') + ) + if (hit && fault.effect === 'throw') throw new Error('simulated Redis write failure') + /** The lock expires under a live controller, and nobody else holds it. */ + if (hit && fault.effect === 'lose_lease') await getRedisClient()!.del(lease.key) + return actual.appendEvents(...args) + }, + } +}) + +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 { eq } from 'drizzle-orm' +import { NextRequest } from 'next/server' +import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis' +import { + MAX_RECOVERY_ATTEMPTS, + RECOVERY_BUDGET_RESET_MS, +} from '@/lib/mothership/request/lifecycle/controller-ownership' +import { isTerminalStreamStatus } from '@/lib/mothership/request/session' +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 { + StreamPersistenceFailedError, + StreamRecoveryExhaustedError, +} from '@/lib/mothership/request/session/turn-failure' +import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' + +const userId = generateId() +const workspaceId = generateId() +const chatIds: string[] = [] + +function redis() { + const client = getRedisClient() + if (!client) throw new Error('The integration suite requires TEST_REDIS_URL') + return client +} + +/** A run whose controller died after persisting two text frames, its lock long expired. */ +async function orphanedRun( + options: { + status?: 'active' | 'paused_waiting_for_tool' + recoveryBackoff?: Record + } = {} +) { + const chatId = generateId() + chatIds.push(chatId) + const streamId = generateId() + const runId = generateId() + const request = { message: 'Summarize', 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, + status: options.status ?? 'active', + toolExecutionVersion: 2, + requestContext: { + requestId: generateId(), + controllerToken: `dead\n${generateId()}`, + recovery: { + kind: 'interactive_stream', + request, + goRoute: '/api/mothership', + clientToolPickupExpected: false, + }, + ...(options.recoveryBackoff ? { recoveryBackoff: options.recoveryBackoff } : {}), + }, + }) + await appendEvents( + [1, 2].map((seq) => + createEvent({ + streamId, + cursor: String(seq), + seq, + requestId: generateId(), + type: 'text', + payload: { channel: 'assistant', text: `part ${seq} ` }, + }) + ), + { streamId } + ) + const frame = (seq: number, type: string, payload: unknown) => ({ + v: 1, + type, + seq, + ts: new Date().toISOString(), + stream: { streamId, chatId }, + payload, + }) + worker.frames = [ + frame(1, 'tool', { + phase: 'call', + toolCallId: 'go-call', + toolName: 'search_online', + executor: 'go', + mode: 'sync', + arguments: { query: 'logs' }, + }), + frame(2, 'tool', { + phase: 'result', + toolCallId: 'go-call', + toolName: 'search_online', + executor: 'go', + mode: 'sync', + success: true, + output: { results: [] }, + }), + frame(3, 'complete', { status: 'complete', textLength: 14 }), + ] + return { chatId, streamId, runId } +} + +async function storedRun(runId: string) { + const [run] = await db + .select({ + status: copilotRuns.status, + error: copilotRuns.error, + requestContext: copilotRuns.requestContext, + }) + .from(copilotRuns) + .where(eq(copilotRuns.id, runId)) + return { + ...run, + recoveryBackoff: (run.requestContext as { recoveryBackoff?: Record }) + .recoveryBackoff, + } +} + +/** + * Holds `tails` browser-like reconnect tails open until the run is terminal and its chat + * lock released, or `windowMs` passes, and returns the worker requests made meanwhile. + */ +async function reconnect( + { chatId, streamId, runId }: { chatId: string; streamId: string; runId: string }, + { tails = 1, windowMs }: { tails?: number; windowMs: number } +) { + const startedAt = Date.now() + const controllers = Array.from({ length: tails }, () => new AbortController()) + const drains = controllers.map(async (controller) => { + const response = await streamGET( + new NextRequest( + `http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=2`, + { signal: controller.signal } + ), + { params: Promise.resolve({}) } + ) + /** The route never observes the request signal under test, so the reader cancels. */ + await response.body?.pipeTo(new WritableStream(), { signal: controller.signal }).catch(() => {}) + }) + while (Date.now() - startedAt < windowMs) { + const settled = + isTerminalStreamStatus((await storedRun(runId)).status) && + !(await redis().get(chatStreamLockKey(chatId))) + if (settled) break + await sleep(100) + } + for (const controller of controllers) controller.abort() + await Promise.allSettled(drains) + for (let i = 0; i < 50 && (await redis().get(chatStreamLockKey(chatId))); i++) await sleep(100) + const requests = worker.posts.filter((post) => post.at >= startedAt) + return { + startedAt, + turns: requests.filter((post) => post.path === '/api/mothership'), + stops: requests.filter((post) => post.path === '/api/streams/explicit-abort'), + } +} + +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))('reconnecting to an orphaned Chat run', () => { + beforeAll(async () => { + const now = new Date() + await db.insert(user).values({ + id: userId, + name: 'Recovery storm fixture', + email: `${userId}@recovery-storm.test`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(workspace).values({ + id: workspaceId, + name: 'Recovery storm fixture', + ownerId: userId, + billedAccountUserId: userId, + }) + await db.insert(permissions).values({ + id: generateId(), + userId, + entityType: 'workspace', + entityId: workspaceId, + permissionType: 'admin', + }) + authMockFns.mockGetSession.mockResolvedValue({ + user: { id: userId }, + session: { id: generateId() }, + }) + }) + + afterAll(async () => { + for (const chatId of chatIds) { + await db.delete(copilotMessages).where(eq(copilotMessages.chatId, chatId)) + await db.delete(copilotRuns).where(eq(copilotRuns.chatId, chatId)) + await db.delete(copilotChats).where(eq(copilotChats.id, chatId)) + } + await db.delete(permissions).where(eq(permissions.userId, userId)) + await db.delete(workspace).where(eq(workspace.id, workspaceId)) + await db.delete(user).where(eq(user.id, userId)) + await closeRedisConnection() + }) + + it.each(['active', 'paused_waiting_for_tool'] as const)( + 'takes over an orphaned %s run at once and completes it', + async (status) => { + worker.mode = 'frames' + const run = await orphanedRun({ status }) + + const { startedAt, turns } = await reconnect(run, { windowMs: 10_000 }) + + expect(turns).toHaveLength(1) + expect(turns[0].at - startedAt).toBeLessThan(1_000) + const stored = await storedRun(run.runId) + expect(stored.status).toBe('complete') + expect(stored.recoveryBackoff).toMatchObject({ attempts: 1 }) + }, + 30_000 + ) + + it.each([ + { frames: 'any_tool', tails: 1 }, + { frames: 'any_tool', tails: 3 }, + { frames: 'tool_result', tails: 1 }, + { frames: 'tool_result', tails: 3 }, + ] as const)( + 'ends the run after one worker request when a $frames frame cannot be persisted ($tails tails)', + async ({ frames, tails }) => { + worker.mode = 'frames' + const run = await orphanedRun() + faults.append = { frames, effect: 'throw' } + try { + const { turns, stops } = await reconnect(run, { tails, windowMs: 15_000 }) + + expect(turns).toHaveLength(1) + expect(stops).toHaveLength(1) + expect(await storedRun(run.runId)).toMatchObject({ + status: 'error', + error: new StreamPersistenceFailedError(undefined).userMessage, + }) + } finally { + faults.append = undefined + } + }, + 60_000 + ) + + it('gives up on a run whose recovered controllers keep losing their lease', async () => { + worker.mode = 'frames' + const run = await orphanedRun() + faults.append = { frames: 'any_tool', effect: 'lose_lease' } + try { + const { turns, stops } = await reconnect(run, { tails: 3, windowMs: 60_000 }) + + expect(turns).toHaveLength(MAX_RECOVERY_ATTEMPTS) + /** Each takeover waits at least the jittered floor of the one before it. */ + turns.slice(1).forEach((turn, i) => { + expect(turn.at - turns[i].at).toBeGreaterThanOrEqual(0.7 * 1_000 * 2 ** i) + }) + expect(stops).toHaveLength(1) + expect(await storedRun(run.runId)).toMatchObject({ + status: 'error', + error: new StreamRecoveryExhaustedError().userMessage, + recoveryBackoff: { attempts: MAX_RECOVERY_ATTEMPTS + 1 }, + }) + } finally { + faults.append = undefined + } + }, 120_000) + + it('waits out the backoff of a recent takeover, then takes over and completes the run', async () => { + worker.mode = 'frames' + const now = Date.now() + const notBefore = now + 3_000 + const run = await orphanedRun({ + recoveryBackoff: { attempts: 2, claimedAt: now - 1_000, notBefore }, + }) + + const { turns } = await reconnect(run, { windowMs: 15_000 }) + + expect(turns).toHaveLength(1) + expect(turns[0].at).toBeGreaterThanOrEqual(notBefore) + const stored = await storedRun(run.runId) + expect(stored.status).toBe('complete') + expect(stored.recoveryBackoff).toMatchObject({ attempts: 3 }) + }, 30_000) + + it('starts a fresh budget once the previous takeover is older than the reset window', async () => { + worker.mode = 'frames' + const claimedAt = Date.now() - RECOVERY_BUDGET_RESET_MS - 1_000 + const run = await orphanedRun({ + recoveryBackoff: { attempts: MAX_RECOVERY_ATTEMPTS, claimedAt, notBefore: claimedAt }, + }) + + const { startedAt, turns } = await reconnect(run, { windowMs: 10_000 }) + + expect(turns).toHaveLength(1) + expect(turns[0].at - startedAt).toBeLessThan(1_000) + const stored = await storedRun(run.runId) + expect(stored.status).toBe('complete') + expect(stored.recoveryBackoff).toMatchObject({ attempts: 1 }) + }, 30_000) + + it('never takes over a parked run whose controller still holds the chat lock', async () => { + worker.mode = 'frames' + const run = await orphanedRun({ status: 'paused_waiting_for_tool' }) + const liveController = `${run.streamId}\n${generateId()}` + await redis().set(chatStreamLockKey(run.chatId), liveController, 'EX', 60) + try { + const { turns } = await reconnect(run, { tails: 3, windowMs: 3_000 }) + + expect(turns).toHaveLength(0) + const stored = await storedRun(run.runId) + expect(stored.status).toBe('paused_waiting_for_tool') + expect(stored.recoveryBackoff).toBeUndefined() + } finally { + await redis().del(chatStreamLockKey(run.chatId)) + } + }, 30_000) + + it.each(['json500', 'drop'] as const)( + 'ends the run after its bounded retries when the worker fails every leg (%s)', + async (mode) => { + worker.mode = mode + const run = await orphanedRun() + try { + const { turns } = await reconnect(run, { windowMs: 20_000 }) + + expect(turns).toHaveLength(4) + expect((await storedRun(run.runId)).status).toBe('error') + } finally { + worker.mode = 'frames' + } + }, + 60_000 + ) +}) 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 376286c0b57..992ecbda154 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -147,6 +147,7 @@ import { StreamReplayBudgetExhaustedError, } from '@/lib/mothership/request/session/replay-budget' import { STREAM_STRING_PREVIEW_UNITS } from '@/lib/mothership/request/session/replay-compaction' +import { STREAM_PERSISTENCE_FAILED_CODE } from '@/lib/mothership/request/session/turn-failure' import type { StreamEvent } from '@/lib/mothership/request/session/types' import { StreamWriter } from '@/lib/mothership/request/session/writer' import type { StreamingContext } from '@/lib/mothership/request/types' @@ -658,7 +659,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge } ) - it('leaves a run it handed off recoverable after a transient append failure', async () => { + it('ends a turn whose append fails while it still holds the lease, and stops the worker', async () => { let eventsKey = '' const { streamId, runId, frames } = await runTurn( [ @@ -673,11 +674,37 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge } ) - expect(frames.map((frame) => frame.type)).toEqual(['session']) - expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeGreaterThan(300) + expect(frames.map((frame) => [frame.type, frame.payload.code ?? frame.payload.status])).toEqual( + [ + ['session', undefined], + ['error', STREAM_PERSISTENCE_FAILED_CODE], + ['complete', 'error'], + ] + ) const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) - expect(stored.status).toBe('active') - expect(worker.abortRequests).toEqual([]) + expect(stored.status).toBe('error') + expect(worker.abortRequests).toEqual([expect.objectContaining({ messageId: streamId })]) + expect(worker.runs[0].dispatched).toEqual([]) + expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + }) + + it('settles a finished turn whose lease became unreadable instead of handing it off', async () => { + const { runId, frames } = await runTurn([ + text('Done.'), + async () => { + // A lock key of the wrong type makes every read of the lease fail, as a Redis error would. + await redis().del(chatStreamLockKey(chatId)) + await redis().hset(chatStreamLockKey(chatId), 'corrupt', '1') + }, + ]) + + try { + 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') + } finally { + await redis().del(chatStreamLockKey(chatId)) + } }) it('still cleans up a finished turn whose terminal events could not be published', async () => { @@ -1038,7 +1065,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge ) } - it('marks its run terminal even when the final events cannot be published', async () => { + it('marks its run terminal even when the final events cannot be persisted', async () => { const controllerToken = `owner\n${generateId()}` const { streamId, runId } = await pausedRun(controllerToken) const lease = { key: chatStreamLockKey(generateId()), value: controllerToken } @@ -1047,10 +1074,8 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge await redis().set(`mothership_stream:${streamId}:events`, 'not a sorted set', 'EX', 60) const publisher = new StreamWriter({ streamId, requestId: generateId(), lease }) - const failure = await finalizeAsError(publisher, runId).catch((error: unknown) => error) + await finalizeAsError(publisher, runId) - expect(failure).toBeInstanceOf(Error) - expect(failure).not.toBeInstanceOf(StreamControllerSupersededError) const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) expect(stored.status).toBe('error') await redis().del(lease.key) diff --git a/apps/sim/lib/mothership/request/session/replay-budget.ts b/apps/sim/lib/mothership/request/session/replay-budget.ts index c64a6203abf..81d6c0f70b4 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.ts @@ -1,4 +1,5 @@ import type { RedisBudgetRefusal } from '@/lib/core/redis/byte-budget.server' +import { StreamTurnFailure } from '@/lib/mothership/request/session/turn-failure' /** Run-error code for a turn stopped because its replay buffer refused a write. */ export const REPLAY_BUDGET_EXHAUSTED_CODE = 'replay_budget_exhausted' @@ -11,10 +12,9 @@ const HOURLY_LIMIT_MESSAGE = /** * The replay buffer refused an event a leased controller had to persist before - * delivering it. Unlike {@link StreamControllerSupersededError} this is not a - * handoff: no replacement can persist the same event either, so the turn ends. + * delivering it. No replacement can persist the same event either, so the turn ends. */ -export class StreamReplayBudgetExhaustedError extends Error { +export class StreamReplayBudgetExhaustedError extends StreamTurnFailure { readonly code = REPLAY_BUDGET_EXHAUSTED_CODE constructor(readonly refusal: RedisBudgetRefusal) { @@ -29,8 +29,3 @@ export class StreamReplayBudgetExhaustedError extends Error { : STREAM_LIMIT_MESSAGE } } - -/** The refusal an abort reason or thrown value carries, if it is one. */ -export function replayRefusal(value: unknown): StreamReplayBudgetExhaustedError | undefined { - return value instanceof StreamReplayBudgetExhaustedError ? value : undefined -} diff --git a/apps/sim/lib/mothership/request/session/turn-failure.ts b/apps/sim/lib/mothership/request/session/turn-failure.ts new file mode 100644 index 00000000000..aa92ed76848 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/turn-failure.ts @@ -0,0 +1,53 @@ +/** + * A failure that ends the turn as an error instead of handing it to a successor. + * Only a controller that provably lost its chat lease hands off + * ({@link StreamControllerSupersededError}); a replacement would meet any other failure + * again, re-POSTing the same run to the worker on every reconnect poll. + */ +export abstract class StreamTurnFailure extends Error { + abstract readonly code: string + /** What the user is told. */ + abstract get userMessage(): string +} + +/** The turn failure an abort reason or thrown value carries, if it is one. */ +export function turnFailure(value: unknown): StreamTurnFailure | undefined { + return value instanceof StreamTurnFailure ? value : undefined +} + +/** Run-error code for a turn whose controller could not persist or verify its stream. */ +export const STREAM_PERSISTENCE_FAILED_CODE = 'stream_persistence_failed' + +/** Run-error code for a run whose controllers kept dying before it could finish. */ +export const STREAM_RECOVERY_EXHAUSTED_CODE = 'stream_recovery_exhausted' + +/** + * The controller could not record an event to the replay buffer, or could not read its + * own lease, for a reason other than losing it (a Redis error that outlasted the retries). + */ +export class StreamPersistenceFailedError extends StreamTurnFailure { + readonly code = STREAM_PERSISTENCE_FAILED_CODE + + constructor(cause: unknown) { + super('Stream persistence failed', { cause }) + this.name = 'StreamPersistenceFailedError' + } + + get userMessage(): string { + return 'This response was stopped because it could not be saved while streaming. The work it already completed has been saved — send a message to continue from there.' + } +} + +/** Recovery took over the run too many times without a controller staying alive. */ +export class StreamRecoveryExhaustedError extends StreamTurnFailure { + readonly code = STREAM_RECOVERY_EXHAUSTED_CODE + + constructor() { + super('Stream recovery exhausted') + this.name = 'StreamRecoveryExhaustedError' + } + + get userMessage(): string { + return 'This response was stopped because it was interrupted repeatedly. The work it already completed has been saved — send a message to continue from there.' + } +} diff --git a/apps/sim/lib/mothership/request/session/writer.test.ts b/apps/sim/lib/mothership/request/session/writer.test.ts index 1eab0e0a636..143a075fbca 100644 --- a/apps/sim/lib/mothership/request/session/writer.test.ts +++ b/apps/sim/lib/mothership/request/session/writer.test.ts @@ -15,6 +15,7 @@ vi.mock('@/lib/mothership/request/session/buffer', () => ({ import { StreamControllerSupersededError } from '@/lib/mothership/request/session/controller-lease' import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' +import { StreamPersistenceFailedError } from '@/lib/mothership/request/session/turn-failure' import { StreamWriter } from '@/lib/mothership/request/session/writer' function decodeChunk(value: Uint8Array): string { @@ -63,7 +64,7 @@ describe('StreamWriter', () => { }) it('does not deliver an event from a controller whose durable append was rejected', async () => { - appendEvents.mockRejectedValueOnce(new Error('ownership lost')) + appendEvents.mockRejectedValueOnce(new StreamControllerSupersededError()) const writer = new StreamWriter({ streamId: 'stream-1', requestId: 'req-1', @@ -73,59 +74,74 @@ describe('StreamWriter', () => { writer.attach(controller as unknown as ReadableStreamDefaultController) await expect( writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'stale' } }) - ).rejects.toThrow('ownership lost') + ).rejects.toBeInstanceOf(StreamControllerSupersededError) expect(controller.enqueue).not.toHaveBeenCalled() - await expect(writer.close()).rejects.toThrow('ownership lost') + await expect(writer.close()).rejects.toBeInstanceOf(StreamControllerSupersededError) expect(controller.close).toHaveBeenCalledOnce() }) - it('ends delivery with a budget error, not a handoff, and still delivers the terminal verdict', async () => { - appendEvents.mockResolvedValueOnce({ - persisted: false, - refusal: { - resource: 'owner_redis_bytes', - currentBytes: 32, - limitBytes: 32, - attemptedBytes: 1, - }, - }) - const writer = new StreamWriter({ - streamId: 'stream-1', - requestId: 'req-1', - userId: 'user-1', - lease: { key: 'chat-lock', value: 'current-controller' }, - }) - const delivered: string[] = [] - const controller = { - enqueue: (value: Uint8Array) => delivered.push(decodeChunk(value)), - close: () => {}, + it.each([ + { + failure: 'a budget refusal', + expected: StreamReplayBudgetExhaustedError, + append: () => + appendEvents.mockResolvedValueOnce({ + persisted: false, + refusal: { + resource: 'owner_redis_bytes', + currentBytes: 32, + limitBytes: 32, + attemptedBytes: 1, + }, + }), + }, + { + failure: 'a failed append', + expected: StreamPersistenceFailedError, + append: () => appendEvents.mockRejectedValueOnce(new Error('connection reset')), + }, + ])( + 'ends delivery after $failure, not a handoff, and still delivers the terminal verdict', + async ({ expected, append }) => { + append() + const writer = new StreamWriter({ + streamId: 'stream-1', + requestId: 'req-1', + userId: 'user-1', + lease: { key: 'chat-lock', value: 'current-controller' }, + }) + const delivered: string[] = [] + const controller = { + enqueue: (value: Uint8Array) => delivered.push(decodeChunk(value)), + close: () => {}, + } + writer.attach(controller as unknown as ReadableStreamDefaultController) + + const refused = writer.publish({ + type: 'text', + payload: { channel: 'assistant', text: 'unsaved' }, + }) + await expect(refused).rejects.toBeInstanceOf(expected) + await expect(refused).rejects.not.toBeInstanceOf(StreamControllerSupersededError) + await expect( + writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'later' } }) + ).rejects.toBeInstanceOf(expected) + expect(delivered).toEqual([]) + await expect(writer.flush()).resolves.toBeUndefined() + + await writer.publish({ + type: 'error', + payload: { message: 'stopped', code: 'replay_budget_exhausted' }, + }) + await writer.publish({ type: 'complete', payload: { status: 'error' } }) + await expect(writer.close()).resolves.toBeUndefined() + expect(delivered.map((frame) => JSON.parse(frame.replace(/^data: /, '')).type)).toEqual([ + 'error', + 'complete', + ]) + expect(writer.sawComplete).toBe(true) } - writer.attach(controller as unknown as ReadableStreamDefaultController) - - const refused = writer.publish({ - type: 'text', - payload: { channel: 'assistant', text: 'unsaved' }, - }) - await expect(refused).rejects.toBeInstanceOf(StreamReplayBudgetExhaustedError) - await expect(refused).rejects.not.toBeInstanceOf(StreamControllerSupersededError) - await expect( - writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'later' } }) - ).rejects.toBeInstanceOf(StreamReplayBudgetExhaustedError) - expect(delivered).toEqual([]) - await expect(writer.flush()).resolves.toBeUndefined() - - await writer.publish({ - type: 'error', - payload: { message: 'stopped', code: 'replay_budget_exhausted' }, - }) - await writer.publish({ type: 'complete', payload: { status: 'error' } }) - await expect(writer.close()).resolves.toBeUndefined() - expect(delivered.map((frame) => JSON.parse(frame.replace(/^data: /, '')).type)).toEqual([ - 'error', - 'complete', - ]) - expect(writer.sawComplete).toBe(true) - }) + ) it('enqueues before persistence completes and flushes pending writes on close', async () => { let releasePersist!: () => void diff --git a/apps/sim/lib/mothership/request/session/writer.ts b/apps/sim/lib/mothership/request/session/writer.ts index 38889bd84a5..2fdc6be5139 100644 --- a/apps/sim/lib/mothership/request/session/writer.ts +++ b/apps/sim/lib/mothership/request/session/writer.ts @@ -4,11 +4,18 @@ import { encodeSSEComment } from '@/lib/core/utils/sse' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { appendEvents } from '@/lib/mothership/request/session/buffer' import type { PersistedStreamEventEnvelope } from '@/lib/mothership/request/session/contract' -import type { ChatStreamLease } from '@/lib/mothership/request/session/controller-lease' +import { + type ChatStreamLease, + StreamControllerSupersededError, +} from '@/lib/mothership/request/session/controller-lease' import { createEvent } from '@/lib/mothership/request/session/event' import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' import { compactStreamEvent } from '@/lib/mothership/request/session/replay-compaction' import { encodeSSEEnvelope } from '@/lib/mothership/request/session/sse' +import { + StreamPersistenceFailedError, + StreamTurnFailure, +} from '@/lib/mothership/request/session/turn-failure' import type { StreamEvent } from '@/lib/mothership/request/session/types' const logger = createLogger('StreamWriter') @@ -28,7 +35,7 @@ export interface StreamWriterOptions { initialSeq?: number } -/** The turn's closing verdict, which a refused writer still owes its live client. */ +/** The turn's closing verdict, which a failed writer still owes its live client. */ function isTurnTerminalEvent(event: StreamEvent): boolean { return ( event.type === MothershipStreamV1EventType.error || @@ -36,8 +43,8 @@ function isTurnTerminalEvent(event: StreamEvent): boolean { ) } -function ignoreReplayBudgetRefusal(error: unknown): void { - if (!(error instanceof StreamReplayBudgetExhaustedError)) throw error +function ignoreTurnFailure(error: unknown): void { + if (!(error instanceof StreamTurnFailure)) throw error } /** Result used when the soft stop is already latched, so no further append is attempted. */ @@ -62,7 +69,7 @@ export class StreamWriter { private pendingEnvelopes: PersistedStreamEventEnvelope[] = [] private persistenceTail: Promise = Promise.resolve() private lastPersistenceError: Error | null = null - private replayBudgetError: StreamReplayBudgetExhaustedError | null = null + private turnFailure: StreamTurnFailure | null = null private readonly lease?: ChatStreamLease constructor(options: StreamWriterOptions) { @@ -91,8 +98,9 @@ export class StreamWriter { } /** - * The replay buffer stopped accepting writes because this stream exhausted its byte - * budget. Leased delivery stops before the refused event; unleased delivery continues. + * The replay buffer stopped accepting writes: this stream exhausted its byte budget, or + * a leased append failed. Leased delivery stops before the failed event; unleased + * delivery continues. */ get persistenceStopped(): boolean { return this._persistenceStopped @@ -137,10 +145,11 @@ export class StreamWriter { * compacted copy, while the caller keeps the full event for dispatch. * * A leased writer persists before delivering. When the buffer refuses, the - * publish rejects with {@link StreamReplayBudgetExhaustedError}, and every later - * publish rejects the same way — except the turn's terminal `error`/`complete`, - * which are delivered unpersisted: the run row records that terminal state, and - * a reconnect replays it from there. + * publish rejects with {@link StreamReplayBudgetExhaustedError}; when the append + * fails for any reason but a lost lease, with {@link StreamPersistenceFailedError}. + * Every later publish rejects the same way — except the turn's terminal + * `error`/`complete`, which are delivered unpersisted: the run row records that + * terminal state, and a reconnect replays it from there. */ publish(event: StreamEvent): void | Promise { const envelope = this.createEnvelope(compactStreamEvent(event)) @@ -149,18 +158,24 @@ export class StreamWriter { // A replacement must see every event the browser has received. Fence // persistence before delivery, and before dispatching the event's tool. const persistThenDeliver = async () => { - if (!this.replayBudgetError) { - const result = await appendEvents( - [envelope], - { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, - lease - ) - if (!result.persisted) { + if (!this.turnFailure) { + try { + const result = await appendEvents( + [envelope], + { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, + lease + ) + if (!result.persisted) { + this._persistenceStopped = true + this.turnFailure = new StreamReplayBudgetExhaustedError(result.refusal) + } + } catch (error) { + if (error instanceof StreamControllerSupersededError) throw error this._persistenceStopped = true - this.replayBudgetError = new StreamReplayBudgetExhaustedError(result.refusal) + this.turnFailure = new StreamPersistenceFailedError(error) } } - if (this.replayBudgetError && !isTurnTerminalEvent(event)) throw this.replayBudgetError + if (this.turnFailure && !isTurnTerminalEvent(event)) throw this.turnFailure this.enqueue(envelope) if (event.type === MothershipStreamV1EventType.complete) this._sawComplete = true } @@ -169,7 +184,7 @@ export class StreamWriter { the append by one microtask behind the previous delivery. */ const delivery = this.persistenceTail.then(persistThenDeliver, (error: unknown) => { - ignoreReplayBudgetRefusal(error) + ignoreTurnFailure(error) return persistThenDeliver() }) this.persistenceTail = delivery @@ -188,9 +203,9 @@ export class StreamWriter { async flush(): Promise { this.flushPendingPersistence() - // A refusal belongs to the publish it refused; later publishes consult - // `replayBudgetError`, and the finalizer's flush must not rethrow it. - await this.persistenceTail.catch(ignoreReplayBudgetRefusal) + // A turn failure belongs to the publish it failed; later publishes consult + // `turnFailure`, and the finalizer's flush must not rethrow it. + await this.persistenceTail.catch(ignoreTurnFailure) if (this.lastPersistenceError) { const error = this.lastPersistenceError this.lastPersistenceError = null From d8bf6ccfd951606fd6512f6293cd3b15b947d8f3 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 1 Oct 2026 13:39:29 -0700 Subject: [PATCH 2/3] fix(mothership): hand off a failed replay append again, bounded by the recovery budget Ending the turn on any leased append error that was not a lost lease turned a short Redis blip (a failover READONLY or a connection reset that outlasts the append retries) into a failed Chat turn and a worker stop. Before, that turn was handed off and a successor completed it. - A failed append that is not a budget refusal hands off again, as on staging. The per-run recovery budget bounds a persistent failure: MAX_RECOVERY_ATTEMPTS takeovers with backoff, and then the run ends as stream_recovery_exhausted. - Drop StreamPersistenceFailedError and stream_persistence_failed. The writer only latches budget refusals again, and finalizeAsError surfaces a failed terminal append as it did before. - Keep the change that an unreadable lease is not a hand-off. Its test now fails a lease read alone: a corrupted lock key also fails the fenced append, which is now a hand-off. - Storm suite: a single failed append hands off and the next controller completes the turn. A persistent append failure exhausts the budget, the same as a lease lost with no successor. --- .../lib/mothership/request/lifecycle/start.ts | 27 ++--- .../session/recovery-storm.integration.ts | 88 +++++++------- .../session/replay-budget.integration.ts | 60 +++++---- .../request/session/turn-failure.ts | 27 +---- .../mothership/request/session/writer.test.ts | 114 ++++++++---------- .../lib/mothership/request/session/writer.ts | 63 ++++------ 6 files changed, 169 insertions(+), 210 deletions(-) diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index a817cbf7ce2..9db811b9f7b 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -62,11 +62,7 @@ import { } from '@/lib/mothership/request/session/controller-lease' import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' import { SSE_RESPONSE_HEADERS } from '@/lib/mothership/request/session/sse' -import { - StreamPersistenceFailedError, - type StreamTurnFailure, - turnFailure, -} from '@/lib/mothership/request/session/turn-failure' +import { type StreamTurnFailure, turnFailure } from '@/lib/mothership/request/session/turn-failure' import { TraceCollector } from '@/lib/mothership/request/trace' import type { OrchestratorResult } from '@/lib/mothership/request/types' import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url' @@ -448,17 +444,18 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS await publisher.publish(event) } catch (error) { /* - Only a lost lease hands the turn to a successor. A refusal, or a - Redis error that outlasted its retries, is terminal: leaving it - recoverable made each replacement re-receive and fail the same - event, re-POSTing the run on every reconnect poll. + A refused write is a terminal failure, not a handoff: leaving it + recoverable made each replacement re-receive and re-refuse the same + event. Any other failure to persist means this controller can no + longer prove ownership of the replay, so a successor takes over; the + run's recovery budget bounds how often that repeats. */ - const failure = - error instanceof StreamControllerSupersededError - ? error - : (turnFailure(error) ?? new StreamPersistenceFailedError(error)) - if (!abortController.signal.aborted) abortController.abort(failure) - throw failure + if (!abortController.signal.aborted) { + abortController.abort( + turnFailure(error) ?? new StreamControllerSupersededError() + ) + } + throw error } }, onAbortObserved: (reason) => { diff --git a/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts b/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts index 8b4deb26b12..ba0b88968a2 100644 --- a/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts +++ b/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts @@ -2,8 +2,8 @@ * How often reconnecting to an orphaned Chat run re-POSTs it to the worker, against real * Redis and PostgreSQL. The production reconnect route, stream recovery, chat lifecycle * and finalization run unmodified; a local HTTP server stands in for the worker. Faults - * are injected at one seam, the leased replay append: an append that fails for a reason - * other than a lost lease, or a lease lost with no successor to take over. + * are injected at one seam, the leased replay append: an append that fails once or every + * time, or a lease lost with no successor to take over. */ import { authMock, authMockFns } from '@sim/testing/mocks/auth.mock' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' @@ -15,7 +15,10 @@ const { redisUrl, inheritedEnv, worker, faults } = await vi.hoisted(async () => /** Which leased appends fail, and how; `undefined` lets every append through. */ append: undefined as | undefined - | { frames: 'any_tool' | 'tool_result'; effect: 'throw' | 'lose_lease' }, + | { + frames: 'any_tool' | 'tool_result' + effect: 'throw' | 'throw_once' | 'lose_lease' + }, } const worker = { posts: [] as Array<{ path: string; at: number }>, @@ -76,6 +79,11 @@ vi.mock('@/lib/mothership/request/session/buffer', async (importOriginal) => { (envelope.payload as { phase?: string }).phase === 'result') ) if (hit && fault.effect === 'throw') throw new Error('simulated Redis write failure') + /** A Redis blip that outlasts the append retries, after which Redis is healthy again. */ + if (hit && fault.effect === 'throw_once') { + faults.append = undefined + throw new Error('simulated transient Redis write failure') + } /** The lock expires under a live controller, and nobody else holds it. */ if (hit && fault.effect === 'lose_lease') await getRedisClient()!.del(lease.key) return actual.appendEvents(...args) @@ -105,10 +113,7 @@ import { isTerminalStreamStatus } from '@/lib/mothership/request/session' 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 { - StreamPersistenceFailedError, - StreamRecoveryExhaustedError, -} from '@/lib/mothership/request/session/turn-failure' +import { StreamRecoveryExhaustedError } from '@/lib/mothership/request/session/turn-failure' import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' const userId = generateId() @@ -331,56 +336,57 @@ describe.runIf(Boolean(redisUrl))('reconnecting to an orphaned Chat run', () => 30_000 ) + it('hands off a turn whose append fails once, and the next controller completes it', async () => { + worker.mode = 'frames' + const run = await orphanedRun() + faults.append = { frames: 'any_tool', effect: 'throw_once' } + try { + const { turns, stops } = await reconnect(run, { windowMs: 15_000 }) + + expect(turns).toHaveLength(2) + expect(stops).toHaveLength(0) + expect(await storedRun(run.runId)).toMatchObject({ + status: 'complete', + recoveryBackoff: { attempts: 2 }, + }) + } finally { + faults.append = undefined + } + }, 60_000) + it.each([ - { frames: 'any_tool', tails: 1 }, - { frames: 'any_tool', tails: 3 }, - { frames: 'tool_result', tails: 1 }, - { frames: 'tool_result', tails: 3 }, + { frames: 'any_tool', effect: 'throw', tails: 1 }, + { frames: 'any_tool', effect: 'throw', tails: 3 }, + { frames: 'tool_result', effect: 'throw', tails: 1 }, + { frames: 'tool_result', effect: 'throw', tails: 3 }, + { frames: 'any_tool', effect: 'lose_lease', tails: 3 }, ] as const)( - 'ends the run after one worker request when a $frames frame cannot be persisted ($tails tails)', - async ({ frames, tails }) => { + 'gives up on a run whose recovered controllers keep failing ($effect on $frames, $tails tails)', + async ({ frames, effect, tails }) => { worker.mode = 'frames' const run = await orphanedRun() - faults.append = { frames, effect: 'throw' } + faults.append = { frames, effect } try { - const { turns, stops } = await reconnect(run, { tails, windowMs: 15_000 }) + const { turns, stops } = await reconnect(run, { tails, windowMs: 60_000 }) - expect(turns).toHaveLength(1) + expect(turns).toHaveLength(MAX_RECOVERY_ATTEMPTS) + /** Each takeover waits at least the jittered floor of the one before it. */ + turns.slice(1).forEach((turn, i) => { + expect(turn.at - turns[i].at).toBeGreaterThanOrEqual(0.7 * 1_000 * 2 ** i) + }) expect(stops).toHaveLength(1) expect(await storedRun(run.runId)).toMatchObject({ status: 'error', - error: new StreamPersistenceFailedError(undefined).userMessage, + error: new StreamRecoveryExhaustedError().userMessage, + recoveryBackoff: { attempts: MAX_RECOVERY_ATTEMPTS + 1 }, }) } finally { faults.append = undefined } }, - 60_000 + 120_000 ) - it('gives up on a run whose recovered controllers keep losing their lease', async () => { - worker.mode = 'frames' - const run = await orphanedRun() - faults.append = { frames: 'any_tool', effect: 'lose_lease' } - try { - const { turns, stops } = await reconnect(run, { tails: 3, windowMs: 60_000 }) - - expect(turns).toHaveLength(MAX_RECOVERY_ATTEMPTS) - /** Each takeover waits at least the jittered floor of the one before it. */ - turns.slice(1).forEach((turn, i) => { - expect(turn.at - turns[i].at).toBeGreaterThanOrEqual(0.7 * 1_000 * 2 ** i) - }) - expect(stops).toHaveLength(1) - expect(await storedRun(run.runId)).toMatchObject({ - status: 'error', - error: new StreamRecoveryExhaustedError().userMessage, - recoveryBackoff: { attempts: MAX_RECOVERY_ATTEMPTS + 1 }, - }) - } finally { - faults.append = undefined - } - }, 120_000) - it('waits out the backoff of a recent takeover, then takes over and completes the run', async () => { worker.mode = 'frames' const now = Date.now() 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 992ecbda154..6eab2cdd23b 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -17,6 +17,8 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { 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[] }, + /** How many coming reads of the controller's own lease fail, as a Redis error would. */ + leaseReadFailures: 0, } const replayRequests: Array> = [] const server = createHttpServer(async (request, response) => { @@ -72,6 +74,20 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { }) vi.mock('@/lib/auth', () => authMock) +vi.mock('@/lib/mothership/request/session/controller-lease', async (importOriginal) => { + const actual = + await importOriginal() + return { + ...actual, + assertChatStreamLease: async (...args: Parameters) => { + if (worker.hooks.leaseReadFailures > 0) { + worker.hooks.leaseReadFailures-- + throw new Error('simulated Redis read failure') + } + return actual.assertChatStreamLease(...args) + }, + } +}) vi.mock('@/lib/mothership/request/lifecycle/run', () => ({ /** * Stands in for the worker leg: forwards each scripted event to the controller's @@ -147,7 +163,6 @@ import { StreamReplayBudgetExhaustedError, } from '@/lib/mothership/request/session/replay-budget' import { STREAM_STRING_PREVIEW_UNITS } from '@/lib/mothership/request/session/replay-compaction' -import { STREAM_PERSISTENCE_FAILED_CODE } from '@/lib/mothership/request/session/turn-failure' import type { StreamEvent } from '@/lib/mothership/request/session/types' import { StreamWriter } from '@/lib/mothership/request/session/writer' import type { StreamingContext } from '@/lib/mothership/request/types' @@ -659,7 +674,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge } ) - it('ends a turn whose append fails while it still holds the lease, and stops the worker', async () => { + it('leaves a run it handed off recoverable after a transient append failure', async () => { let eventsKey = '' const { streamId, runId, frames } = await runTurn( [ @@ -674,36 +689,27 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge } ) - expect(frames.map((frame) => [frame.type, frame.payload.code ?? frame.payload.status])).toEqual( - [ - ['session', undefined], - ['error', STREAM_PERSISTENCE_FAILED_CODE], - ['complete', 'error'], - ] - ) + expect(frames.map((frame) => frame.type)).toEqual(['session']) + expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeGreaterThan(300) const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) - expect(stored.status).toBe('error') - expect(worker.abortRequests).toEqual([expect.objectContaining({ messageId: streamId })]) - expect(worker.runs[0].dispatched).toEqual([]) - expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + expect(stored.status).toBe('active') + expect(worker.abortRequests).toEqual([]) }) - it('settles a finished turn whose lease became unreadable instead of handing it off', async () => { - const { runId, frames } = await runTurn([ - text('Done.'), - async () => { - // A lock key of the wrong type makes every read of the lease fail, as a Redis error would. - await redis().del(chatStreamLockKey(chatId)) - await redis().hset(chatStreamLockKey(chatId), 'corrupt', '1') - }, - ]) - + it('settles a finished turn whose lease could not be read instead of handing it off', async () => { try { + const { runId, frames } = await runTurn([ + text('Done.'), + async () => { + worker.hooks.leaseReadFailures = 1 + }, + ]) + 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') } finally { - await redis().del(chatStreamLockKey(chatId)) + worker.hooks.leaseReadFailures = 0 } }) @@ -1065,7 +1071,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge ) } - it('marks its run terminal even when the final events cannot be persisted', async () => { + it('marks its run terminal even when the final events cannot be published', async () => { const controllerToken = `owner\n${generateId()}` const { streamId, runId } = await pausedRun(controllerToken) const lease = { key: chatStreamLockKey(generateId()), value: controllerToken } @@ -1074,8 +1080,10 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge await redis().set(`mothership_stream:${streamId}:events`, 'not a sorted set', 'EX', 60) const publisher = new StreamWriter({ streamId, requestId: generateId(), lease }) - await finalizeAsError(publisher, runId) + const failure = await finalizeAsError(publisher, runId).catch((error: unknown) => error) + expect(failure).toBeInstanceOf(Error) + expect(failure).not.toBeInstanceOf(StreamControllerSupersededError) const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) expect(stored.status).toBe('error') await redis().del(lease.key) diff --git a/apps/sim/lib/mothership/request/session/turn-failure.ts b/apps/sim/lib/mothership/request/session/turn-failure.ts index aa92ed76848..4976615786a 100644 --- a/apps/sim/lib/mothership/request/session/turn-failure.ts +++ b/apps/sim/lib/mothership/request/session/turn-failure.ts @@ -1,8 +1,7 @@ /** - * A failure that ends the turn as an error instead of handing it to a successor. - * Only a controller that provably lost its chat lease hands off - * ({@link StreamControllerSupersededError}); a replacement would meet any other failure - * again, re-POSTing the same run to the worker on every reconnect poll. + * A failure that ends the turn as an error instead of handing it to a successor + * ({@link StreamControllerSupersededError}): a replacement would meet it again, + * re-POSTing the same run to the worker on every reconnect poll. */ export abstract class StreamTurnFailure extends Error { abstract readonly code: string @@ -15,29 +14,9 @@ export function turnFailure(value: unknown): StreamTurnFailure | undefined { return value instanceof StreamTurnFailure ? value : undefined } -/** Run-error code for a turn whose controller could not persist or verify its stream. */ -export const STREAM_PERSISTENCE_FAILED_CODE = 'stream_persistence_failed' - /** Run-error code for a run whose controllers kept dying before it could finish. */ export const STREAM_RECOVERY_EXHAUSTED_CODE = 'stream_recovery_exhausted' -/** - * The controller could not record an event to the replay buffer, or could not read its - * own lease, for a reason other than losing it (a Redis error that outlasted the retries). - */ -export class StreamPersistenceFailedError extends StreamTurnFailure { - readonly code = STREAM_PERSISTENCE_FAILED_CODE - - constructor(cause: unknown) { - super('Stream persistence failed', { cause }) - this.name = 'StreamPersistenceFailedError' - } - - get userMessage(): string { - return 'This response was stopped because it could not be saved while streaming. The work it already completed has been saved — send a message to continue from there.' - } -} - /** Recovery took over the run too many times without a controller staying alive. */ export class StreamRecoveryExhaustedError extends StreamTurnFailure { readonly code = STREAM_RECOVERY_EXHAUSTED_CODE diff --git a/apps/sim/lib/mothership/request/session/writer.test.ts b/apps/sim/lib/mothership/request/session/writer.test.ts index 143a075fbca..1eab0e0a636 100644 --- a/apps/sim/lib/mothership/request/session/writer.test.ts +++ b/apps/sim/lib/mothership/request/session/writer.test.ts @@ -15,7 +15,6 @@ vi.mock('@/lib/mothership/request/session/buffer', () => ({ import { StreamControllerSupersededError } from '@/lib/mothership/request/session/controller-lease' import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' -import { StreamPersistenceFailedError } from '@/lib/mothership/request/session/turn-failure' import { StreamWriter } from '@/lib/mothership/request/session/writer' function decodeChunk(value: Uint8Array): string { @@ -64,7 +63,7 @@ describe('StreamWriter', () => { }) it('does not deliver an event from a controller whose durable append was rejected', async () => { - appendEvents.mockRejectedValueOnce(new StreamControllerSupersededError()) + appendEvents.mockRejectedValueOnce(new Error('ownership lost')) const writer = new StreamWriter({ streamId: 'stream-1', requestId: 'req-1', @@ -74,74 +73,59 @@ describe('StreamWriter', () => { writer.attach(controller as unknown as ReadableStreamDefaultController) await expect( writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'stale' } }) - ).rejects.toBeInstanceOf(StreamControllerSupersededError) + ).rejects.toThrow('ownership lost') expect(controller.enqueue).not.toHaveBeenCalled() - await expect(writer.close()).rejects.toBeInstanceOf(StreamControllerSupersededError) + await expect(writer.close()).rejects.toThrow('ownership lost') expect(controller.close).toHaveBeenCalledOnce() }) - it.each([ - { - failure: 'a budget refusal', - expected: StreamReplayBudgetExhaustedError, - append: () => - appendEvents.mockResolvedValueOnce({ - persisted: false, - refusal: { - resource: 'owner_redis_bytes', - currentBytes: 32, - limitBytes: 32, - attemptedBytes: 1, - }, - }), - }, - { - failure: 'a failed append', - expected: StreamPersistenceFailedError, - append: () => appendEvents.mockRejectedValueOnce(new Error('connection reset')), - }, - ])( - 'ends delivery after $failure, not a handoff, and still delivers the terminal verdict', - async ({ expected, append }) => { - append() - const writer = new StreamWriter({ - streamId: 'stream-1', - requestId: 'req-1', - userId: 'user-1', - lease: { key: 'chat-lock', value: 'current-controller' }, - }) - const delivered: string[] = [] - const controller = { - enqueue: (value: Uint8Array) => delivered.push(decodeChunk(value)), - close: () => {}, - } - writer.attach(controller as unknown as ReadableStreamDefaultController) - - const refused = writer.publish({ - type: 'text', - payload: { channel: 'assistant', text: 'unsaved' }, - }) - await expect(refused).rejects.toBeInstanceOf(expected) - await expect(refused).rejects.not.toBeInstanceOf(StreamControllerSupersededError) - await expect( - writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'later' } }) - ).rejects.toBeInstanceOf(expected) - expect(delivered).toEqual([]) - await expect(writer.flush()).resolves.toBeUndefined() - - await writer.publish({ - type: 'error', - payload: { message: 'stopped', code: 'replay_budget_exhausted' }, - }) - await writer.publish({ type: 'complete', payload: { status: 'error' } }) - await expect(writer.close()).resolves.toBeUndefined() - expect(delivered.map((frame) => JSON.parse(frame.replace(/^data: /, '')).type)).toEqual([ - 'error', - 'complete', - ]) - expect(writer.sawComplete).toBe(true) + it('ends delivery with a budget error, not a handoff, and still delivers the terminal verdict', async () => { + appendEvents.mockResolvedValueOnce({ + persisted: false, + refusal: { + resource: 'owner_redis_bytes', + currentBytes: 32, + limitBytes: 32, + attemptedBytes: 1, + }, + }) + const writer = new StreamWriter({ + streamId: 'stream-1', + requestId: 'req-1', + userId: 'user-1', + lease: { key: 'chat-lock', value: 'current-controller' }, + }) + const delivered: string[] = [] + const controller = { + enqueue: (value: Uint8Array) => delivered.push(decodeChunk(value)), + close: () => {}, } - ) + writer.attach(controller as unknown as ReadableStreamDefaultController) + + const refused = writer.publish({ + type: 'text', + payload: { channel: 'assistant', text: 'unsaved' }, + }) + await expect(refused).rejects.toBeInstanceOf(StreamReplayBudgetExhaustedError) + await expect(refused).rejects.not.toBeInstanceOf(StreamControllerSupersededError) + await expect( + writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'later' } }) + ).rejects.toBeInstanceOf(StreamReplayBudgetExhaustedError) + expect(delivered).toEqual([]) + await expect(writer.flush()).resolves.toBeUndefined() + + await writer.publish({ + type: 'error', + payload: { message: 'stopped', code: 'replay_budget_exhausted' }, + }) + await writer.publish({ type: 'complete', payload: { status: 'error' } }) + await expect(writer.close()).resolves.toBeUndefined() + expect(delivered.map((frame) => JSON.parse(frame.replace(/^data: /, '')).type)).toEqual([ + 'error', + 'complete', + ]) + expect(writer.sawComplete).toBe(true) + }) it('enqueues before persistence completes and flushes pending writes on close', async () => { let releasePersist!: () => void diff --git a/apps/sim/lib/mothership/request/session/writer.ts b/apps/sim/lib/mothership/request/session/writer.ts index 2fdc6be5139..38889bd84a5 100644 --- a/apps/sim/lib/mothership/request/session/writer.ts +++ b/apps/sim/lib/mothership/request/session/writer.ts @@ -4,18 +4,11 @@ import { encodeSSEComment } from '@/lib/core/utils/sse' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { appendEvents } from '@/lib/mothership/request/session/buffer' import type { PersistedStreamEventEnvelope } from '@/lib/mothership/request/session/contract' -import { - type ChatStreamLease, - StreamControllerSupersededError, -} from '@/lib/mothership/request/session/controller-lease' +import type { ChatStreamLease } from '@/lib/mothership/request/session/controller-lease' import { createEvent } from '@/lib/mothership/request/session/event' import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' import { compactStreamEvent } from '@/lib/mothership/request/session/replay-compaction' import { encodeSSEEnvelope } from '@/lib/mothership/request/session/sse' -import { - StreamPersistenceFailedError, - StreamTurnFailure, -} from '@/lib/mothership/request/session/turn-failure' import type { StreamEvent } from '@/lib/mothership/request/session/types' const logger = createLogger('StreamWriter') @@ -35,7 +28,7 @@ export interface StreamWriterOptions { initialSeq?: number } -/** The turn's closing verdict, which a failed writer still owes its live client. */ +/** The turn's closing verdict, which a refused writer still owes its live client. */ function isTurnTerminalEvent(event: StreamEvent): boolean { return ( event.type === MothershipStreamV1EventType.error || @@ -43,8 +36,8 @@ function isTurnTerminalEvent(event: StreamEvent): boolean { ) } -function ignoreTurnFailure(error: unknown): void { - if (!(error instanceof StreamTurnFailure)) throw error +function ignoreReplayBudgetRefusal(error: unknown): void { + if (!(error instanceof StreamReplayBudgetExhaustedError)) throw error } /** Result used when the soft stop is already latched, so no further append is attempted. */ @@ -69,7 +62,7 @@ export class StreamWriter { private pendingEnvelopes: PersistedStreamEventEnvelope[] = [] private persistenceTail: Promise = Promise.resolve() private lastPersistenceError: Error | null = null - private turnFailure: StreamTurnFailure | null = null + private replayBudgetError: StreamReplayBudgetExhaustedError | null = null private readonly lease?: ChatStreamLease constructor(options: StreamWriterOptions) { @@ -98,9 +91,8 @@ export class StreamWriter { } /** - * The replay buffer stopped accepting writes: this stream exhausted its byte budget, or - * a leased append failed. Leased delivery stops before the failed event; unleased - * delivery continues. + * The replay buffer stopped accepting writes because this stream exhausted its byte + * budget. Leased delivery stops before the refused event; unleased delivery continues. */ get persistenceStopped(): boolean { return this._persistenceStopped @@ -145,11 +137,10 @@ export class StreamWriter { * compacted copy, while the caller keeps the full event for dispatch. * * A leased writer persists before delivering. When the buffer refuses, the - * publish rejects with {@link StreamReplayBudgetExhaustedError}; when the append - * fails for any reason but a lost lease, with {@link StreamPersistenceFailedError}. - * Every later publish rejects the same way — except the turn's terminal - * `error`/`complete`, which are delivered unpersisted: the run row records that - * terminal state, and a reconnect replays it from there. + * publish rejects with {@link StreamReplayBudgetExhaustedError}, and every later + * publish rejects the same way — except the turn's terminal `error`/`complete`, + * which are delivered unpersisted: the run row records that terminal state, and + * a reconnect replays it from there. */ publish(event: StreamEvent): void | Promise { const envelope = this.createEnvelope(compactStreamEvent(event)) @@ -158,24 +149,18 @@ export class StreamWriter { // A replacement must see every event the browser has received. Fence // persistence before delivery, and before dispatching the event's tool. const persistThenDeliver = async () => { - if (!this.turnFailure) { - try { - const result = await appendEvents( - [envelope], - { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, - lease - ) - if (!result.persisted) { - this._persistenceStopped = true - this.turnFailure = new StreamReplayBudgetExhaustedError(result.refusal) - } - } catch (error) { - if (error instanceof StreamControllerSupersededError) throw error + if (!this.replayBudgetError) { + const result = await appendEvents( + [envelope], + { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, + lease + ) + if (!result.persisted) { this._persistenceStopped = true - this.turnFailure = new StreamPersistenceFailedError(error) + this.replayBudgetError = new StreamReplayBudgetExhaustedError(result.refusal) } } - if (this.turnFailure && !isTurnTerminalEvent(event)) throw this.turnFailure + if (this.replayBudgetError && !isTurnTerminalEvent(event)) throw this.replayBudgetError this.enqueue(envelope) if (event.type === MothershipStreamV1EventType.complete) this._sawComplete = true } @@ -184,7 +169,7 @@ export class StreamWriter { the append by one microtask behind the previous delivery. */ const delivery = this.persistenceTail.then(persistThenDeliver, (error: unknown) => { - ignoreTurnFailure(error) + ignoreReplayBudgetRefusal(error) return persistThenDeliver() }) this.persistenceTail = delivery @@ -203,9 +188,9 @@ export class StreamWriter { async flush(): Promise { this.flushPendingPersistence() - // A turn failure belongs to the publish it failed; later publishes consult - // `turnFailure`, and the finalizer's flush must not rethrow it. - await this.persistenceTail.catch(ignoreTurnFailure) + // A refusal belongs to the publish it refused; later publishes consult + // `replayBudgetError`, and the finalizer's flush must not rethrow it. + await this.persistenceTail.catch(ignoreReplayBudgetRefusal) if (this.lastPersistenceError) { const error = this.lastPersistenceError this.lastPersistenceError = null From 6c7bf90693db2847e58f378073f714828e9b48db Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Thu, 1 Oct 2026 15:09:49 -0700 Subject: [PATCH 3/3] fix(mothership): claim a recovering run only once its takeover can start Recovery claimed the run before reading its replay, resolving billing and checking access. A takeover that then failed released the lock but had already spent an attempt and refreshed the run, so an exhausted run whose takeover kept failing there stayed active, and the refreshed row kept the orphan sweep away. Claim last: a takeover that fails before it starts leaves the run untouched, and an exhausted claim always reaches the terminal path. --- .../request/application/recover-stream.ts | 43 +++++++++-------- .../session/recovery-storm.integration.ts | 46 ++++++++++++++++++- 2 files changed, 68 insertions(+), 21 deletions(-) diff --git a/apps/sim/lib/mothership/request/application/recover-stream.ts b/apps/sim/lib/mothership/request/application/recover-stream.ts index 95c3e8673bd..a13c63715fe 100644 --- a/apps/sim/lib/mothership/request/application/recover-stream.ts +++ b/apps/sim/lib/mothership/request/application/recover-stream.ts @@ -90,25 +90,6 @@ export const readChatStream = defineAuthorizedChatUseCase({ if (!(await acquirePendingChatStream(chatId, run.streamId, 0))) return run const lease = getLocalChatStreamLease(chatId, run.streamId)! try { - await assertChatStreamLease(lease) - if ( - !(await claimRunController({ - runId: run.id, - chatId, - previousToken: saved.controllerToken, - token: lease.value, - recoveryBackoff: plan.backoff, - })) - ) { - await releasePendingChatStream(chatId, run.streamId, lease) - return (await getLatestRunForStream(run.streamId, userId)) ?? run - } - logger.info('Claimed stream run for recovery', { - runId: run.id, - streamId: run.streamId, - attempt: plan.backoff.attempts, - exhausted: plan.kind === 'exhausted', - }) if (isHosted && !config.data.billingAdmission) throw new OrchestrationError( 'forbidden', @@ -145,6 +126,30 @@ export const readChatStream = defineAuthorizedChatUseCase({ const recoveredEvents = ringIntact ? events : [] const lastEvent = recoveredEvents.at(-1) const resumeSeq = lastEvent ? lastEvent.seq : ((await getLatestSeq(run.streamId)) ?? 0) + /** + * Claim last: everything before it can fail without touching the run, so a takeover + * that cannot start neither spends the recovery budget nor refreshes the run, and + * an exhausted claim always reaches the terminal path below. + */ + await assertChatStreamLease(lease) + if ( + !(await claimRunController({ + runId: run.id, + chatId, + previousToken: saved.controllerToken, + token: lease.value, + recoveryBackoff: plan.backoff, + })) + ) { + await releasePendingChatStream(chatId, run.streamId, lease) + return (await getLatestRunForStream(run.streamId, userId)) ?? run + } + logger.info('Claimed stream run for recovery', { + runId: run.id, + streamId: run.streamId, + attempt: plan.backoff.attempts, + exhausted: plan.kind === 'exhausted', + }) const requestId = typeof saved?.requestId === 'string' ? saved.requestId : generateId() const completion = { chatId, diff --git a/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts b/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts index ba0b88968a2..bd7f34718cf 100644 --- a/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts +++ b/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts @@ -2,8 +2,9 @@ * How often reconnecting to an orphaned Chat run re-POSTs it to the worker, against real * Redis and PostgreSQL. The production reconnect route, stream recovery, chat lifecycle * and finalization run unmodified; a local HTTP server stands in for the worker. Faults - * are injected at one seam, the leased replay append: an append that fails once or every - * time, or a lease lost with no successor to take over. + * are injected at two seams: the leased replay append (an append that fails once or every + * time, or a lease lost with no successor to take over) and a recovering controller's read + * of the replay ring. */ import { authMock, authMockFns } from '@sim/testing/mocks/auth.mock' import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' @@ -19,6 +20,8 @@ const { redisUrl, inheritedEnv, worker, faults } = await vi.hoisted(async () => frames: 'any_tool' | 'tool_result' effect: 'throw' | 'throw_once' | 'lose_lease' }, + /** Fails a recovering controller's read of the whole replay ring. */ + recoveryRead: false, } const worker = { posts: [] as Array<{ path: string; at: number }>, @@ -88,6 +91,12 @@ vi.mock('@/lib/mothership/request/session/buffer', async (importOriginal) => { if (hit && fault.effect === 'lose_lease') await getRedisClient()!.del(lease.key) return actual.appendEvents(...args) }, + readEvents: async (...args: Parameters) => { + const [, afterCursor] = args + if (faults.recoveryRead && afterCursor === '0') + throw new Error('simulated Redis read failure') + return actual.readEvents(...args) + }, } }) @@ -420,6 +429,39 @@ describe.runIf(Boolean(redisUrl))('reconnecting to an orphaned Chat run', () => expect(stored.recoveryBackoff).toMatchObject({ attempts: 1 }) }, 30_000) + it('leaves the run untouched when a takeover fails before it starts, then ends it once exhausted', async () => { + worker.mode = 'frames' + const claimedAt = Date.now() - 1_000 + const recoveryBackoff = { attempts: MAX_RECOVERY_ATTEMPTS, claimedAt, notBefore: claimedAt } + const run = await orphanedRun({ recoveryBackoff }) + const { controllerToken } = (await storedRun(run.runId)).requestContext as { + controllerToken: string + } + faults.recoveryRead = true + try { + const { turns } = await reconnect(run, { windowMs: 2_000 }) + + expect(turns).toHaveLength(0) + expect(await storedRun(run.runId)).toMatchObject({ + status: 'active', + requestContext: { controllerToken }, + recoveryBackoff, + }) + } finally { + faults.recoveryRead = false + } + + const { turns, stops } = await reconnect(run, { windowMs: 10_000 }) + + expect(turns).toHaveLength(0) + expect(stops).toHaveLength(1) + expect(await storedRun(run.runId)).toMatchObject({ + status: 'error', + error: new StreamRecoveryExhaustedError().userMessage, + recoveryBackoff: { attempts: MAX_RECOVERY_ATTEMPTS + 1 }, + }) + }, 30_000) + it('never takes over a parked run whose controller still holds the chat lock', async () => { worker.mode = 'frames' const run = await orphanedRun({ status: 'paused_waiting_for_tool' })