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..a13c63715fe 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,21 +85,11 @@ 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 { - await assertChatStreamLease(lease) - if ( - !(await claimRunController({ - runId: run.id, - chatId, - previousToken: saved.controllerToken, - token: lease.value, - })) - ) { - await releasePendingChatStream(chatId, run.streamId, lease) - return (await getLatestRunForStream(run.streamId, userId)) ?? run - } if (isHosted && !config.data.billingAdmission) throw new OrchestrationError( 'forbidden', @@ -132,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, @@ -163,6 +181,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..9db811b9f7b 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -60,11 +60,9 @@ 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 { 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 +93,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 +133,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 +226,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 +249,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 +278,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 +292,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), }) @@ -427,11 +447,12 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS 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. + longer prove ownership of the replay, so a successor takes over; the + run's recovery budget bounds how often that repeats. */ if (!abortController.signal.aborted) { abortController.abort( - replayRefusal(error) ?? new StreamControllerSupersededError() + turnFailure(error) ?? new StreamControllerSupersededError() ) } throw error @@ -445,11 +466,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 +479,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 +493,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..bd7f34718cf --- /dev/null +++ b/apps/sim/lib/mothership/request/session/recovery-storm.integration.ts @@ -0,0 +1,498 @@ +/** + * 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 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' + +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' | '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 }>, + 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') + /** 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) + }, + readEvents: async (...args: Parameters) => { + const [, afterCursor] = args + if (faults.recoveryRead && afterCursor === '0') + throw new Error('simulated Redis read failure') + return actual.readEvents(...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 { 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('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', 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)( + '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 } + try { + const { turns, stops } = await reconnect(run, { tails, 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('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' }) + 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..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 @@ -680,6 +696,23 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge expect(worker.abortRequests).toEqual([]) }) + 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 { + worker.hooks.leaseReadFailures = 0 + } + }) + it('still cleans up a finished turn whose terminal events could not be published', async () => { let eventsKey = '' const { streamId, runId } = await runTurn( 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..4976615786a --- /dev/null +++ b/apps/sim/lib/mothership/request/session/turn-failure.ts @@ -0,0 +1,32 @@ +/** + * 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 + /** 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 run whose controllers kept dying before it could finish. */ +export const STREAM_RECOVERY_EXHAUSTED_CODE = 'stream_recovery_exhausted' + +/** 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.' + } +}