From c920082042b38c5615c8f5e1e4dd2b2519fde1bd Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 18:42:12 -0700 Subject: [PATCH 01/18] fix(mothership): bound replay frames and end refused turns instead of handing them off A leased Chat stream persists every event to its Redis replay buffer before delivering or dispatching it. When the buffer refused a write (a frame over the 1 MiB single-write ceiling, or a stream past its 32 MiB budget), the writer threw a generic error, the controller classified it as a takeover and released its lock without finalizing, and every reconnect started a recovery controller that re-received and re-refused the same event until the run deadline. The refused tool call was never dispatched, so the worker waited. Refusals: - The writer throws StreamReplayBudgetExhaustedError and rejects later publishes the same way, but still delivers the turn's terminal error and complete events unpersisted; flush no longer rethrows the refusal. - The controller aborts with that error rather than a supersession, the lifecycle classifies it as an error (not a user cancel), the run is finalized as an error with code replay_budget_exhausted, and the worker is told to stop. Reconnects then see a terminal run and start no controller. - Only a genuine StreamControllerSupersededError still hands the run off. Compaction, applied only to the copy the writer delivers and persists: - Payloads over 256 KiB have long string leaves cut to an 8 KiB head in place, then medium strings to 500 characters, and only then the largest unprotected subtrees replaced by { omitted, bytes }, targeting 128 KiB. A payload-level streamTruncation marker lists every changed field. - Identity fields, UI-read keys, arguments of client-executable calls, file preview events, and generate_api_key results are never compacted; such a frame that stays too large is refused and ends the turn cleanly. - Oversized assistant text splits into contiguous exact pieces, and each tool call forwards only the head of its argument deltas. --- apps/sim/lib/mothership/request/go/parser.ts | 4 +- .../mothership/request/lifecycle/finalize.ts | 1 + .../mothership/request/lifecycle/run.test.ts | 61 +++ .../lib/mothership/request/lifecycle/run.ts | 35 +- .../lib/mothership/request/lifecycle/start.ts | 95 +++- .../mothership/request/session/omission.ts | 50 ++ .../session/replay-budget.integration.ts | 467 ++++++++++++++++++ .../request/session/replay-budget.ts | 21 + .../request/session/replay-compaction.test.ts | 263 ++++++++++ .../request/session/replay-compaction.ts | 332 +++++++++++++ .../mothership/request/session/writer.test.ts | 77 ++- .../lib/mothership/request/session/writer.ts | 102 +++- apps/sim/lib/mothership/request/types.ts | 2 + 13 files changed, 1481 insertions(+), 29 deletions(-) create mode 100644 apps/sim/lib/mothership/request/session/omission.ts create mode 100644 apps/sim/lib/mothership/request/session/replay-budget.integration.ts create mode 100644 apps/sim/lib/mothership/request/session/replay-budget.ts create mode 100644 apps/sim/lib/mothership/request/session/replay-compaction.test.ts create mode 100644 apps/sim/lib/mothership/request/session/replay-compaction.ts diff --git a/apps/sim/lib/mothership/request/go/parser.ts b/apps/sim/lib/mothership/request/go/parser.ts index c17ce6ae9e7..d37a71685e0 100644 --- a/apps/sim/lib/mothership/request/go/parser.ts +++ b/apps/sim/lib/mothership/request/go/parser.ts @@ -2,6 +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' const logger = createLogger('CopilotSseParser') @@ -49,7 +50,8 @@ export async function processSSEStream( } catch (error) { if ( error instanceof FatalSseEventError || - error instanceof StreamControllerSupersededError + error instanceof StreamControllerSupersededError || + error instanceof StreamReplayBudgetExhaustedError ) throw error logger.warn('Failed to handle SSE event', { diff --git a/apps/sim/lib/mothership/request/lifecycle/finalize.ts b/apps/sim/lib/mothership/request/lifecycle/finalize.ts index 2aaaa04c3ff..d57648c6c42 100644 --- a/apps/sim/lib/mothership/request/lifecycle/finalize.ts +++ b/apps/sim/lib/mothership/request/lifecycle/finalize.ts @@ -146,6 +146,7 @@ async function handleError( message: errorMessage, error: errorMessage, displayMessage: errorMessage, + ...(result.errorCode ? { code: result.errorCode } : {}), data: { displayMessage: errorMessage }, }, }) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 38cd3f1faf6..ce362e0f001 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -206,6 +206,11 @@ import { StreamEndedWithoutTerminalError, } from '@/lib/mothership/request/go/stream' import { runCopilotLifecycle } from '@/lib/mothership/request/lifecycle/run' +import { + REPLAY_BUDGET_EXHAUSTED_CODE, + REPLAY_BUDGET_EXHAUSTED_MESSAGE, + StreamReplayBudgetExhaustedError, +} from '@/lib/mothership/request/session/replay-budget' import { executeToolAndReport } from '@/lib/mothership/request/tools/executor' const mockExecuteAppTool = toolsMockFns.mockExecuteTool @@ -1947,6 +1952,62 @@ describe('runCopilotLifecycle', () => { ) }) + it.each([ + ['throws the refusal', true], + ['returns after the abort', false], + ])( + 'ends a turn stopped by a refused replay write as an error, not a user cancel, when the stream %s', + async (_label, throws) => { + const abortController = new AbortController() + mockRunStreamLoop.mockImplementationOnce( + async ( + _fetchUrl: string, + _fetchOptions: RequestInit, + context: StreamingContext + ): Promise => { + context.accumulatedContent = 'partial answer' + const refusal = new StreamReplayBudgetExhaustedError({ + resource: 'owner_redis_bytes', + currentBytes: 32 * 1024 * 1024, + limitBytes: 32 * 1024 * 1024, + attemptedBytes: 512, + }) + abortController.abort(refusal) + context.wasAborted = true + if (throws) throw refusal + } + ) + + const result = await runCopilotLifecycle( + { message: 'hello', messageId: 'stream-1' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'exec-1', + runId: 'run-1', + abortSignal: abortController.signal, + executionContext: { + userId: 'user-1', + workflowId: '', + workspaceId: 'ws-1', + chatId: 'chat-1', + }, + } + ) + + expect(result).toEqual( + expect.objectContaining({ + success: false, + cancelled: false, + content: 'partial answer', + error: REPLAY_BUDGET_EXHAUSTED_MESSAGE, + errorCode: REPLAY_BUDGET_EXHAUSTED_CODE, + }) + ) + } + ) + it('returns the cancelled result when cancelled completion persistence fails', async () => { const abortController = new AbortController() const onComplete = vi.fn().mockRejectedValue(new Error('db unavailable')) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index 110a4bbf45f..62ef1d5ccca 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -53,6 +53,11 @@ 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 { + REPLAY_BUDGET_EXHAUSTED_CODE, + REPLAY_BUDGET_EXHAUSTED_MESSAGE, + StreamReplayBudgetExhaustedError, +} from '@/lib/mothership/request/session/replay-budget' import { getToolCallTerminalData, requireToolCallStateResult, @@ -534,16 +539,24 @@ 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 replayBudgetExhausted = stoppedByReplayBudget(lifecycleOptions.abortSignal) + if (replayBudgetExhausted && !context.errors.includes(REPLAY_BUDGET_EXHAUSTED_MESSAGE)) { + context.errors.push(REPLAY_BUDGET_EXHAUSTED_MESSAGE) + } // 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 = - context.completionStatus === MothershipStreamV1CompletionStatus.cancelled || - context.wasAborted || - (lifecycleOptions.abortSignal?.aborted ?? false) + !replayBudgetExhausted && + (context.completionStatus === MothershipStreamV1CompletionStatus.cancelled || + context.wasAborted || + (lifecycleOptions.abortSignal?.aborted ?? false)) const succeeded = + !replayBudgetExhausted && !turnWasAborted && (backendFinishedTurn || (!context.completionStatus && context.errors.length === 0)) @@ -564,6 +577,9 @@ export async function runCopilotLifecycle( toolCalls: buildToolCallSummaries(context), chatId: context.chatId, requestId: context.requestId, + ...(replayBudgetExhausted + ? { error: REPLAY_BUDGET_EXHAUSTED_MESSAGE, errorCode: REPLAY_BUDGET_EXHAUSTED_CODE } + : {}), errors: !succeeded && context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, @@ -601,7 +617,11 @@ export async function runCopilotLifecycle( // partial content can be appended. // Return `cancelled: true` so upstream classification stays // consistent with the success-path cancel result. - const wasCancelled = lifecycleOptions.abortSignal?.aborted ?? false + const replayBudgetExhausted = + error instanceof StreamReplayBudgetExhaustedError || + stoppedByReplayBudget(lifecycleOptions.abortSignal) + const wasCancelled = + !replayBudgetExhausted && (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 @@ -616,7 +636,8 @@ export async function runCopilotLifecycle( toolCalls: buildToolCallSummaries(context), chatId: context.chatId, requestId: context.requestId, - error: err.message, + error: replayBudgetExhausted ? REPLAY_BUDGET_EXHAUSTED_MESSAGE : err.message, + ...(replayBudgetExhausted ? { errorCode: REPLAY_BUDGET_EXHAUSTED_CODE } : {}), errors: context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, @@ -1660,6 +1681,10 @@ async function withEnterpriseByokKey( return byokApiKey ? { ...refreshed, byokApiKey } : refreshed } +function stoppedByReplayBudget(signal: AbortSignal | undefined): boolean { + return signal?.reason instanceof StreamReplayBudgetExhaustedError +} + function isAborted(options: CopilotLifecycleOptions, context: StreamingContext): boolean { return !!(options.abortSignal?.aborted || context.wasAborted) } diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 9efcbc019ed..4e74c81c9b5 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -59,8 +59,14 @@ import { assertChatStreamLease, StreamControllerSupersededError, } from '@/lib/mothership/request/session/controller-lease' +import { + REPLAY_BUDGET_EXHAUSTED_CODE, + REPLAY_BUDGET_EXHAUSTED_MESSAGE, + StreamReplayBudgetExhaustedError, +} from '@/lib/mothership/request/session/replay-budget' import { SSE_RESPONSE_HEADERS } from '@/lib/mothership/request/session/sse' import { TraceCollector } from '@/lib/mothership/request/trace' +import type { OrchestratorResult } from '@/lib/mothership/request/types' import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url' export { SSE_RESPONSE_HEADERS } @@ -204,6 +210,52 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS const collector = new TraceCollector() + /** + * 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. + */ + const finalizeAfterReplayRefusal = async ( + refusal: StreamReplayBudgetExhaustedError, + result?: OrchestratorResult + ) => { + logger.warn(`[${requestId}] Stream replay budget exhausted; ending the turn`, { + streamId, + resource: refusal.refusal.resource, + attemptedBytes: refusal.refusal.attemptedBytes, + currentBytes: refusal.refusal.currentBytes, + limitBytes: refusal.refusal.limitBytes, + }) + await finalizeStream( + { + content: '', + contentBlocks: [], + toolCalls: [], + ...result, + success: false, + cancelled: false, + error: REPLAY_BUDGET_EXHAUSTED_MESSAGE, + errorCode: REPLAY_BUDGET_EXHAUSTED_CODE, + }, + publisher, + runId, + RequestTraceV1Outcome.error, + requestId + ) + try { + const { requestExplicitStreamAbort } = await import( + '@/lib/mothership/request/session/explicit-abort' + ) + await requestExplicitStreamAbort({ streamId, userId, chatId }) + } catch (error) { + logger.warn(`[${requestId}] Worker stop after replay refusal was not delivered`, { + streamId, + error: getErrorMessage(error), + }) + } + } + return new ReadableStream({ async start(controller) { publisher.attach(controller) @@ -348,7 +400,19 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS try { await publisher.publish(event) } catch (error) { - abortController.abort(new StreamControllerSupersededError()) + /* + 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. + */ + if (!abortController.signal.aborted) { + abortController.abort( + error instanceof StreamReplayBudgetExhaustedError + ? error + : new StreamControllerSupersededError() + ) + } throw error } }, @@ -360,17 +424,27 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS }) lifecycleResult = result + const replayRefusal = + abortController.signal.reason instanceof StreamReplayBudgetExhaustedError + ? abortController.signal.reason + : undefined // A completed result wins a late Stop; passive disconnection never cancels. outcome = result.success ? RequestTraceV1Outcome.success - : result.cancelled || abortController.signal.aborted - ? RequestTraceV1Outcome.cancelled - : RequestTraceV1Outcome.error + : replayRefusal + ? RequestTraceV1Outcome.error + : result.cancelled || abortController.signal.aborted + ? RequestTraceV1Outcome.cancelled + : RequestTraceV1Outcome.error if (outcome === RequestTraceV1Outcome.cancelled) { cancelReason = recordCancelled() } await assertControllerOwnership() - await finalizeStream(result, publisher, runId, outcome, requestId) + if (replayRefusal && !result.success) { + await finalizeAfterReplayRefusal(replayRefusal, result) + } else { + await finalizeStream(result, publisher, runId, outcome, requestId) + } } catch (error) { if ( error instanceof StreamControllerSupersededError || @@ -380,6 +454,17 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS return } await assertControllerOwnership() + const replayRefusal = + error instanceof StreamReplayBudgetExhaustedError + ? error + : abortController.signal.reason instanceof StreamReplayBudgetExhaustedError + ? abortController.signal.reason + : undefined + if (replayRefusal) { + outcome = RequestTraceV1Outcome.error + await finalizeAfterReplayRefusal(replayRefusal) + return + } const wasCancelled = abortController.signal.aborted outcome = wasCancelled ? RequestTraceV1Outcome.cancelled : RequestTraceV1Outcome.error if (outcome === RequestTraceV1Outcome.cancelled) { diff --git a/apps/sim/lib/mothership/request/session/omission.ts b/apps/sim/lib/mothership/request/session/omission.ts new file mode 100644 index 00000000000..fd798b13e74 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/omission.ts @@ -0,0 +1,50 @@ +import { isRecordLike } from '@sim/utils/object' + +/** + * Stand-in for a subtree the bounded chat stream could not carry even after its + * long strings were shortened. The full value still reaches the model and the + * durable tool records; only the live and replayed UI stream carries this. + */ +export interface OmittedStreamValue { + omitted: true + bytes: number +} + +/** One field shortened or omitted by stream compaction, addressed by JSON pointer. */ +export interface StreamTruncatedField { + path: string + /** Full UTF-8 size of the original value. */ + bytes: number + /** UTF-8 size kept in the stream; absent when the value was omitted. */ + previewBytes?: number +} + +/** Payload key of the {@link StreamTruncationMarker} on a compacted event. */ +export const STREAM_TRUNCATION_KEY = 'streamTruncation' + +/** + * Payload-level marker naming every field compaction changed. It sits beside the + * compacted fields rather than inside them, so `arguments` and `output` keep the + * shape the UI reads. + */ +export interface StreamTruncationMarker { + /** UTF-8 size of the payload before compaction. */ + bytes: number + fields: StreamTruncatedField[] +} + +export function isOmittedStreamValue(value: unknown): value is OmittedStreamValue { + return isRecordLike(value) && value.omitted === true && typeof value.bytes === 'number' +} + +/** Formats a byte count the way the UI states an omitted size ("1.6 MB"). */ +export function formatByteSize(bytes: number): string { + if (bytes < 1024) return `${bytes} B` + if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB` + return `${(bytes / (1024 * 1024)).toFixed(1)} MB` +} + +/** Replaces omitted subtrees with a sentence stating their size, for display. */ +export function describeOmittedStreamValues(_key: string, value: unknown): unknown { + return isOmittedStreamValue(value) ? `Too large to show (${formatByteSize(value.bytes)})` : value +} diff --git a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts new file mode 100644 index 00000000000..b07688ff92e --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -0,0 +1,467 @@ +/** + * The leased Chat stream writer and its replay-budget failure path against real + * Redis and PostgreSQL. A local HTTP server stands in for the worker's abort + * endpoint, and a scripted lifecycle stands in for the worker's event stream; + * everything between them — the writer, the Redis append script, the chat lock, + * run finalization, the reconnect route, and stream recovery — is production code. + */ +import { authMock, authMockFns } from '@sim/testing/mocks/auth.mock' +import { toRecord } from '@sim/utils/object' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' + +const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { + const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') + const { createServer: createHttpServer } = await import('node:http') + const abortRequests: Array> = [] + const server = createHttpServer(async (request, response) => { + let body = '' + for await (const chunk of request) body += chunk + if (request.url === '/api/streams/explicit-abort') abortRequests.push(JSON.parse(body)) + response.writeHead(200, { 'content-type': 'application/json' }) + response.end(JSON.stringify({ settled: true })) + }) + 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: { + server, + abortRequests, + /** Events the scripted worker streams through the controller's sink, in order. */ + script: [] as unknown[], + /** Controller lifecycles started, and what each sink call threw. */ + runs: [] as Array<{ dispatched: unknown[]; sinkErrors: unknown[] }>, + }, + } +}) + +vi.mock('@/lib/auth', () => authMock) +vi.mock('@/lib/mothership/request/lifecycle/run', () => ({ + /** + * Stands in for the worker leg: forwards each scripted event to the controller's + * sink and dispatches it only once the sink accepted it, as the real loop does. + */ + runCopilotLifecycle: async ( + _payload: unknown, + options: { onEvent?: (event: unknown) => Promise; abortSignal?: AbortSignal } + ) => { + const run = { dispatched: [] as unknown[], sinkErrors: [] as unknown[] } + worker.runs.push(run) + for (const event of worker.script) { + try { + await options.onEvent?.(event) + } catch (error) { + run.sinkErrors.push(error) + break + } + run.dispatched.push(event) + } + return { + success: run.sinkErrors.length === 0, + cancelled: options.abortSignal?.aborted ?? false, + content: '', + contentBlocks: [], + toolCalls: [], + } + }, +})) + +import { db } from '@sim/db' +import { copilotChats, copilotRuns, permissions, user, workspace } from '@sim/db/schema' +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 { getRedisBudgetKeys, getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' +import { readChatStream } from '@/lib/mothership/request/application/recover-stream' +import { createStreamingContext } from '@/lib/mothership/request/context/request-context' +import { restoreStreamingContext } from '@/lib/mothership/request/context/restore' +import { createSSEStream } from '@/lib/mothership/request/lifecycle/start' +import { acquirePendingChatStream } from '@/lib/mothership/request/session/abort' +import { readEvents } from '@/lib/mothership/request/session/buffer' +import { + type ChatStreamLease, + chatStreamLockKey, + StreamControllerSupersededError, +} from '@/lib/mothership/request/session/controller-lease' +import { eventToStreamEvent } from '@/lib/mothership/request/session/event' +import { STREAM_TRUNCATION_KEY } from '@/lib/mothership/request/session/omission' +import { + REPLAY_BUDGET_EXHAUSTED_CODE, + REPLAY_BUDGET_EXHAUSTED_MESSAGE, + StreamReplayBudgetExhaustedError, +} from '@/lib/mothership/request/session/replay-budget' +import { STREAM_STRING_PREVIEW_UNITS } from '@/lib/mothership/request/session/replay-compaction' +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' +import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' + +const MB = 1024 * 1024 + +function redis() { + const client = getRedisClient() + if (!client) throw new Error('The integration suite requires TEST_REDIS_URL') + return client +} + +/** A writer that owns a real chat lock, with the frames it delivers to its client. */ +async function leasedWriter(options: { lease?: ChatStreamLease; ownLock?: boolean } = {}) { + const streamId = generateId() + const lease = options.lease ?? { + key: chatStreamLockKey(generateId()), + value: `${streamId}\n${generateId()}`, + } + if (options.ownLock !== false) await redis().set(lease.key, lease.value, 'EX', 60) + const writer = new StreamWriter({ + streamId, + requestId: generateId(), + userId: generateId(), + lease, + }) + const delivered: string[] = [] + const client = new ReadableStream({ + start: (controller) => writer.attach(controller), + }) + const received = client.pipeTo( + new WritableStream({ + write: (chunk) => { + delivered.push( + new TextDecoder() + .decode(chunk) + .replace(/^data: /, '') + .trim() + ) + }, + }) + ) + /** The frames the client received, once the writer has closed its stream. */ + const frames = async () => { + await received + return delivered + } + return { streamId, writer, frames } +} + +async function storedMembers(streamId: string): Promise { + return redis().zrange(`mothership_stream:${streamId}:events`, 0, -1) +} + +function toolCall( + toolCallId: string, + toolName: string, + args: Record, + executor: 'sim' | 'go' | 'client' = 'sim' +): StreamEvent { + return { + type: 'tool', + payload: { toolCallId, toolName, executor, mode: 'async', phase: 'call', arguments: args }, + } +} + +function text(value: string): StreamEvent { + return { type: 'text', payload: { channel: 'assistant', text: value } } +} + +afterAll(async () => { + await closeRedisConnection() + await new Promise((resolve) => worker.server.close(() => resolve())) + for (const [key, value] of Object.entries(inheritedEnv)) { + if (value === undefined) delete process.env[key] + else process.env[key] = value + } +}) + +describe.runIf(Boolean(redisUrl))('leased Chat stream writer with Redis', () => { + it('delivers and persists the same bounded copy of a 2 MB tool call, leaving the dispatched event whole', async () => { + const { streamId, writer, frames } = await leasedWriter() + const stdout = 'o'.repeat(2 * MB) + const call = toolCall('call-large', 'cli_logs_get', { + activity: { id: 'activity-1', title: 'Reading logs' }, + command: 'logs get --run latest', + stdout, + }) + + await writer.publish(text('before ')) + await writer.publish(call) + await writer.publish(text('after')) + await writer.close() + + expect(await frames()).toEqual(await storedMembers(streamId)) + const replayed = await readEvents(streamId, '0') + expect(replayed.map((envelope) => envelope.seq)).toEqual([1, 2, 3]) + const payload = toRecord(replayed[1].payload) + expect(payload).toMatchObject({ + toolCallId: 'call-large', + toolName: 'cli_logs_get', + phase: 'call', + arguments: { + activity: { id: 'activity-1', title: 'Reading logs' }, + command: 'logs get --run latest', + stdout: stdout.slice(0, STREAM_STRING_PREVIEW_UNITS), + }, + [STREAM_TRUNCATION_KEY]: { + fields: [ + { path: '/arguments/stdout', bytes: 2 * MB, previewBytes: STREAM_STRING_PREVIEW_UNITS }, + ], + }, + }) + expect(toRecord(call.payload).arguments).toMatchObject({ stdout }) + }) + + it('publishes a bounded preview of a committed 1.6 MB tool result and keeps persisting later events', async () => { + const { streamId, writer, frames } = await leasedWriter() + const resources = [{ type: 'file', id: 'file-1', title: 'report.csv' }] + + await writer.publish({ + type: 'tool', + payload: { + toolCallId: 'call-result', + toolName: 'run_code', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + status: 'success', + output: { stdout: 'r'.repeat(1.6 * MB), exitCode: 0, resources }, + }, + }) + await writer.publish(text('The report is ready.')) + await writer.close() + + expect(await frames()).toEqual(await storedMembers(streamId)) + const replayed = await readEvents(streamId, '0') + expect(replayed.map((envelope) => envelope.type)).toEqual(['tool', 'text']) + expect(toRecord(replayed[0].payload)).toMatchObject({ + success: true, + status: 'success', + output: { exitCode: 0, resources, stdout: 'r'.repeat(STREAM_STRING_PREVIEW_UNITS) }, + }) + expect(writer.persistenceStopped).toBe(false) + }) + + it('restores the same tool calls, text, and activity from a compacted buffer as from the full events, without dispatching', async () => { + const { streamId, writer } = await leasedWriter() + const workflowArgs = { + workflowId: generateId(), + input: { rows: 'w'.repeat(400 * 1024) }, + } + const events: StreamEvent[] = [ + text('Checking the logs. '), + toolCall('call-logs', 'cli_logs_get', { + activity: { id: 'activity-logs', title: 'Reading logs' }, + command: 'logs get', + stdout: 'l'.repeat(2 * MB), + }), + { + type: 'tool', + payload: { + toolCallId: 'call-logs', + toolName: 'cli_logs_get', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + status: 'success', + output: { stdout: 'l'.repeat(1.6 * MB), exitCode: 0 }, + }, + }, + toolCall('call-workflow', 'run_workflow', workflowArgs, 'client'), + text('Done.'), + ] + for (const event of events) await writer.publish(event) + await writer.close() + + const restore = async (source: readonly StreamEvent[]): Promise => { + const context = createStreamingContext({ messageId: streamId }) + await restoreStreamingContext(source, context, { userId: generateId(), workflowId: '' }) + return context + } + const summary = (context: StreamingContext) => ({ + text: context.accumulatedContent, + calls: [...context.toolCalls.values()].map((call) => ({ + id: call.id, + name: call.name, + status: call.status, + activity: call.params?.activity, + })), + pending: context.pendingToolPromises.size, + }) + const fromBuffer = await restore((await readEvents(streamId, '0')).map(eventToStreamEvent)) + const fromFull = await restore(events) + + expect(summary(fromBuffer)).toEqual(summary(fromFull)) + expect(summary(fromBuffer).pending).toBe(0) + expect(JSON.stringify(fromBuffer.toolCalls.get('call-workflow')?.params)).toBe( + JSON.stringify(workflowArgs) + ) + }) + + it('still refuses a controller whose lease was taken over', async () => { + const lease = { key: chatStreamLockKey(generateId()), value: `stale\n${generateId()}` } + await redis().set(lease.key, `successor\n${generateId()}`, 'EX', 60) + const { writer, frames } = await leasedWriter({ lease, ownLock: false }) + + const published = writer.publish(text('from the stale controller')) + + await expect(published).rejects.toBeInstanceOf(StreamControllerSupersededError) + await expect(writer.close()).rejects.toBeInstanceOf(StreamControllerSupersededError) + expect(await frames()).toEqual([]) + }) +}) + +describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budget', () => { + const userId = generateId() + const workspaceId = generateId() + const chatId = generateId() + const streamId = generateId() + const runId = generateId() + + beforeAll(async () => { + const now = new Date() + await db.insert(user).values({ + id: userId, + name: 'Replay budget fixture', + email: `${userId}@replay-budget.test`, + emailVerified: true, + createdAt: now, + updatedAt: now, + }) + await db.insert(workspace).values({ + id: workspaceId, + name: 'Replay budget fixture', + ownerId: userId, + billedAccountUserId: userId, + }) + await db.insert(permissions).values({ + id: generateId(), + userId, + entityType: 'workspace', + entityId: workspaceId, + permissionType: 'admin', + }) + await db.insert(copilotChats).values({ id: chatId, userId, workspaceId, type: 'mothership' }) + }) + + afterAll(async () => { + 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)) + }) + + it('ends as an error, marks its run terminal, stops the worker, and starts no recovery controller', async () => { + expect(await acquirePendingChatStream(chatId, streamId, 0)).toBe(true) + const lease = { + key: chatStreamLockKey(chatId), + value: (await redis().get(chatStreamLockKey(chatId)))!, + } + const request = { + message: 'Summarize the latest logs', + userId, + messageId: streamId, + chatId, + workspaceId, + } + const [run] = await db + .insert(copilotRuns) + .values({ + id: runId, + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + requestContext: { + requestId: generateId(), + controllerToken: lease.value, + recovery: { + kind: 'interactive_stream', + request, + goRoute: '/api/mothership', + clientToolPickupExpected: false, + }, + }, + }) + .returning() + const { maxOwnerBytes } = getRedisBudgetLimits('copilot_stream') + const [ownerBudgetKey] = getRedisBudgetKeys({ kind: 'copilot_stream', id: streamId }) + /** Room for the turn's opening session frame, not for the worker's tool call. */ + await redis().set(ownerBudgetKey, String(maxOwnerBytes - 512), 'EX', 3600) + worker.script = [ + toolCall('call-refused', 'cli_blocks_get', { command: `blocks get ${'b'.repeat(1024)}` }), + text('never delivered'), + ] + + const response = createSSEStream({ + requestPayload: request, + userId, + streamId, + executionId: run.executionId, + runId, + chatId, + currentChat: { title: 'Existing title' }, + message: request.message, + titleModel: '', + requestId: generateId(), + workspaceId, + admittedRun: run, + orchestrateOptions: { userId, workspaceId, chatId, runId, interactive: true }, + }) + const frames = (await new Response(response).text()) + .split('\n\n') + .filter((frame) => frame.startsWith('data: ')) + .map((frame) => JSON.parse(frame.slice('data: '.length))) + + const [controllerRun] = worker.runs + expect(controllerRun.dispatched).toEqual([]) + expect(controllerRun.sinkErrors[0]).toBeInstanceOf(StreamReplayBudgetExhaustedError) + expect(controllerRun.sinkErrors[0]).not.toBeInstanceOf(StreamControllerSupersededError) + expect(frames.map((frame) => [frame.type, frame.payload.code ?? frame.payload.status])).toEqual( + [ + ['session', undefined], + ['error', REPLAY_BUDGET_EXHAUSTED_CODE], + ['complete', 'error'], + ] + ) + expect(frames[1].payload.message).toBe(REPLAY_BUDGET_EXHAUSTED_MESSAGE) + + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored).toMatchObject({ status: 'error', error: REPLAY_BUDGET_EXHAUSTED_MESSAGE }) + expect(worker.abortRequests).toEqual([expect.objectContaining({ messageId: streamId })]) + expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + + authMockFns.mockGetSession.mockResolvedValue({ + user: { id: userId }, + session: { id: generateId() }, + }) + const reconnect = await streamGET( + new NextRequest(`http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=0`), + { params: Promise.resolve({}) } + ) + const replayed = (await reconnect.text()) + .split('\n\n') + .filter((frame) => frame.startsWith('data: ')) + .map((frame) => JSON.parse(frame.slice('data: '.length))) + expect(replayed.map((frame) => frame.type)).toEqual(['session', 'error', 'complete']) + expect(replayed[1].payload.message).toBe(REPLAY_BUDGET_EXHAUSTED_MESSAGE) + expect(replayed[2].payload.status).toBe('error') + + const recovered = await readChatStream.execute({ + principal: { kind: 'session', userId, sessionId: generateId() }, + input: { streamId }, + }) + expect(recovered?.status).toBe('error') + expect(worker.runs).toHaveLength(1) + expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + }) +}) diff --git a/apps/sim/lib/mothership/request/session/replay-budget.ts b/apps/sim/lib/mothership/request/session/replay-budget.ts new file mode 100644 index 00000000000..9c66d26e756 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-budget.ts @@ -0,0 +1,21 @@ +import type { RedisBudgetRefusal } from '@/lib/core/redis/byte-budget.server' + +/** Run-error code for a turn stopped because its replay buffer refused a write. */ +export const REPLAY_BUDGET_EXHAUSTED_CODE = 'replay_budget_exhausted' + +export const REPLAY_BUDGET_EXHAUSTED_MESSAGE = + 'This turn produced more output than a single response can stream, so it was stopped. The work it already completed has been saved — send a message to continue from there.' + +/** + * 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. + */ +export class StreamReplayBudgetExhaustedError extends Error { + readonly code = REPLAY_BUDGET_EXHAUSTED_CODE + + constructor(readonly refusal: RedisBudgetRefusal) { + super('Stream replay byte budget exhausted') + this.name = 'StreamReplayBudgetExhaustedError' + } +} diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts new file mode 100644 index 00000000000..414377a8543 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -0,0 +1,263 @@ +/** + * Failure modes of stream compaction, each a way a compacted replay frame would + * break the UI, restore, or the byte budget: + * - a bulky string survives, so the frame stays unpersistable; + * - a structured field the UI reads (exit code, resources, citations, cancel + * reason, activity) is truncated or dropped; + * - a whole `output` or `arguments` object is replaced, changing the shape the + * UI reads; + * - arguments the browser executes from, a file preview, or a one-time API key + * are altered; + * - many medium strings keep the frame over budget; + * - split assistant text no longer concatenates to the original receipt; + * - the caller's event is mutated, so dispatch sees the compacted copy. + */ +import { toRecord } from '@sim/utils/object' +import { describe, expect, it } from 'vitest' +import { + isOmittedStreamValue, + STREAM_TRUNCATION_KEY, + type StreamTruncationMarker, +} from '@/lib/mothership/request/session/omission' +import { + compactStreamEventForReplay, + STREAM_EVENT_COMPACTION_THRESHOLD_BYTES, + STREAM_SHORT_STRING_PREVIEW_UNITS, + STREAM_STRING_PREVIEW_UNITS, +} from '@/lib/mothership/request/session/replay-compaction' +import type { StreamEvent } from '@/lib/mothership/request/session/types' + +const MB = 1024 * 1024 + +function payloadOf(event: StreamEvent): Record { + return toRecord(event.payload) +} + +function payloadBytes(event: StreamEvent): number { + return Buffer.byteLength(JSON.stringify(event.payload)) +} + +function marker(event: StreamEvent): StreamTruncationMarker { + return payloadOf(event)[STREAM_TRUNCATION_KEY] as StreamTruncationMarker +} + +describe('compactStreamEventForReplay', () => { + it('truncates only the bulky strings of a Sim tool result and keeps its structured fields', () => { + const citations = [{ index: 1, title: 'Runbook', url: 'https://docs.example/runbook' }] + const observations = [{ kind: 'image', mediaType: 'image/png', name: 'chart.png' }] + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'call-1', + toolName: 'run_code', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + status: 'success', + output: { + stdout: 'o'.repeat(2 * MB), + stderr: 'e'.repeat(300 * 1024), + exitCode: 0, + resources: [{ type: 'file', id: 'file-1', title: 'out.csv' }], + sinkError: 'sink unavailable', + observations, + citations, + }, + }, + } + + const [compacted] = compactStreamEventForReplay(event) + const output = payloadOf(compacted).output as Record + + expect(payloadBytes(compacted)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) + expect(output.stdout).toBe('o'.repeat(STREAM_STRING_PREVIEW_UNITS)) + expect(output.stderr).toBe('e'.repeat(STREAM_STRING_PREVIEW_UNITS)) + expect(output.exitCode).toBe(0) + expect(output.resources).toEqual([{ type: 'file', id: 'file-1', title: 'out.csv' }]) + expect(output.sinkError).toBe('sink unavailable') + expect(output.observations).toEqual(observations) + expect(output.citations).toEqual(citations) + expect(marker(compacted).fields).toEqual([ + { path: '/output/stdout', bytes: 2 * MB, previewBytes: STREAM_STRING_PREVIEW_UNITS }, + { path: '/output/stderr', bytes: 300 * 1024, previewBytes: STREAM_STRING_PREVIEW_UNITS }, + ]) + expect(payloadOf(compacted)).toMatchObject({ + toolCallId: 'call-1', + toolName: 'run_code', + phase: 'result', + success: true, + status: 'success', + }) + }) + + it('keeps retrieval result metadata and cuts each result to a short preview when many medium strings remain', () => { + const results = Array.from({ length: 80 }, (_, index) => ({ + id: `doc-${index}`, + title: `Document ${index}`, + url: `https://docs.example/${index}`, + score: 0.5, + content: 'c'.repeat(6 * 1024), + })) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'call-2', + toolName: 'search_documentation', + executor: 'go', + mode: 'sync', + phase: 'result', + success: true, + output: { data: { results } }, + }, + } + + const [compacted] = compactStreamEventForReplay(event) + const kept = ( + (payloadOf(compacted).output as Record).data as { + results: Array> + } + ).results + + expect(payloadBytes(compacted)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) + expect(kept).toHaveLength(80) + expect(kept[79]).toEqual({ + id: 'doc-79', + title: 'Document 79', + url: 'https://docs.example/79', + score: 0.5, + content: 'c'.repeat(STREAM_SHORT_STRING_PREVIEW_UNITS), + }) + }) + + it('omits a subtree only as a last resort, keeping the arguments object and the keys the UI reads', () => { + const rows = Array.from({ length: 40_000 }, (_, index) => ({ id: index, name: 'row' })) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'call-3', + toolName: 'cli_tables_rows_query', + executor: 'sim', + mode: 'async', + phase: 'call', + arguments: { + activity: { id: 'a-1', title: 'Querying rows' }, + operation: 'query', + description: 'd'.repeat(10_000), + format: 'json', + rows, + }, + }, + } + + const [compacted] = compactStreamEventForReplay(event) + const args = payloadOf(compacted).arguments as Record + + expect(payloadBytes(compacted)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) + expect(args.activity).toEqual({ id: 'a-1', title: 'Querying rows' }) + expect(args.operation).toBe('query') + expect(args.description).toBe('d'.repeat(10_000)) + expect(args.format).toBe('json') + expect(isOmittedStreamValue(args.rows)).toBe(true) + expect(marker(compacted).fields).toEqual([ + { path: '/arguments/rows', bytes: Buffer.byteLength(JSON.stringify(rows)) }, + ]) + }) + + it.each([ + ['a client-executable workflow tool', 'run_workflow', 'sim'], + ['a client-executed tool', 'custom_client_tool', 'client'], + ['a user-local VFS read', 'read', 'go'], + ['a terminal command', 'terminal', 'go'], + ])('leaves the arguments of %s whole', (_label, toolName, executor) => { + const args = { + path: 'user-local/notes.md', + operation: 'run', + input: { body: 'b'.repeat(400 * 1024) }, + } + const event = { + type: 'tool', + payload: { + toolCallId: 'call-4', + toolName, + executor, + mode: 'async', + phase: 'call', + arguments: args, + }, + } as StreamEvent + + const [compacted] = compactStreamEventForReplay(event) + + expect(JSON.stringify(payloadOf(compacted).arguments)).toBe(JSON.stringify(args)) + }) + + it('never compacts a file preview or a generated API key', () => { + const preview: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'call-5', + toolName: 'prepare_file_edit', + previewPhase: 'file_preview_content', + content: 'p'.repeat(400 * 1024), + contentMode: 'snapshot', + previewVersion: 3, + fileName: 'notes.md', + }, + } + const apiKey: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'call-6', + toolName: 'generate_api_key', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + output: { key: 'k'.repeat(400 * 1024) }, + }, + } + + expect(compactStreamEventForReplay(preview)).toEqual([preview]) + expect(compactStreamEventForReplay(apiKey)).toEqual([apiKey]) + }) + + it('splits oversized assistant text into contiguous events that reassemble exactly', () => { + const text = `${'a'.repeat(300 * 1024)}😀${'b'.repeat(300 * 1024)}` + const event: StreamEvent = { + type: 'text', + payload: { channel: 'assistant', text, textOffset: 17 }, + } + + const pieces = compactStreamEventForReplay(event) + + expect(pieces.length).toBeGreaterThan(1) + expect(pieces.map((piece) => payloadOf(piece).text).join('')).toBe(text) + let offset = 17 + for (const piece of pieces) { + expect(payloadOf(piece).textOffset).toBe(offset) + expect(payloadBytes(piece)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) + offset += (payloadOf(piece).text as string).length + } + }) + + it('compacts a copy and leaves the caller’s event whole for dispatch', () => { + const stdout = 's'.repeat(MB) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'call-7', + toolName: 'cli_logs_get', + executor: 'sim', + mode: 'async', + phase: 'call', + arguments: { stdout }, + }, + } + + compactStreamEventForReplay(event) + + expect((payloadOf(event).arguments as Record).stdout).toBe(stdout) + expect(STREAM_TRUNCATION_KEY in payloadOf(event)).toBe(false) + }) +}) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts new file mode 100644 index 00000000000..862dfaff5cf --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -0,0 +1,332 @@ +import { isCurrentBrowserToolName } from '@sim/browser-protocol' +import { isTerminalToolName } from '@sim/terminal-protocol' +import { isRecordLike } from '@sim/utils/object' +import { + MothershipStreamV1EventType, + MothershipStreamV1ToolExecutor, + MothershipStreamV1ToolPhase, +} from '@/lib/mothership/generated/mothership-stream-v1' +import { GenerateApiKey, TOOL_CATALOG } from '@/lib/mothership/generated/tool-catalog-v1' +import { isNativeFileTool, isUserLocalVfsToolCall } from '@/lib/mothership/tools/local-filesystem' +import { + type OmittedStreamValue, + STREAM_TRUNCATION_KEY, + type StreamTruncatedField, + type StreamTruncationMarker, +} from './omission' +import type { StreamEvent } from './types' + +/** + * An event whose payload serializes larger than this is compacted before it is + * persisted and delivered. Well under the replay buffer's 1 MiB single-write + * ceiling, so a compacted event is persistable. + */ +export const STREAM_EVENT_COMPACTION_THRESHOLD_BYTES = 256 * 1024 + +/** + * Compaction stops once the payload fits this, leaving room in the per-stream + * budget for a turn of many large tool results. + */ +export const STREAM_EVENT_COMPACTION_TARGET_BYTES = 128 * 1024 + +/** UTF-16 units kept at the head of a long string leaf. */ +export const STREAM_STRING_PREVIEW_UNITS = 8 * 1024 + +/** + * Second pass for payloads made large by many medium strings (retrieval + * results): each is cut to this many units, keeping its sibling metadata. + */ +export const STREAM_SHORT_STRING_PREVIEW_UNITS = 500 + +/** + * UTF-16 units of `args_delta` forwarded per tool call. The deltas only feed + * a progressive title and a live preview of the arguments; the call frame and + * the result supersede them, so the rest is not forwarded. + */ +export const TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS = 64 * 1024 + +/** + * UTF-16 units per piece of a split assistant text. JSON escapes a unit to at + * most six bytes, so a piece always serializes under the threshold. + */ +const SPLIT_TEXT_MAX_UNITS = Math.floor(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES / 8) + +/** Payload fields that restore, dedupe, and the UI key on; never compacted. */ +const IDENTITY_PAYLOAD_KEYS: ReadonlySet = new Set([ + 'activityDescription', + 'activityReceipt', + 'agent', + 'channel', + 'error', + 'event', + 'execName', + 'executor', + 'kind', + 'mode', + 'op', + 'partial', + 'phase', + 'replay', + 'resource', + 'status', + 'success', + 'targetWorkspaceId', + 'textLength', + 'textOffset', + 'toolCallId', + 'toolName', + 'ui', + 'workspaceId', +]) + +/** + * Keys the UI reads from arguments and outputs (activity labels, targets, + * cancellation). Kept whole at any depth. + */ +const PROTECTED_NESTED_KEYS: ReadonlySet = new Set([ + 'activity', + 'cancelledByUser', + 'description', + 'elementId', + 'fileName', + 'operation', + 'path', + 'pattern', + 'reason', + 'terminalId', + 'timeoutMs', + 'title', + 'toolId', + 'url', + 'workflowId', +]) + +type JsonRecord = Record +type JsonContainer = JsonRecord | unknown[] + +function utf8Bytes(value: string): number { + return Buffer.byteLength(value, 'utf8') +} + +function jsonBytes(value: unknown): number { + return utf8Bytes(JSON.stringify(value) ?? '') +} + +/** Cuts at most `maxUnits` UTF-16 units without splitting a surrogate pair. */ +function headUnits(text: string, maxUnits: number): string { + if (text.length <= maxUnits) return text + const code = text.charCodeAt(maxUnits - 1) + return text.slice(0, code >= 0xd800 && code <= 0xdbff ? maxUnits - 1 : maxUnits) +} + +function splitText(text: string): string[] { + const pieces: string[] = [] + let rest = text + while (rest.length > 0) { + const piece = headUnits(rest, SPLIT_TEXT_MAX_UNITS) + pieces.push(piece) + rest = rest.slice(piece.length) + } + return pieces +} + +function pointer(base: string, key: string | number): string { + return `${base}/${String(key).replaceAll('~', '~0').replaceAll('/', '~1')}` +} + +function entriesOf(container: JsonContainer): Array<[string | number, unknown]> { + return Array.isArray(container) + ? container.map((value, index) => [index, value]) + : Object.entries(container) +} + +function setEntry(container: JsonContainer, key: string | number, value: unknown): void { + if (Array.isArray(container)) container[Number(key)] = value + else container[String(key)] = value +} + +function isProtectedKey(key: string | number): boolean { + return typeof key === 'string' && PROTECTED_NESTED_KEYS.has(key) +} + +function containsProtectedKey(value: unknown): boolean { + if (Array.isArray(value)) return value.some(containsProtectedKey) + if (!isRecordLike(value)) return false + return Object.entries(value).some( + ([key, field]) => PROTECTED_NESTED_KEYS.has(key) || containsProtectedKey(field) + ) +} + +interface CompactionState { + payload: JsonRecord + roots: ReadonlySet + fields: Map +} + +/** Shortens every string leaf longer than `maxUnits` to its head, in place. */ +function shortenStrings( + state: CompactionState, + container: JsonContainer, + base: string, + maxUnits: number +): void { + for (const [key, value] of entriesOf(container)) { + if (base === '' ? !state.roots.has(String(key)) : isProtectedKey(key)) continue + const path = pointer(base, key) + if (typeof value === 'string') { + if (value.length <= maxUnits) continue + const head = headUnits(value, maxUnits) + const bytes = state.fields.get(path)?.bytes ?? utf8Bytes(value) + state.fields.set(path, { path, bytes, previewBytes: utf8Bytes(head) }) + setEntry(container, key, head) + } else if (Array.isArray(value) || isRecordLike(value)) { + shortenStrings(state, value, path, maxUnits) + } + } +} + +interface Candidate { + container: JsonContainer + key: string | number + path: string + bytes: number +} + +function collectCandidates( + state: CompactionState, + container: JsonContainer, + base: string, + out: Candidate[] +): void { + for (const [key, value] of entriesOf(container)) { + if (base === '' ? !state.roots.has(String(key)) : isProtectedKey(key)) continue + const path = pointer(base, key) + if (Array.isArray(value) || isRecordLike(value)) { + if (!containsProtectedKey(value)) { + out.push({ container, key, path, bytes: jsonBytes(value) }) + } + collectCandidates(state, value, path, out) + } else if (typeof value === 'string') { + out.push({ container, key, path, bytes: jsonBytes(value) }) + } + } +} + +/** Last resort: replaces the largest unprotected subtrees with stubs until the payload fits. */ +function omitLargestSubtrees(state: CompactionState): void { + const candidates: Candidate[] = [] + collectCandidates(state, state.payload, '', candidates) + candidates.sort((a, b) => b.bytes - a.bytes) + const omitted: string[] = [] + for (const candidate of candidates) { + if (jsonBytes(state.payload) <= STREAM_EVENT_COMPACTION_TARGET_BYTES) return + if (omitted.some((path) => candidate.path.startsWith(`${path}/`))) continue + const stub: OmittedStreamValue = { + omitted: true, + bytes: state.fields.get(candidate.path)?.bytes ?? candidate.bytes, + } + setEntry(candidate.container, candidate.key, stub) + for (const path of state.fields.keys()) { + if (path.startsWith(`${candidate.path}/`)) state.fields.delete(path) + } + state.fields.set(candidate.path, { path: candidate.path, bytes: stub.bytes }) + omitted.push(candidate.path) + } +} + +/** + * The browser may start these tools from the call frame's arguments, so the + * arguments are never compacted; a frame too large with them whole is refused. + */ +function isClientExecutableCall(payload: JsonRecord): boolean { + const toolName = typeof payload.toolName === 'string' ? payload.toolName : '' + const args = isRecordLike(payload.arguments) ? payload.arguments : undefined + const catalogEntry = Object.hasOwn(TOOL_CATALOG, toolName) ? TOOL_CATALOG[toolName] : undefined + return ( + payload.executor === MothershipStreamV1ToolExecutor.client || + (isRecordLike(payload.ui) && payload.ui.clientExecutable === true) || + catalogEntry?.route === 'client' || + catalogEntry?.clientExecutable === true || + isNativeFileTool(toolName) || + isUserLocalVfsToolCall(toolName, args) || + isCurrentBrowserToolName(toolName) || + isTerminalToolName(toolName) + ) +} + +function isNeverCompacted(payload: JsonRecord): boolean { + return ( + 'previewPhase' in payload || + (payload.phase === MothershipStreamV1ToolPhase.result && payload.toolName === GenerateApiKey.id) + ) +} + +function withPayload(event: StreamEvent, payload: JsonRecord): StreamEvent { + return { ...event, payload } as StreamEvent +} + +/** + * Bounds one outgoing stream event so the replay buffer can persist it. Applied + * only by the stream writer, to the copy it delivers and persists; the caller + * keeps the full event for dispatch, the async tool row, and approval. + * + * Returns the event unchanged when its payload is within + * {@link STREAM_EVENT_COMPACTION_THRESHOLD_BYTES}. Otherwise: + * - assistant text splits into contiguous events whose concatenation is exact; + * - long string leaves are cut to their head in place, so `arguments` and + * `output` keep their shape; a second pass cuts medium strings shorter; + * - only if the payload is still too large are the largest unprotected subtrees + * replaced by an {@link OmittedStreamValue}. + * Every changed field is listed in a payload-level {@link StreamTruncationMarker}. + * + * File previews, `generate_api_key` results, identity fields, and the arguments + * of client-executable calls are never compacted; such an event that stays too + * large is refused by the buffer and ends the turn with an error. + */ +export function compactStreamEventForReplay(event: StreamEvent): StreamEvent[] { + const original = event.payload + if (!isRecordLike(original)) return [event] + const originalBytes = jsonBytes(original) + if (originalBytes <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES || isNeverCompacted(original)) { + return [event] + } + + if (event.type === MothershipStreamV1EventType.text) { + const { text, textOffset } = event.payload + let offset = textOffset + return splitText(text).map((piece) => { + const split = withPayload(event, { + ...original, + text: piece, + ...(offset !== undefined ? { textOffset: offset } : {}), + }) + if (offset !== undefined) offset += piece.length + return split + }) + } + + const payload: JsonRecord = structuredClone(original) + const protectsArguments = + payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutableCall(payload) + const roots = new Set( + Object.keys(payload).filter( + (key) => !IDENTITY_PAYLOAD_KEYS.has(key) && !(protectsArguments && key === 'arguments') + ) + ) + const state: CompactionState = { payload, roots, fields: new Map() } + + shortenStrings(state, payload, '', STREAM_STRING_PREVIEW_UNITS) + if (jsonBytes(payload) > STREAM_EVENT_COMPACTION_TARGET_BYTES) { + shortenStrings(state, payload, '', STREAM_SHORT_STRING_PREVIEW_UNITS) + } + if (jsonBytes(payload) > STREAM_EVENT_COMPACTION_TARGET_BYTES) { + omitLargestSubtrees(state) + } + if (state.fields.size === 0) return [event] + + const marker: StreamTruncationMarker = { + bytes: originalBytes, + fields: [...state.fields.values()], + } + return [withPayload(event, { ...payload, [STREAM_TRUNCATION_KEY]: marker })] +} diff --git a/apps/sim/lib/mothership/request/session/writer.test.ts b/apps/sim/lib/mothership/request/session/writer.test.ts index 6fa5c4361a4..935aa4b33bc 100644 --- a/apps/sim/lib/mothership/request/session/writer.test.ts +++ b/apps/sim/lib/mothership/request/session/writer.test.ts @@ -13,6 +13,9 @@ vi.mock('@/lib/mothership/request/session/buffer', () => ({ appendEvents, })) +import { StreamControllerSupersededError } from '@/lib/mothership/request/session/controller-lease' +import { StreamReplayBudgetExhaustedError } from '@/lib/mothership/request/session/replay-budget' +import { TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS } from '@/lib/mothership/request/session/replay-compaction' import { StreamWriter } from '@/lib/mothership/request/session/writer' function decodeChunk(value: Uint8Array): string { @@ -77,7 +80,7 @@ describe('StreamWriter', () => { expect(controller.close).toHaveBeenCalledOnce() }) - it('does not deliver unreplayable events when an owned controller exhausts its byte budget', async () => { + it('ends delivery with a budget error, not a handoff, and still delivers the terminal verdict', async () => { appendEvents.mockResolvedValueOnce({ persisted: false, refusal: { @@ -93,14 +96,76 @@ describe('StreamWriter', () => { userId: 'user-1', lease: { key: 'chat-lock', value: 'current-controller' }, }) - const controller = { enqueue: vi.fn(), close: vi.fn() } + 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: 'unsaved' } }) - ).rejects.toThrow('Stream replay byte budget exhausted') - expect(controller.enqueue).not.toHaveBeenCalled() + writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'later' } }) + ).rejects.toBeInstanceOf(StreamReplayBudgetExhaustedError) + expect(delivered).toEqual([]) expect(writer.persistenceStopped).toBe(true) - await expect(writer.close()).rejects.toThrow('Stream replay byte budget exhausted') + 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("forwards only the head of a tool call's argument deltas", async () => { + const writer = new StreamWriter({ streamId: 'stream-1', requestId: 'req-1' }) + const delivered: string[] = [] + writer.attach({ + enqueue: (value: Uint8Array) => delivered.push(decodeChunk(value)), + close: () => {}, + } as unknown as ReadableStreamDefaultController) + const delta = (toolCallId: string, argumentsDelta: string): StreamEvent => ({ + type: 'tool', + payload: { + toolCallId, + toolName: 'sim_cli', + executor: 'go', + mode: 'async', + phase: 'args_delta', + argumentsDelta, + }, + }) + const streamed = `{"argv":["logs","get"],"stdin":"${'x'.repeat(100_000)}"}` + + for (let index = 0; index < streamed.length; index += 30_000) { + await writer.publish(delta('call-1', streamed.slice(index, index + 30_000))) + } + await writer.publish(delta('call-2', '{"argv":["blocks"]}')) + await writer.close() + + const forwarded = delivered.map((frame) => JSON.parse(frame.replace(/^data: /, '')).payload) + expect( + forwarded + .filter((payload) => payload.toolCallId === 'call-1') + .map((payload) => payload.argumentsDelta) + .join('') + ).toBe(streamed.slice(0, TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS)) + expect(forwarded.at(-1)).toMatchObject({ + toolCallId: 'call-2', + argumentsDelta: '{"argv":["blocks"]}', + }) }) it('enqueues before persistence completes and flushes pending writes on close', async () => { diff --git a/apps/sim/lib/mothership/request/session/writer.ts b/apps/sim/lib/mothership/request/session/writer.ts index 05eb280b809..ab422fca718 100644 --- a/apps/sim/lib/mothership/request/session/writer.ts +++ b/apps/sim/lib/mothership/request/session/writer.ts @@ -1,11 +1,19 @@ import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { encodeSSEComment } from '@/lib/core/utils/sse' -import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' +import { + MothershipStreamV1EventType, + MothershipStreamV1ToolPhase, +} from '@/lib/mothership/generated/mothership-stream-v1' import { appendEvents } from './buffer' import type { PersistedStreamEventEnvelope } from './contract' import type { ChatStreamLease } from './controller-lease' import { createEvent } from './event' +import { StreamReplayBudgetExhaustedError } from './replay-budget' +import { + compactStreamEventForReplay, + TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS, +} from './replay-compaction' import { encodeSSEEnvelope } from './sse' import type { StreamEvent } from './types' @@ -26,6 +34,18 @@ export interface StreamWriterOptions { initialSeq?: number } +/** The turn's closing verdict, which a refused writer still owes its live client. */ +function isTurnTerminalEvent(event: StreamEvent): boolean { + return ( + event.type === MothershipStreamV1EventType.error || + event.type === MothershipStreamV1EventType.complete + ) +} + +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. */ const PERSISTENCE_ALREADY_STOPPED = { persisted: true } as const @@ -48,6 +68,8 @@ export class StreamWriter { private pendingEnvelopes: PersistedStreamEventEnvelope[] = [] private persistenceTail: Promise = Promise.resolve() private lastPersistenceError: Error | null = null + private replayBudgetError: StreamReplayBudgetExhaustedError | null = null + private readonly forwardedArgsDeltaUnits = new Map() private readonly lease?: ChatStreamLease constructor(options: StreamWriterOptions) { @@ -77,7 +99,8 @@ 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. + * budget. Leased delivery stops before the refused event (only the terminal verdict + * still reaches the client); unleased delivery continues. */ get persistenceStopped(): boolean { return this._persistenceStopped @@ -114,29 +137,57 @@ export class StreamWriter { } } + /** + * Delivers an event and records it for replay. Oversized events are compacted + * first ({@link compactStreamEventForReplay}); the compacted copy is what both + * the client and the replay buffer receive, 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(event: StreamEvent): void | Promise { - const envelope = this.createEnvelope(event) + const forwarded = this.limitArgsDelta(event) + if (!forwarded) return + const envelopes = compactStreamEventForReplay(forwarded).map((compacted) => + this.createEnvelope(compacted) + ) if (this.lease) { + const lease = this.lease // A replacement must see every event the browser has received. Fence // persistence before delivery, and before dispatching the event's tool. - const delivery = this.persistenceTail.then(async () => { + const persistThenDeliver = async () => { + if (this.replayBudgetError) { + if (!isTurnTerminalEvent(event)) throw this.replayBudgetError + this.deliver(envelopes, event) + return + } const result = await appendEvents( - [envelope], + envelopes, { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, - this.lease + lease ) if (!result.persisted) { this._persistenceStopped = true - throw new Error('Stream replay byte budget exhausted') + this.replayBudgetError = new StreamReplayBudgetExhaustedError(result.refusal) + if (!isTurnTerminalEvent(event)) throw this.replayBudgetError } - this.enqueue(envelope) - if (event.type === MothershipStreamV1EventType.complete) this._sawComplete = true + this.deliver(envelopes, event) + } + const delivery = this.persistenceTail.then(persistThenDeliver, (error: unknown) => { + ignoreReplayBudgetRefusal(error) + return persistThenDeliver() }) this.persistenceTail = delivery return delivery } - this.enqueue(envelope) - this.queuePersistence(envelope) + for (const envelope of envelopes) { + this.enqueue(envelope) + this.queuePersistence(envelope) + } if (event.type === MothershipStreamV1EventType.complete) { this._sawComplete = true } @@ -148,7 +199,9 @@ export class StreamWriter { async flush(): Promise { this.flushPendingPersistence() - await this.persistenceTail + // 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 @@ -171,6 +224,31 @@ export class StreamWriter { } } + /** + * Forwards the head of each tool call's argument deltas and drops the rest: + * they only feed a progressive title, and the call frame supersedes them. + */ + private limitArgsDelta(event: StreamEvent): StreamEvent | undefined { + if (event.type !== MothershipStreamV1EventType.tool || !('phase' in event.payload)) { + return event + } + const payload = event.payload + if (payload.phase !== MothershipStreamV1ToolPhase.args_delta) return event + const forwarded = this.forwardedArgsDeltaUnits.get(payload.toolCallId) ?? 0 + const remaining = TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS - forwarded + if (remaining <= 0) return undefined + const argumentsDelta = payload.argumentsDelta.slice(0, remaining) + this.forwardedArgsDeltaUnits.set(payload.toolCallId, forwarded + argumentsDelta.length) + return argumentsDelta === payload.argumentsDelta + ? event + : { ...event, payload: { ...payload, argumentsDelta } } + } + + private deliver(envelopes: PersistedStreamEventEnvelope[], event: StreamEvent): void { + for (const envelope of envelopes) this.enqueue(envelope) + if (event.type === MothershipStreamV1EventType.complete) this._sawComplete = true + } + private enqueue(envelope: PersistedStreamEventEnvelope): void { if (this._clientDisconnected || !this.controller) return try { diff --git a/apps/sim/lib/mothership/request/types.ts b/apps/sim/lib/mothership/request/types.ts index 527b111b13c..1c11e6b7d4c 100644 --- a/apps/sim/lib/mothership/request/types.ts +++ b/apps/sim/lib/mothership/request/types.ts @@ -298,6 +298,8 @@ export interface OrchestratorResult { chatId?: string requestId?: string error?: string + /** Machine-readable cause of `error`, forwarded as the stream error event's `code`. */ + errorCode?: string errors?: string[] usage?: { prompt: number; completion: number } cost?: { input: number; output: number; total: number } From 38cc7f9eb16be14ae59046456a7ae0e10cf78119 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 18:42:14 -0700 Subject: [PATCH 02/18] fix(mothership): render stream-omitted tool values honestly A call frame whose whole argument object was omitted by stream compaction no longer replaces a tool node's arguments, and the permission card and generic tool output render omitted values as "Too large to show (N KB)" instead of raw marker JSON. --- .../agent-group/tool-permission-card.tsx | 3 ++- .../generic-resource-content.tsx | 3 ++- .../home/hooks/stream/turn-model.test.ts | 19 +++++++++++++++++++ .../home/hooks/stream/turn-model.ts | 5 ++++- 4 files changed, 27 insertions(+), 3 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx index 0af0622c44d..3426050bcc8 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx @@ -22,6 +22,7 @@ import { } from '@/components/ui/activity-status' import { requestJson } from '@/lib/api/client/request' import { copilotToolPermissionContract } from '@/lib/api/contracts/copilot' +import { describeOmittedStreamValues } from '@/lib/mothership/request/session/omission' import { generalSettingsKeys } from '@/hooks/queries/current-user-data' import { useToolPermissionStore } from '@/stores/tool-permission/store' @@ -59,7 +60,7 @@ function isGoneError(error: unknown): boolean { function argsPreview(params: Record | undefined): string | undefined { if (!params || Object.keys(params).length === 0) return undefined try { - const json = JSON.stringify(params, null, 2) + const json = JSON.stringify(params, describeOmittedStreamValues, 2) return json.length > 600 ? `${json.slice(0, 600)}…` : json } catch { return undefined diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx index d763ae269f7..2ad127f6be3 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx @@ -2,6 +2,7 @@ import { useEffect, useRef } from 'react' import { PillsRing } from '@sim/emcn' +import { describeOmittedStreamValues } from '@/lib/mothership/request/session/omission' import { getToolStatusDisplayTitle } from '@/lib/mothership/tools/tool-display' import type { GenericResourceData } from '@/app/workspace/[workspaceId]/home/types' @@ -59,7 +60,7 @@ export function GenericResourceContent({ data }: GenericResourceContentProps) { 
               {typeof entry.result.output === 'string'
                 ? entry.result.output
-                : JSON.stringify(entry.result.output, null, 2)}
+                : JSON.stringify(entry.result.output, describeOmittedStreamValues, 2)}
)} {entry.result?.error && ( diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts index eb8bdc88674..27b47d03b1c 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts @@ -136,6 +136,25 @@ describe('reduceEvent — tool lifecycle', () => { ]) expect(tool(m, 'tc-1').status).toBe('cancelled') }) + + it('keeps the arguments it has when a compacted call frame omitted the whole argument object', () => { + const m = apply([ + envelope(1, 'tool', { + phase: 'call', + toolCallId: 'tc-1', + toolName: 'cli_tables_rows_query', + status: 'generating', + arguments: { operation: 'query' }, + }), + envelope(2, 'tool', { + phase: 'call', + toolCallId: 'tc-1', + toolName: 'cli_tables_rows_query', + arguments: { omitted: true, bytes: 3_000_000 }, + }), + ]) + expect(tool(m, 'tc-1').args).toEqual({ operation: 'query' }) + }) }) describe('reduceEvent — subagent lifecycle', () => { diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts index 283b7e0451e..386031fb9c2 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts @@ -13,6 +13,7 @@ import { import { ToolActivity } from '@/lib/mothership/generated/protocol' import { CallIntegrationTool } from '@/lib/mothership/generated/tool-catalog-v1' import type { PersistedStreamEventEnvelope } from '@/lib/mothership/request/session/contract' +import { isOmittedStreamValue } from '@/lib/mothership/request/session/omission' import type { TaskBlockInfo } from '@/lib/mothership/request/types' import { extractStreamingStringArgument } from '@/lib/mothership/tools/streaming-args' import { @@ -538,7 +539,9 @@ export function reduceEvent(model: TurnModel, envelope: PersistedStreamEventEnve // back into an ordinary running row without waiting for the result. node.status = 'running' } - if (isRecordLike(payload.arguments)) { + // Stream compaction may have omitted the whole argument object; its stub + // is not arguments, so the node keeps whatever it already had. + if (isRecordLike(payload.arguments) && !isOmittedStreamValue(payload.arguments)) { node.args = payload.arguments const activity = ToolActivity.safeParse(payload.arguments.activity) if (!node.activity && activity.success) node.activity = activity.data From a0fe73fd60d5181cf07e3e96b52a59172e2fff69 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 19:14:43 -0700 Subject: [PATCH 03/18] refactor(mothership): replace replay compaction with one string-truncation pass MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Compaction is now a single linear walk: when a payload's strings could serialize past 256 KiB, string leaves longer than 8K units are cut to their head with an inline "…[truncated, N total]" note, copying only what changes. Identity and UI-read keys, file previews, generate_api_key results, and the arguments of calls the browser executes are never cut; anything still over the write ceiling is refused and ends the turn. The subtree omission, second cut pass, truncation marker, text splitting, argument-delta cap, and the UI stub handling are removed. One predicate, isClientExecutedToolCall, now names the tool calls the browser starts from the call frame; both the stream compaction and the client dispatch use it, and the workflow tool names move beside it. Also fixes three defects in the refusal path: - A superseded controller that is refused an oversized frame no longer expires its successor's stream or clears its abort marker; teardown keys on whether this controller handed off, not on the abort reason. - A refusal after a user Stop stays a cancellation: the abort reason decides. - The per-user hourly ceiling gets its own message instead of telling the user to continue immediately. --- apps/sim/app/api/copilot/confirm/route.ts | 2 +- .../agent-group/tool-permission-card.tsx | 3 +- .../generic-resource-content.tsx | 3 +- .../home/hooks/message-reconcile.ts | 2 +- .../hooks/stream/handle-tool-event.test.ts | 3 - .../home/hooks/stream/handle-tool-event.ts | 71 +--- .../home/hooks/stream/turn-model.test.ts | 19 - .../home/hooks/stream/turn-model.ts | 5 +- .../[workspaceId]/home/hooks/use-chat.ts | 2 +- .../lib/mothership/request/handlers/tool.ts | 6 +- .../mothership/request/lifecycle/run.test.ts | 78 ++-- .../lib/mothership/request/lifecycle/run.ts | 34 +- .../lib/mothership/request/lifecycle/start.ts | 47 ++- .../mothership/request/session/omission.ts | 50 --- .../session/replay-budget.integration.ts | 134 ++++--- .../request/session/replay-budget.ts | 19 +- .../request/session/replay-compaction.test.ts | 222 ++++------- .../request/session/replay-compaction.ts | 365 ++++-------------- .../mothership/request/session/writer.test.ts | 40 -- .../lib/mothership/request/session/writer.ts | 88 ++--- .../mothership/tools/client-executed-tools.ts | 32 ++ .../lib/mothership/tools/workflow-tools.ts | 14 +- 22 files changed, 416 insertions(+), 823 deletions(-) delete mode 100644 apps/sim/lib/mothership/request/session/omission.ts create mode 100644 apps/sim/lib/mothership/tools/client-executed-tools.ts diff --git a/apps/sim/app/api/copilot/confirm/route.ts b/apps/sim/app/api/copilot/confirm/route.ts index b34ba42a32b..430b38b0b71 100644 --- a/apps/sim/app/api/copilot/confirm/route.ts +++ b/apps/sim/app/api/copilot/confirm/route.ts @@ -42,13 +42,13 @@ import { retainSealedClientToolContext, sealClientToolCompletion, } from '@/lib/mothership/request/tools/client-completion-seal.server' +import { isWorkflowToolName } from '@/lib/mothership/tools/client-executed-tools' import { createStructuralWorkflowToolCompletionData, getWorkflowToolCompletionExecutionId, getWorkflowToolCompletionMessage, getWorkflowToolConfirmationStatus, getWorkflowToolLaunchError, - isWorkflowToolName, resolveWorkflowToolTargetId, WORKFLOW_EXECUTION_BUSY, type WorkflowToolLaunchError, diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx index 3426050bcc8..0af0622c44d 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/message-content/components/agent-group/tool-permission-card.tsx @@ -22,7 +22,6 @@ import { } from '@/components/ui/activity-status' import { requestJson } from '@/lib/api/client/request' import { copilotToolPermissionContract } from '@/lib/api/contracts/copilot' -import { describeOmittedStreamValues } from '@/lib/mothership/request/session/omission' import { generalSettingsKeys } from '@/hooks/queries/current-user-data' import { useToolPermissionStore } from '@/stores/tool-permission/store' @@ -60,7 +59,7 @@ function isGoneError(error: unknown): boolean { function argsPreview(params: Record | undefined): string | undefined { if (!params || Object.keys(params).length === 0) return undefined try { - const json = JSON.stringify(params, describeOmittedStreamValues, 2) + const json = JSON.stringify(params, null, 2) return json.length > 600 ? `${json.slice(0, 600)}…` : json } catch { return undefined diff --git a/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx b/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx index 2ad127f6be3..d763ae269f7 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-content/components/generic-resource-content/generic-resource-content.tsx @@ -2,7 +2,6 @@ import { useEffect, useRef } from 'react' import { PillsRing } from '@sim/emcn' -import { describeOmittedStreamValues } from '@/lib/mothership/request/session/omission' import { getToolStatusDisplayTitle } from '@/lib/mothership/tools/tool-display' import type { GenericResourceData } from '@/app/workspace/[workspaceId]/home/types' @@ -60,7 +59,7 @@ export function GenericResourceContent({ data }: GenericResourceContentProps) { 
               {typeof entry.result.output === 'string'
                 ? entry.result.output
-                : JSON.stringify(entry.result.output, describeOmittedStreamValues, 2)}
+                : JSON.stringify(entry.result.output, null, 2)}
)} {entry.result?.error && ( diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts index dccc95ad1a9..5245f98bc1e 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts @@ -18,7 +18,7 @@ import { MothershipStreamV1ToolPhase, } from '@/lib/mothership/generated/mothership-stream-v1' import type { StreamBatchEvent } from '@/lib/mothership/request/session/types' -import { isWorkflowToolName } from '@/lib/mothership/tools/workflow-tools' +import { isWorkflowToolName } from '@/lib/mothership/tools/client-executed-tools' import type { ContentBlock } from '@/app/workspace/[workspaceId]/home/types' import type { MothershipChatHistory } from '@/hooks/queries/mothership-chats' import { isZeroStreamCursor } from './stream-protocol' diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.test.ts index b5a13580024..e42181d47c2 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.test.ts @@ -4,9 +4,6 @@ vi.mock('@/lib/mothership/resources/extraction', () => ({ isResourceToolName: vi.fn(() => false), extractResourcesFromToolResult: vi.fn(() => []), })) -vi.mock('@/lib/mothership/tools/workflow-tools', () => ({ - isWorkflowToolName: vi.fn(() => false), -})) vi.mock( '@/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-registry', () => ({ invalidateResourceQueries: vi.fn() }) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts index 0fb9877fa43..1d92bcf8ddb 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/handle-tool-event.ts @@ -14,8 +14,10 @@ import { extractResourcesFromToolResult, isResourceToolName, } from '@/lib/mothership/resources/extraction' -import { isNativeFileTool, isUserLocalVfsToolCall } from '@/lib/mothership/tools/local-filesystem' -import { isWorkflowToolName } from '@/lib/mothership/tools/workflow-tools' +import { + isClientExecutedToolCall, + isWorkflowToolName, +} from '@/lib/mothership/tools/client-executed-tools' import { invalidateResourceQueries } from '@/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-registry' import type { StreamLoopContext } from '@/app/workspace/[workspaceId]/home/hooks/stream/stream-context' import { @@ -188,56 +190,23 @@ export function handleToolEvent(ctx: StreamLoopContext, parsed: ToolEvent): void const name = payload.toolName const isPartial = payload.partial === true || payload.status === MothershipStreamV1ToolStatus.generating - if (isWorkflowToolName(name) && !isPartial) { - const shouldStartWorkflowTool = - !deps.options.suppressedWorkflowToolStartIds?.has(rawId) && - node?.kind === 'tool' && - node.status === 'running' && - !node.result - if (shouldStartWorkflowTool) { - const args = payload.arguments as Record | undefined + const args = payload.arguments as Record | undefined + const shouldStartClientTool = + isClientExecutedToolCall(name, args) && + !isPartial && + !deps.options.suppressedWorkflowToolStartIds?.has(rawId) && + node?.kind === 'tool' && + node.status === 'running' && + !node.result + if (shouldStartClientTool) { + if (isWorkflowToolName(name)) { deps.startClientWorkflowTool(rawId, name, args ?? {}) - } - } - const localFilesystemArgs = payload.arguments as Record | undefined - if ((isNativeFileTool(name) || isUserLocalVfsToolCall(name, localFilesystemArgs)) && !isPartial) { - const shouldStartLocalFilesystemTool = - !deps.options.suppressedWorkflowToolStartIds?.has(rawId) && - node?.kind === 'tool' && - node.status === 'running' && - !node.result - if (shouldStartLocalFilesystemTool) { - deps.startClientLocalFilesystemTool(rawId, name, localFilesystemArgs ?? {}) - } - } - if (isCurrentBrowserToolName(name) && !isPartial) { - const shouldStartBrowserTool = - !deps.options.suppressedWorkflowToolStartIds?.has(rawId) && - node?.kind === 'tool' && - node.status === 'running' && - !node.result - if (shouldStartBrowserTool) { - deps.startClientBrowserTool( - rawId, - name, - (payload.arguments as Record | undefined) ?? {}, - parsed.ts - ) - } - } - if (isTerminalToolName(name) && !isPartial) { - const shouldStartTerminalTool = - !deps.options.suppressedWorkflowToolStartIds?.has(rawId) && - node?.kind === 'tool' && - node.status === 'running' && - !node.result - if (shouldStartTerminalTool) { - deps.startClientTerminalTool( - rawId, - name, - (payload.arguments as Record | undefined) ?? {}, - parsed.ts - ) + } else if (isCurrentBrowserToolName(name)) { + deps.startClientBrowserTool(rawId, name, args ?? {}, parsed.ts) + } else if (isTerminalToolName(name)) { + deps.startClientTerminalTool(rawId, name, args ?? {}, parsed.ts) + } else { + deps.startClientLocalFilesystemTool(rawId, name, args ?? {}) } } ops.flush() diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts index 27b47d03b1c..eb8bdc88674 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.test.ts @@ -136,25 +136,6 @@ describe('reduceEvent — tool lifecycle', () => { ]) expect(tool(m, 'tc-1').status).toBe('cancelled') }) - - it('keeps the arguments it has when a compacted call frame omitted the whole argument object', () => { - const m = apply([ - envelope(1, 'tool', { - phase: 'call', - toolCallId: 'tc-1', - toolName: 'cli_tables_rows_query', - status: 'generating', - arguments: { operation: 'query' }, - }), - envelope(2, 'tool', { - phase: 'call', - toolCallId: 'tc-1', - toolName: 'cli_tables_rows_query', - arguments: { omitted: true, bytes: 3_000_000 }, - }), - ]) - expect(tool(m, 'tc-1').args).toEqual({ operation: 'query' }) - }) }) describe('reduceEvent — subagent lifecycle', () => { diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts index 386031fb9c2..283b7e0451e 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/turn-model.ts @@ -13,7 +13,6 @@ import { import { ToolActivity } from '@/lib/mothership/generated/protocol' import { CallIntegrationTool } from '@/lib/mothership/generated/tool-catalog-v1' import type { PersistedStreamEventEnvelope } from '@/lib/mothership/request/session/contract' -import { isOmittedStreamValue } from '@/lib/mothership/request/session/omission' import type { TaskBlockInfo } from '@/lib/mothership/request/types' import { extractStreamingStringArgument } from '@/lib/mothership/tools/streaming-args' import { @@ -539,9 +538,7 @@ export function reduceEvent(model: TurnModel, envelope: PersistedStreamEventEnve // back into an ordinary running row without waiting for the result. node.status = 'running' } - // Stream compaction may have omitted the whole argument object; its stub - // is not arguments, so the node keeps whatever it already had. - if (isRecordLike(payload.arguments) && !isOmittedStreamValue(payload.arguments)) { + if (isRecordLike(payload.arguments)) { node.args = payload.arguments const activity = ToolActivity.safeParse(payload.arguments.activity) if (!node.activity && activity.success) node.activity = activity.data diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index 2885e0fd8f1..cbac7f0d951 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -82,8 +82,8 @@ import { } from '@/lib/mothership/tools/client/run-tool-execution' import { executeTerminalToolOnClient } from '@/lib/mothership/tools/client/terminal-tool-execution' import { setCurrentChatTraceparent } from '@/lib/mothership/tools/client/trace-context' +import { isWorkflowToolName } from '@/lib/mothership/tools/client-executed-tools' import { isNativeFileTool, isUserLocalVfsToolCall } from '@/lib/mothership/tools/local-filesystem' -import { isWorkflowToolName } from '@/lib/mothership/tools/workflow-tools' import { initTerminalTransport } from '@/lib/terminal/transport' import { getQueryClient } from '@/app/_shell/providers/get-query-client' import { chatUrl } from '@/app/workspace/[workspaceId]/home/hooks/chat-url' diff --git a/apps/sim/lib/mothership/request/handlers/tool.ts b/apps/sim/lib/mothership/request/handlers/tool.ts index 1e81d1bb07c..deadb956753 100644 --- a/apps/sim/lib/mothership/request/handlers/tool.ts +++ b/apps/sim/lib/mothership/request/handlers/tool.ts @@ -48,6 +48,7 @@ import type { } from '@/lib/mothership/request/types' import { getToolEntry, isSimExecuted } from '@/lib/mothership/tool-executor' import { isToolHiddenInUi } from '@/lib/mothership/tools/client/hidden-tools' +import { isWorkflowToolName } from '@/lib/mothership/tools/client-executed-tools' import { isUserLocalVfsToolCall } from '@/lib/mothership/tools/local-filesystem' import { extractStreamingStringArgument } from '@/lib/mothership/tools/streaming-args' import { readToolActivity } from '@/lib/mothership/tools/tool-activity' @@ -56,10 +57,7 @@ import { normalizeToolActivityDescription, refineStreamingCliToolName, } from '@/lib/mothership/tools/tool-display' -import { - isWorkflowToolName, - resolveWorkflowToolTargetId, -} from '@/lib/mothership/tools/workflow-tools' +import { resolveWorkflowToolTargetId } from '@/lib/mothership/tools/workflow-tools' import { getBlockByToolName } from '@/blocks/registry' import { abortPendingToolIfStreamDead, diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index ce362e0f001..e3747581b41 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -208,7 +208,6 @@ import { import { runCopilotLifecycle } from '@/lib/mothership/request/lifecycle/run' import { REPLAY_BUDGET_EXHAUSTED_CODE, - REPLAY_BUDGET_EXHAUSTED_MESSAGE, StreamReplayBudgetExhaustedError, } from '@/lib/mothership/request/session/replay-budget' import { executeToolAndReport } from '@/lib/mothership/request/tools/executor' @@ -1952,6 +1951,34 @@ describe('runCopilotLifecycle', () => { ) }) + const ownerRefusal = () => + new StreamReplayBudgetExhaustedError({ + resource: 'owner_redis_bytes', + currentBytes: 32 * 1024 * 1024, + limitBytes: 32 * 1024 * 1024, + attemptedBytes: 512, + }) + + function runWithStreamAbort(abortController: AbortController) { + return runCopilotLifecycle( + { message: 'hello', messageId: 'stream-1' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'exec-1', + runId: 'run-1', + abortSignal: abortController.signal, + executionContext: { + userId: 'user-1', + workflowId: '', + workspaceId: 'ws-1', + chatId: 'chat-1', + }, + } + ) + } + it.each([ ['throws the refusal', true], ['returns after the abort', false], @@ -1959,55 +1986,46 @@ describe('runCopilotLifecycle', () => { 'ends a turn stopped by a refused replay write as an error, not a user cancel, when the stream %s', async (_label, throws) => { const abortController = new AbortController() + const refusal = ownerRefusal() mockRunStreamLoop.mockImplementationOnce( - async ( - _fetchUrl: string, - _fetchOptions: RequestInit, - context: StreamingContext - ): Promise => { + async (_url: string, _init: RequestInit, context: StreamingContext): Promise => { context.accumulatedContent = 'partial answer' - const refusal = new StreamReplayBudgetExhaustedError({ - resource: 'owner_redis_bytes', - currentBytes: 32 * 1024 * 1024, - limitBytes: 32 * 1024 * 1024, - attemptedBytes: 512, - }) abortController.abort(refusal) context.wasAborted = true if (throws) throw refusal } ) - const result = await runCopilotLifecycle( - { message: 'hello', messageId: 'stream-1' }, - { - userId: 'user-1', - workspaceId: 'ws-1', - chatId: 'chat-1', - executionId: 'exec-1', - runId: 'run-1', - abortSignal: abortController.signal, - executionContext: { - userId: 'user-1', - workflowId: '', - workspaceId: 'ws-1', - chatId: 'chat-1', - }, - } - ) + const result = await runWithStreamAbort(abortController) expect(result).toEqual( expect.objectContaining({ success: false, cancelled: false, content: 'partial answer', - error: REPLAY_BUDGET_EXHAUSTED_MESSAGE, + error: refusal.userMessage, errorCode: REPLAY_BUDGET_EXHAUSTED_CODE, }) ) } ) + it('keeps a Stop a cancellation when a replay refusal follows it', async () => { + const abortController = new AbortController() + mockRunStreamLoop.mockImplementationOnce( + async (_url: string, _init: RequestInit, context: StreamingContext): Promise => { + abortController.abort('user_stop:abortActiveStream') + context.wasAborted = true + throw ownerRefusal() + } + ) + + const result = await runWithStreamAbort(abortController) + + expect(result).toEqual(expect.objectContaining({ success: false, cancelled: true })) + expect(result.errorCode).toBeUndefined() + }) + it('returns the cancelled result when cancelled completion persistence fails', async () => { const abortController = new AbortController() const onComplete = vi.fn().mockRejectedValue(new Error('db unavailable')) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index 62ef1d5ccca..b4a6b6f9765 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -53,11 +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 { - REPLAY_BUDGET_EXHAUSTED_CODE, - REPLAY_BUDGET_EXHAUSTED_MESSAGE, - StreamReplayBudgetExhaustedError, -} from '@/lib/mothership/request/session/replay-budget' +import { replayRefusal } from '@/lib/mothership/request/session/replay-budget' import { getToolCallTerminalData, requireToolCallStateResult, @@ -541,22 +537,19 @@ export async function runCopilotLifecycle( 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 replayBudgetExhausted = stoppedByReplayBudget(lifecycleOptions.abortSignal) - if (replayBudgetExhausted && !context.errors.includes(REPLAY_BUDGET_EXHAUSTED_MESSAGE)) { - context.errors.push(REPLAY_BUDGET_EXHAUSTED_MESSAGE) - } + const refusal = replayRefusal(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 = - !replayBudgetExhausted && + !refusal && (context.completionStatus === MothershipStreamV1CompletionStatus.cancelled || context.wasAborted || (lifecycleOptions.abortSignal?.aborted ?? false)) const succeeded = - !replayBudgetExhausted && + !refusal && !turnWasAborted && (backendFinishedTurn || (!context.completionStatus && context.errors.length === 0)) @@ -577,9 +570,7 @@ export async function runCopilotLifecycle( toolCalls: buildToolCallSummaries(context), chatId: context.chatId, requestId: context.requestId, - ...(replayBudgetExhausted - ? { error: REPLAY_BUDGET_EXHAUSTED_MESSAGE, errorCode: REPLAY_BUDGET_EXHAUSTED_CODE } - : {}), + ...(refusal ? { error: refusal.userMessage, errorCode: refusal.code } : {}), errors: !succeeded && context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, @@ -617,11 +608,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 replayBudgetExhausted = - error instanceof StreamReplayBudgetExhaustedError || - stoppedByReplayBudget(lifecycleOptions.abortSignal) - const wasCancelled = - !replayBudgetExhausted && (lifecycleOptions.abortSignal?.aborted ?? false) + const refusal = replayRefusal(lifecycleOptions.abortSignal?.reason) + const wasCancelled = !refusal && (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 @@ -636,8 +624,8 @@ export async function runCopilotLifecycle( toolCalls: buildToolCallSummaries(context), chatId: context.chatId, requestId: context.requestId, - error: replayBudgetExhausted ? REPLAY_BUDGET_EXHAUSTED_MESSAGE : err.message, - ...(replayBudgetExhausted ? { errorCode: REPLAY_BUDGET_EXHAUSTED_CODE } : {}), + error: refusal?.userMessage ?? err.message, + ...(refusal ? { errorCode: refusal.code } : {}), errors: context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, @@ -1681,10 +1669,6 @@ async function withEnterpriseByokKey( return byokApiKey ? { ...refreshed, byokApiKey } : refreshed } -function stoppedByReplayBudget(signal: AbortSignal | undefined): boolean { - return signal?.reason instanceof StreamReplayBudgetExhaustedError -} - function isAborted(options: CopilotLifecycleOptions, context: StreamingContext): boolean { return !!(options.abortSignal?.aborted || context.wasAborted) } diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 4e74c81c9b5..7d173f461d2 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -60,9 +60,8 @@ import { StreamControllerSupersededError, } from '@/lib/mothership/request/session/controller-lease' import { - REPLAY_BUDGET_EXHAUSTED_CODE, - REPLAY_BUDGET_EXHAUSTED_MESSAGE, - StreamReplayBudgetExhaustedError, + replayRefusal, + type StreamReplayBudgetExhaustedError, } from '@/lib/mothership/request/session/replay-budget' import { SSE_RESPONSE_HEADERS } from '@/lib/mothership/request/session/sse' import { TraceCollector } from '@/lib/mothership/request/trace' @@ -209,6 +208,14 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS } const collector = new TraceCollector() + let handedOff = false + + /** + * 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. + */ + const refusalOf = (thrown?: unknown) => + replayRefusal(abortController.signal.aborted ? abortController.signal.reason : thrown) /** * A refused replay write ends the turn as an error: every replacement would be @@ -235,8 +242,8 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS ...result, success: false, cancelled: false, - error: REPLAY_BUDGET_EXHAUSTED_MESSAGE, - errorCode: REPLAY_BUDGET_EXHAUSTED_CODE, + error: refusal.userMessage, + errorCode: refusal.code, }, publisher, runId, @@ -408,9 +415,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS */ if (!abortController.signal.aborted) { abortController.abort( - error instanceof StreamReplayBudgetExhaustedError - ? error - : new StreamControllerSupersededError() + replayRefusal(error) ?? new StreamControllerSupersededError() ) } throw error @@ -424,14 +429,11 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS }) lifecycleResult = result - const replayRefusal = - abortController.signal.reason instanceof StreamReplayBudgetExhaustedError - ? abortController.signal.reason - : undefined + const refusal = refusalOf() // A completed result wins a late Stop; passive disconnection never cancels. outcome = result.success ? RequestTraceV1Outcome.success - : replayRefusal + : refusal ? RequestTraceV1Outcome.error : result.cancelled || abortController.signal.aborted ? RequestTraceV1Outcome.cancelled @@ -440,8 +442,8 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS cancelReason = recordCancelled() } await assertControllerOwnership() - if (replayRefusal && !result.success) { - await finalizeAfterReplayRefusal(replayRefusal, result) + if (refusal && !result.success) { + await finalizeAfterReplayRefusal(refusal, result) } else { await finalizeStream(result, publisher, runId, outcome, requestId) } @@ -451,18 +453,14 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS abortController.signal.reason instanceof StreamControllerSupersededError ) { logger.info('Stream controller handed off; leaving its run recoverable', { streamId }) + handedOff = true return } await assertControllerOwnership() - const replayRefusal = - error instanceof StreamReplayBudgetExhaustedError - ? error - : abortController.signal.reason instanceof StreamReplayBudgetExhaustedError - ? abortController.signal.reason - : undefined - if (replayRefusal) { + const refusal = refusalOf(error) + if (refusal) { outcome = RequestTraceV1Outcome.error - await finalizeAfterReplayRefusal(replayRefusal) + await finalizeAfterReplayRefusal(refusal) return } const wasCancelled = abortController.signal.aborted @@ -512,7 +510,8 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS await releasePendingChatStream(chatId, streamId, lease) } processResourcesReleased = true - if (!(abortController.signal.reason instanceof StreamControllerSupersededError)) { + // A superseded controller must not expire or clear its successor's stream. + if (!handedOff) { await scheduleBufferCleanup(streamId) await scheduleFilePreviewSessionCleanup(streamId) await cleanupAbortMarker(streamId) diff --git a/apps/sim/lib/mothership/request/session/omission.ts b/apps/sim/lib/mothership/request/session/omission.ts deleted file mode 100644 index fd798b13e74..00000000000 --- a/apps/sim/lib/mothership/request/session/omission.ts +++ /dev/null @@ -1,50 +0,0 @@ -import { isRecordLike } from '@sim/utils/object' - -/** - * Stand-in for a subtree the bounded chat stream could not carry even after its - * long strings were shortened. The full value still reaches the model and the - * durable tool records; only the live and replayed UI stream carries this. - */ -export interface OmittedStreamValue { - omitted: true - bytes: number -} - -/** One field shortened or omitted by stream compaction, addressed by JSON pointer. */ -export interface StreamTruncatedField { - path: string - /** Full UTF-8 size of the original value. */ - bytes: number - /** UTF-8 size kept in the stream; absent when the value was omitted. */ - previewBytes?: number -} - -/** Payload key of the {@link StreamTruncationMarker} on a compacted event. */ -export const STREAM_TRUNCATION_KEY = 'streamTruncation' - -/** - * Payload-level marker naming every field compaction changed. It sits beside the - * compacted fields rather than inside them, so `arguments` and `output` keep the - * shape the UI reads. - */ -export interface StreamTruncationMarker { - /** UTF-8 size of the payload before compaction. */ - bytes: number - fields: StreamTruncatedField[] -} - -export function isOmittedStreamValue(value: unknown): value is OmittedStreamValue { - return isRecordLike(value) && value.omitted === true && typeof value.bytes === 'number' -} - -/** Formats a byte count the way the UI states an omitted size ("1.6 MB"). */ -export function formatByteSize(bytes: number): string { - if (bytes < 1024) return `${bytes} B` - if (bytes < 1024 * 1024) return `${(bytes / 1024).toFixed(1)} KB` - return `${(bytes / (1024 * 1024)).toFixed(1)} MB` -} - -/** Replaces omitted subtrees with a sentence stating their size, for display. */ -export function describeOmittedStreamValues(_key: string, value: unknown): unknown { - return isOmittedStreamValue(value) ? `Too large to show (${formatByteSize(value.bytes)})` : value -} 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 b07688ff92e..2f2500a4bf4 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -37,6 +37,7 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { server, abortRequests, /** Events the scripted worker streams through the controller's sink, in order. */ + /** Events, or steps to run between them, that the scripted worker streams in order. */ script: [] as unknown[], /** Controller lifecycles started, and what each sink call threw. */ runs: [] as Array<{ dispatched: unknown[]; sinkErrors: unknown[] }>, @@ -57,6 +58,10 @@ vi.mock('@/lib/mothership/request/lifecycle/run', () => ({ const run = { dispatched: [] as unknown[], sinkErrors: [] as unknown[] } worker.runs.push(run) for (const event of worker.script) { + if (typeof event === 'function') { + await event() + continue + } try { await options.onEvent?.(event) } catch (error) { @@ -77,6 +82,7 @@ vi.mock('@/lib/mothership/request/lifecycle/run', () => ({ import { db } from '@sim/db' import { copilotChats, copilotRuns, permissions, user, workspace } from '@sim/db/schema' +import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { eq } from 'drizzle-orm' import { NextRequest } from 'next/server' @@ -94,10 +100,8 @@ import { StreamControllerSupersededError, } from '@/lib/mothership/request/session/controller-lease' import { eventToStreamEvent } from '@/lib/mothership/request/session/event' -import { STREAM_TRUNCATION_KEY } from '@/lib/mothership/request/session/omission' import { REPLAY_BUDGET_EXHAUSTED_CODE, - REPLAY_BUDGET_EXHAUSTED_MESSAGE, StreamReplayBudgetExhaustedError, } from '@/lib/mothership/request/session/replay-budget' import { STREAM_STRING_PREVIEW_UNITS } from '@/lib/mothership/request/session/replay-compaction' @@ -168,6 +172,13 @@ function toolCall( } } +function dataFrames(body: string) { + return body + .split('\n\n') + .filter((frame) => frame.startsWith('data: ')) + .map((frame) => JSON.parse(frame.slice('data: '.length))) +} + function text(value: string): StreamEvent { return { type: 'text', payload: { channel: 'assistant', text: value } } } @@ -207,12 +218,7 @@ describe.runIf(Boolean(redisUrl))('leased Chat stream writer with Redis', () => arguments: { activity: { id: 'activity-1', title: 'Reading logs' }, command: 'logs get --run latest', - stdout: stdout.slice(0, STREAM_STRING_PREVIEW_UNITS), - }, - [STREAM_TRUNCATION_KEY]: { - fields: [ - { path: '/arguments/stdout', bytes: 2 * MB, previewBytes: STREAM_STRING_PREVIEW_UNITS }, - ], + stdout: `${stdout.slice(0, STREAM_STRING_PREVIEW_UNITS)}…[truncated, 2 MB total]`, }, }) expect(toRecord(call.payload).arguments).toMatchObject({ stdout }) @@ -244,7 +250,11 @@ describe.runIf(Boolean(redisUrl))('leased Chat stream writer with Redis', () => expect(toRecord(replayed[0].payload)).toMatchObject({ success: true, status: 'success', - output: { exitCode: 0, resources, stdout: 'r'.repeat(STREAM_STRING_PREVIEW_UNITS) }, + output: { + exitCode: 0, + resources, + stdout: `${'r'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 1.6 MB total]`, + }, }) expect(writer.persistenceStopped).toBe(false) }) @@ -323,8 +333,12 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge const userId = generateId() const workspaceId = generateId() const chatId = generateId() - const streamId = generateId() - const runId = generateId() + const ownerMessage = new StreamReplayBudgetExhaustedError({ + resource: 'owner_redis_bytes', + currentBytes: 0, + limitBytes: 0, + attemptedBytes: 0, + }).userMessage beforeAll(async () => { const now = new Date() @@ -350,6 +364,10 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge permissionType: 'admin', }) await db.insert(copilotChats).values({ id: chatId, userId, workspaceId, type: 'mothership' }) + authMockFns.mockGetSession.mockResolvedValue({ + user: { id: userId }, + session: { id: generateId() }, + }) }) afterAll(async () => { @@ -359,14 +377,14 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge await db.delete(user).where(eq(user.id, userId)) }) - it('ends as an error, marks its run terminal, stops the worker, and starts no recovery controller', async () => { + /** Admits a turn the way the chat POST does, runs its controller, and returns its frames. */ + async function runTurn(script: unknown[], prepare?: (streamId: string) => Promise) { + const streamId = generateId() + const runId = generateId() expect(await acquirePendingChatStream(chatId, streamId, 0)).toBe(true) - const lease = { - key: chatStreamLockKey(chatId), - value: (await redis().get(chatStreamLockKey(chatId)))!, - } + const controllerToken = (await redis().get(chatStreamLockKey(chatId)))! const request = { - message: 'Summarize the latest logs', + message: 'Summarize the logs', userId, messageId: streamId, chatId, @@ -383,7 +401,7 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge streamId, requestContext: { requestId: generateId(), - controllerToken: lease.value, + controllerToken, recovery: { kind: 'interactive_stream', request, @@ -393,15 +411,10 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge }, }) .returning() - const { maxOwnerBytes } = getRedisBudgetLimits('copilot_stream') - const [ownerBudgetKey] = getRedisBudgetKeys({ kind: 'copilot_stream', id: streamId }) - /** Room for the turn's opening session frame, not for the worker's tool call. */ - await redis().set(ownerBudgetKey, String(maxOwnerBytes - 512), 'EX', 3600) - worker.script = [ - toolCall('call-refused', 'cli_blocks_get', { command: `blocks get ${'b'.repeat(1024)}` }), - text('never delivered'), - ] - + await prepare?.(streamId) + worker.runs.length = 0 + worker.abortRequests.length = 0 + worker.script = script const response = createSSEStream({ requestPayload: request, userId, @@ -417,15 +430,23 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge admittedRun: run, orchestrateOptions: { userId, workspaceId, chatId, runId, interactive: true }, }) - const frames = (await new Response(response).text()) - .split('\n\n') - .filter((frame) => frame.startsWith('data: ')) - .map((frame) => JSON.parse(frame.slice('data: '.length))) + return { streamId, runId, frames: dataFrames(await new Response(response).text()) } + } + + it('ends as an error, marks its run terminal, stops the worker, and starts no recovery controller', async () => { + const { streamId, runId, frames } = await runTurn( + [ + toolCall('call-refused', 'cli_blocks_get', { command: `blocks get ${'b'.repeat(1024)}` }), + text('never delivered'), + ], + async (streamId) => { + const { maxOwnerBytes } = getRedisBudgetLimits('copilot_stream') + const [ownerBudgetKey] = getRedisBudgetKeys({ kind: 'copilot_stream', id: streamId }) + /** Room for the turn's opening session frame, not for the worker's tool call. */ + await redis().set(ownerBudgetKey, String(maxOwnerBytes - 512), 'EX', 3600) + } + ) - const [controllerRun] = worker.runs - expect(controllerRun.dispatched).toEqual([]) - expect(controllerRun.sinkErrors[0]).toBeInstanceOf(StreamReplayBudgetExhaustedError) - expect(controllerRun.sinkErrors[0]).not.toBeInstanceOf(StreamControllerSupersededError) expect(frames.map((frame) => [frame.type, frame.payload.code ?? frame.payload.status])).toEqual( [ ['session', undefined], @@ -433,27 +454,19 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge ['complete', 'error'], ] ) - expect(frames[1].payload.message).toBe(REPLAY_BUDGET_EXHAUSTED_MESSAGE) - + expect(frames[1].payload.message).toBe(ownerMessage) const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) - expect(stored).toMatchObject({ status: 'error', error: REPLAY_BUDGET_EXHAUSTED_MESSAGE }) + expect(stored).toMatchObject({ status: 'error', error: ownerMessage }) expect(worker.abortRequests).toEqual([expect.objectContaining({ messageId: streamId })]) expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() - authMockFns.mockGetSession.mockResolvedValue({ - user: { id: userId }, - session: { id: generateId() }, - }) const reconnect = await streamGET( new NextRequest(`http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=0`), { params: Promise.resolve({}) } ) - const replayed = (await reconnect.text()) - .split('\n\n') - .filter((frame) => frame.startsWith('data: ')) - .map((frame) => JSON.parse(frame.slice('data: '.length))) + const replayed = dataFrames(await reconnect.text()) expect(replayed.map((frame) => frame.type)).toEqual(['session', 'error', 'complete']) - expect(replayed[1].payload.message).toBe(REPLAY_BUDGET_EXHAUSTED_MESSAGE) + expect(replayed[1].payload.message).toBe(ownerMessage) expect(replayed[2].payload.status).toBe('error') const recovered = await readChatStream.execute({ @@ -463,5 +476,32 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge expect(recovered?.status).toBe('error') expect(worker.runs).toHaveLength(1) expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + + const [controllerRun] = worker.runs + expect(controllerRun.dispatched).toEqual([]) + expect(controllerRun.sinkErrors[0]).toBeInstanceOf(StreamReplayBudgetExhaustedError) + }) + + it("leaves its successor's stream untouched when a superseded controller is refused an oversized frame", async () => { + const successorToken = `successor\n${generateId()}` + const { streamId, runId, frames } = await runTurn([ + async () => { + await redis().set(chatStreamLockKey(chatId), successorToken, 'EX', 60) + }, + toolCall('call-oversized', 'run_workflow', { + workflowId: generateId(), + input: 'w'.repeat(1.5 * MB), + }), + ]) + + expect(frames.map((frame) => frame.type)).toEqual(['session']) + /** The client stream closes before the controller's teardown; let teardown finish. */ + await sleep(500) + expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeGreaterThan(300) + expect(await redis().get(chatStreamLockKey(chatId))).toBe(successorToken) + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('active') + expect(worker.abortRequests).toEqual([]) + await redis().del(chatStreamLockKey(chatId)) }) }) diff --git a/apps/sim/lib/mothership/request/session/replay-budget.ts b/apps/sim/lib/mothership/request/session/replay-budget.ts index 9c66d26e756..c64a6203abf 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.ts @@ -3,8 +3,11 @@ import type { RedisBudgetRefusal } from '@/lib/core/redis/byte-budget.server' /** Run-error code for a turn stopped because its replay buffer refused a write. */ export const REPLAY_BUDGET_EXHAUSTED_CODE = 'replay_budget_exhausted' -export const REPLAY_BUDGET_EXHAUSTED_MESSAGE = - 'This turn produced more output than a single response can stream, so it was stopped. The work it already completed has been saved — send a message to continue from there.' +const STREAM_LIMIT_MESSAGE = + 'This response produced more output than a single response can stream, so it was stopped. The work it already completed has been saved — send a message to continue from there.' + +const HOURLY_LIMIT_MESSAGE = + 'Your recent responses streamed more output than the hourly limit allows, so this one was stopped. The work it already completed has been saved — you can continue once the limit resets within the hour.' /** * The replay buffer refused an event a leased controller had to persist before @@ -18,4 +21,16 @@ export class StreamReplayBudgetExhaustedError extends Error { super('Stream replay byte budget exhausted') this.name = 'StreamReplayBudgetExhaustedError' } + + /** What the user is told; the per-user ceiling is a fixed hourly window. */ + get userMessage(): string { + return this.refusal.resource === 'user_redis_bytes' + ? HOURLY_LIMIT_MESSAGE + : 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/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index 414377a8543..425b914c5c0 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -2,27 +2,18 @@ * Failure modes of stream compaction, each a way a compacted replay frame would * break the UI, restore, or the byte budget: * - a bulky string survives, so the frame stays unpersistable; - * - a structured field the UI reads (exit code, resources, citations, cancel - * reason, activity) is truncated or dropped; - * - a whole `output` or `arguments` object is replaced, changing the shape the - * UI reads; + * - a field the UI reads (exit code, resources, citations, cancel reason, + * activity, identity) is cut or dropped, or an object changes shape; * - arguments the browser executes from, a file preview, or a one-time API key * are altered; - * - many medium strings keep the frame over budget; - * - split assistant text no longer concatenates to the original receipt; - * - the caller's event is mutated, so dispatch sees the compacted copy. + * - the caller's event is mutated, so dispatch sees the compacted copy; + * - a cut splits a surrogate pair into invalid text. */ import { toRecord } from '@sim/utils/object' import { describe, expect, it } from 'vitest' import { - isOmittedStreamValue, - STREAM_TRUNCATION_KEY, - type StreamTruncationMarker, -} from '@/lib/mothership/request/session/omission' -import { - compactStreamEventForReplay, + compactStreamEvent, STREAM_EVENT_COMPACTION_THRESHOLD_BYTES, - STREAM_SHORT_STRING_PREVIEW_UNITS, STREAM_STRING_PREVIEW_UNITS, } from '@/lib/mothership/request/session/replay-compaction' import type { StreamEvent } from '@/lib/mothership/request/session/types' @@ -33,18 +24,10 @@ function payloadOf(event: StreamEvent): Record { return toRecord(event.payload) } -function payloadBytes(event: StreamEvent): number { - return Buffer.byteLength(JSON.stringify(event.payload)) -} - -function marker(event: StreamEvent): StreamTruncationMarker { - return payloadOf(event)[STREAM_TRUNCATION_KEY] as StreamTruncationMarker -} - -describe('compactStreamEventForReplay', () => { - it('truncates only the bulky strings of a Sim tool result and keeps its structured fields', () => { +describe('compactStreamEvent', () => { + it('cuts only the bulky strings of a tool result and keeps every other field and shape', () => { const citations = [{ index: 1, title: 'Runbook', url: 'https://docs.example/runbook' }] - const observations = [{ kind: 'image', mediaType: 'image/png', name: 'chart.png' }] + const resources = [{ type: 'file', id: 'file-1', title: 'out.csv' }] const event: StreamEvent = { type: 'tool', payload: { @@ -59,29 +42,25 @@ describe('compactStreamEventForReplay', () => { stdout: 'o'.repeat(2 * MB), stderr: 'e'.repeat(300 * 1024), exitCode: 0, - resources: [{ type: 'file', id: 'file-1', title: 'out.csv' }], - sinkError: 'sink unavailable', - observations, + resources, citations, + reason: 'r'.repeat(20_000), }, }, } - const [compacted] = compactStreamEventForReplay(event) - const output = payloadOf(compacted).output as Record - - expect(payloadBytes(compacted)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) - expect(output.stdout).toBe('o'.repeat(STREAM_STRING_PREVIEW_UNITS)) - expect(output.stderr).toBe('e'.repeat(STREAM_STRING_PREVIEW_UNITS)) - expect(output.exitCode).toBe(0) - expect(output.resources).toEqual([{ type: 'file', id: 'file-1', title: 'out.csv' }]) - expect(output.sinkError).toBe('sink unavailable') - expect(output.observations).toEqual(observations) - expect(output.citations).toEqual(citations) - expect(marker(compacted).fields).toEqual([ - { path: '/output/stdout', bytes: 2 * MB, previewBytes: STREAM_STRING_PREVIEW_UNITS }, - { path: '/output/stderr', bytes: 300 * 1024, previewBytes: STREAM_STRING_PREVIEW_UNITS }, - ]) + const compacted = compactStreamEvent(event) + const output = toRecord(payloadOf(compacted).output) + + expect(Buffer.byteLength(JSON.stringify(compacted.payload))).toBeLessThan( + STREAM_EVENT_COMPACTION_THRESHOLD_BYTES + ) + expect(output.stdout).toBe(`${'o'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 2 MB total]`) + expect(output.stderr).toBe( + `${'e'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 300 KB total]` + ) + expect(output).toMatchObject({ exitCode: 0, resources, citations, reason: 'r'.repeat(20_000) }) + expect(Object.keys(output)).toEqual(Object.keys(toRecord(payloadOf(event).output))) expect(payloadOf(compacted)).toMatchObject({ toolCallId: 'call-1', toolName: 'run_code', @@ -91,112 +70,52 @@ describe('compactStreamEventForReplay', () => { }) }) - it('keeps retrieval result metadata and cuts each result to a short preview when many medium strings remain', () => { - const results = Array.from({ length: 80 }, (_, index) => ({ - id: `doc-${index}`, - title: `Document ${index}`, - url: `https://docs.example/${index}`, - score: 0.5, - content: 'c'.repeat(6 * 1024), - })) - const event: StreamEvent = { + it.each([ + ['a workflow run', 'run_workflow', { workflowId: 'wf-1' }], + ['a user-local VFS read', 'read', { path: 'user-local/notes.md' }], + ['a terminal command', 'terminal', { operation: 'run' }], + ['a browser action', 'browser_click', { elementId: 'e-1' }], + ])('leaves the arguments of %s whole', (_label, toolName, identity) => { + const args = { ...identity, input: 'b'.repeat(400 * 1024) } + const event = { type: 'tool', payload: { - toolCallId: 'call-2', - toolName: 'search_documentation', + toolCallId: 'c', + toolName, executor: 'go', - mode: 'sync', - phase: 'result', - success: true, - output: { data: { results } }, - }, - } - - const [compacted] = compactStreamEventForReplay(event) - const kept = ( - (payloadOf(compacted).output as Record).data as { - results: Array> - } - ).results - - expect(payloadBytes(compacted)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) - expect(kept).toHaveLength(80) - expect(kept[79]).toEqual({ - id: 'doc-79', - title: 'Document 79', - url: 'https://docs.example/79', - score: 0.5, - content: 'c'.repeat(STREAM_SHORT_STRING_PREVIEW_UNITS), - }) - }) - - it('omits a subtree only as a last resort, keeping the arguments object and the keys the UI reads', () => { - const rows = Array.from({ length: 40_000 }, (_, index) => ({ id: index, name: 'row' })) - const event: StreamEvent = { - type: 'tool', - payload: { - toolCallId: 'call-3', - toolName: 'cli_tables_rows_query', - executor: 'sim', mode: 'async', phase: 'call', - arguments: { - activity: { id: 'a-1', title: 'Querying rows' }, - operation: 'query', - description: 'd'.repeat(10_000), - format: 'json', - rows, - }, + arguments: args, }, - } - - const [compacted] = compactStreamEventForReplay(event) - const args = payloadOf(compacted).arguments as Record + } as StreamEvent - expect(payloadBytes(compacted)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) - expect(args.activity).toEqual({ id: 'a-1', title: 'Querying rows' }) - expect(args.operation).toBe('query') - expect(args.description).toBe('d'.repeat(10_000)) - expect(args.format).toBe('json') - expect(isOmittedStreamValue(args.rows)).toBe(true) - expect(marker(compacted).fields).toEqual([ - { path: '/arguments/rows', bytes: Buffer.byteLength(JSON.stringify(rows)) }, - ]) + expect(compactStreamEvent(event)).toBe(event) }) - it.each([ - ['a client-executable workflow tool', 'run_workflow', 'sim'], - ['a client-executed tool', 'custom_client_tool', 'client'], - ['a user-local VFS read', 'read', 'go'], - ['a terminal command', 'terminal', 'go'], - ])('leaves the arguments of %s whole', (_label, toolName, executor) => { - const args = { - path: 'user-local/notes.md', - operation: 'run', - input: { body: 'b'.repeat(400 * 1024) }, - } - const event = { + it('cuts the arguments of a call the browser does not execute', () => { + const event: StreamEvent = { type: 'tool', payload: { - toolCallId: 'call-4', - toolName, - executor, + toolCallId: 'c', + toolName: 'sim_cli', + executor: 'go', mode: 'async', phase: 'call', - arguments: args, + arguments: { activity: { id: 'a', title: 'Reading logs' }, stdin: 'x'.repeat(MB) }, }, - } as StreamEvent + } - const [compacted] = compactStreamEventForReplay(event) + const args = toRecord(payloadOf(compactStreamEvent(event)).arguments) - expect(JSON.stringify(payloadOf(compacted).arguments)).toBe(JSON.stringify(args)) + expect(args.activity).toEqual({ id: 'a', title: 'Reading logs' }) + expect(args.stdin).toBe(`${'x'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 1 MB total]`) }) it('never compacts a file preview or a generated API key', () => { const preview: StreamEvent = { type: 'tool', payload: { - toolCallId: 'call-5', + toolCallId: 'c', toolName: 'prepare_file_edit', previewPhase: 'file_preview_content', content: 'p'.repeat(400 * 1024), @@ -208,7 +127,7 @@ describe('compactStreamEventForReplay', () => { const apiKey: StreamEvent = { type: 'tool', payload: { - toolCallId: 'call-6', + toolCallId: 'c', toolName: 'generate_api_key', executor: 'sim', mode: 'async', @@ -218,46 +137,45 @@ describe('compactStreamEventForReplay', () => { }, } - expect(compactStreamEventForReplay(preview)).toEqual([preview]) - expect(compactStreamEventForReplay(apiKey)).toEqual([apiKey]) + expect(compactStreamEvent(preview)).toBe(preview) + expect(compactStreamEvent(apiKey)).toBe(apiKey) }) - it('splits oversized assistant text into contiguous events that reassemble exactly', () => { - const text = `${'a'.repeat(300 * 1024)}😀${'b'.repeat(300 * 1024)}` + it('compacts a copy and leaves the caller’s event whole for dispatch', () => { + const stdout = 's'.repeat(MB) const event: StreamEvent = { - type: 'text', - payload: { channel: 'assistant', text, textOffset: 17 }, + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_logs_get', + executor: 'sim', + mode: 'async', + phase: 'call', + arguments: { stdout }, + }, } - const pieces = compactStreamEventForReplay(event) - - expect(pieces.length).toBeGreaterThan(1) - expect(pieces.map((piece) => payloadOf(piece).text).join('')).toBe(text) - let offset = 17 - for (const piece of pieces) { - expect(payloadOf(piece).textOffset).toBe(offset) - expect(payloadBytes(piece)).toBeLessThanOrEqual(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) - offset += (payloadOf(piece).text as string).length - } + expect(compactStreamEvent(event)).not.toBe(event) + expect(toRecord(payloadOf(event).arguments).stdout).toBe(stdout) }) - it('compacts a copy and leaves the caller’s event whole for dispatch', () => { - const stdout = 's'.repeat(MB) + it('does not split a surrogate pair at the cut', () => { + const text = `${'a'.repeat(STREAM_STRING_PREVIEW_UNITS - 1)}😀${'b'.repeat(MB)}` const event: StreamEvent = { type: 'tool', payload: { - toolCallId: 'call-7', + toolCallId: 'c', toolName: 'cli_logs_get', executor: 'sim', mode: 'async', - phase: 'call', - arguments: { stdout }, + phase: 'result', + success: true, + output: { text }, }, } - compactStreamEventForReplay(event) + const cut = toRecord(payloadOf(compactStreamEvent(event)).output).text as string - expect((payloadOf(event).arguments as Record).stdout).toBe(stdout) - expect(STREAM_TRUNCATION_KEY in payloadOf(event)).toBe(false) + expect(cut.startsWith(`${'a'.repeat(STREAM_STRING_PREVIEW_UNITS - 1)}…`)).toBe(true) }) }) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index 862dfaff5cf..ce265209271 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -1,332 +1,119 @@ -import { isCurrentBrowserToolName } from '@sim/browser-protocol' -import { isTerminalToolName } from '@sim/terminal-protocol' -import { isRecordLike } from '@sim/utils/object' -import { - MothershipStreamV1EventType, - MothershipStreamV1ToolExecutor, - MothershipStreamV1ToolPhase, -} from '@/lib/mothership/generated/mothership-stream-v1' -import { GenerateApiKey, TOOL_CATALOG } from '@/lib/mothership/generated/tool-catalog-v1' -import { isNativeFileTool, isUserLocalVfsToolCall } from '@/lib/mothership/tools/local-filesystem' -import { - type OmittedStreamValue, - STREAM_TRUNCATION_KEY, - type StreamTruncatedField, - type StreamTruncationMarker, -} from './omission' +import { isRecordLike, toRecordOrNull } from '@sim/utils/object' +import { MothershipStreamV1ToolPhase } from '@/lib/mothership/generated/mothership-stream-v1' +import { isClientExecutedToolCall } from '@/lib/mothership/tools/client-executed-tools' +import { formatFileSize } from '@/lib/uploads/utils/file-utils' import type { StreamEvent } from './types' /** - * An event whose payload serializes larger than this is compacted before it is - * persisted and delivered. Well under the replay buffer's 1 MiB single-write - * ceiling, so a compacted event is persistable. + * Payloads whose strings could serialize past this are compacted before they are + * persisted and delivered; well under the replay buffer's 1 MiB write ceiling. */ export const STREAM_EVENT_COMPACTION_THRESHOLD_BYTES = 256 * 1024 -/** - * Compaction stops once the payload fits this, leaving room in the per-stream - * budget for a turn of many large tool results. - */ -export const STREAM_EVENT_COMPACTION_TARGET_BYTES = 128 * 1024 - -/** UTF-16 units kept at the head of a long string leaf. */ +/** UTF-16 units kept at the head of a long string. */ export const STREAM_STRING_PREVIEW_UNITS = 8 * 1024 -/** - * Second pass for payloads made large by many medium strings (retrieval - * results): each is cut to this many units, keeping its sibling metadata. - */ -export const STREAM_SHORT_STRING_PREVIEW_UNITS = 500 - -/** - * UTF-16 units of `args_delta` forwarded per tool call. The deltas only feed - * a progressive title and a live preview of the arguments; the call frame and - * the result supersede them, so the rest is not forwarded. - */ -export const TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS = 64 * 1024 - -/** - * UTF-16 units per piece of a split assistant text. JSON escapes a unit to at - * most six bytes, so a piece always serializes under the threshold. - */ -const SPLIT_TEXT_MAX_UNITS = Math.floor(STREAM_EVENT_COMPACTION_THRESHOLD_BYTES / 8) - -/** Payload fields that restore, dedupe, and the UI key on; never compacted. */ -const IDENTITY_PAYLOAD_KEYS: ReadonlySet = new Set([ +/** Identity fields and the keys the UI reads from arguments and outputs; kept whole at any depth. */ +const PRESERVED_KEYS: ReadonlySet = new Set([ + 'activity', 'activityDescription', - 'activityReceipt', 'agent', + 'cancelledByUser', 'channel', + 'description', + 'elementId', 'error', 'event', 'execName', 'executor', + 'fileName', 'kind', 'mode', 'op', + 'operation', 'partial', + 'path', + 'pattern', 'phase', + 'reason', 'replay', 'resource', 'status', 'success', 'targetWorkspaceId', - 'textLength', - 'textOffset', - 'toolCallId', - 'toolName', - 'ui', - 'workspaceId', -]) - -/** - * Keys the UI reads from arguments and outputs (activity labels, targets, - * cancellation). Kept whole at any depth. - */ -const PROTECTED_NESTED_KEYS: ReadonlySet = new Set([ - 'activity', - 'cancelledByUser', - 'description', - 'elementId', - 'fileName', - 'operation', - 'path', - 'pattern', - 'reason', 'terminalId', 'timeoutMs', 'title', + 'toolCallId', 'toolId', + 'toolName', + 'ui', 'url', 'workflowId', + 'workspaceId', ]) -type JsonRecord = Record -type JsonContainer = JsonRecord | unknown[] - -function utf8Bytes(value: string): number { - return Buffer.byteLength(value, 'utf8') -} - -function jsonBytes(value: unknown): number { - return utf8Bytes(JSON.stringify(value) ?? '') -} - -/** Cuts at most `maxUnits` UTF-16 units without splitting a surrogate pair. */ -function headUnits(text: string, maxUnits: number): string { - if (text.length <= maxUnits) return text - const code = text.charCodeAt(maxUnits - 1) - return text.slice(0, code >= 0xd800 && code <= 0xdbff ? maxUnits - 1 : maxUnits) -} - -function splitText(text: string): string[] { - const pieces: string[] = [] - let rest = text - while (rest.length > 0) { - const piece = headUnits(rest, SPLIT_TEXT_MAX_UNITS) - pieces.push(piece) - rest = rest.slice(piece.length) +/** Its live result is the only place the plaintext key reaches the browser. */ +const GENERATE_API_KEY_TOOL = 'generate_api_key' + +function stringUnits(value: unknown): number { + if (typeof value === 'string') return value.length + if (Array.isArray(value)) return value.reduce((sum, item) => sum + stringUnits(item), 0) + if (!isRecordLike(value)) return 0 + let sum = 0 + for (const key in value) sum += key.length + stringUnits(value[key]) + return sum +} + +/** Cuts long strings to a head with a readable size note, copying only what changes. */ +function truncateStrings(value: unknown, skipKeys: ReadonlySet): unknown { + if (typeof value === 'string') { + if (value.length <= STREAM_STRING_PREVIEW_UNITS) return value + const end = STREAM_STRING_PREVIEW_UNITS + const lastUnit = value.charCodeAt(end - 1) + const cut = lastUnit >= 0xd800 && lastUnit <= 0xdbff ? end - 1 : end + const size = formatFileSize(Buffer.byteLength(value, 'utf8')) + return `${value.slice(0, cut)}…[truncated, ${size} total]` } - return pieces -} - -function pointer(base: string, key: string | number): string { - return `${base}/${String(key).replaceAll('~', '~0').replaceAll('/', '~1')}` -} - -function entriesOf(container: JsonContainer): Array<[string | number, unknown]> { - return Array.isArray(container) - ? container.map((value, index) => [index, value]) - : Object.entries(container) -} - -function setEntry(container: JsonContainer, key: string | number, value: unknown): void { - if (Array.isArray(container)) container[Number(key)] = value - else container[String(key)] = value -} - -function isProtectedKey(key: string | number): boolean { - return typeof key === 'string' && PROTECTED_NESTED_KEYS.has(key) -} - -function containsProtectedKey(value: unknown): boolean { - if (Array.isArray(value)) return value.some(containsProtectedKey) - if (!isRecordLike(value)) return false - return Object.entries(value).some( - ([key, field]) => PROTECTED_NESTED_KEYS.has(key) || containsProtectedKey(field) - ) -} - -interface CompactionState { - payload: JsonRecord - roots: ReadonlySet - fields: Map -} - -/** Shortens every string leaf longer than `maxUnits` to its head, in place. */ -function shortenStrings( - state: CompactionState, - container: JsonContainer, - base: string, - maxUnits: number -): void { - for (const [key, value] of entriesOf(container)) { - if (base === '' ? !state.roots.has(String(key)) : isProtectedKey(key)) continue - const path = pointer(base, key) - if (typeof value === 'string') { - if (value.length <= maxUnits) continue - const head = headUnits(value, maxUnits) - const bytes = state.fields.get(path)?.bytes ?? utf8Bytes(value) - state.fields.set(path, { path, bytes, previewBytes: utf8Bytes(head) }) - setEntry(container, key, head) - } else if (Array.isArray(value) || isRecordLike(value)) { - shortenStrings(state, value, path, maxUnits) - } - } -} - -interface Candidate { - container: JsonContainer - key: string | number - path: string - bytes: number -} - -function collectCandidates( - state: CompactionState, - container: JsonContainer, - base: string, - out: Candidate[] -): void { - for (const [key, value] of entriesOf(container)) { - if (base === '' ? !state.roots.has(String(key)) : isProtectedKey(key)) continue - const path = pointer(base, key) - if (Array.isArray(value) || isRecordLike(value)) { - if (!containsProtectedKey(value)) { - out.push({ container, key, path, bytes: jsonBytes(value) }) - } - collectCandidates(state, value, path, out) - } else if (typeof value === 'string') { - out.push({ container, key, path, bytes: jsonBytes(value) }) - } + if (Array.isArray(value)) { + let copy: unknown[] | undefined + value.forEach((item, index) => { + const next = truncateStrings(item, PRESERVED_KEYS) + if (next !== item) (copy ??= [...value])[index] = next + }) + return copy ?? value } -} - -/** Last resort: replaces the largest unprotected subtrees with stubs until the payload fits. */ -function omitLargestSubtrees(state: CompactionState): void { - const candidates: Candidate[] = [] - collectCandidates(state, state.payload, '', candidates) - candidates.sort((a, b) => b.bytes - a.bytes) - const omitted: string[] = [] - for (const candidate of candidates) { - if (jsonBytes(state.payload) <= STREAM_EVENT_COMPACTION_TARGET_BYTES) return - if (omitted.some((path) => candidate.path.startsWith(`${path}/`))) continue - const stub: OmittedStreamValue = { - omitted: true, - bytes: state.fields.get(candidate.path)?.bytes ?? candidate.bytes, - } - setEntry(candidate.container, candidate.key, stub) - for (const path of state.fields.keys()) { - if (path.startsWith(`${candidate.path}/`)) state.fields.delete(path) - } - state.fields.set(candidate.path, { path: candidate.path, bytes: stub.bytes }) - omitted.push(candidate.path) + if (!isRecordLike(value)) return value + let copy: Record | undefined + for (const [key, field] of Object.entries(value)) { + if (skipKeys.has(key)) continue + const next = truncateStrings(field, PRESERVED_KEYS) + if (next !== field) (copy ??= { ...value })[key] = next } + return copy ?? value } /** - * The browser may start these tools from the call frame's arguments, so the - * arguments are never compacted; a frame too large with them whole is refused. + * Bounds an outgoing stream event so the replay buffer can persist it. Applied + * only to the copy the writer delivers and persists; the caller keeps the full + * event for dispatch. Long strings are cut to their head in place, so every + * object keeps its shape. File previews, `generate_api_key` results, preserved + * keys, and the arguments of calls the browser executes are never cut; an event + * still too large is refused by the buffer, which ends the turn with an error. */ -function isClientExecutableCall(payload: JsonRecord): boolean { +export function compactStreamEvent(event: StreamEvent): StreamEvent { + const payload = toRecordOrNull(event.payload) + if (!payload || 'previewPhase' in payload) return event + if (stringUnits(payload) * 3 <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) return event const toolName = typeof payload.toolName === 'string' ? payload.toolName : '' - const args = isRecordLike(payload.arguments) ? payload.arguments : undefined - const catalogEntry = Object.hasOwn(TOOL_CATALOG, toolName) ? TOOL_CATALOG[toolName] : undefined - return ( - payload.executor === MothershipStreamV1ToolExecutor.client || - (isRecordLike(payload.ui) && payload.ui.clientExecutable === true) || - catalogEntry?.route === 'client' || - catalogEntry?.clientExecutable === true || - isNativeFileTool(toolName) || - isUserLocalVfsToolCall(toolName, args) || - isCurrentBrowserToolName(toolName) || - isTerminalToolName(toolName) - ) -} - -function isNeverCompacted(payload: JsonRecord): boolean { - return ( - 'previewPhase' in payload || - (payload.phase === MothershipStreamV1ToolPhase.result && payload.toolName === GenerateApiKey.id) - ) -} - -function withPayload(event: StreamEvent, payload: JsonRecord): StreamEvent { - return { ...event, payload } as StreamEvent -} - -/** - * Bounds one outgoing stream event so the replay buffer can persist it. Applied - * only by the stream writer, to the copy it delivers and persists; the caller - * keeps the full event for dispatch, the async tool row, and approval. - * - * Returns the event unchanged when its payload is within - * {@link STREAM_EVENT_COMPACTION_THRESHOLD_BYTES}. Otherwise: - * - assistant text splits into contiguous events whose concatenation is exact; - * - long string leaves are cut to their head in place, so `arguments` and - * `output` keep their shape; a second pass cuts medium strings shorter; - * - only if the payload is still too large are the largest unprotected subtrees - * replaced by an {@link OmittedStreamValue}. - * Every changed field is listed in a payload-level {@link StreamTruncationMarker}. - * - * File previews, `generate_api_key` results, identity fields, and the arguments - * of client-executable calls are never compacted; such an event that stays too - * large is refused by the buffer and ends the turn with an error. - */ -export function compactStreamEventForReplay(event: StreamEvent): StreamEvent[] { - const original = event.payload - if (!isRecordLike(original)) return [event] - const originalBytes = jsonBytes(original) - if (originalBytes <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES || isNeverCompacted(original)) { - return [event] - } - - if (event.type === MothershipStreamV1EventType.text) { - const { text, textOffset } = event.payload - let offset = textOffset - return splitText(text).map((piece) => { - const split = withPayload(event, { - ...original, - text: piece, - ...(offset !== undefined ? { textOffset: offset } : {}), - }) - if (offset !== undefined) offset += piece.length - return split - }) + if (payload.phase === MothershipStreamV1ToolPhase.result && toolName === GENERATE_API_KEY_TOOL) { + return event } - - const payload: JsonRecord = structuredClone(original) - const protectsArguments = - payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutableCall(payload) - const roots = new Set( - Object.keys(payload).filter( - (key) => !IDENTITY_PAYLOAD_KEYS.has(key) && !(protectsArguments && key === 'arguments') - ) - ) - const state: CompactionState = { payload, roots, fields: new Map() } - - shortenStrings(state, payload, '', STREAM_STRING_PREVIEW_UNITS) - if (jsonBytes(payload) > STREAM_EVENT_COMPACTION_TARGET_BYTES) { - shortenStrings(state, payload, '', STREAM_SHORT_STRING_PREVIEW_UNITS) - } - if (jsonBytes(payload) > STREAM_EVENT_COMPACTION_TARGET_BYTES) { - omitLargestSubtrees(state) - } - if (state.fields.size === 0) return [event] - - const marker: StreamTruncationMarker = { - bytes: originalBytes, - fields: [...state.fields.values()], - } - return [withPayload(event, { ...payload, [STREAM_TRUNCATION_KEY]: marker })] + const args = isRecordLike(payload.arguments) ? payload.arguments : undefined + const skipKeys = + payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutedToolCall(toolName, args) + ? new Set([...PRESERVED_KEYS, 'arguments']) + : PRESERVED_KEYS + const compacted = truncateStrings(payload, skipKeys) + return compacted === payload ? event : ({ ...event, payload: compacted } as StreamEvent) } diff --git a/apps/sim/lib/mothership/request/session/writer.test.ts b/apps/sim/lib/mothership/request/session/writer.test.ts index 935aa4b33bc..6da1ac447c8 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 { TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS } from '@/lib/mothership/request/session/replay-compaction' import { StreamWriter } from '@/lib/mothership/request/session/writer' function decodeChunk(value: Uint8Array): string { @@ -129,45 +128,6 @@ describe('StreamWriter', () => { expect(writer.sawComplete).toBe(true) }) - it("forwards only the head of a tool call's argument deltas", async () => { - const writer = new StreamWriter({ streamId: 'stream-1', requestId: 'req-1' }) - const delivered: string[] = [] - writer.attach({ - enqueue: (value: Uint8Array) => delivered.push(decodeChunk(value)), - close: () => {}, - } as unknown as ReadableStreamDefaultController) - const delta = (toolCallId: string, argumentsDelta: string): StreamEvent => ({ - type: 'tool', - payload: { - toolCallId, - toolName: 'sim_cli', - executor: 'go', - mode: 'async', - phase: 'args_delta', - argumentsDelta, - }, - }) - const streamed = `{"argv":["logs","get"],"stdin":"${'x'.repeat(100_000)}"}` - - for (let index = 0; index < streamed.length; index += 30_000) { - await writer.publish(delta('call-1', streamed.slice(index, index + 30_000))) - } - await writer.publish(delta('call-2', '{"argv":["blocks"]}')) - await writer.close() - - const forwarded = delivered.map((frame) => JSON.parse(frame.replace(/^data: /, '')).payload) - expect( - forwarded - .filter((payload) => payload.toolCallId === 'call-1') - .map((payload) => payload.argumentsDelta) - .join('') - ).toBe(streamed.slice(0, TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS)) - expect(forwarded.at(-1)).toMatchObject({ - toolCallId: 'call-2', - argumentsDelta: '{"argv":["blocks"]}', - }) - }) - it('enqueues before persistence completes and flushes pending writes on close', async () => { let releasePersist!: () => void appendEvents.mockImplementation( diff --git a/apps/sim/lib/mothership/request/session/writer.ts b/apps/sim/lib/mothership/request/session/writer.ts index ab422fca718..a06fac713e9 100644 --- a/apps/sim/lib/mothership/request/session/writer.ts +++ b/apps/sim/lib/mothership/request/session/writer.ts @@ -1,19 +1,13 @@ import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { encodeSSEComment } from '@/lib/core/utils/sse' -import { - MothershipStreamV1EventType, - MothershipStreamV1ToolPhase, -} from '@/lib/mothership/generated/mothership-stream-v1' +import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { appendEvents } from './buffer' import type { PersistedStreamEventEnvelope } from './contract' import type { ChatStreamLease } from './controller-lease' import { createEvent } from './event' import { StreamReplayBudgetExhaustedError } from './replay-budget' -import { - compactStreamEventForReplay, - TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS, -} from './replay-compaction' +import { compactStreamEvent } from './replay-compaction' import { encodeSSEEnvelope } from './sse' import type { StreamEvent } from './types' @@ -69,7 +63,6 @@ export class StreamWriter { private persistenceTail: Promise = Promise.resolve() private lastPersistenceError: Error | null = null private replayBudgetError: StreamReplayBudgetExhaustedError | null = null - private readonly forwardedArgsDeltaUnits = new Map() private readonly lease?: ChatStreamLease constructor(options: StreamWriterOptions) { @@ -103,7 +96,7 @@ export class StreamWriter { * still reaches the client); unleased delivery continues. */ get persistenceStopped(): boolean { - return this._persistenceStopped + return this._persistenceStopped || this.replayBudgetError !== null } updateRequestId(id: string): void { @@ -138,10 +131,9 @@ export class StreamWriter { } /** - * Delivers an event and records it for replay. Oversized events are compacted - * first ({@link compactStreamEventForReplay}); the compacted copy is what both - * the client and the replay buffer receive, while the caller keeps the full - * event for dispatch. + * Delivers an event and records it for replay. Oversized strings are cut first + * ({@link compactStreamEvent}); the client and the replay buffer receive the + * same 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 @@ -150,33 +142,30 @@ export class StreamWriter { * a reconnect replays it from there. */ publish(event: StreamEvent): void | Promise { - const forwarded = this.limitArgsDelta(event) - if (!forwarded) return - const envelopes = compactStreamEventForReplay(forwarded).map((compacted) => - this.createEnvelope(compacted) - ) + const envelope = this.createEnvelope(compactStreamEvent(event)) if (this.lease) { const lease = this.lease // 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) { - if (!isTurnTerminalEvent(event)) throw this.replayBudgetError - this.deliver(envelopes, event) - return + if (!this.replayBudgetError) { + const result = await appendEvents( + [envelope], + { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, + lease + ) + if (!result.persisted) { + this.replayBudgetError = new StreamReplayBudgetExhaustedError(result.refusal) + } } - const result = await appendEvents( - envelopes, - { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, - lease - ) - if (!result.persisted) { - this._persistenceStopped = true - this.replayBudgetError = new StreamReplayBudgetExhaustedError(result.refusal) - if (!isTurnTerminalEvent(event)) throw this.replayBudgetError - } - this.deliver(envelopes, event) + if (this.replayBudgetError && !isTurnTerminalEvent(event)) throw this.replayBudgetError + this.enqueue(envelope) + if (event.type === MothershipStreamV1EventType.complete) this._sawComplete = true } + /* + Two-argument `then` rather than `.catch().then()`: an extra hop would delay + the append by one microtask behind the previous delivery. + */ const delivery = this.persistenceTail.then(persistThenDeliver, (error: unknown) => { ignoreReplayBudgetRefusal(error) return persistThenDeliver() @@ -184,10 +173,8 @@ export class StreamWriter { this.persistenceTail = delivery return delivery } - for (const envelope of envelopes) { - this.enqueue(envelope) - this.queuePersistence(envelope) - } + this.enqueue(envelope) + this.queuePersistence(envelope) if (event.type === MothershipStreamV1EventType.complete) { this._sawComplete = true } @@ -224,31 +211,6 @@ export class StreamWriter { } } - /** - * Forwards the head of each tool call's argument deltas and drops the rest: - * they only feed a progressive title, and the call frame supersedes them. - */ - private limitArgsDelta(event: StreamEvent): StreamEvent | undefined { - if (event.type !== MothershipStreamV1EventType.tool || !('phase' in event.payload)) { - return event - } - const payload = event.payload - if (payload.phase !== MothershipStreamV1ToolPhase.args_delta) return event - const forwarded = this.forwardedArgsDeltaUnits.get(payload.toolCallId) ?? 0 - const remaining = TOOL_ARGS_DELTA_FORWARD_LIMIT_UNITS - forwarded - if (remaining <= 0) return undefined - const argumentsDelta = payload.argumentsDelta.slice(0, remaining) - this.forwardedArgsDeltaUnits.set(payload.toolCallId, forwarded + argumentsDelta.length) - return argumentsDelta === payload.argumentsDelta - ? event - : { ...event, payload: { ...payload, argumentsDelta } } - } - - private deliver(envelopes: PersistedStreamEventEnvelope[], event: StreamEvent): void { - for (const envelope of envelopes) this.enqueue(envelope) - if (event.type === MothershipStreamV1EventType.complete) this._sawComplete = true - } - private enqueue(envelope: PersistedStreamEventEnvelope): void { if (this._clientDisconnected || !this.controller) return try { diff --git a/apps/sim/lib/mothership/tools/client-executed-tools.ts b/apps/sim/lib/mothership/tools/client-executed-tools.ts new file mode 100644 index 00000000000..63c506bdff3 --- /dev/null +++ b/apps/sim/lib/mothership/tools/client-executed-tools.ts @@ -0,0 +1,32 @@ +import { isCurrentBrowserToolName } from '@sim/browser-protocol' +import { isTerminalToolName } from '@sim/terminal-protocol' +import { isNativeFileTool, isUserLocalVfsToolCall } from '@/lib/mothership/tools/local-filesystem' + +const WORKFLOW_TOOL_NAMES = new Set([ + 'run_workflow', + 'run_workflow_until_block', + 'run_block', + 'run_from_block', +]) + +export function isWorkflowToolName(name: string): boolean { + return WORKFLOW_TOOL_NAMES.has(name) +} + +/** + * Tool calls the browser starts from the call frame's own arguments: workflow + * runs, local file access, browser actions, and terminal commands. The stream + * must deliver those arguments exactly as the model sent them. + */ +export function isClientExecutedToolCall( + name: string, + args: Record | undefined +): boolean { + return ( + isWorkflowToolName(name) || + isNativeFileTool(name) || + isUserLocalVfsToolCall(name, args) || + isCurrentBrowserToolName(name) || + isTerminalToolName(name) + ) +} diff --git a/apps/sim/lib/mothership/tools/workflow-tools.ts b/apps/sim/lib/mothership/tools/workflow-tools.ts index 167d0abb0d1..666c111d2fd 100644 --- a/apps/sim/lib/mothership/tools/workflow-tools.ts +++ b/apps/sim/lib/mothership/tools/workflow-tools.ts @@ -7,15 +7,7 @@ import { isWorkflowToolExecutionClaimable, } from '@/lib/mothership/async-runs/lifecycle' import { COPILOT_WORKFLOW_EXECUTION_CONFLICT_CODE } from '@/lib/mothership/constants' - -const WORKFLOW_TOOL_NAMES = [ - 'run_workflow', - 'run_workflow_until_block', - 'run_block', - 'run_from_block', -] as const - -const WORKFLOW_TOOL_NAME_SET = new Set(WORKFLOW_TOOL_NAMES) +import { isWorkflowToolName } from '@/lib/mothership/tools/client-executed-tools' export const ASYNC_WORKFLOW_DEPLOYMENT_ERRORS = { missing: { @@ -142,10 +134,6 @@ export function classifyWorkflowToolBinding(params: { return { ok: true } } -export function isWorkflowToolName(name: string): boolean { - return WORKFLOW_TOOL_NAME_SET.has(name) -} - /** Resolves the workflow target from immutable tool arguments, then the owning Copilot run. */ export function resolveWorkflowToolTargetId( args: unknown, From 9b4f7c8090453b595631d94817a5f3d4345fac21 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 19:21:26 -0700 Subject: [PATCH 04/18] fix(mothership): stop stale approval cards, raw backend bodies, and errored-turn spinners - A replayed call frame is history and never gates anything, so Sim now clears the worker's awaiting_approval stamp on it as it already does for live frames. Replays no longer render Allow/Skip cards when approvals are off. - Stop settles every unfinished tool row (pending, executing, or awaiting approval) through one shared predicate, isUnsettledToolState, on the persisted, live, and pending-message paths. - An errored turn settles its unfinished tool rows as errored when it is persisted, so they no longer reload as spinners. - A failed backend response no longer puts the upstream body in the error shown to users: 5xx and non-JSON bodies say the service is temporarily unavailable, 4xx uses the worker's displayMessage or a generic message. The status and body stay on the error for logs. --- .../home/hooks/message-reconcile.test.ts | 28 ++++++++++ .../home/hooks/message-reconcile.ts | 12 +++-- .../home/hooks/stream/stream-helpers.test.ts | 18 +++++++ .../home/hooks/stream/stream-helpers.ts | 3 +- .../[workspaceId]/home/hooks/use-chat.ts | 11 ++-- .../mothership/chat/persisted-message.test.ts | 23 ++++++++ .../lib/mothership/chat/persisted-message.ts | 50 +++++++++++------- .../lib/mothership/request/go/stream.test.ts | 52 +++++++++++++++++++ apps/sim/lib/mothership/request/go/stream.ts | 32 ++++++++++-- .../request/handlers/handlers.test.ts | 29 +++++++++++ .../lib/mothership/request/handlers/tool.ts | 6 ++- 11 files changed, 230 insertions(+), 34 deletions(-) create mode 100644 apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.test.ts create mode 100644 apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.test.ts new file mode 100644 index 00000000000..6a1a931cd34 --- /dev/null +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from 'vitest' +import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message' +import { markMessageStopped } from '@/app/workspace/[workspaceId]/home/hooks/message-reconcile' + +describe('markMessageStopped', () => { + it.each(['executing', 'pending', 'awaiting_approval'] as const)( + 'settles a %s tool row as stopped', + (state) => { + const message: PersistedMessage = { + id: 'assistant-1', + role: 'assistant', + content: '', + timestamp: '2026-09-29T00:00:00.000Z', + contentBlocks: [ + { + type: 'tool', + endedAt: 1, + toolCall: { id: 'call-1', name: 'gmail_read_v2', state }, + }, + ], + } + + const stopped = markMessageStopped(message) + + expect(stopped.contentBlocks?.[0].toolCall?.state).toBe('cancelled') + } + ) +}) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts index 5245f98bc1e..8686d925660 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts @@ -8,7 +8,11 @@ import { isBrowserToolName } from '@sim/browser-protocol' import type { PersistedContentBlock } from '@/lib/api/contracts/copilot-messages' import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message' -import { normalizeMessage, withBlockTiming } from '@/lib/mothership/chat/persisted-message' +import { + isUnsettledToolState, + normalizeMessage, + withBlockTiming, +} from '@/lib/mothership/chat/persisted-message' import { MothershipStreamV1CompletionStatus, MothershipStreamV1EventType, @@ -142,8 +146,8 @@ export function buildAssistantSnapshotMessage(params: { } export function markMessageStopped(message: PersistedMessage): PersistedMessage { - const hasExecutingTool = message.contentBlocks?.some( - (block) => block.toolCall?.state === 'executing' + const hasExecutingTool = message.contentBlocks?.some((block) => + isUnsettledToolState(block.toolCall?.state) ) const hasOpenBlock = message.contentBlocks?.some((block) => block.endedAt === undefined) if (!hasExecutingTool && !hasOpenBlock) { @@ -153,7 +157,7 @@ export function markMessageStopped(message: PersistedMessage): PersistedMessage const stopTs = Date.now() const nextBlocks = (message.contentBlocks ?? []).map((block) => { const stamped = block.endedAt === undefined ? { ...block, endedAt: stopTs } : block - if (stamped.toolCall?.state !== 'executing') { + if (!stamped.toolCall || !isUnsettledToolState(stamped.toolCall.state)) { return stamped } return { diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts new file mode 100644 index 00000000000..4a1b6d61b36 --- /dev/null +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts @@ -0,0 +1,18 @@ +import { describe, expect, it } from 'vitest' +import { finalizeResidualToolCalls } from '@/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers' +import type { ContentBlock } from '@/app/workspace/[workspaceId]/home/types' + +describe('finalizeResidualToolCalls', () => { + it.each(['executing', 'awaiting_approval'] as const)( + 'settles a %s tool row with the turn outcome on Stop', + (status) => { + const blocks: ContentBlock[] = [ + { type: 'tool_call', toolCall: { id: 'call-1', name: 'gmail_read_v2', status } }, + ] + + finalizeResidualToolCalls(blocks, 'cancelled') + + expect(blocks[0].toolCall?.status).toBe('cancelled') + } + ) +}) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts index 99d455dd7b1..4ed30860902 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts @@ -1,5 +1,6 @@ import { createLogger } from '@sim/logger' import { isRecordLike } from '@sim/utils/object' +import { isUnsettledToolState } from '@/lib/mothership/chat/persisted-message' import { CallIntegrationTool, CreateEmptyFile, @@ -96,7 +97,7 @@ export function finalizeResidualToolCalls( continue } const tc = block.toolCall - if (!tc || tc.status !== ToolCallStatus.executing) continue + if (!tc || !isUnsettledToolState(tc.status)) continue tc.status = propagated if (propagated === ToolCallStatus.cancelled) { tc.displayTitle = 'Stopped by user' diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index cbac7f0d951..ac540bdad0d 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -45,9 +45,10 @@ import { import { getMothershipAttachmentPreviewUrl } from '@/lib/mothership/chat/attachment-preview' import { toDisplayMessage } from '@/lib/mothership/chat/display-message' import { getLiveAssistantMessageId } from '@/lib/mothership/chat/live-message-id' -import type { - PersistedFileAttachment, - PersistedMessage, +import { + isUnsettledToolState, + type PersistedFileAttachment, + type PersistedMessage, } from '@/lib/mothership/chat/persisted-message' import { type RevealedSimKeysByMessage, @@ -4370,8 +4371,8 @@ export function useChat( } else { setPendingMessages((prev) => prev.map((msg) => { - const hasExecutingTool = msg.contentBlocks?.some( - (block) => block.toolCall?.status === 'executing' + const hasExecutingTool = msg.contentBlocks?.some((block) => + isUnsettledToolState(block.toolCall?.status) ) const hasOpenBlock = msg.contentBlocks?.some((block) => block.endedAt === undefined) if (!hasExecutingTool && !hasOpenBlock) { diff --git a/apps/sim/lib/mothership/chat/persisted-message.test.ts b/apps/sim/lib/mothership/chat/persisted-message.test.ts index 78861c7169c..160b76348e9 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.test.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.test.ts @@ -761,3 +761,26 @@ describe('stripToolResultOutput', () => { expect(JSON.stringify(blocks)).not.toContain('file contents') }) }) + +describe('buildPersistedAssistantMessage on an errored turn', () => { + it.each(['pending', 'executing'] as const)( + 'settles a %s tool row as errored so it does not reload as a spinner', + (status) => { + const persisted = buildPersistedAssistantMessage({ + success: false, + error: 'The agent service is temporarily unavailable. Please try again.', + content: '', + toolCalls: [], + contentBlocks: [ + { + type: 'tool_call', + timestamp: 1, + toolCall: { id: 'call-1', name: 'gmail_read_v2', status }, + }, + ], + }) + + expect(persisted.contentBlocks?.[0].toolCall?.state).toBe('error') + } + ) +}) diff --git a/apps/sim/lib/mothership/chat/persisted-message.ts b/apps/sim/lib/mothership/chat/persisted-message.ts index 2e15cabb9da..9424e2dce95 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.ts @@ -336,8 +336,9 @@ export function buildPersistedAssistantMessage( 'An unexpected error occurred while processing the response.' ) normalized.contentBlocks = [ - ...(normalized.contentBlocks ?? - (message.content + ...(normalized.contentBlocks + ? settleUnfinishedToolCalls(normalized.contentBlocks, 'error') + : message.content ? [ { type: MothershipStreamV1EventType.text, @@ -345,7 +346,7 @@ export function buildPersistedAssistantMessage( content: message.content, }, ] - : [])), + : []), { type: MothershipStreamV1EventType.error, content: buildMothershipErrorTag({ message: error }), @@ -357,23 +358,34 @@ export function buildPersistedAssistantMessage( return message } -export function withStoppedContentBlock(message: PersistedMessage): PersistedMessage { - const contentBlocks = (message.contentBlocks ?? []).map( - (block): PersistedContentBlock => - block.toolCall && - (block.toolCall.state === 'executing' || - block.toolCall.state === 'pending' || - block.toolCall.state === 'awaiting_approval') - ? { - ...block, - toolCall: { - ...block.toolCall, - state: 'cancelled', - display: { title: 'Stopped by user' }, - }, - } - : block +const UNSETTLED_TOOL_STATES: ReadonlySet = new Set([ + 'pending', + 'executing', + 'awaiting_approval', +]) + +/** A tool row that has not finished: waiting to run, running, or awaiting a decision. */ +export function isUnsettledToolState(state: string | undefined): boolean { + return state !== undefined && UNSETTLED_TOOL_STATES.has(state) +} + +/** Settles every unfinished tool row at a turn terminal so none reloads as a spinner. */ +function settleUnfinishedToolCalls( + blocks: PersistedContentBlock[], + state: 'cancelled' | 'error', + display?: { title: string } +): PersistedContentBlock[] { + return blocks.map((block) => + block.toolCall && isUnsettledToolState(block.toolCall.state) + ? { ...block, toolCall: { ...block.toolCall, state, ...(display ? { display } : {}) } } + : block ) +} + +export function withStoppedContentBlock(message: PersistedMessage): PersistedMessage { + const contentBlocks = settleUnfinishedToolCalls(message.contentBlocks ?? [], 'cancelled', { + title: 'Stopped by user', + }) const hasAssistantText = contentBlocks.some( (block) => block.type === MothershipStreamV1EventType.text && diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index f412a194479..dc04e1b2392 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -78,8 +78,10 @@ vi.mock('@/lib/mothership/tools/server/files/file-preview', async () => { } }) +import { buildPersistedAssistantMessage } from '@/lib/mothership/chat/persisted-message' import { buildPreviewContentUpdate, + CopilotBackendError, decodeJsonStringPrefix, extractEditContent, runStreamLoop, @@ -250,6 +252,56 @@ describe('copilot go stream helpers', () => { expect(context.streamComplete).toBe(true) }) + it.each([ + [ + 'an HTML gateway page', + 502, + 'text/html', + '502 Bad Gatewaynginx', + 'The agent service is temporarily unavailable. Please try again.', + ], + [ + 'an internal JSON error', + 400, + 'application/json', + '{"error":"Bad Request","message":"userId required for internal API key"}', + 'The agent service could not process this request.', + ], + [ + 'a worker display message', + 409, + 'application/json', + '{"error":"conflict","displayMessage":"This chat is already answering another message."}', + 'This chat is already answering another message.', + ], + ])( + 'never shows the user the raw body of %s', + async (_label, status, contentType, body, userMessage) => { + vi.mocked(fetch).mockResolvedValueOnce( + new Response(body, { status, headers: { 'Content-Type': contentType } }) + ) + + const error = await runStreamLoop( + 'https://example.com/api/mothership', + {}, + createStreamingContext(), + turnScopedExecContext(), + { timeout: 1000 } + ).catch((thrown: unknown) => thrown) + + expect(error).toBeInstanceOf(CopilotBackendError) + expect(error).toMatchObject({ message: userMessage, status, body }) + const persisted = buildPersistedAssistantMessage({ + success: false, + error: (error as Error).message, + content: '', + contentBlocks: [], + toolCalls: [], + }) + expect(JSON.stringify(persisted)).not.toMatch(/ { const identity = createProviderToolCallIdentity('exhausted-identity-run') identity.retainedBytes = PROVIDER_TOOL_CALL_IDENTITY_LIMITS.maxRetainedBytes diff --git a/apps/sim/lib/mothership/request/go/stream.ts b/apps/sim/lib/mothership/request/go/stream.ts index 91dba0de262..01769002178 100644 --- a/apps/sim/lib/mothership/request/go/stream.ts +++ b/apps/sim/lib/mothership/request/go/stream.ts @@ -1,6 +1,7 @@ import { type Context, SpanStatusCode } from '@opentelemetry/api' import { createLogger } from '@sim/logger' import { getErrorMessage } from '@sim/utils/errors' +import { toRecordOrNull } from '@sim/utils/object' import { ORCHESTRATION_TIMEOUT_MS } from '@/lib/mothership/constants' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { CopilotSseCloseReason } from '@/lib/mothership/generated/trace-attribute-values-v1' @@ -65,6 +66,29 @@ export class CopilotBackendError extends Error { } } +const BACKEND_UNAVAILABLE_MESSAGE = + 'The agent service is temporarily unavailable. Please try again.' +const BACKEND_REJECTED_MESSAGE = 'The agent service could not process this request.' + +/** + * What the user is told about a failed backend response. The body is upstream + * detail (a proxy's HTML page, an internal validation error) and stays on the + * error for logs; only a worker-supplied `displayMessage` is meant for users. + */ +function backendErrorMessage(status: number, body: string): string { + let parsed: unknown + try { + parsed = JSON.parse(body) + } catch { + return BACKEND_UNAVAILABLE_MESSAGE + } + if (status >= 500) return BACKEND_UNAVAILABLE_MESSAGE + const displayMessage = toRecordOrNull(parsed)?.displayMessage + return typeof displayMessage === 'string' && displayMessage.trim() + ? displayMessage.trim() + : BACKEND_REJECTED_MESSAGE +} + export class BillingLimitError extends Error { constructor(public readonly userId: string) { super('Usage limit reached') @@ -175,10 +199,10 @@ export async function runStreamLoop( throw new BillingLimitError(execContext.userId) } - throw new CopilotBackendError( - `Copilot backend error (${response.status}): ${errorText || response.statusText}`, - { status: response.status, body: errorText || response.statusText } - ) + throw new CopilotBackendError(backendErrorMessage(response.status, errorText), { + status: response.status, + body: errorText || response.statusText, + }) } if (!response.body) { diff --git a/apps/sim/lib/mothership/request/handlers/handlers.test.ts b/apps/sim/lib/mothership/request/handlers/handlers.test.ts index d8c93d82444..59ed69ba3f3 100644 --- a/apps/sim/lib/mothership/request/handlers/handlers.test.ts +++ b/apps/sim/lib/mothership/request/handlers/handlers.test.ts @@ -402,6 +402,35 @@ describe('sse-handlers tool lifecycle', () => { expect(upsertAsyncToolCall).not.toHaveBeenCalled() }) + it('clears a Go-stamped approval frame on a replayed call, which never gates anything', async () => { + toolRequiresApproval.mockReturnValue(false) + context.runId = 'run-1' + context.toolPermissions = { + enabled: false, + autoAllowed: new Set(), + autoAllowPermitted: true, + } + + const event = { + type: MothershipStreamV1EventType.tool, + payload: { + toolCallId: 'gmail-2', + toolName: 'gmail_read_v2', + arguments: {}, + executor: MothershipStreamV1ToolExecutor.sim, + mode: MothershipStreamV1ToolMode.async, + phase: MothershipStreamV1ToolPhase.call, + status: 'awaiting_approval', + replay: true, + }, + } as unknown as StreamEvent + + await prePersistClientExecutableToolCall(event, context, {}) + + expect((event.payload as { status?: string }).status).toBeUndefined() + expect(upsertAsyncToolCall).not.toHaveBeenCalled() + }) + it('clears a Go-stamped approval frame on an internal tool', async () => { toolRequiresApproval.mockReturnValue(true) context.runId = 'run-1' diff --git a/apps/sim/lib/mothership/request/handlers/tool.ts b/apps/sim/lib/mothership/request/handlers/tool.ts index deadb956753..6ee96dcdf87 100644 --- a/apps/sim/lib/mothership/request/handlers/tool.ts +++ b/apps/sim/lib/mothership/request/handlers/tool.ts @@ -214,7 +214,11 @@ export async function prePersistClientExecutableToolCall( if (!isToolCallStreamEvent(event)) return const data = event.payload - if (data.replay) return + // A replay is history and never gates anything, so it cannot carry a prompt. + if (data.replay) { + if (data.status === TOOL_AWAITING_APPROVAL_STATUS) data.status = undefined + return + } const isGenerating = data.status === TOOL_CALL_STATUS.generating const isPartial = data.partial === true || isGenerating if (isPartial) return From 9491b149792ed1703dee61e83dc380301e5e8a16 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 19:43:13 -0700 Subject: [PATCH 05/18] fix(mothership): never forward an approval stamp on a frame that cannot be gated The worker stamps awaiting_approval on resolved integration calls and relies on Sim to keep or clear it. Sim cleared it on live complete frames that were not gated, but a partial frame returned before that decision and kept the stamp. Only a live, complete call frame can be held behind a prompt, so every replayed or partial frame now drops the stamp at the one point each frame passes through before it is persisted and delivered. With approvals off, no frame, replay, snapshot, or persisted row can carry the stamp to the client. --- .../request/handlers/handlers.test.ts | 78 +++++++++++++++++++ .../lib/mothership/request/handlers/tool.ts | 13 ++-- 2 files changed, 86 insertions(+), 5 deletions(-) diff --git a/apps/sim/lib/mothership/request/handlers/handlers.test.ts b/apps/sim/lib/mothership/request/handlers/handlers.test.ts index 59ed69ba3f3..290b818306b 100644 --- a/apps/sim/lib/mothership/request/handlers/handlers.test.ts +++ b/apps/sim/lib/mothership/request/handlers/handlers.test.ts @@ -73,12 +73,17 @@ import { sseHandlers, subAgentHandlers, } from '@/lib/mothership/request/handlers' +import { createEvent } from '@/lib/mothership/request/session/event' import { shouldSkipToolCallEvent } from '@/lib/mothership/request/sse-utils' import type { ExecutionContext, StreamEvent, StreamingContext, } from '@/lib/mothership/request/types' +import { + createTurnModel, + reduceEvent, +} from '@/app/workspace/[workspaceId]/home/hooks/stream/turn-model' import { ResolvedSecretTraceRegistry } from '@/executor/utils/resolved-secret-trace-registry' const { @@ -431,6 +436,79 @@ describe('sse-handlers tool lifecycle', () => { expect(upsertAsyncToolCall).not.toHaveBeenCalled() }) + it('clears a Go-stamped approval frame on a partial call, which is never held', async () => { + toolRequiresApproval.mockReturnValue(true) + context.runId = 'run-1' + context.toolPermissions = { + enabled: true, + autoAllowed: new Set(), + autoAllowPermitted: true, + } + + const event = { + type: MothershipStreamV1EventType.tool, + payload: { + toolCallId: 'gmail-3', + toolName: 'gmail_read_v2', + arguments: {}, + executor: MothershipStreamV1ToolExecutor.sim, + mode: MothershipStreamV1ToolMode.async, + phase: MothershipStreamV1ToolPhase.call, + status: 'awaiting_approval', + partial: true, + }, + } as unknown as StreamEvent + + await prePersistClientExecutableToolCall(event, context, {}) + + expect((event.payload as { status?: string }).status).toBeUndefined() + }) + + it.each([ + ['a live', {}], + ['a replayed', { replay: true }], + ['a partial', { partial: true }], + ])( + 'renders %s stamped call as an ordinary row when approvals are off', + async (_label, variant) => { + toolRequiresApproval.mockReturnValue(false) + context.runId = 'run-1' + context.toolPermissions = { + enabled: false, + autoAllowed: new Set(), + autoAllowPermitted: true, + } + const event = { + type: MothershipStreamV1EventType.tool, + payload: { + toolCallId: 'gmail-4', + toolName: 'gmail_read_v2', + arguments: {}, + executor: MothershipStreamV1ToolExecutor.sim, + mode: MothershipStreamV1ToolMode.async, + phase: MothershipStreamV1ToolPhase.call, + status: 'awaiting_approval', + ...variant, + }, + } as unknown as StreamEvent + + await prePersistClientExecutableToolCall(event, context, {}) + const model = reduceEvent( + createTurnModel(), + createEvent({ + streamId: 'stream-1', + cursor: '1', + seq: 1, + requestId: 'request-1', + type: event.type, + payload: event.payload, + } as Parameters[0]) + ) + + expect(model.nodes.get('gmail-4')?.status).toBe('running') + } + ) + it('clears a Go-stamped approval frame on an internal tool', async () => { toolRequiresApproval.mockReturnValue(true) context.runId = 'run-1' diff --git a/apps/sim/lib/mothership/request/handlers/tool.ts b/apps/sim/lib/mothership/request/handlers/tool.ts index 6ee96dcdf87..357b7a8b9b2 100644 --- a/apps/sim/lib/mothership/request/handlers/tool.ts +++ b/apps/sim/lib/mothership/request/handlers/tool.ts @@ -214,14 +214,17 @@ export async function prePersistClientExecutableToolCall( if (!isToolCallStreamEvent(event)) return const data = event.payload - // A replay is history and never gates anything, so it cannot carry a prompt. - if (data.replay) { + const isGenerating = data.status === TOOL_CALL_STATUS.generating + const isPartial = data.partial === true || isGenerating + /* + Only a live, complete call frame can be held behind a prompt, and only when + the gate below decides so. A replay is history and a partial frame is not yet + a call, so neither may carry the worker's approval stamp to the client. + */ + if (data.replay || isPartial) { if (data.status === TOOL_AWAITING_APPROVAL_STATUS) data.status = undefined return } - const isGenerating = data.status === TOOL_CALL_STATUS.generating - const isPartial = data.partial === true || isGenerating - if (isPartial) return const ui = getToolCallUI(data) const catalogEntry = getToolEntry(data.toolName) From bafd6aa5559779d2abcdee37380ac855d147515e Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 19:43:14 -0700 Subject: [PATCH 06/18] fix(mothership): keep retrying an unreachable worker through a task replacement A worker task replacement leaves the load balancer answering 502/504 for about a minute. Stream retries stopped after three attempts, about four seconds, and ended the turn with an error. While no worker answers (a 5xx or no connection), retries now continue with the existing backoff for up to two minutes from the first failure, still bounded by the run deadline and by Stop. Each attempt re-sends the same message identity, which the worker treats as a reattach rather than a second run. Streams that end without a terminal keep the three-retry limit. --- .../mothership/request/lifecycle/run.test.ts | 96 ++++++++++++++++++- .../request/lifecycle/stream-retry.test.ts | 33 ++++--- .../request/lifecycle/stream-retry.ts | 32 +++++-- 3 files changed, 138 insertions(+), 23 deletions(-) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index e3747581b41..498b90df99a 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -21,7 +21,7 @@ import { workspaceFileSecretProvenanceMockFns, } from '@sim/testing/mocks/workspace-file-secret-provenance.mock' import { generateId } from '@sim/utils/id' -import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' +import { afterAll, afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { scopeProviderToolCallId } from '@/lib/mothership/request/go/tool-call-identity' import { handleBillingLimitResponse } from '@/lib/mothership/request/tools/billing' import type { ExecutionContext, StreamingContext } from '@/lib/mothership/request/types' @@ -2877,6 +2877,100 @@ describe('runCopilotLifecycle', () => { ) }) + describe('when the worker task is being replaced', () => { + const replacementMs = 70_000 + + function workerUnavailableFor( + unavailableMs: number, + attempts: Array<{ at: number; body: string }> + ) { + const start = Date.now() + mockRunStreamLoop.mockImplementation( + async (_url: string, init: RequestInit, context: StreamingContext): Promise => { + attempts.push({ at: Date.now() - start, body: String(init.body) }) + if (Date.now() - start < unavailableMs) { + throw new CopilotBackendError('The agent service is temporarily unavailable.', { + status: 502, + }) + } + context.streamComplete = true + context.completionStatus = MothershipStreamV1CompletionStatus.complete + } + ) + } + + function send(abortSignal?: AbortSignal) { + return runCopilotLifecycle( + { message: 'hello', messageId: 'stream-replacement' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'exec-1', + runId: 'run-1', + ...(abortSignal ? { abortSignal } : {}), + executionContext: { + userId: 'user-1', + workflowId: '', + workspaceId: 'ws-1', + chatId: 'chat-1', + }, + } + ) + } + + afterEach(() => { + mockRunStreamLoop.mockReset() + vi.useRealTimers() + }) + + it('retries the same send through the replacement and completes one run', async () => { + vi.useFakeTimers() + const attempts: Array<{ at: number; body: string }> = [] + workerUnavailableFor(replacementMs, attempts) + + const pending = send() + await vi.advanceTimersByTimeAsync(replacementMs + 10_000) + const result = await pending + + expect(result).toEqual(expect.objectContaining({ success: true, cancelled: false })) + expect(attempts.at(-1)?.at).toBeGreaterThanOrEqual(replacementMs) + expect(new Set(attempts.map((attempt) => JSON.parse(attempt.body).messageId))).toEqual( + new Set(['stream-replacement']) + ) + }) + + it('gives up once the recovery window has passed', async () => { + vi.useFakeTimers() + const attempts: Array<{ at: number; body: string }> = [] + workerUnavailableFor(Number.POSITIVE_INFINITY, attempts) + + const pending = send() + await vi.advanceTimersByTimeAsync(10 * 60_000) + const result = await pending + + expect(result).toEqual(expect.objectContaining({ success: false, cancelled: false })) + expect(attempts.at(-1)?.at).toBeGreaterThanOrEqual(90_000) + expect(attempts.at(-1)?.at).toBeLessThanOrEqual(130_000) + }) + + it('stops waiting as soon as the user stops the turn', async () => { + vi.useFakeTimers() + const attempts: Array<{ at: number; body: string }> = [] + workerUnavailableFor(Number.POSITIVE_INFINITY, attempts) + const stop = new AbortController() + + const pending = send(stop.signal) + await vi.advanceTimersByTimeAsync(1) + stop.abort('user_stop:abortActiveStream') + await vi.advanceTimersByTimeAsync(1) + const result = await pending + + expect(result).toEqual(expect.objectContaining({ success: false, cancelled: true })) + expect(attempts).toHaveLength(1) + }) + }) + it('retries an interrupted resume with the same identity and keeps prior content', async () => { const executionContext: ExecutionContext = { userId: 'user-1', diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index ed75d7106e1..94a5a795e57 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -8,12 +8,9 @@ import { StreamRetryWindow } from '@/lib/mothership/request/lifecycle/stream-ret afterEach(() => vi.useRealTimers()) describe('stream recovery budget', () => { - it.each([ - new TypeError('fetch failed'), - new StreamEndedWithoutTerminalError('/api/mothership'), - new CopilotBackendError('Unavailable', { status: 503 }), - ])('stops after three retries despite a long task budget: %s', (error) => { + it('stops an ended-without-terminal stream after three retries despite a long task budget', () => { vi.useFakeTimers() + const error = new StreamEndedWithoutTerminalError('/api/mothership') const retry = new StreamRetryWindow() for (let index = 0; index < 3; index++) { const delay = retry.nextDelay(error) @@ -25,15 +22,23 @@ describe('stream recovery budget', () => { expect(retry.remainingMs()).toBeGreaterThan(3_500_000) }) - it('bounds the recovery period from the first failure without shortening healthy work', () => { - vi.useFakeTimers() - const retry = new StreamRetryWindow() - vi.advanceTimersByTime(600_000) - expect(retry.nextDelay(new TypeError('fetch failed'))).not.toBeNull() - vi.advanceTimersByTime(30_000) - expect(retry.nextDelay(new TypeError('fetch failed'))).toBeNull() - expect(retry.remainingMs()).toBe(2_970_000) - }) + it.each([new TypeError('fetch failed'), new CopilotBackendError('Unavailable', { status: 502 })])( + 'keeps retrying an unreachable worker for two minutes from the first failure: %s', + (error) => { + vi.useFakeTimers() + const retry = new StreamRetryWindow() + vi.advanceTimersByTime(600_000) + const firstFailure = Date.now() + for (;;) { + const delay = retry.nextDelay(error) + if (delay === null) break + vi.advanceTimersByTime(delay) + } + expect(Date.now() - firstFailure).toBeGreaterThan(110_000) + expect(Date.now() - firstFailure).toBeLessThanOrEqual(120_000) + expect(retry.remainingMs()).toBeGreaterThan(2_800_000) + } + ) it('never extends the original execution deadline', () => { vi.useFakeTimers() diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts index 98e382edb41..87101978c5b 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -8,11 +8,19 @@ import { const MAX_STREAM_RETRIES = 3 const STREAM_RECOVERY_WINDOW_MS = 30_000 +/** + * While the worker cannot be reached at all (a 5xx from the load balancer, or no + * connection), retries continue for this long from the first failure: long + * enough to outlast a worker task replacement (about 70 s of 502/504). Every + * attempt re-sends the same message identity, which the worker treats as a + * reattach, never a second run. + */ +const WORKER_REPLACEMENT_WINDOW_MS = 120_000 /** Recovery is bounded independently of the healthy run's execution budget. */ export class StreamRetryWindow { private readonly deadline: number - private recoveryDeadline?: number + private firstFailureAt?: number attempt = 0 constructor(timeoutMs = ORCHESTRATION_TIMEOUT_MS) { @@ -28,23 +36,31 @@ export class StreamRetryWindow { nextDelay(error: unknown, signal?: AbortSignal): number | null { if (signal?.aborted || !isRetryableStreamError(error)) return null - this.recoveryDeadline ??= Date.now() + STREAM_RECOVERY_WINDOW_MS - if (this.attempt >= MAX_STREAM_RETRIES) return null + this.firstFailureAt ??= Date.now() + const unreachable = isWorkerUnreachable(error) + if (!unreachable && this.attempt >= MAX_STREAM_RETRIES) return null + const recoveryDeadline = + this.firstFailureAt + (unreachable ? WORKER_REPLACEMENT_WINDOW_MS : STREAM_RECOVERY_WINDOW_MS) const delay = backoffWithJitter(this.attempt + 1, null, { baseMs: 250, maxMs: 5_000 }) - if (Date.now() + delay >= Math.min(this.deadline, this.recoveryDeadline)) return null + if (Date.now() + delay >= Math.min(this.deadline, recoveryDeadline)) return null this.attempt++ return delay } } +/** No worker answered: the load balancer's 5xx, or no connection at all. */ +function isWorkerUnreachable(error: unknown): boolean { + if (error instanceof CopilotBackendError) { + return error.status !== undefined && error.status >= 500 + } + return error instanceof TypeError +} + /** Initial sends and resumes both replay one durable identity after an ambiguous response. */ function isRetryableStreamError(error: unknown): boolean { if (error instanceof Error && error.name === 'AbortError') return false if (error instanceof StreamEndedWithoutTerminalError || error instanceof StreamContinuityError) { return true } - if (error instanceof CopilotBackendError) { - return error.status !== undefined && error.status >= 500 - } - return error instanceof TypeError + return isWorkerUnreachable(error) } From 61388bf44d9d132baabef3b69ca386c6e7a79a89 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 19:53:18 -0700 Subject: [PATCH 07/18] fix(mothership): settle runs whose terminal events fail, and tighten compaction and cleanup - A controller that lost its lock while handling an ordinary failure no longer expires its successor's stream or clears its abort marker: teardown skips cleanup when this controller handed off or was superseded. - The terminal run status is now written even when publishing the terminal events fails. The write stays scoped to this controller's token and never replaces a terminal status, so a successor's run is left alone. Before, such a failure left the run non-terminal with nothing to settle it. - A failed backend response is described by its status alone: 5xx says the service is temporarily unavailable, anything else that the request could not be processed. - Compaction no longer exempts whole subtrees by key name. Every string the UI reads for identity, status, titles, or targets is far shorter than the cut, and long text keeps its head, so only client-executed call arguments, file previews, and generate_api_key results are left whole. --- .../lib/mothership/request/go/stream.test.ts | 14 +- apps/sim/lib/mothership/request/go/stream.ts | 19 +-- .../mothership/request/lifecycle/finalize.ts | 151 ++++++++++-------- .../lib/mothership/request/lifecycle/start.ts | 8 +- .../session/replay-budget.integration.ts | 113 +++++++++++-- .../request/session/replay-compaction.test.ts | 38 ++++- .../request/session/replay-compaction.ts | 62 ++----- 7 files changed, 254 insertions(+), 151 deletions(-) diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index dc04e1b2392..539e5c2dcbb 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -268,11 +268,11 @@ describe('copilot go stream helpers', () => { 'The agent service could not process this request.', ], [ - 'a worker display message', - 409, - 'application/json', - '{"error":"conflict","displayMessage":"This chat is already answering another message."}', - 'This chat is already answering another message.', + 'a plain-text request rejection', + 400, + 'text/plain', + 'Invalid request body', + 'The agent service could not process this request.', ], ])( 'never shows the user the raw body of %s', @@ -298,7 +298,9 @@ describe('copilot go stream helpers', () => { contentBlocks: [], toolCalls: [], }) - expect(JSON.stringify(persisted)).not.toMatch(/= 500) return BACKEND_UNAVAILABLE_MESSAGE - const displayMessage = toRecordOrNull(parsed)?.displayMessage - return typeof displayMessage === 'string' && displayMessage.trim() - ? displayMessage.trim() - : BACKEND_REJECTED_MESSAGE +function backendErrorMessage(status: number): string { + return status >= 500 ? BACKEND_UNAVAILABLE_MESSAGE : BACKEND_REJECTED_MESSAGE } export class BillingLimitError extends Error { @@ -199,7 +188,7 @@ export async function runStreamLoop( throw new BillingLimitError(execContext.userId) } - throw new CopilotBackendError(backendErrorMessage(response.status, errorText), { + throw new CopilotBackendError(backendErrorMessage(response.status), { status: response.status, body: errorText || response.statusText, }) diff --git a/apps/sim/lib/mothership/request/lifecycle/finalize.ts b/apps/sim/lib/mothership/request/lifecycle/finalize.ts index d57648c6c42..e77838a66a1 100644 --- a/apps/sim/lib/mothership/request/lifecycle/finalize.ts +++ b/apps/sim/lib/mothership/request/lifecycle/finalize.ts @@ -86,27 +86,29 @@ async function handleAborted( toolCallCount, blockCount, }) - if (!publisher.sawComplete) { - const partialContent = result.content || undefined - await publisher.publish({ - type: MothershipStreamV1EventType.complete, - payload: { - status: MothershipStreamV1CompletionStatus.cancelled, - ...(partialContent ? { partialContent } : {}), - ...(partialContentLen ? { partialContentLen } : {}), - ...(toolCallCount ? { toolCallCount } : {}), - }, - }) - } - await publisher.flush() - await loggedRunStatusUpdate( - runId, - MothershipStreamV1CompletionStatus.cancelled, - requestId, - { - completedAt: new Date(), + await publishThenSettle( + publisher, + async () => { + if (publisher.sawComplete) return + const partialContent = result.content || undefined + await publisher.publish({ + type: MothershipStreamV1EventType.complete, + payload: { + status: MothershipStreamV1CompletionStatus.cancelled, + ...(partialContent ? { partialContent } : {}), + ...(partialContentLen ? { partialContentLen } : {}), + ...(toolCallCount ? { toolCallCount } : {}), + }, + }) }, - publisher.controllerToken + () => + loggedRunStatusUpdate( + runId, + MothershipStreamV1CompletionStatus.cancelled, + requestId, + { completedAt: new Date() }, + publisher.controllerToken + ) ) } @@ -140,37 +142,38 @@ async function handleError( // Surface the real error (Go already classifies provider errors like // "overloaded" into a friendly displayMessage). Don't clobber it with a // generic string. - await publisher.publish({ - type: MothershipStreamV1EventType.error, - payload: { - message: errorMessage, - error: errorMessage, - displayMessage: errorMessage, - ...(result.errorCode ? { code: result.errorCode } : {}), - data: { displayMessage: errorMessage }, - }, - }) - if (!publisher.sawComplete) { - await publisher.publish({ - type: MothershipStreamV1EventType.complete, - payload: { - status: MothershipStreamV1CompletionStatus.error, - ...(partialContent ? { partialContent } : {}), - ...(partialContentLen ? { partialContentLen } : {}), - ...(toolCallCount ? { toolCallCount } : {}), - }, - }) - } - await publisher.flush() - await loggedRunStatusUpdate( - runId, - MothershipStreamV1CompletionStatus.error, - requestId, - { - completedAt: new Date(), - error: errorMessage, + await publishThenSettle( + publisher, + async () => { + await publisher.publish({ + type: MothershipStreamV1EventType.error, + payload: { + message: errorMessage, + error: errorMessage, + displayMessage: errorMessage, + ...(result.errorCode ? { code: result.errorCode } : {}), + data: { displayMessage: errorMessage }, + }, + }) + if (publisher.sawComplete) return + await publisher.publish({ + type: MothershipStreamV1EventType.complete, + payload: { + status: MothershipStreamV1CompletionStatus.error, + ...(partialContent ? { partialContent } : {}), + ...(partialContentLen ? { partialContentLen } : {}), + ...(toolCallCount ? { toolCallCount } : {}), + }, + }) }, - publisher.controllerToken + () => + loggedRunStatusUpdate( + runId, + MothershipStreamV1CompletionStatus.error, + requestId, + { completedAt: new Date(), error: errorMessage }, + publisher.controllerToken + ) ) } @@ -179,24 +182,44 @@ async function handleSuccess( runId: string, requestId: string ): Promise { - if (!publisher.sawComplete) { - await publisher.publish({ - type: MothershipStreamV1EventType.complete, - payload: { status: MothershipStreamV1CompletionStatus.complete }, - }) - } - await publisher.flush() - await loggedRunStatusUpdate( - runId, - MothershipStreamV1CompletionStatus.complete, - requestId, - { - completedAt: new Date(), + await publishThenSettle( + publisher, + async () => { + if (publisher.sawComplete) return + await publisher.publish({ + type: MothershipStreamV1EventType.complete, + payload: { status: MothershipStreamV1CompletionStatus.complete }, + }) }, - publisher.controllerToken + () => + loggedRunStatusUpdate( + runId, + MothershipStreamV1CompletionStatus.complete, + requestId, + { completedAt: new Date() }, + publisher.controllerToken + ) ) } +/** + * Publishes the terminal events, then records the run's terminal status even if + * publishing failed: nothing else settles a run, and the status write is scoped + * to this controller's token, so it never settles a run a successor has claimed. + */ +async function publishThenSettle( + publisher: StreamWriter, + publish: () => Promise, + settle: () => Promise +): Promise { + try { + await publish() + await publisher.flush() + } finally { + await settle() + } +} + async function loggedRunStatusUpdate( runId: string, status: Parameters[1], diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 7d173f461d2..358ea64cc01 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -510,8 +510,12 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS await releasePendingChatStream(chatId, streamId, lease) } processResourcesReleased = true - // A superseded controller must not expire or clear its successor's stream. - if (!handedOff) { + // A superseded controller must not expire or clear its successor's stream, + // whether it handed off or lost ownership while handling another failure. + if ( + !handedOff && + !(abortController.signal.reason instanceof StreamControllerSupersededError) + ) { await scheduleBufferCleanup(streamId) await scheduleFilePreviewSessionCleanup(streamId) await cleanupAbortMarker(streamId) 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 2f2500a4bf4..d08c2cdceeb 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -91,6 +91,7 @@ import { getRedisBudgetKeys, getRedisBudgetLimits } from '@/lib/core/redis/byte- import { readChatStream } from '@/lib/mothership/request/application/recover-stream' import { createStreamingContext } from '@/lib/mothership/request/context/request-context' import { restoreStreamingContext } from '@/lib/mothership/request/context/restore' +import { finalizeStream } from '@/lib/mothership/request/lifecycle/finalize' import { createSSEStream } from '@/lib/mothership/request/lifecycle/start' import { acquirePendingChatStream } from '@/lib/mothership/request/session/abort' import { readEvents } from '@/lib/mothership/request/session/buffer' @@ -482,26 +483,108 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge expect(controllerRun.sinkErrors[0]).toBeInstanceOf(StreamReplayBudgetExhaustedError) }) - it("leaves its successor's stream untouched when a superseded controller is refused an oversized frame", async () => { - const successorToken = `successor\n${generateId()}` - const { streamId, runId, frames } = await runTurn([ - async () => { - await redis().set(chatStreamLockKey(chatId), successorToken, 'EX', 60) - }, + it.each([ + [ + 'is refused an oversized frame', toolCall('call-oversized', 'run_workflow', { workflowId: generateId(), input: 'w'.repeat(1.5 * MB), }), - ]) + ], + [ + 'fails its worker leg', + async () => { + throw new Error('worker leg failed') + }, + ], + ])( + "leaves its successor's stream untouched when a superseded controller %s", + async (_label, step) => { + const successorToken = `successor\n${generateId()}` + const { streamId, runId, frames } = await runTurn([ + async () => { + await redis().set(chatStreamLockKey(chatId), successorToken, 'EX', 60) + }, + step, + ]) + + expect(frames.map((frame) => frame.type)).toEqual(['session']) + /** The client stream closes before the controller's teardown; let teardown finish. */ + await sleep(500) + expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeGreaterThan(300) + expect(await redis().get(chatStreamLockKey(chatId))).toBe(successorToken) + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('active') + expect(worker.abortRequests).toEqual([]) + await redis().del(chatStreamLockKey(chatId)) + } + ) + + it('marks its run terminal even when the final events cannot be published', async () => { + const streamId = generateId() + const runId = generateId() + const controllerToken = `${streamId}\n${generateId()}` + await db.insert(copilotRuns).values({ + id: runId, + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + status: 'paused_waiting_for_tool', + requestContext: { requestId: generateId(), controllerToken }, + }) + /** No chat lock holds this lease, so publishing the terminal events throws. */ + const lease = { key: chatStreamLockKey(generateId()), value: controllerToken } + const publisher = new StreamWriter({ streamId, requestId: generateId(), lease }) + + await expect( + finalizeStream( + { + success: false, + error: 'The agent service is temporarily unavailable. Please try again.', + content: '', + contentBlocks: [], + toolCalls: [], + }, + publisher, + runId, + 'error', + generateId() + ) + ).rejects.toBeInstanceOf(StreamControllerSupersededError) + + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('error') + }) + + it('does not settle a run another controller has claimed', async () => { + const streamId = generateId() + const runId = generateId() + await db.insert(copilotRuns).values({ + id: runId, + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + status: 'paused_waiting_for_tool', + requestContext: { requestId: generateId(), controllerToken: `successor\n${generateId()}` }, + }) + const lease = { key: chatStreamLockKey(generateId()), value: `stale\n${generateId()}` } + const publisher = new StreamWriter({ streamId, requestId: generateId(), lease }) + + await expect( + finalizeStream( + { success: false, error: 'failed', content: '', contentBlocks: [], toolCalls: [] }, + publisher, + runId, + 'error', + generateId() + ) + ).rejects.toBeInstanceOf(StreamControllerSupersededError) - expect(frames.map((frame) => frame.type)).toEqual(['session']) - /** The client stream closes before the controller's teardown; let teardown finish. */ - await sleep(500) - expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeGreaterThan(300) - expect(await redis().get(chatStreamLockKey(chatId))).toBe(successorToken) const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) - expect(stored.status).toBe('active') - expect(worker.abortRequests).toEqual([]) - await redis().del(chatStreamLockKey(chatId)) + expect(stored.status).toBe('paused_waiting_for_tool') }) }) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index 425b914c5c0..2ea67e71bb0 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -44,7 +44,7 @@ describe('compactStreamEvent', () => { exitCode: 0, resources, citations, - reason: 'r'.repeat(20_000), + reason: 'user_cancelled', }, }, } @@ -59,7 +59,7 @@ describe('compactStreamEvent', () => { expect(output.stderr).toBe( `${'e'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 300 KB total]` ) - expect(output).toMatchObject({ exitCode: 0, resources, citations, reason: 'r'.repeat(20_000) }) + expect(output).toMatchObject({ exitCode: 0, resources, citations, reason: 'user_cancelled' }) expect(Object.keys(output)).toEqual(Object.keys(toRecord(payloadOf(event).output))) expect(payloadOf(compacted)).toMatchObject({ toolCallId: 'call-1', @@ -70,6 +70,40 @@ describe('compactStreamEvent', () => { }) }) + it('cuts long text even under keys the UI reads, keeping its head', () => { + const rows = Array.from({ length: 3 }, (_, index) => ({ + id: `row-${index}`, + description: `row ${index} ${'d'.repeat(700 * 1024)}`, + })) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_tables_rows_query', + executor: 'sim', + mode: 'async', + phase: 'result', + success: false, + status: 'error', + error: `Query failed: ${'e'.repeat(2 * MB)}`, + output: { error: `Query failed: ${'e'.repeat(2 * MB)}`, rows }, + }, + } + + const compacted = payloadOf(compactStreamEvent(event)) + const output = toRecord(compacted.output) + + expect(Buffer.byteLength(JSON.stringify(compacted))).toBeLessThan( + STREAM_EVENT_COMPACTION_THRESHOLD_BYTES + ) + expect(compacted.error).toMatch(/^Query failed: e+…\[truncated, 2 MB total\]$/) + expect(output.error).toMatch(/^Query failed: e+…\[truncated, 2 MB total\]$/) + expect((output.rows as Array<{ id: string; description: string }>)[2]).toEqual({ + id: 'row-2', + description: `row 2 ${'d'.repeat(STREAM_STRING_PREVIEW_UNITS - 6)}…[truncated, 700 KB total]`, + }) + }) + it.each([ ['a workflow run', 'run_workflow', { workflowId: 'wf-1' }], ['a user-local VFS read', 'read', { path: 'user-local/notes.md' }], diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index ce265209271..3bf1d03fb29 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -13,45 +13,8 @@ export const STREAM_EVENT_COMPACTION_THRESHOLD_BYTES = 256 * 1024 /** UTF-16 units kept at the head of a long string. */ export const STREAM_STRING_PREVIEW_UNITS = 8 * 1024 -/** Identity fields and the keys the UI reads from arguments and outputs; kept whole at any depth. */ -const PRESERVED_KEYS: ReadonlySet = new Set([ - 'activity', - 'activityDescription', - 'agent', - 'cancelledByUser', - 'channel', - 'description', - 'elementId', - 'error', - 'event', - 'execName', - 'executor', - 'fileName', - 'kind', - 'mode', - 'op', - 'operation', - 'partial', - 'path', - 'pattern', - 'phase', - 'reason', - 'replay', - 'resource', - 'status', - 'success', - 'targetWorkspaceId', - 'terminalId', - 'timeoutMs', - 'title', - 'toolCallId', - 'toolId', - 'toolName', - 'ui', - 'url', - 'workflowId', - 'workspaceId', -]) +const NO_KEYS: ReadonlySet = new Set() +const ARGUMENTS_KEY: ReadonlySet = new Set(['arguments']) /** Its live result is the only place the plaintext key reaches the browser. */ const GENERATE_API_KEY_TOOL = 'generate_api_key' @@ -65,8 +28,13 @@ function stringUnits(value: unknown): number { return sum } -/** Cuts long strings to a head with a readable size note, copying only what changes. */ -function truncateStrings(value: unknown, skipKeys: ReadonlySet): unknown { +/** + * Cuts long strings to a head with a readable size note, copying only what + * changes. Every string the UI reads for identity, status, titles, or targets is + * far shorter than the cut, and text fields keep their head, so nothing but the + * named top-level keys is exempt. + */ +function truncateStrings(value: unknown, skipKeys: ReadonlySet = NO_KEYS): unknown { if (typeof value === 'string') { if (value.length <= STREAM_STRING_PREVIEW_UNITS) return value const end = STREAM_STRING_PREVIEW_UNITS @@ -78,7 +46,7 @@ function truncateStrings(value: unknown, skipKeys: ReadonlySet): unknown if (Array.isArray(value)) { let copy: unknown[] | undefined value.forEach((item, index) => { - const next = truncateStrings(item, PRESERVED_KEYS) + const next = truncateStrings(item) if (next !== item) (copy ??= [...value])[index] = next }) return copy ?? value @@ -87,7 +55,7 @@ function truncateStrings(value: unknown, skipKeys: ReadonlySet): unknown let copy: Record | undefined for (const [key, field] of Object.entries(value)) { if (skipKeys.has(key)) continue - const next = truncateStrings(field, PRESERVED_KEYS) + const next = truncateStrings(field) if (next !== field) (copy ??= { ...value })[key] = next } return copy ?? value @@ -97,8 +65,8 @@ function truncateStrings(value: unknown, skipKeys: ReadonlySet): unknown * Bounds an outgoing stream event so the replay buffer can persist it. Applied * only to the copy the writer delivers and persists; the caller keeps the full * event for dispatch. Long strings are cut to their head in place, so every - * object keeps its shape. File previews, `generate_api_key` results, preserved - * keys, and the arguments of calls the browser executes are never cut; an event + * object keeps its shape. File previews, `generate_api_key` results, and the + * arguments of calls the browser executes are never cut; an event * still too large is refused by the buffer, which ends the turn with an error. */ export function compactStreamEvent(event: StreamEvent): StreamEvent { @@ -112,8 +80,8 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { const args = isRecordLike(payload.arguments) ? payload.arguments : undefined const skipKeys = payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutedToolCall(toolName, args) - ? new Set([...PRESERVED_KEYS, 'arguments']) - : PRESERVED_KEYS + ? ARGUMENTS_KEY + : NO_KEYS const compacted = truncateStrings(payload, skipKeys) return compacted === payload ? event : ({ ...event, payload: compacted } as StreamEvent) } From 7fc6db9aa887fa186982df7c5c8c5c6beb755182 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 21:20:27 -0700 Subject: [PATCH 08/18] fix(mothership): settle tool rows left unfinished by any finished turn - A completed turn now persists its unfinished tool rows settled as the live view settled them at complete, so they no longer reload as spinners. - A stored message's unfinished tool rows render as interrupted, which settles rows persisted before these fixes; the live streaming message is unchanged. - The client's error finalize settles pending and awaiting-approval rows too: finalizeResidualToolCalls reports whether it settled anything, and that decides whether the snapshot is rewritten. --- .../home/hooks/stream/stream-helpers.test.ts | 12 ++++++++ .../home/hooks/stream/stream-helpers.ts | 14 +++++---- .../[workspaceId]/home/hooks/use-chat.ts | 3 +- .../mothership/chat/display-message.test.ts | 29 ++++++++++++++++++ .../lib/mothership/chat/display-message.ts | 30 ++++++++++++------- .../mothership/chat/persisted-message.test.ts | 18 +++++++++++ .../lib/mothership/chat/persisted-message.ts | 6 +++- 7 files changed, 93 insertions(+), 19 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts index 4a1b6d61b36..2d0e905bb31 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts @@ -15,4 +15,16 @@ describe('finalizeResidualToolCalls', () => { expect(blocks[0].toolCall?.status).toBe('cancelled') } ) + + it('reports whether any tool row was left to settle', () => { + const open: ContentBlock[] = [ + { type: 'tool_call', toolCall: { id: 'call-1', name: 'read', status: 'awaiting_approval' } }, + ] + const settled: ContentBlock[] = [ + { type: 'tool_call', toolCall: { id: 'call-2', name: 'read', status: 'success' } }, + ] + + expect(finalizeResidualToolCalls(open, 'error')).toBe(true) + expect(finalizeResidualToolCalls(settled, 'error')).toBe(false) + }) }) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts index 4ed30860902..3067ba815ba 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts @@ -69,17 +69,17 @@ export function asPayloadRecord(value: unknown): StreamPayload | undefined { } /** - * Settles any tool row still `executing` at a turn terminal by propagating the - * turn's outcome — the deterministic replacement for the old `interrupted` - * invention. A clean `complete` means the turn succeeded, so a straggler is - * settled `success` (with explicit tool/span terminals from the backend there - * are normally none); a stop settles `cancelled`; an error settles `error`. + * Settles every unfinished tool row (running, pending, or awaiting approval) at + * a turn terminal by propagating the turn's outcome: a clean `complete` settles + * a straggler `success`, a stop `cancelled`, an error `error`. Also closes any + * open subagent lane. Returns whether any tool row was settled. */ export function finalizeResidualToolCalls( blocks: ContentBlock[], turnTerminal: 'complete' | 'cancelled' | 'error' -): void { +): boolean { const endedAt = Date.now() + let settled = false const propagated = turnTerminal === 'cancelled' ? ToolCallStatus.cancelled @@ -98,6 +98,7 @@ export function finalizeResidualToolCalls( } const tc = block.toolCall if (!tc || !isUnsettledToolState(tc.status)) continue + settled = true tc.status = propagated if (propagated === ToolCallStatus.cancelled) { tc.displayTitle = 'Stopped by user' @@ -106,6 +107,7 @@ export function finalizeResidualToolCalls( block.endedAt = endedAt } } + return settled } function stringParam(value: unknown): string | undefined { diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index ac540bdad0d..d6b514fc733 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -3228,8 +3228,7 @@ export function useChat( const isError = !!options?.error if (isError) { const blocks = streamingBlocksRef.current - if (blocks.some((block) => block.toolCall?.status === 'executing')) { - finalizeResidualToolCalls(blocks, 'error') + if (finalizeResidualToolCalls(blocks, 'error')) { const assistantId = activeTurnRef.current?.assistantMessageId ?? (streamIdRef.current ? getLiveAssistantMessageId(streamIdRef.current) : undefined) diff --git a/apps/sim/lib/mothership/chat/display-message.test.ts b/apps/sim/lib/mothership/chat/display-message.test.ts index 85f34a64992..efd2f9d9ba9 100644 --- a/apps/sim/lib/mothership/chat/display-message.test.ts +++ b/apps/sim/lib/mothership/chat/display-message.test.ts @@ -1,6 +1,18 @@ import { describe, expect, it } from 'vitest' +import { getLiveAssistantMessageId } from '@/lib/mothership/chat/live-message-id' +import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message' import { toDisplayMessage } from './display-message' +function storedAssistant(id: string, state: string): PersistedMessage { + return { + id, + role: 'assistant', + content: '', + timestamp: '2026-09-29T00:00:00.000Z', + contentBlocks: [{ type: 'tool', toolCall: { id: 'call-1', name: 'read', state } }], + } as PersistedMessage +} + describe('display-message', () => { it('maps canonical tool, subagent text, and cancelled complete blocks to display blocks', () => { const display = toDisplayMessage({ @@ -195,4 +207,21 @@ describe('display-message', () => { }, ]) }) + + it.each(['pending', 'executing', 'awaiting_approval'])( + 'shows a %s row of a stored message as interrupted, not running', + (state) => { + const display = toDisplayMessage(storedAssistant('assistant-1', state)) + + expect(display.contentBlocks?.[0].toolCall?.status).toBe('interrupted') + } + ) + + it('keeps a running row of the live message running', () => { + const display = toDisplayMessage( + storedAssistant(getLiveAssistantMessageId('stream-1'), 'executing') + ) + + expect(display.contentBlocks?.[0].toolCall?.status).toBe('executing') + }) }) diff --git a/apps/sim/lib/mothership/chat/display-message.ts b/apps/sim/lib/mothership/chat/display-message.ts index fbe3d568e3c..58155f701e8 100644 --- a/apps/sim/lib/mothership/chat/display-message.ts +++ b/apps/sim/lib/mothership/chat/display-message.ts @@ -18,8 +18,9 @@ import { ToolCallStatus, } from '@/app/workspace/[workspaceId]/home/types' import { getMothershipAttachmentPreviewUrl } from './attachment-preview' +import { isLiveAssistantMessageId } from './live-message-id' import type { PersistedMessage } from './persisted-message' -import { withBlockTiming } from './persisted-message' +import { isUnsettledToolState, withBlockTiming } from './persisted-message' const STATE_TO_STATUS: Record = { [MothershipStreamV1ToolOutcome.success]: ToolCallStatus.success, @@ -35,11 +36,15 @@ const STATE_TO_STATUS: Record = { awaiting_approval: ToolCallStatus.awaiting_approval, } -function toToolCallInfo(block: PersistedContentBlock): ToolCallInfo | undefined { +function toToolCallInfo(block: PersistedContentBlock, stored: boolean): ToolCallInfo | undefined { const tc = block.toolCall if (!tc) return undefined if (isToolHiddenInUi(tc.name)) return undefined - const status: ToolCallStatus = STATE_TO_STATUS[tc.state] ?? ToolCallStatus.error + /** A stored turn has ended; a row it left unfinished did not finish. */ + const status: ToolCallStatus = + stored && isUnsettledToolState(tc.state) + ? ToolCallStatus.interrupted + : (STATE_TO_STATUS[tc.state] ?? ToolCallStatus.error) const activityDescription = normalizeToolActivityDescription(tc.activityDescription) return { id: tc.id, @@ -53,8 +58,8 @@ function toToolCallInfo(block: PersistedContentBlock): ToolCallInfo | undefined } } -function toDisplayBlock(block: PersistedContentBlock): ContentBlock | undefined { - const displayed = toDisplayBlockBody(block) +function toDisplayBlock(block: PersistedContentBlock, stored: boolean): ContentBlock | undefined { + const displayed = toDisplayBlockBody(block, stored) if (!displayed) return undefined if (block.parentToolCallId && displayed.parentToolCallId === undefined) { displayed.parentToolCallId = block.parentToolCallId @@ -68,7 +73,10 @@ function toDisplayBlock(block: PersistedContentBlock): ContentBlock | undefined return withBlockTiming(displayed, block) } -function toDisplayBlockBody(block: PersistedContentBlock): ContentBlock | undefined { +function toDisplayBlockBody( + block: PersistedContentBlock, + stored: boolean +): ContentBlock | undefined { switch (block.type) { case 'task': return block.task ? { type: ContentBlockType.task, task: block.task } : undefined @@ -91,9 +99,10 @@ function toDisplayBlockBody(block: PersistedContentBlock): ContentBlock | undefi return { type: ContentBlockType.thinking, content: block.content } } return { type: ContentBlockType.text, content: block.content } - case MothershipStreamV1EventType.tool: - if (!toToolCallInfo(block)) return undefined - return { type: ContentBlockType.tool_call, toolCall: toToolCallInfo(block) } + case MothershipStreamV1EventType.tool: { + const toolCall = toToolCallInfo(block, stored) + return toolCall ? { type: ContentBlockType.tool_call, toolCall } : undefined + } case MothershipStreamV1EventType.span: if (block.lifecycle === MothershipStreamV1SpanLifecycleEvent.end) { return { @@ -218,8 +227,9 @@ export function toDisplayMessage(msg: PersistedMessage): ChatMessage { } if (msg.contentBlocks && msg.contentBlocks.length > 0) { + const stored = msg.role === 'assistant' && !isLiveAssistantMessageId(msg.id) const displayBlocks = msg.contentBlocks - .map(toDisplayBlock) + .map((block) => toDisplayBlock(block, stored)) .filter((block): block is ContentBlock => !!block) display.contentBlocks = foldFileWriteBlocks(displayBlocks) } diff --git a/apps/sim/lib/mothership/chat/persisted-message.test.ts b/apps/sim/lib/mothership/chat/persisted-message.test.ts index 160b76348e9..9d08415e68a 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.test.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.test.ts @@ -784,3 +784,21 @@ describe('buildPersistedAssistantMessage on an errored turn', () => { } ) }) + +describe('buildPersistedAssistantMessage on a completed turn', () => { + it.each(['pending', 'executing'] as const)( + 'settles a %s tool row the way the live view settled it', + (status) => { + const persisted = buildPersistedAssistantMessage({ + success: true, + content: 'Done.', + toolCalls: [], + contentBlocks: [ + { type: 'tool_call', timestamp: 1, toolCall: { id: 'call-1', name: 'read', status } }, + ], + }) + + expect(persisted.contentBlocks?.[0].toolCall?.state).toBe('success') + } + ) +}) diff --git a/apps/sim/lib/mothership/chat/persisted-message.ts b/apps/sim/lib/mothership/chat/persisted-message.ts index 9424e2dce95..2d0481cda1e 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.ts @@ -355,6 +355,10 @@ export function buildPersistedAssistantMessage( return normalized } + // A completed turn settles its stragglers as the live view did at `complete`. + if (result.success && message.contentBlocks) { + message.contentBlocks = settleUnfinishedToolCalls(message.contentBlocks, 'success') + } return message } @@ -372,7 +376,7 @@ export function isUnsettledToolState(state: string | undefined): boolean { /** Settles every unfinished tool row at a turn terminal so none reloads as a spinner. */ function settleUnfinishedToolCalls( blocks: PersistedContentBlock[], - state: 'cancelled' | 'error', + state: 'success' | 'cancelled' | 'error', display?: { title: string } ): PersistedContentBlock[] { return blocks.map((block) => From 177ac847eaeabc7dfa313c60460934d25ec0469e Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 21:20:27 -0700 Subject: [PATCH 09/18] fix(mothership): harden replay compaction, stream retries, teardown, and backend errors - Assistant text is never compacted: its length is part of the text receipt a replacement controller and the worker check. When cutting strings alone cannot bound an event, long arrays keep their first 100 items and a count. - Only a gateway error page (502/503/504 without a JSON body) or a failed connection counts as an unreachable worker for the two-minute window; a JSON 5xx from a reachable worker keeps the three-retry limit. A retry that streams again resets the budget for the next outage. - A controller cleans up its stream's buffer and abort marker only while it still holds the chat lock, checked just before cleanup and before releasing the lock, instead of inferring ownership from flags. - A controller superseded while publishing its terminal events leaves the run to its successor; other publish failures still settle the run. - A 4xx shows the worker's own short, plain reason (rate limit, BYOK, plan mode, model selection, version skew) and hides internal identifiers. - File preview snapshots are spaced by size, so a large file's patch preview can no longer fill the stream's replay budget; small files keep their pace. --- .../request/go/file-preview-adapter.test.ts | 67 ++++++++- .../request/go/file-preview-adapter.ts | 16 ++- .../lib/mothership/request/go/stream.test.ts | 44 ++++-- apps/sim/lib/mothership/request/go/stream.ts | 36 ++++- .../request/handlers/handlers.test.ts | 29 ---- .../mothership/request/lifecycle/finalize.ts | 11 +- .../mothership/request/lifecycle/run.test.ts | 88 ++++++++++++ .../lib/mothership/request/lifecycle/run.ts | 11 +- .../request/lifecycle/start.test.ts | 13 +- .../lib/mothership/request/lifecycle/start.ts | 24 ++-- .../request/lifecycle/stream-retry.test.ts | 65 +++++++-- .../request/lifecycle/stream-retry.ts | 33 ++++- .../request/session/controller-lease.ts | 10 ++ .../session/replay-budget.integration.ts | 131 ++++++++++++++---- .../request/session/replay-compaction.test.ts | 38 +++++ .../request/session/replay-compaction.ts | 92 ++++++++---- .../mothership/request/session/writer.test.ts | 1 - .../lib/mothership/request/session/writer.ts | 6 +- 18 files changed, 571 insertions(+), 144 deletions(-) diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts index af13587cad4..ba5c6e5e8c0 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts @@ -1,5 +1,5 @@ import { flushMicrotasks } from '@sim/testing/helpers/async' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { MothershipStreamV1EventType, MothershipStreamV1ToolExecutor, @@ -163,6 +163,71 @@ describe('processFilePreviewStreamEvent — preview content emission', () => { }) }) +describe('processFilePreviewStreamEvent — preview byte rate', () => { + const execContext: ExecutionContext = { + userId: 'user-1', + workflowId: 'workflow-1', + workspaceId: 'workspace-1', + chatId: 'chat-1', + messageId: 'msg-1', + } + + afterEach(() => vi.useRealTimers()) + + it('streams a 300 KB patch for 10 s well under the stream budget and ends on its final content', async () => { + vi.useFakeTimers() + const base = `ANCHOR\n${'line of existing file content\n'.repeat(10_000)}` + const edit = { strategy: 'anchored', mode: 'insert_after', anchor: 'ANCHOR' } + peekFileIntentMock.mockResolvedValue({ existingContent: base, edit }) + const state = createFilePreviewAdapterState() + const intent = { + ...makeIntent({ operation: 'patch', fileId: 'file-big', fileName: 'big.md' }), + edit, + } + const payloads: Array> = [] + const drive = async (streamEvent: StreamEvent) => { + const context = createStreamingContext() + context.activeFileIntents.set('', intent) + await processFilePreviewStreamEvent({ + streamId: STREAM_ID, + streamEvent, + context, + execContext, + options: { + onEvent: (event) => { + payloads.push((event as { payload: Record }).payload) + }, + }, + state, + }) + } + + let streamed = '' + await drive(editContentDelta('{"content":"')) + for (let tick = 0; tick < 625; tick++) { + vi.advanceTimersByTime(16) + streamed += `word${tick} ` + await drive(editContentDelta(`word${tick} `)) + } + await drive( + toolEvent({ + toolCallId: EDIT_TOOL_CALL_ID, + toolName: 'apply_file_edit', + phase: MothershipStreamV1ToolPhase.result, + success: true, + }) + ) + + const contents = payloads.filter((payload) => payload.previewPhase === 'file_preview_content') + const streamedBytes = contents.reduce( + (sum, payload) => sum + Buffer.byteLength(String(payload.content)), + 0 + ) + expect(streamedBytes).toBeLessThan(8 * 1024 * 1024) + expect(contents.at(-1)?.content).toBe(base.replace('ANCHOR\n', `ANCHOR\n${streamed}\n`)) + }) +}) + /** * The adapter runs while the tool-call frame is still on the wire, so the * execution context it receives is turn-scoped and carries no `toolCallId`. diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts index 925ebe4dbcc..df31af488af 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts @@ -56,6 +56,17 @@ type ParsedWorkspaceFileArgs = { const PATCH_PREVIEW_SNAPSHOT_INTERVAL_MS = 80 const DELTA_PREVIEW_CHECKPOINT_INTERVAL_MS = 1000 +/** + * Full snapshots of one file's preview are spaced so they stream at most about + * this many characters per second. Small files keep the base intervals; a large + * file's snapshots slow down instead of filling the stream's replay budget. + */ +const PREVIEW_SNAPSHOT_CHARS_PER_SECOND = 256 * 1024 + +/** The minimum gap between full snapshots of a preview this long. */ +function snapshotIntervalMs(baseMs: number, previewText: string): number { + return Math.max(baseMs, (previewText.length / PREVIEW_SNAPSHOT_CHARS_PER_SECOND) * 1000) +} function asJsonRecord(value: unknown): JsonRecord | undefined { return isRecordLike(value) ? (value as JsonRecord) : undefined @@ -326,7 +337,7 @@ export function buildPreviewContentUpdate( previousText.length === 0 || !nextText.startsWith(previousText) || operation === 'patch' || - now - lastSnapshotAt >= DELTA_PREVIEW_CHECKPOINT_INTERVAL_MS + now - lastSnapshotAt >= snapshotIntervalMs(DELTA_PREVIEW_CHECKPOINT_INTERVAL_MS, nextText) if (shouldForceSnapshot) { return { @@ -691,7 +702,8 @@ export async function processFilePreviewStreamEvent(input: { if ( nextSession.operation === 'patch' && - now - currentPreview.lastSnapshotAt < PATCH_PREVIEW_SNAPSHOT_INTERVAL_MS + now - currentPreview.lastSnapshotAt < + snapshotIntervalMs(PATCH_PREVIEW_SNAPSHOT_INTERVAL_MS, nextSession.previewText) ) { filePreviewState.set(editIntent.toolCallId, { session: nextSession, diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index 539e5c2dcbb..2a0719bf5a6 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -78,7 +78,6 @@ vi.mock('@/lib/mothership/tools/server/files/file-preview', async () => { } }) -import { buildPersistedAssistantMessage } from '@/lib/mothership/chat/persisted-message' import { buildPreviewContentUpdate, CopilotBackendError, @@ -261,7 +260,36 @@ describe('copilot go stream helpers', () => { 'The agent service is temporarily unavailable. Please try again.', ], [ - 'an internal JSON error', + "the worker's internal error", + 500, + 'application/json', + '{"error":"Internal error"}', + 'The agent service is temporarily unavailable. Please try again.', + ], + ['a rate limit', 429, 'application/json', '{"error":"Too many requests"}', 'Too many requests'], + [ + 'an enterprise-only request', + 403, + 'application/json', + '{"error":"Enterprise BYOK required"}', + 'Enterprise BYOK required', + ], + [ + 'a model selection problem', + 400, + 'application/json', + '{"error":"This workspace uses an Anthropic API key. Select Opus 5.5 to continue."}', + 'This workspace uses an Anthropic API key. Select Opus 5.5 to continue.', + ], + [ + 'protocol skew', + 426, + 'application/json', + '{"error":"protocol_version_mismatch","expected":3,"got":2,"message":"This Sim build speaks a different mothership protocol version. Update the older side."}', + 'This Sim build speaks a different mothership protocol version. Update the older side.', + ], + [ + 'an internal validation detail', 400, 'application/json', '{"error":"Bad Request","message":"userId required for internal API key"}', @@ -275,7 +303,7 @@ describe('copilot go stream helpers', () => { 'The agent service could not process this request.', ], ])( - 'never shows the user the raw body of %s', + 'tells the user about %s without the raw body', async (_label, status, contentType, body, userMessage) => { vi.mocked(fetch).mockResolvedValueOnce( new Response(body, { status, headers: { 'Content-Type': contentType } }) @@ -291,16 +319,6 @@ describe('copilot go stream helpers', () => { expect(error).toBeInstanceOf(CopilotBackendError) expect(error).toMatchObject({ message: userMessage, status, body }) - const persisted = buildPersistedAssistantMessage({ - success: false, - error: (error as Error).message, - content: '', - contentBlocks: [], - toolCalls: [], - }) - expect(JSON.stringify(persisted)).not.toMatch( - /= 500 ? BACKEND_UNAVAILABLE_MESSAGE : BACKEND_REJECTED_MESSAGE +function userFacingRejection(value: unknown): string | undefined { + if (typeof value !== 'string') return undefined + const message = value.trim() + if (!message || message.length > 200 || /[<\n]/.test(message)) return undefined + if (/\b\w*[a-z][A-Z]\w*\b|\b\w+_\w+\b/.test(message)) return undefined + return message +} + +/** + * What the user is told about a failed backend response. A 5xx or a gateway page + * is upstream detail and stays on the error for logs; a 4xx may carry the + * worker's own reason, which is shown when it is safe to. + */ +function backendErrorMessage(status: number, body: string): string { + if (status >= 500) return BACKEND_UNAVAILABLE_MESSAGE + let parsed: unknown + try { + parsed = JSON.parse(body) + } catch { + return BACKEND_REJECTED_MESSAGE + } + // The worker puts its reason in `error`, or a code there and the reason in `message`. + const record = toRecordOrNull(parsed) + const reason = record && 'message' in record ? record.message : record?.error + return userFacingRejection(reason) ?? BACKEND_REJECTED_MESSAGE } export class BillingLimitError extends Error { @@ -188,7 +212,7 @@ export async function runStreamLoop( throw new BillingLimitError(execContext.userId) } - throw new CopilotBackendError(backendErrorMessage(response.status), { + throw new CopilotBackendError(backendErrorMessage(response.status, errorText), { status: response.status, body: errorText || response.statusText, }) diff --git a/apps/sim/lib/mothership/request/handlers/handlers.test.ts b/apps/sim/lib/mothership/request/handlers/handlers.test.ts index 290b818306b..88a75b6e976 100644 --- a/apps/sim/lib/mothership/request/handlers/handlers.test.ts +++ b/apps/sim/lib/mothership/request/handlers/handlers.test.ts @@ -407,35 +407,6 @@ describe('sse-handlers tool lifecycle', () => { expect(upsertAsyncToolCall).not.toHaveBeenCalled() }) - it('clears a Go-stamped approval frame on a replayed call, which never gates anything', async () => { - toolRequiresApproval.mockReturnValue(false) - context.runId = 'run-1' - context.toolPermissions = { - enabled: false, - autoAllowed: new Set(), - autoAllowPermitted: true, - } - - const event = { - type: MothershipStreamV1EventType.tool, - payload: { - toolCallId: 'gmail-2', - toolName: 'gmail_read_v2', - arguments: {}, - executor: MothershipStreamV1ToolExecutor.sim, - mode: MothershipStreamV1ToolMode.async, - phase: MothershipStreamV1ToolPhase.call, - status: 'awaiting_approval', - replay: true, - }, - } as unknown as StreamEvent - - await prePersistClientExecutableToolCall(event, context, {}) - - expect((event.payload as { status?: string }).status).toBeUndefined() - expect(upsertAsyncToolCall).not.toHaveBeenCalled() - }) - it('clears a Go-stamped approval frame on a partial call, which is never held', async () => { toolRequiresApproval.mockReturnValue(true) context.runId = 'run-1' diff --git a/apps/sim/lib/mothership/request/lifecycle/finalize.ts b/apps/sim/lib/mothership/request/lifecycle/finalize.ts index e77838a66a1..2bbab2c5167 100644 --- a/apps/sim/lib/mothership/request/lifecycle/finalize.ts +++ b/apps/sim/lib/mothership/request/lifecycle/finalize.ts @@ -14,6 +14,7 @@ import { CopilotFinalizeOutcome } from '@/lib/mothership/generated/trace-attribu import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1' import type { StreamWriter } from '@/lib/mothership/request/session' +import { StreamControllerSupersededError } from '@/lib/mothership/request/session/controller-lease' import type { OrchestratorResult } from '@/lib/mothership/request/types' const logger = createLogger('CopilotStreamFinalize') @@ -204,8 +205,8 @@ async function handleSuccess( /** * Publishes the terminal events, then records the run's terminal status even if - * publishing failed: nothing else settles a run, and the status write is scoped - * to this controller's token, so it never settles a run a successor has claimed. + * publishing failed, since nothing else settles a run. A controller superseded + * while publishing leaves the run to its successor instead. */ async function publishThenSettle( publisher: StreamWriter, @@ -215,9 +216,11 @@ async function publishThenSettle( try { await publish() await publisher.flush() - } finally { - await settle() + } catch (error) { + if (!(error instanceof StreamControllerSupersededError)) await settle() + throw error } + await settle() } async function loggedRunStatusUpdate( diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 498b90df99a..f22d22d4054 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -196,6 +196,7 @@ vi.mock('@/lib/mothership/request/enterprise-byok', () => ({ resolveEnterpriseByokKey: mockResolveEnterpriseByokKey, })) +import { buildPersistedAssistantMessage } from '@/lib/mothership/chat/persisted-message' import { MothershipStreamV1CompletionStatus, MothershipStreamV1ToolOutcome, @@ -2877,6 +2878,93 @@ describe('runCopilotLifecycle', () => { ) }) + it('persists a failed backend response without its upstream body', async () => { + const body = '502 Bad Gatewaynginx' + mockRunStreamLoop.mockRejectedValueOnce( + Object.assign( + new CopilotBackendError('The agent service is temporarily unavailable. Please try again.', { + status: 400, + }), + { body } + ) + ) + let persisted: unknown + await runCopilotLifecycle( + { message: 'hello', messageId: 'stream-backend-error' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'exec-1', + runId: 'run-1', + executionContext: { + userId: 'user-1', + workflowId: '', + workspaceId: 'ws-1', + chatId: 'chat-1', + }, + onError: async (_error, result) => { + persisted = result && buildPersistedAssistantMessage(result) + }, + } + ) + + expect(JSON.stringify(persisted)).toContain('The agent service is temporarily unavailable.') + expect(JSON.stringify(persisted)).not.toMatch(/ { + vi.useFakeTimers() + try { + let attempts = 0 + mockRunStreamLoop.mockImplementation( + async ( + _url: string, + _init: RequestInit, + context: StreamingContext, + _exec: ExecutionContext, + options: { onEvent?: (event: unknown) => Promise } + ): Promise => { + attempts++ + if (attempts < 6) { + await options.onEvent?.({ + type: 'text', + payload: { channel: 'assistant', text: `part ${attempts} ` }, + }) + context.errors.push(STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE) + throw new StreamEndedWithoutTerminalError('/api/mothership') + } + context.streamComplete = true + context.completionStatus = MothershipStreamV1CompletionStatus.complete + } + ) + + const pending = runCopilotLifecycle( + { message: 'hello', messageId: 'stream-repeated-outages' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'exec-1', + runId: 'run-1', + executionContext: { + userId: 'user-1', + workflowId: '', + workspaceId: 'ws-1', + chatId: 'chat-1', + }, + } + ) + await vi.advanceTimersByTimeAsync(60_000) + + expect(await pending).toEqual(expect.objectContaining({ success: true })) + expect(attempts).toBe(6) + } finally { + mockRunStreamLoop.mockReset() + vi.useRealTimers() + } + }) + describe('when the worker task is being replaced', () => { const replacementMs = 70_000 diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index b4a6b6f9765..550c21a1e8e 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -845,6 +845,14 @@ async function runResumeLegWithRetry( hostedBillingRequest?: AttributedBillingRequestEnvelope ): Promise { const retry = new StreamRetryWindow(options.timeout) + /** A leg that streams again has recovered; a later outage gets its own budget. */ + const legOptions: CopilotLifecycleOptions = { + ...options, + onEvent: async (event) => { + retry.recovered() + await options.onEvent?.(event) + }, + } for (;;) { options.abortSignal?.throwIfAborted() const errorsBeforeAttempt = leg.errors.length @@ -859,7 +867,7 @@ async function runResumeLegWithRetry( }, leg, execContext, - { ...options, timeout: retry.remainingMs() } + { ...legOptions, timeout: retry.remainingMs() } ) return } catch (error) { @@ -1137,6 +1145,7 @@ async function runCheckpointLoop( has an HTTP buffer worth a per-event macrotask flush. */ flushAfterEvent: options.flushAfterEvent ?? Boolean(callerOnEvent), onEvent: async (event: StreamEvent) => { + retry?.recovered() if ( event.type === MothershipStreamV1EventType.run && event.payload.kind === MothershipStreamV1RunKind.checkpoint_pause && diff --git a/apps/sim/lib/mothership/request/lifecycle/start.test.ts b/apps/sim/lib/mothership/request/lifecycle/start.test.ts index 57a54b64b0a..dcc2c6937d3 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.test.ts @@ -80,6 +80,7 @@ vi.mock('@/lib/mothership/request/session/abort', () => ({ vi.mock('@/lib/mothership/request/session/controller-lease', async (original) => ({ ...(await original()), assertChatStreamLease: vi.fn().mockResolvedValue(undefined), + holdsChatStreamLease: vi.fn().mockResolvedValue(true), })) vi.mock('@/lib/billing/core/billing-attribution', () => billingAttributionMock) @@ -468,12 +469,14 @@ describe('createSSEStream terminal error handling', () => { ) expect(appendEvent).not.toHaveBeenCalledWith(expect.objectContaining({ type: 'error' })) expect(unregisterActiveStream).toHaveBeenCalledWith('stream-1', expect.any(AbortController)) - expect(releasePendingChatStream).toHaveBeenCalledWith( - '11111111-1111-4111-8111-111111111111', - 'stream-1', - expect.objectContaining({ value: 'stream-1\ncontroller' }) + await vi.waitFor(() => + expect(releasePendingChatStream).toHaveBeenCalledWith( + '11111111-1111-4111-8111-111111111111', + 'stream-1', + expect.objectContaining({ value: 'stream-1\ncontroller' }) + ) ) - await vi.waitFor(() => expect(scheduleBufferCleanup).toHaveBeenCalledWith('stream-1')) + expect(scheduleBufferCleanup).toHaveBeenCalledWith('stream-1') }) it('names an untitled chat on the next accepted turn after an initial Stop, then leaves its title alone', async () => { diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 358ea64cc01..5bfe2039526 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -57,6 +57,7 @@ import { getLocalChatStreamLease } from '@/lib/mothership/request/session/abort' import { AbortReason } from '@/lib/mothership/request/session/abort-reason' import { assertChatStreamLease, + holdsChatStreamLease, StreamControllerSupersededError, } from '@/lib/mothership/request/session/controller-lease' import { @@ -208,7 +209,6 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS } const collector = new TraceCollector() - let handedOff = false /** * The replay refusal that ends this turn. Once the controller is aborted its @@ -453,7 +453,6 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS abortController.signal.reason instanceof StreamControllerSupersededError ) { logger.info('Stream controller handed off; leaving its run recoverable', { streamId }) - handedOff = true return } await assertControllerOwnership() @@ -506,20 +505,21 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS }) } unregisterActiveStream(streamId, abortController) - if (chatId) { - await releasePendingChatStream(chatId, streamId, lease) - } - processResourcesReleased = true - // A superseded controller must not expire or clear its successor's stream, - // whether it handed off or lost ownership while handling another failure. - if ( - !handedOff && - !(abortController.signal.reason instanceof StreamControllerSupersededError) - ) { + /* + The stream's buffer and abort marker belong to whoever holds the chat + lock now. Clean them up only while this controller still holds it, and + before releasing it, so a successor's stream is never expired and its + Stop is never cleared. + */ + if (!chatId || (lease && (await holdsChatStreamLease(lease)))) { await scheduleBufferCleanup(streamId) await scheduleFilePreviewSessionCleanup(streamId) await cleanupAbortMarker(streamId) } + if (chatId) { + await releasePendingChatStream(chatId, streamId, lease) + } + processResourcesReleased = true rootOutcome = outcome if (lifecycleResult?.usage) { diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index 94a5a795e57..c557edd270b 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -22,23 +22,60 @@ describe('stream recovery budget', () => { expect(retry.remainingMs()).toBeGreaterThan(3_500_000) }) - it.each([new TypeError('fetch failed'), new CopilotBackendError('Unavailable', { status: 502 })])( - 'keeps retrying an unreachable worker for two minutes from the first failure: %s', - (error) => { - vi.useFakeTimers() - const retry = new StreamRetryWindow() - vi.advanceTimersByTime(600_000) - const firstFailure = Date.now() - for (;;) { + it.each([ + new CopilotBackendError('Unavailable', { status: 500, body: '{"error":"Internal error"}' }), + new CopilotBackendError('Unavailable', { + status: 503, + body: '{"error":"Account admission is unavailable"}', + }), + ])('gives a reachable worker that answers with a 5xx only three retries: %s', (error) => { + vi.useFakeTimers() + const retry = new StreamRetryWindow() + for (let index = 0; index < 3; index++) { + const delay = retry.nextDelay(error) + expect(delay).not.toBeNull() + vi.advanceTimersByTime(delay ?? 0) + } + expect(retry.nextDelay(error)).toBeNull() + }) + + it('starts a fresh recovery budget once a retry reconnects', () => { + vi.useFakeTimers() + const error = new StreamEndedWithoutTerminalError('/api/mothership') + const retry = new StreamRetryWindow() + for (let outage = 0; outage < 2; outage++) { + for (let index = 0; index < 3; index++) { const delay = retry.nextDelay(error) - if (delay === null) break - vi.advanceTimersByTime(delay) + expect(delay).not.toBeNull() + vi.advanceTimersByTime(delay ?? 0) } - expect(Date.now() - firstFailure).toBeGreaterThan(110_000) - expect(Date.now() - firstFailure).toBeLessThanOrEqual(120_000) - expect(retry.remainingMs()).toBeGreaterThan(2_800_000) + retry.recovered() + vi.advanceTimersByTime(60_000) + } + expect(retry.nextDelay(error)).not.toBeNull() + }) + + it.each([ + new TypeError('fetch failed'), + new CopilotBackendError('Unavailable', { status: 502 }), + new CopilotBackendError('Unavailable', { + status: 504, + body: '

504 Gateway Time-out

', + }), + ])('keeps retrying an unreachable worker for two minutes from the first failure: %s', (error) => { + vi.useFakeTimers() + const retry = new StreamRetryWindow() + vi.advanceTimersByTime(600_000) + const firstFailure = Date.now() + for (;;) { + const delay = retry.nextDelay(error) + if (delay === null) break + vi.advanceTimersByTime(delay) } - ) + expect(Date.now() - firstFailure).toBeGreaterThan(110_000) + expect(Date.now() - firstFailure).toBeLessThanOrEqual(120_000) + expect(retry.remainingMs()).toBeGreaterThan(2_800_000) + }) it('never extends the original execution deadline', () => { vi.useFakeTimers() diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts index 87101978c5b..e9235b6c11f 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -8,8 +8,10 @@ import { const MAX_STREAM_RETRIES = 3 const STREAM_RECOVERY_WINDOW_MS = 30_000 +/** The load balancer's answers while no worker task is registered behind it. */ +const GATEWAY_STATUSES: ReadonlySet = new Set([502, 503, 504]) /** - * While the worker cannot be reached at all (a 5xx from the load balancer, or no + * While the worker cannot be reached at all (a gateway error page, or no * connection), retries continue for this long from the first failure: long * enough to outlast a worker task replacement (about 70 s of 502/504). Every * attempt re-sends the same message identity, which the worker treats as a @@ -34,6 +36,12 @@ export class StreamRetryWindow { return remaining } + /** A retry reconnected and made progress; a later outage gets a fresh budget. */ + recovered(): void { + this.attempt = 0 + this.firstFailureAt = undefined + } + nextDelay(error: unknown, signal?: AbortSignal): number | null { if (signal?.aborted || !isRetryableStreamError(error)) return null this.firstFailureAt ??= Date.now() @@ -48,19 +56,36 @@ export class StreamRetryWindow { } } -/** No worker answered: the load balancer's 5xx, or no connection at all. */ +/** + * No worker answered: a gateway error page from the load balancer, or no + * connection at all. A JSON 5xx comes from a reachable worker, so it gets the + * short recovery budget instead. + */ function isWorkerUnreachable(error: unknown): boolean { if (error instanceof CopilotBackendError) { - return error.status !== undefined && error.status >= 500 + return error.status !== undefined && GATEWAY_STATUSES.has(error.status) && !isJson(error.body) } return error instanceof TypeError } +function isJson(body: string | undefined): boolean { + if (!body) return false + try { + JSON.parse(body) + return true + } catch { + return false + } +} + /** Initial sends and resumes both replay one durable identity after an ambiguous response. */ function isRetryableStreamError(error: unknown): boolean { if (error instanceof Error && error.name === 'AbortError') return false if (error instanceof StreamEndedWithoutTerminalError || error instanceof StreamContinuityError) { return true } - return isWorkerUnreachable(error) + if (error instanceof CopilotBackendError) { + return error.status !== undefined && error.status >= 500 + } + return error instanceof TypeError } diff --git a/apps/sim/lib/mothership/request/session/controller-lease.ts b/apps/sim/lib/mothership/request/session/controller-lease.ts index 109b2b2a58d..b7d89647a33 100644 --- a/apps/sim/lib/mothership/request/session/controller-lease.ts +++ b/apps/sim/lib/mothership/request/session/controller-lease.ts @@ -27,3 +27,13 @@ export async function assertChatStreamLease(lease: ChatStreamLease): Promise { + try { + await assertChatStreamLease(lease) + return true + } catch { + return false + } +} 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 d08c2cdceeb..be71ec56412 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -13,10 +13,14 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { const { readTestRedisUrl } = await import('@sim/db/testing/test-infrastructure') const { createServer: createHttpServer } = await import('node:http') const abortRequests: Array> = [] + const hooks = { onAbort: undefined as (() => Promise) | undefined } const server = createHttpServer(async (request, response) => { let body = '' for await (const chunk of request) body += chunk - if (request.url === '/api/streams/explicit-abort') abortRequests.push(JSON.parse(body)) + if (request.url === '/api/streams/explicit-abort') { + abortRequests.push(JSON.parse(body)) + await hooks.onAbort?.() + } response.writeHead(200, { 'content-type': 'application/json' }) response.end(JSON.stringify({ settled: true })) }) @@ -36,7 +40,7 @@ const { redisUrl, inheritedEnv, worker } = await vi.hoisted(async () => { worker: { server, abortRequests, - /** Events the scripted worker streams through the controller's sink, in order. */ + hooks, /** Events, or steps to run between them, that the scripted worker streams in order. */ script: [] as unknown[], /** Controller lifecycles started, and what each sink call threw. */ @@ -80,14 +84,17 @@ vi.mock('@/lib/mothership/request/lifecycle/run', () => ({ }, })) +import { trace } from '@opentelemetry/api' +import { BasicTracerProvider } from '@opentelemetry/sdk-trace-base' import { db } from '@sim/db' import { copilotChats, copilotRuns, permissions, user, workspace } from '@sim/db/schema' -import { sleep } from '@sim/utils/helpers' import { generateId } from '@sim/utils/id' import { eq } from 'drizzle-orm' import { NextRequest } from 'next/server' import { closeRedisConnection, getRedisClient } from '@/lib/core/config/redis' import { getRedisBudgetKeys, getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' +import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' +import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1' import { readChatStream } from '@/lib/mothership/request/application/recover-stream' import { createStreamingContext } from '@/lib/mothership/request/context/request-context' import { restoreStreamingContext } from '@/lib/mothership/request/context/restore' @@ -257,7 +264,6 @@ describe.runIf(Boolean(redisUrl))('leased Chat stream writer with Redis', () => stdout: `${'r'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 1.6 MB total]`, }, }) - expect(writer.persistenceStopped).toBe(false) }) it('restores the same tool calls, text, and activity from a compacted buffer as from the full events, without dispatching', async () => { @@ -379,8 +385,33 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge }) /** Admits a turn the way the chat POST does, runs its controller, and returns its frames. */ + /** Resolves when a controller's root span ends, the last step of its teardown. */ + const teardowns = new Map void>() + beforeAll(() => { + trace.disable() + trace.setGlobalTracerProvider( + new BasicTracerProvider({ + spanProcessors: [ + { + onStart: () => {}, + onEnd: (span) => { + const streamId = span.attributes[TraceAttr.StreamId] + if (span.name === TraceSpan.GenAiAgentExecute && typeof streamId === 'string') { + teardowns.get(streamId)?.() + } + }, + forceFlush: async () => {}, + shutdown: async () => {}, + }, + ], + }) + ) + }) + afterAll(() => trace.disable()) + async function runTurn(script: unknown[], prepare?: (streamId: string) => Promise) { const streamId = generateId() + const teardown = new Promise((resolve) => teardowns.set(streamId, resolve)) const runId = generateId() expect(await acquirePendingChatStream(chatId, streamId, 0)).toBe(true) const controllerToken = (await redis().get(chatStreamLockKey(chatId)))! @@ -431,7 +462,9 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge admittedRun: run, orchestrateOptions: { userId, workspaceId, chatId, runId, interactive: true }, }) - return { streamId, runId, frames: dataFrames(await new Response(response).text()) } + const frames = dataFrames(await new Response(response).text()) + await teardown + return { streamId, runId, frames } } it('ends as an error, marks its run terminal, stops the worker, and starts no recovery controller', async () => { @@ -509,8 +542,6 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge ]) expect(frames.map((frame) => frame.type)).toEqual(['session']) - /** The client stream closes before the controller's teardown; let teardown finish. */ - await sleep(500) expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeGreaterThan(300) expect(await redis().get(chatStreamLockKey(chatId))).toBe(successorToken) const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) @@ -520,10 +551,33 @@ 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("leaves its successor's stream untouched when the lease is lost while ending a refused turn", async () => { + const successorToken = `successor\n${generateId()}` + worker.hooks.onAbort = async () => { + await redis().set(chatStreamLockKey(chatId), successorToken, 'EX', 60) + } + try { + const { streamId, frames } = await runTurn( + [toolCall('call-refused', 'cli_blocks_get', { command: `blocks get ${'b'.repeat(1024)}` })], + async (streamId) => { + const { maxOwnerBytes } = getRedisBudgetLimits('copilot_stream') + const [ownerBudgetKey] = getRedisBudgetKeys({ kind: 'copilot_stream', id: streamId }) + await redis().set(ownerBudgetKey, String(maxOwnerBytes - 512), 'EX', 3600) + } + ) + + expect(frames.map((frame) => frame.type)).toEqual(['session', 'error', 'complete']) + expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeGreaterThan(300) + expect(await redis().get(chatStreamLockKey(chatId))).toBe(successorToken) + } finally { + worker.hooks.onAbort = undefined + await redis().del(chatStreamLockKey(chatId)) + } + }) + + async function pausedRun(controllerToken: string) { const streamId = generateId() const runId = generateId() - const controllerToken = `${streamId}\n${generateId()}` await db.insert(copilotRuns).values({ id: runId, executionId: generateId(), @@ -534,28 +588,55 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge status: 'paused_waiting_for_tool', requestContext: { requestId: generateId(), controllerToken }, }) - /** No chat lock holds this lease, so publishing the terminal events throws. */ + return { streamId, runId } + } + + async function finalizeAsError(publisher: StreamWriter, runId: string) { + return finalizeStream( + { + success: false, + error: 'The agent service is temporarily unavailable. Please try again.', + content: '', + contentBlocks: [], + toolCalls: [], + }, + publisher, + runId, + 'error', + generateId() + ) + } + + 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 } + await redis().set(lease.key, lease.value, 'EX', 60) + /** A corrupt buffer key makes every append fail while the lease is still held. */ + await redis().set(`mothership_stream:${streamId}:events`, 'not a sorted set', 'EX', 60) const publisher = new StreamWriter({ streamId, requestId: generateId(), lease }) - await expect( - finalizeStream( - { - success: false, - error: 'The agent service is temporarily unavailable. Please try again.', - content: '', - contentBlocks: [], - toolCalls: [], - }, - publisher, - runId, - 'error', - generateId() - ) - ).rejects.toBeInstanceOf(StreamControllerSupersededError) + 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) + }) + + it('leaves its run for a successor when it lost the lease while publishing', async () => { + const controllerToken = `stale\n${generateId()}` + const { streamId, runId } = await pausedRun(controllerToken) + const lease = { key: chatStreamLockKey(generateId()), value: controllerToken } + const publisher = new StreamWriter({ streamId, requestId: generateId(), lease }) + + await expect(finalizeAsError(publisher, runId)).rejects.toBeInstanceOf( + StreamControllerSupersededError + ) + + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('paused_waiting_for_tool') }) it('does not settle a run another controller has claimed', async () => { diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index 2ea67e71bb0..89021177a89 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -212,4 +212,42 @@ describe('compactStreamEvent', () => { expect(cut.startsWith(`${'a'.repeat(STREAM_STRING_PREVIEW_UNITS - 1)}…`)).toBe(true) }) + + it('never cuts assistant text, whose length is part of the text receipt', () => { + const event: StreamEvent = { + type: 'text', + payload: { channel: 'assistant', text: 'a'.repeat(MB), textOffset: 0 }, + } + + expect(compactStreamEvent(event)).toBe(event) + }) + + it('keeps the head of a long array of short items when strings alone cannot bound it', () => { + const rows = Array.from({ length: 40_000 }, (_, index) => ({ id: index, name: 'row' })) + const citations = [{ index: 1, title: 'Runbook', url: 'https://docs.example/runbook' }] + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_tables_rows_query', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + output: { rows, data: { results: citations } }, + }, + } + + const compacted = compactStreamEvent(event) + const output = toRecord(payloadOf(compacted).output) + const kept = output.rows as unknown[] + + expect(Buffer.byteLength(JSON.stringify(compacted.payload))).toBeLessThan( + STREAM_EVENT_COMPACTION_THRESHOLD_BYTES + ) + expect(kept.slice(0, 2)).toEqual([rows[0], rows[1]]) + expect(kept.at(-1)).toMatch(/^…\[truncated, \d+ more items\]$/) + expect(output.data).toEqual({ results: citations }) + expect(toRecord(payloadOf(event).output).rows).toHaveLength(40_000) + }) }) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index 3bf1d03fb29..10e699802b4 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -1,5 +1,8 @@ import { isRecordLike, toRecordOrNull } from '@sim/utils/object' -import { MothershipStreamV1ToolPhase } from '@/lib/mothership/generated/mothership-stream-v1' +import { + MothershipStreamV1EventType, + MothershipStreamV1ToolPhase, +} from '@/lib/mothership/generated/mothership-stream-v1' import { isClientExecutedToolCall } from '@/lib/mothership/tools/client-executed-tools' import { formatFileSize } from '@/lib/uploads/utils/file-utils' import type { StreamEvent } from './types' @@ -19,12 +22,17 @@ const ARGUMENTS_KEY: ReadonlySet = new Set(['arguments']) /** Its live result is the only place the plaintext key reaches the browser. */ const GENERATE_API_KEY_TOOL = 'generate_api_key' -function stringUnits(value: unknown): number { - if (typeof value === 'string') return value.length - if (Array.isArray(value)) return value.reduce((sum, item) => sum + stringUnits(item), 0) - if (!isRecordLike(value)) return 0 - let sum = 0 - for (const key in value) sum += key.length + stringUnits(value[key]) +/** Items kept at the head of an array that string cuts alone could not bound. */ +export const STREAM_ARRAY_HEAD_ITEMS = 100 + +/** An upper bound on a value's serialized size, without serializing it. */ +function estimateBytes(value: unknown): number { + if (typeof value === 'string') return value.length * 3 + 2 + if (Array.isArray(value)) + return value.reduce((sum, item) => sum + estimateBytes(item) + 1, 2) + if (!isRecordLike(value)) return 24 + let sum = 2 + for (const key in value) sum += key.length * 3 + 4 + estimateBytes(value[key]) return sum } @@ -35,27 +43,56 @@ function stringUnits(value: unknown): number { * named top-level keys is exempt. */ function truncateStrings(value: unknown, skipKeys: ReadonlySet = NO_KEYS): unknown { - if (typeof value === 'string') { - if (value.length <= STREAM_STRING_PREVIEW_UNITS) return value + return mapLeaves(value, skipKeys, (leaf) => { + if (typeof leaf !== 'string' || leaf.length <= STREAM_STRING_PREVIEW_UNITS) return leaf const end = STREAM_STRING_PREVIEW_UNITS - const lastUnit = value.charCodeAt(end - 1) + const lastUnit = leaf.charCodeAt(end - 1) const cut = lastUnit >= 0xd800 && lastUnit <= 0xdbff ? end - 1 : end - const size = formatFileSize(Buffer.byteLength(value, 'utf8')) - return `${value.slice(0, cut)}…[truncated, ${size} total]` - } + const size = formatFileSize(Buffer.byteLength(leaf, 'utf8')) + return `${leaf.slice(0, cut)}…[truncated, ${size} total]` + }) +} + +/** Keeps the head of every long array, with a note of how many items were dropped. */ +function trimArrays(value: unknown, skipKeys: ReadonlySet): unknown { + return mapLeaves( + value, + skipKeys, + (leaf) => leaf, + (items) => + items.length <= STREAM_ARRAY_HEAD_ITEMS + ? items + : [ + ...items.slice(0, STREAM_ARRAY_HEAD_ITEMS), + `…[truncated, ${items.length - STREAM_ARRAY_HEAD_ITEMS} more items]`, + ] + ) +} + +/** + * Rebuilds a value with `leaf` applied to every non-container and `array` to + * every array, copying only what changes. `skipKeys` exempts top-level fields. + */ +function mapLeaves( + value: unknown, + skipKeys: ReadonlySet, + leaf: (value: unknown) => unknown, + array: (items: unknown[]) => unknown[] = (items) => items +): unknown { if (Array.isArray(value)) { - let copy: unknown[] | undefined - value.forEach((item, index) => { - const next = truncateStrings(item) - if (next !== item) (copy ??= [...value])[index] = next + const items = array(value) + let copy: unknown[] | undefined = items === value ? undefined : [...items] + items.forEach((item, index) => { + const next = mapLeaves(item, NO_KEYS, leaf, array) + if (next !== item) (copy ??= [...items])[index] = next }) return copy ?? value } - if (!isRecordLike(value)) return value + if (!isRecordLike(value)) return leaf(value) let copy: Record | undefined for (const [key, field] of Object.entries(value)) { if (skipKeys.has(key)) continue - const next = truncateStrings(field) + const next = mapLeaves(field, NO_KEYS, leaf, array) if (next !== field) (copy ??= { ...value })[key] = next } return copy ?? value @@ -65,14 +102,18 @@ function truncateStrings(value: unknown, skipKeys: ReadonlySet = NO_KEYS * Bounds an outgoing stream event so the replay buffer can persist it. Applied * only to the copy the writer delivers and persists; the caller keeps the full * event for dispatch. Long strings are cut to their head in place, so every - * object keeps its shape. File previews, `generate_api_key` results, and the - * arguments of calls the browser executes are never cut; an event + * object keeps its shape; if that is not enough, long arrays keep their head. + * Assistant text, file previews, `generate_api_key` results, and the arguments + * of calls the browser executes are never cut; an event * still too large is refused by the buffer, which ends the turn with an error. */ export function compactStreamEvent(event: StreamEvent): StreamEvent { const payload = toRecordOrNull(event.payload) - if (!payload || 'previewPhase' in payload) return event - if (stringUnits(payload) * 3 <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) return event + // Text length is part of the receipt the worker and a replacement check. + if (!payload || event.type === MothershipStreamV1EventType.text || 'previewPhase' in payload) { + return event + } + if (estimateBytes(payload) <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) return event const toolName = typeof payload.toolName === 'string' ? payload.toolName : '' if (payload.phase === MothershipStreamV1ToolPhase.result && toolName === GENERATE_API_KEY_TOOL) { return event @@ -82,6 +123,9 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutedToolCall(toolName, args) ? ARGUMENTS_KEY : NO_KEYS - const compacted = truncateStrings(payload, skipKeys) + let compacted = truncateStrings(payload, skipKeys) + if (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { + compacted = trimArrays(compacted, skipKeys) + } return compacted === payload ? event : ({ ...event, payload: compacted } as StreamEvent) } diff --git a/apps/sim/lib/mothership/request/session/writer.test.ts b/apps/sim/lib/mothership/request/session/writer.test.ts index 6da1ac447c8..1eab0e0a636 100644 --- a/apps/sim/lib/mothership/request/session/writer.test.ts +++ b/apps/sim/lib/mothership/request/session/writer.test.ts @@ -112,7 +112,6 @@ describe('StreamWriter', () => { writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'later' } }) ).rejects.toBeInstanceOf(StreamReplayBudgetExhaustedError) expect(delivered).toEqual([]) - expect(writer.persistenceStopped).toBe(true) await expect(writer.flush()).resolves.toBeUndefined() await writer.publish({ diff --git a/apps/sim/lib/mothership/request/session/writer.ts b/apps/sim/lib/mothership/request/session/writer.ts index a06fac713e9..400d3f0dfce 100644 --- a/apps/sim/lib/mothership/request/session/writer.ts +++ b/apps/sim/lib/mothership/request/session/writer.ts @@ -92,11 +92,10 @@ export class StreamWriter { /** * The replay buffer stopped accepting writes because this stream exhausted its byte - * budget. Leased delivery stops before the refused event (only the terminal verdict - * still reaches the client); unleased delivery continues. + * budget. Leased delivery stops before the refused event; unleased delivery continues. */ get persistenceStopped(): boolean { - return this._persistenceStopped || this.replayBudgetError !== null + return this._persistenceStopped } updateRequestId(id: string): void { @@ -155,6 +154,7 @@ export class StreamWriter { lease ) if (!result.persisted) { + this._persistenceStopped = true this.replayBudgetError = new StreamReplayBudgetExhaustedError(result.refusal) } } From 89977d00d96ed508657d71e17fcfeb08560e5517 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 21:38:47 -0700 Subject: [PATCH 10/18] fix(mothership): separate retry budgets, cap preview totals, and clean up only ended turns - An unreachable worker and a reachable one now have independent retry budgets. Any answer from the worker restarts only the two-minute unreachable window; the three-retry, 30 s budget for failures of a worker that answered keeps its original semantics, so a failure that repeats after every reattach still stops, and a long outage no longer spends the replacement's retries. - Intermediate file preview content is capped at 8M characters per edit; past it the preview holds until the final snapshot, which is always sent. - A controller cleans up its stream only when it finalized the turn and still holds the lock, so a run it leaves recoverable keeps its buffer and any pending Stop. - Remove the unreachable generate_api_key compaction exemption, avoid a second copy of trimmed arrays, and correct the size-estimate comment. --- .../lib/mothership/chat/display-message.ts | 2 +- .../request/go/file-preview-adapter.test.ts | 40 ++++++++--- .../request/go/file-preview-adapter.ts | 66 +++++++++++-------- .../lib/mothership/request/handlers/tool.ts | 6 +- .../mothership/request/lifecycle/run.test.ts | 21 ++---- .../lib/mothership/request/lifecycle/start.ts | 23 ++++--- .../request/lifecycle/stream-retry.test.ts | 48 ++++++++++++-- .../request/lifecycle/stream-retry.ts | 41 +++++++++--- .../session/replay-budget.integration.ts | 22 +++++++ .../request/session/replay-compaction.test.ts | 15 +---- .../request/session/replay-compaction.ts | 19 +++--- 11 files changed, 199 insertions(+), 104 deletions(-) diff --git a/apps/sim/lib/mothership/chat/display-message.ts b/apps/sim/lib/mothership/chat/display-message.ts index 58155f701e8..56607544997 100644 --- a/apps/sim/lib/mothership/chat/display-message.ts +++ b/apps/sim/lib/mothership/chat/display-message.ts @@ -40,7 +40,7 @@ function toToolCallInfo(block: PersistedContentBlock, stored: boolean): ToolCall const tc = block.toolCall if (!tc) return undefined if (isToolHiddenInUi(tc.name)) return undefined - /** A stored turn has ended; a row it left unfinished did not finish. */ + // A stored turn has ended, so a row it left unfinished did not finish. const status: ToolCallStatus = stored && isUnsettledToolState(tc.state) ? ToolCallStatus.interrupted diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts index ba5c6e5e8c0..68a4ba244ba 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts @@ -25,6 +25,7 @@ import { createStreamingContext } from '@/lib/mothership/request/context/request import { createFilePreviewAdapterState, type FilePreviewAdapterState, + PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS, processFilePreviewStreamEvent, } from '@/lib/mothership/request/go/file-preview-adapter' import { createEvent, eventToStreamEvent } from '@/lib/mothership/request/session' @@ -174,7 +175,7 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { afterEach(() => vi.useRealTimers()) - it('streams a 300 KB patch for 10 s well under the stream budget and ends on its final content', async () => { + async function streamPatch(durationMs: number, tickMs: number) { vi.useFakeTimers() const base = `ANCHOR\n${'line of existing file content\n'.repeat(10_000)}` const edit = { strategy: 'anchored', mode: 'insert_after', anchor: 'ANCHOR' } @@ -204,10 +205,10 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { let streamed = '' await drive(editContentDelta('{"content":"')) - for (let tick = 0; tick < 625; tick++) { - vi.advanceTimersByTime(16) - streamed += `word${tick} ` - await drive(editContentDelta(`word${tick} `)) + for (let elapsed = 0; elapsed < durationMs; elapsed += tickMs) { + vi.advanceTimersByTime(tickMs) + streamed += `word${elapsed} ` + await drive(editContentDelta(`word${elapsed} `)) } await drive( toolEvent({ @@ -219,12 +220,31 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { ) const contents = payloads.filter((payload) => payload.previewPhase === 'file_preview_content') - const streamedBytes = contents.reduce( - (sum, payload) => sum + Buffer.byteLength(String(payload.content)), - 0 - ) + const finalContent = base.replace('ANCHOR\n', `ANCHOR\n${streamed}\n`) + return { + finalContent, + lastContent: contents.at(-1)?.content, + streamedBytes: contents.reduce( + (sum, payload) => sum + Buffer.byteLength(String(payload.content)), + 0 + ), + } + } + + it('streams a 300 KB patch for 10 s well under the stream budget and ends on its final content', async () => { + const { finalContent, lastContent, streamedBytes } = await streamPatch(10_000, 16) + expect(streamedBytes).toBeLessThan(8 * 1024 * 1024) - expect(contents.at(-1)?.content).toBe(base.replace('ANCHOR\n', `ANCHOR\n${streamed}\n`)) + expect(lastContent).toBe(finalContent) + }) + + it('caps a 300 KB patch streaming for 5 minutes and still ends on its final content', async () => { + const { finalContent, lastContent, streamedBytes } = await streamPatch(300_000, 200) + + expect(streamedBytes).toBeLessThanOrEqual( + PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS + Buffer.byteLength(finalContent) + ) + expect(lastContent).toBe(finalContent) }) }) diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts index df31af488af..8ed0f08795e 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts @@ -44,6 +44,8 @@ type FilePreviewStreamState = { session: FilePreviewSession lastEmittedPreviewText: string lastSnapshotAt: number + /** Characters of intermediate preview content streamed so far. */ + streamedChars?: number } type ParsedWorkspaceFileArgs = { @@ -63,6 +65,13 @@ const DELTA_PREVIEW_CHECKPOINT_INTERVAL_MS = 1000 */ const PREVIEW_SNAPSHOT_CHARS_PER_SECOND = 256 * 1024 +/** + * Intermediate preview content streamed per edit is capped at this many + * characters; past it the preview holds still until the final snapshot, which + * is always sent when the edit completes. + */ +export const PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS = 8 * 1024 * 1024 + /** The minimum gap between full snapshots of a preview this long. */ function snapshotIntervalMs(baseMs: number, previewText: string): number { return Math.max(baseMs, (previewText.length / PREVIEW_SNAPSHOT_CHARS_PER_SECOND) * 1000) @@ -706,9 +715,8 @@ export async function processFilePreviewStreamEvent(input: { snapshotIntervalMs(PATCH_PREVIEW_SNAPSHOT_INTERVAL_MS, nextSession.previewText) ) { filePreviewState.set(editIntent.toolCallId, { + ...currentPreview, session: nextSession, - lastEmittedPreviewText: currentPreview.lastEmittedPreviewText, - lastSnapshotAt: currentPreview.lastSnapshotAt, }) } else { const previewUpdate = buildPreviewContentUpdate( @@ -718,33 +726,37 @@ export async function processFilePreviewStreamEvent(input: { now, nextSession.operation ) - - filePreviewState.set(editIntent.toolCallId, { - session: nextSession, - lastEmittedPreviewText: nextSession.previewText, - lastSnapshotAt: previewUpdate.lastSnapshotAt, - }) - - await emitPreviewEvent(streamEvent, options, { - toolCallId: nextSession.toolCallId, - toolName: 'prepare_file_edit', - previewPhase: 'file_preview_content', - content: previewUpdate.content, - contentMode: previewUpdate.contentMode, - previewVersion: nextSession.previewVersion, - fileName: nextSession.fileName, - ...(nextSession.fileId ? { fileId: nextSession.fileId } : {}), - ...(nextSession.targetKind ? { targetKind: nextSession.targetKind } : {}), - ...(nextSession.operation ? { operation: nextSession.operation } : {}), - ...(nextSession.edit ? { edit: nextSession.edit } : {}), - }) + const streamedChars = (currentPreview.streamedChars ?? 0) + previewUpdate.content.length + if (streamedChars > PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS) { + filePreviewState.set(editIntent.toolCallId, { + ...currentPreview, + session: nextSession, + }) + } else { + filePreviewState.set(editIntent.toolCallId, { + session: nextSession, + lastEmittedPreviewText: nextSession.previewText, + lastSnapshotAt: previewUpdate.lastSnapshotAt, + streamedChars, + }) + + await emitPreviewEvent(streamEvent, options, { + toolCallId: nextSession.toolCallId, + toolName: 'prepare_file_edit', + previewPhase: 'file_preview_content', + content: previewUpdate.content, + contentMode: previewUpdate.contentMode, + previewVersion: nextSession.previewVersion, + fileName: nextSession.fileName, + ...(nextSession.fileId ? { fileId: nextSession.fileId } : {}), + ...(nextSession.targetKind ? { targetKind: nextSession.targetKind } : {}), + ...(nextSession.operation ? { operation: nextSession.operation } : {}), + ...(nextSession.edit ? { edit: nextSession.edit } : {}), + }) + } } } else { - filePreviewState.set(editIntent.toolCallId, { - session: currentPreview.session, - lastEmittedPreviewText: currentPreview.lastEmittedPreviewText, - lastSnapshotAt: currentPreview.lastSnapshotAt, - }) + filePreviewState.set(editIntent.toolCallId, currentPreview) } } } diff --git a/apps/sim/lib/mothership/request/handlers/tool.ts b/apps/sim/lib/mothership/request/handlers/tool.ts index 357b7a8b9b2..cb65ad989ef 100644 --- a/apps/sim/lib/mothership/request/handlers/tool.ts +++ b/apps/sim/lib/mothership/request/handlers/tool.ts @@ -216,11 +216,7 @@ export async function prePersistClientExecutableToolCall( const data = event.payload const isGenerating = data.status === TOOL_CALL_STATUS.generating const isPartial = data.partial === true || isGenerating - /* - Only a live, complete call frame can be held behind a prompt, and only when - the gate below decides so. A replay is history and a partial frame is not yet - a call, so neither may carry the worker's approval stamp to the client. - */ + // Only a live, complete call can be held behind a prompt; drop the stamp on any other frame. if (data.replay || isPartial) { if (data.status === TOOL_AWAITING_APPROVAL_STATUS) data.status = undefined return diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index f22d22d4054..011bd49fc86 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -2913,7 +2913,7 @@ describe('runCopilotLifecycle', () => { expect(JSON.stringify(persisted)).not.toMatch(/ { + it('stops a failure that repeats after every reattach at three retries', async () => { vi.useFakeTimers() try { let attempts = 0 @@ -2926,21 +2926,14 @@ describe('runCopilotLifecycle', () => { options: { onEvent?: (event: unknown) => Promise } ): Promise => { attempts++ - if (attempts < 6) { - await options.onEvent?.({ - type: 'text', - payload: { channel: 'assistant', text: `part ${attempts} ` }, - }) - context.errors.push(STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE) - throw new StreamEndedWithoutTerminalError('/api/mothership') - } - context.streamComplete = true - context.completionStatus = MothershipStreamV1CompletionStatus.complete + await options.onEvent?.({ type: 'session', payload: { kind: 'start' } }) + context.errors.push(STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE) + throw new StreamEndedWithoutTerminalError('/api/mothership') } ) const pending = runCopilotLifecycle( - { message: 'hello', messageId: 'stream-repeated-outages' }, + { message: 'hello', messageId: 'stream-repeated-failure' }, { userId: 'user-1', workspaceId: 'ws-1', @@ -2957,8 +2950,8 @@ describe('runCopilotLifecycle', () => { ) await vi.advanceTimersByTimeAsync(60_000) - expect(await pending).toEqual(expect.objectContaining({ success: true })) - expect(attempts).toBe(6) + expect(attempts).toBe(4) + expect(await pending).toEqual(expect.objectContaining({ success: false })) } finally { mockRunStreamLoop.mockReset() vi.useRealTimers() diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 5bfe2039526..05582947bf4 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -217,6 +217,13 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS const refusalOf = (thrown?: unknown) => replayRefusal(abortController.signal.aborted ? abortController.signal.reason : thrown) + /** Set once this controller finalized the turn; only then does it clean up the stream. */ + let turnEnded = false + const endTurn = async (...args: Parameters) => { + await finalizeStream(...args) + turnEnded = true + } + /** * 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 @@ -234,7 +241,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS currentBytes: refusal.refusal.currentBytes, limitBytes: refusal.refusal.limitBytes, }) - await finalizeStream( + await endTurn( { content: '', contentBlocks: [], @@ -330,7 +337,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS outcome = RequestTraceV1Outcome.cancelled abortController.abort(AbortReason.UserStop) cancelReason = recordCancelled() - await finalizeStream( + await endTurn( { success: false, cancelled: true, @@ -445,7 +452,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS if (refusal && !result.success) { await finalizeAfterReplayRefusal(refusal, result) } else { - await finalizeStream(result, publisher, runId, outcome, requestId) + await endTurn(result, publisher, runId, outcome, requestId) } } catch (error) { if ( @@ -485,7 +492,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS toolCalls: [], error: 'An unexpected error occurred while processing the response.', } - await finalizeStream(syntheticResult, publisher, runId, outcome, requestId) + await endTurn(syntheticResult, publisher, runId, outcome, requestId) } finally { collector.endSpan( requestSpan, @@ -507,11 +514,11 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS unregisterActiveStream(streamId, abortController) /* The stream's buffer and abort marker belong to whoever holds the chat - lock now. Clean them up only while this controller still holds it, and - before releasing it, so a successor's stream is never expired and its - Stop is never cleared. + lock now. Clean them up only when this controller ended the turn and + still holds the lock, and before releasing it: a run left recoverable, + or a successor's stream, keeps its buffer and any pending Stop. */ - if (!chatId || (lease && (await holdsChatStreamLease(lease)))) { + if (turnEnded && (!chatId || (lease && (await holdsChatStreamLease(lease))))) { await scheduleBufferCleanup(streamId) await scheduleFilePreviewSessionCleanup(streamId) await cleanupAbortMarker(streamId) diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index c557edd270b..bd4e4e1ac55 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -39,20 +39,58 @@ describe('stream recovery budget', () => { expect(retry.nextDelay(error)).toBeNull() }) - it('starts a fresh recovery budget once a retry reconnects', () => { + it('keeps a reachable failure to three retries however often the worker reattaches', () => { vi.useFakeTimers() const error = new StreamEndedWithoutTerminalError('/api/mothership') const retry = new StreamRetryWindow() + for (let index = 0; index < 3; index++) { + retry.recovered() + const delay = retry.nextDelay(error) + expect(delay).not.toBeNull() + vi.advanceTimersByTime(delay ?? 0) + } + retry.recovered() + expect(retry.nextDelay(error)).toBeNull() + }) + + it('still gives the replacement worker its reachable retries after a long outage', () => { + vi.useFakeTimers() + const retry = new StreamRetryWindow() + const unreachable = new CopilotBackendError('Unavailable', { + status: 502, + body: '502 Bad Gateway', + }) + const start = Date.now() + while (Date.now() - start < 70_000) { + const delay = retry.nextDelay(unreachable) + expect(delay).not.toBeNull() + vi.advanceTimersByTime(delay ?? 0) + } + const reachable = new CopilotBackendError('Unavailable', { + status: 503, + body: '{"error":"Account admission is unavailable"}', + }) + for (let index = 0; index < 3; index++) { + const delay = retry.nextDelay(reachable) + expect(delay).not.toBeNull() + vi.advanceTimersByTime(delay ?? 0) + } + expect(retry.nextDelay(reachable)).toBeNull() + }) + + it('restarts the unreachable window once the worker answers again', () => { + vi.useFakeTimers() + const retry = new StreamRetryWindow() + const unreachable = new TypeError('fetch failed') for (let outage = 0; outage < 2; outage++) { - for (let index = 0; index < 3; index++) { - const delay = retry.nextDelay(error) + const start = Date.now() + while (Date.now() - start < 100_000) { + const delay = retry.nextDelay(unreachable) expect(delay).not.toBeNull() vi.advanceTimersByTime(delay ?? 0) } retry.recovered() - vi.advanceTimersByTime(60_000) } - expect(retry.nextDelay(error)).not.toBeNull() }) it.each([ diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts index e9235b6c11f..d3987af37a2 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -19,10 +19,17 @@ const GATEWAY_STATUSES: ReadonlySet = new Set([502, 503, 504]) */ const WORKER_REPLACEMENT_WINDOW_MS = 120_000 -/** Recovery is bounded independently of the healthy run's execution budget. */ +/** + * Recovery is bounded independently of the healthy run's execution budget, by + * two budgets that never share state: an unreachable worker gets a two-minute + * window from the moment it stopped answering, and any failure of a worker that + * did answer gets the original three retries within 30 s. + */ export class StreamRetryWindow { private readonly deadline: number private firstFailureAt?: number + private firstUnreachableAt?: number + private unreachableAttempt = 0 attempt = 0 constructor(timeoutMs = ORCHESTRATION_TIMEOUT_MS) { @@ -36,24 +43,38 @@ export class StreamRetryWindow { return remaining } - /** A retry reconnected and made progress; a later outage gets a fresh budget. */ + /** The worker answered, so a later loss of it starts a fresh unreachable window. */ recovered(): void { - this.attempt = 0 - this.firstFailureAt = undefined + this.firstUnreachableAt = undefined + this.unreachableAttempt = 0 } nextDelay(error: unknown, signal?: AbortSignal): number | null { if (signal?.aborted || !isRetryableStreamError(error)) return null + if (isWorkerUnreachable(error)) { + this.firstUnreachableAt ??= Date.now() + const delay = backoff(this.unreachableAttempt) + if (!this.fits(delay, this.firstUnreachableAt + WORKER_REPLACEMENT_WINDOW_MS)) return null + this.unreachableAttempt++ + return delay + } + // Any other retryable failure is an answer from the worker. + this.recovered() this.firstFailureAt ??= Date.now() - const unreachable = isWorkerUnreachable(error) - if (!unreachable && this.attempt >= MAX_STREAM_RETRIES) return null - const recoveryDeadline = - this.firstFailureAt + (unreachable ? WORKER_REPLACEMENT_WINDOW_MS : STREAM_RECOVERY_WINDOW_MS) - const delay = backoffWithJitter(this.attempt + 1, null, { baseMs: 250, maxMs: 5_000 }) - if (Date.now() + delay >= Math.min(this.deadline, recoveryDeadline)) return null + if (this.attempt >= MAX_STREAM_RETRIES) return null + const delay = backoff(this.attempt) + if (!this.fits(delay, this.firstFailureAt + STREAM_RECOVERY_WINDOW_MS)) return null this.attempt++ return delay } + + private fits(delay: number, recoveryDeadline: number): boolean { + return Date.now() + delay < Math.min(this.deadline, recoveryDeadline) + } +} + +function backoff(attempt: number): number { + return backoffWithJitter(attempt + 1, null, { baseMs: 250, maxMs: 5_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 be71ec56412..6b9ade70f18 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -551,6 +551,28 @@ 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 () => { + let eventsKey = '' + const { streamId, runId, frames } = await runTurn( + [ + async () => { + // A corrupt buffer key makes the next append fail while the lease is still held. + await redis().set(eventsKey, 'corrupt', 'EX', 3600) + }, + toolCall('call-unsaved', 'cli_blocks_get', { command: 'blocks get' }), + ], + async (id) => { + eventsKey = `mothership_stream:${id}:events` + } + ) + + 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('active') + expect(worker.abortRequests).toEqual([]) + }) + it("leaves its successor's stream untouched when the lease is lost while ending a refused turn", async () => { const successorToken = `successor\n${generateId()}` worker.hooks.onAbort = async () => { diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index 89021177a89..e208e561c6e 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -145,7 +145,7 @@ describe('compactStreamEvent', () => { expect(args.stdin).toBe(`${'x'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 1 MB total]`) }) - it('never compacts a file preview or a generated API key', () => { + it('never compacts a file preview', () => { const preview: StreamEvent = { type: 'tool', payload: { @@ -158,21 +158,8 @@ describe('compactStreamEvent', () => { fileName: 'notes.md', }, } - const apiKey: StreamEvent = { - type: 'tool', - payload: { - toolCallId: 'c', - toolName: 'generate_api_key', - executor: 'sim', - mode: 'async', - phase: 'result', - success: true, - output: { key: 'k'.repeat(400 * 1024) }, - }, - } expect(compactStreamEvent(preview)).toBe(preview) - expect(compactStreamEvent(apiKey)).toBe(apiKey) }) it('compacts a copy and leaves the caller’s event whole for dispatch', () => { diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index 10e699802b4..99f76f979c5 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -19,13 +19,15 @@ export const STREAM_STRING_PREVIEW_UNITS = 8 * 1024 const NO_KEYS: ReadonlySet = new Set() const ARGUMENTS_KEY: ReadonlySet = new Set(['arguments']) -/** Its live result is the only place the plaintext key reaches the browser. */ -const GENERATE_API_KEY_TOOL = 'generate_api_key' - /** Items kept at the head of an array that string cuts alone could not bound. */ export const STREAM_ARRAY_HEAD_ITEMS = 100 -/** An upper bound on a value's serialized size, without serializing it. */ +/** + * A cheap estimate of a value's serialized size, without serializing it. It can + * undercount strings of control characters, which JSON escapes to six bytes + * each, but even then an event under the threshold stays under 6 / 3 × 256 KiB, + * half the replay buffer's 1 MiB write ceiling. + */ function estimateBytes(value: unknown): number { if (typeof value === 'string') return value.length * 3 + 2 if (Array.isArray(value)) @@ -81,7 +83,7 @@ function mapLeaves( ): unknown { if (Array.isArray(value)) { const items = array(value) - let copy: unknown[] | undefined = items === value ? undefined : [...items] + let copy: unknown[] | undefined = items === value ? undefined : items items.forEach((item, index) => { const next = mapLeaves(item, NO_KEYS, leaf, array) if (next !== item) (copy ??= [...items])[index] = next @@ -103,8 +105,8 @@ function mapLeaves( * only to the copy the writer delivers and persists; the caller keeps the full * event for dispatch. Long strings are cut to their head in place, so every * object keeps its shape; if that is not enough, long arrays keep their head. - * Assistant text, file previews, `generate_api_key` results, and the arguments - * of calls the browser executes are never cut; an event + * Assistant text, file previews, and the arguments of calls the browser + * executes are never cut; an event * still too large is refused by the buffer, which ends the turn with an error. */ export function compactStreamEvent(event: StreamEvent): StreamEvent { @@ -115,9 +117,6 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { } if (estimateBytes(payload) <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) return event const toolName = typeof payload.toolName === 'string' ? payload.toolName : '' - if (payload.phase === MothershipStreamV1ToolPhase.result && toolName === GENERATE_API_KEY_TOOL) { - return event - } const args = isRecordLike(payload.arguments) ? payload.arguments : undefined const skipKeys = payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutedToolCall(toolName, args) From 55386c094242265522739938aed7f5d798d088de Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 22:08:32 -0700 Subject: [PATCH 11/18] fix(mothership): clean up turns whose terminal publish failed and cap previews per turn in bytes - A controller that ended its turn but could not publish the terminal events still cleans up its stream: the run is settled either way, so the buffer no longer lingers for its full TTL. Only a superseded controller skips cleanup, and the lease check still protects a successor. - The intermediate file preview cap is now 8 MiB of UTF-8 per stream loop, across all of its edits, instead of characters per edit. Each edit's final snapshot is still always sent. --- .../request/go/file-preview-adapter.test.ts | 51 +++++++++++++------ .../request/go/file-preview-adapter.ts | 20 ++++---- .../lib/mothership/request/lifecycle/start.ts | 15 ++++-- .../session/replay-budget.integration.ts | 22 +++++++- 4 files changed, 79 insertions(+), 29 deletions(-) diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts index 68a4ba244ba..422dc7ba243 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts @@ -25,7 +25,7 @@ import { createStreamingContext } from '@/lib/mothership/request/context/request import { createFilePreviewAdapterState, type FilePreviewAdapterState, - PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS, + PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES, processFilePreviewStreamEvent, } from '@/lib/mothership/request/go/file-preview-adapter' import { createEvent, eventToStreamEvent } from '@/lib/mothership/request/session' @@ -175,15 +175,22 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { afterEach(() => vi.useRealTimers()) - async function streamPatch(durationMs: number, tickMs: number) { + /** Streams one anchored patch of a 300 KB file through `state`, then completes it. */ + async function streamPatch( + durationMs: number, + tickMs: number, + state = createFilePreviewAdapterState(), + edit = 0 + ) { vi.useFakeTimers() const base = `ANCHOR\n${'line of existing file content\n'.repeat(10_000)}` - const edit = { strategy: 'anchored', mode: 'insert_after', anchor: 'ANCHOR' } - peekFileIntentMock.mockResolvedValue({ existingContent: base, edit }) - const state = createFilePreviewAdapterState() + const anchoredEdit = { strategy: 'anchored', mode: 'insert_after', anchor: 'ANCHOR' } + peekFileIntentMock.mockResolvedValue({ existingContent: base, edit: anchoredEdit }) + const editToolCallId = `${EDIT_TOOL_CALL_ID}-${edit}` const intent = { - ...makeIntent({ operation: 'patch', fileId: 'file-big', fileName: 'big.md' }), - edit, + ...makeIntent({ operation: 'patch', fileId: `file-big-${edit}`, fileName: 'big.md' }), + toolCallId: `${WORKSPACE_FILE_TOOL_CALL_ID}-${edit}`, + edit: anchoredEdit, } const payloads: Array> = [] const drive = async (streamEvent: StreamEvent) => { @@ -202,17 +209,24 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { state, }) } + const delta = (argumentsDelta: string) => + toolEvent({ + toolCallId: editToolCallId, + toolName: 'apply_file_edit', + phase: MothershipStreamV1ToolPhase.args_delta, + argumentsDelta, + }) let streamed = '' - await drive(editContentDelta('{"content":"')) + await drive(delta('{"content":"')) for (let elapsed = 0; elapsed < durationMs; elapsed += tickMs) { vi.advanceTimersByTime(tickMs) streamed += `word${elapsed} ` - await drive(editContentDelta(`word${elapsed} `)) + await drive(delta(`word${elapsed} `)) } await drive( toolEvent({ - toolCallId: EDIT_TOOL_CALL_ID, + toolCallId: editToolCallId, toolName: 'apply_file_edit', phase: MothershipStreamV1ToolPhase.result, success: true, @@ -238,13 +252,18 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { expect(lastContent).toBe(finalContent) }) - it('caps a 300 KB patch streaming for 5 minutes and still ends on its final content', async () => { - const { finalContent, lastContent, streamedBytes } = await streamPatch(300_000, 200) + it('caps a turn of four long patches and still ends each on its final content', async () => { + const state = createFilePreviewAdapterState() + let streamedBytes = 0 + let finalBytes = 0 + for (let edit = 0; edit < 4; edit++) { + const result = await streamPatch(120_000, 200, state, edit) + expect(result.lastContent).toBe(result.finalContent) + streamedBytes += result.streamedBytes + finalBytes += Buffer.byteLength(result.finalContent) + } - expect(streamedBytes).toBeLessThanOrEqual( - PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS + Buffer.byteLength(finalContent) - ) - expect(lastContent).toBe(finalContent) + expect(streamedBytes).toBeLessThanOrEqual(PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES + finalBytes) }) }) diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts index 8ed0f08795e..4bcef27d749 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts @@ -44,8 +44,6 @@ type FilePreviewStreamState = { session: FilePreviewSession lastEmittedPreviewText: string lastSnapshotAt: number - /** Characters of intermediate preview content streamed so far. */ - streamedChars?: number } type ParsedWorkspaceFileArgs = { @@ -66,11 +64,11 @@ const DELTA_PREVIEW_CHECKPOINT_INTERVAL_MS = 1000 const PREVIEW_SNAPSHOT_CHARS_PER_SECOND = 256 * 1024 /** - * Intermediate preview content streamed per edit is capped at this many - * characters; past it the preview holds still until the final snapshot, which - * is always sent when the edit completes. + * UTF-8 bytes of intermediate preview content one stream loop may send across + * all of its edits, well under the stream's 32 MiB replay budget. Past it the + * previews hold still; each edit's final snapshot is still sent on completion. */ -export const PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS = 8 * 1024 * 1024 +export const PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES = 8 * 1024 * 1024 /** The minimum gap between full snapshots of a preview this long. */ function snapshotIntervalMs(baseMs: number, previewText: string): number { @@ -366,12 +364,15 @@ export function buildPreviewContentUpdate( export interface FilePreviewAdapterState { editContentState: Map filePreviewState: Map + /** UTF-8 bytes of intermediate preview content sent by this stream loop. */ + intermediateBytes: number } export function createFilePreviewAdapterState(): FilePreviewAdapterState { return { editContentState: new Map(), filePreviewState: new Map(), + intermediateBytes: 0, } } @@ -726,8 +727,9 @@ export async function processFilePreviewStreamEvent(input: { now, nextSession.operation ) - const streamedChars = (currentPreview.streamedChars ?? 0) + previewUpdate.content.length - if (streamedChars > PREVIEW_INTERMEDIATE_SNAPSHOT_CHARS) { + const intermediateBytes = + state.intermediateBytes + Buffer.byteLength(previewUpdate.content, 'utf8') + if (intermediateBytes > PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES) { filePreviewState.set(editIntent.toolCallId, { ...currentPreview, session: nextSession, @@ -737,8 +739,8 @@ export async function processFilePreviewStreamEvent(input: { session: nextSession, lastEmittedPreviewText: nextSession.previewText, lastSnapshotAt: previewUpdate.lastSnapshotAt, - streamedChars, }) + state.intermediateBytes = intermediateBytes await emitPreviewEvent(streamEvent, options, { toolCallId: nextSession.toolCallId, diff --git a/apps/sim/lib/mothership/request/lifecycle/start.ts b/apps/sim/lib/mothership/request/lifecycle/start.ts index 05582947bf4..4be603e2be8 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -217,11 +217,20 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS const refusalOf = (thrown?: unknown) => replayRefusal(abortController.signal.aborted ? abortController.signal.reason : thrown) - /** Set once this controller finalized the turn; only then does it clean up the stream. */ + /** + * Set once this controller ended the turn, even if publishing its terminal + * events failed (the run is settled regardless); only then does it clean up + * the stream. A superseded controller leaves the turn to its successor. + */ let turnEnded = false const endTurn = async (...args: Parameters) => { - await finalizeStream(...args) - turnEnded = true + try { + await finalizeStream(...args) + turnEnded = true + } catch (error) { + turnEnded = !(error instanceof StreamControllerSupersededError) + throw error + } } /** 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 6b9ade70f18..eb128d122c5 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -462,7 +462,8 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge admittedRun: run, orchestrateOptions: { userId, workspaceId, chatId, runId, interactive: true }, }) - const frames = dataFrames(await new Response(response).text()) + /** A controller whose teardown throws errors its client stream; its frames are then moot. */ + const frames = dataFrames(await new Response(response).text().catch(() => '')) await teardown return { streamId, runId, frames } } @@ -573,6 +574,25 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge expect(worker.abortRequests).toEqual([]) }) + it('still cleans up a finished turn whose terminal events could not be published', async () => { + let eventsKey = '' + const { streamId, runId } = await runTurn( + [ + async () => { + // A corrupt buffer key makes the terminal append fail while the lease is held. + await redis().set(eventsKey, 'corrupt', 'EX', 3600) + }, + ], + async (id) => { + eventsKey = `mothership_stream:${id}:events` + } + ) + + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('complete') + expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeLessThanOrEqual(300) + }) + it("leaves its successor's stream untouched when the lease is lost while ending a refused turn", async () => { const successorToken = `successor\n${generateId()}` worker.hooks.onAbort = async () => { From a4debbf3a545fceecaaebfb275ae9072e66f18b9 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 22:42:13 -0700 Subject: [PATCH 12/18] fix(mothership): bound preview frames per write and per turn, and narrow unreachable retries - A file preview frame whose content would not fit one replay write, or that would exceed the turn's preview budget, is skipped instead of sent: the preview holds its last content and the client loads the stored file on completion. A completion too large to send whole drops its tool output. The per-frame limit derives from the replay buffer's single-write ceiling, and the 8 MiB budget now covers every preview frame in the turn, finals included, across all legs. A large file preview no longer ends the turn. - Only a request that fails before any response headers counts as an unreachable worker. A TypeError raised after the response began, including one from handling delivered events, keeps the three-retry budget, so a deterministic failure can no longer retry until the run deadline. - A worker that cannot be reached is reported with the generic "temporarily unavailable" message; the network error stays on the error's cause. - Retry logs and spans count retries from both budgets. --- .../request/context/request-context.ts | 1 + .../request/go/file-preview-adapter.test.ts | 60 +++++--- .../request/go/file-preview-adapter.ts | 136 +++++++++++------- .../lib/mothership/request/go/stream.test.ts | 12 +- apps/sim/lib/mothership/request/go/stream.ts | 14 +- .../request/handlers/handlers.test.ts | 1 + .../mothership/request/lifecycle/run.test.ts | 52 +++++++ .../lib/mothership/request/lifecycle/run.ts | 8 +- .../request/lifecycle/stream-retry.test.ts | 19 ++- .../request/lifecycle/stream-retry.ts | 16 ++- .../session/replay-budget.integration.ts | 76 +++++++++- apps/sim/lib/mothership/request/types.ts | 2 + 12 files changed, 306 insertions(+), 91 deletions(-) diff --git a/apps/sim/lib/mothership/request/context/request-context.ts b/apps/sim/lib/mothership/request/context/request-context.ts index c9e57efdfbb..1e8b5b6b262 100644 --- a/apps/sim/lib/mothership/request/context/request-context.ts +++ b/apps/sim/lib/mothership/request/context/request-context.ts @@ -29,6 +29,7 @@ export function createStreamingContext(overrides?: Partial): S wasAborted: false, errors: [], activeFileIntents: new Map(), + filePreviewBudget: { contentBytes: 0 }, trace: new TraceCollector(), toolPermissions: { enabled: false, autoAllowed: new Set(), autoAllowPermitted: true }, ...overrides, diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts index 422dc7ba243..dc1f0e28832 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts @@ -25,7 +25,8 @@ import { createStreamingContext } from '@/lib/mothership/request/context/request import { createFilePreviewAdapterState, type FilePreviewAdapterState, - PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES, + PREVIEW_FRAME_MAX_BYTES, + PREVIEW_TURN_CONTENT_BYTES, processFilePreviewStreamEvent, } from '@/lib/mothership/request/go/file-preview-adapter' import { createEvent, eventToStreamEvent } from '@/lib/mothership/request/session' @@ -175,15 +176,19 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { afterEach(() => vi.useRealTimers()) - /** Streams one anchored patch of a 300 KB file through `state`, then completes it. */ + /** One turn's adapter state and streaming context, shared by all of its edits. */ + function newTurn() { + return { state: createFilePreviewAdapterState(), context: createStreamingContext() } + } + + /** Streams one anchored patch of a file of `baseLines` lines through `turn`, then completes it. */ async function streamPatch( durationMs: number, tickMs: number, - state = createFilePreviewAdapterState(), - edit = 0 + { turn = newTurn(), edit = 0, baseLines = 10_000 } = {} ) { vi.useFakeTimers() - const base = `ANCHOR\n${'line of existing file content\n'.repeat(10_000)}` + const base = `ANCHOR\n${'line of existing file content\n'.repeat(baseLines)}` const anchoredEdit = { strategy: 'anchored', mode: 'insert_after', anchor: 'ANCHOR' } peekFileIntentMock.mockResolvedValue({ existingContent: base, edit: anchoredEdit }) const editToolCallId = `${EDIT_TOOL_CALL_ID}-${edit}` @@ -194,19 +199,18 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { } const payloads: Array> = [] const drive = async (streamEvent: StreamEvent) => { - const context = createStreamingContext() - context.activeFileIntents.set('', intent) + turn.context.activeFileIntents.set('', intent) await processFilePreviewStreamEvent({ streamId: STREAM_ID, streamEvent, - context, + context: turn.context, execContext, options: { onEvent: (event) => { payloads.push((event as { payload: Record }).payload) }, }, - state, + state: turn.state, }) } const delta = (argumentsDelta: string) => @@ -234,11 +238,12 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { ) const contents = payloads.filter((payload) => payload.previewPhase === 'file_preview_content') - const finalContent = base.replace('ANCHOR\n', `ANCHOR\n${streamed}\n`) return { - finalContent, + payloads, + finalContent: base.replace('ANCHOR\n', `ANCHOR\n${streamed}\n`), lastContent: contents.at(-1)?.content, - streamedBytes: contents.reduce( + completed: payloads.some((payload) => payload.previewPhase === 'file_preview_complete'), + contentBytes: contents.reduce( (sum, payload) => sum + Buffer.byteLength(String(payload.content)), 0 ), @@ -246,24 +251,33 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { } it('streams a 300 KB patch for 10 s well under the stream budget and ends on its final content', async () => { - const { finalContent, lastContent, streamedBytes } = await streamPatch(10_000, 16) + const { finalContent, lastContent, contentBytes } = await streamPatch(10_000, 16) - expect(streamedBytes).toBeLessThan(8 * 1024 * 1024) + expect(contentBytes).toBeLessThan(8 * 1024 * 1024) expect(lastContent).toBe(finalContent) }) - it('caps a turn of four long patches and still ends each on its final content', async () => { - const state = createFilePreviewAdapterState() - let streamedBytes = 0 - let finalBytes = 0 + it('never sends a preview frame over the replay write ceiling, and still completes the edit', async () => { + const { payloads, completed } = await streamPatch(2_000, 200, { baseLines: 45_000 }) + + for (const payload of payloads) { + expect(Buffer.byteLength(JSON.stringify(payload))).toBeLessThanOrEqual( + PREVIEW_FRAME_MAX_BYTES + ) + } + expect(completed).toBe(true) + }) + + it('bounds all preview content in a turn of four long patches, and still completes each', async () => { + const turn = newTurn() + let contentBytes = 0 for (let edit = 0; edit < 4; edit++) { - const result = await streamPatch(120_000, 200, state, edit) - expect(result.lastContent).toBe(result.finalContent) - streamedBytes += result.streamedBytes - finalBytes += Buffer.byteLength(result.finalContent) + const result = await streamPatch(120_000, 200, { turn, edit }) + expect(result.completed).toBe(true) + contentBytes += result.contentBytes } - expect(streamedBytes).toBeLessThanOrEqual(PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES + finalBytes) + expect(contentBytes).toBeLessThanOrEqual(PREVIEW_TURN_CONTENT_BYTES) }) }) diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts index 4bcef27d749..46efed8bd09 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts @@ -1,6 +1,7 @@ import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { isRecordLike } from '@sim/utils/object' +import { getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' import { executeCopilotFileUseCase } from '@/lib/mothership/application/execute-file-use-case' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { @@ -64,11 +65,17 @@ const DELTA_PREVIEW_CHECKPOINT_INTERVAL_MS = 1000 const PREVIEW_SNAPSHOT_CHARS_PER_SECOND = 256 * 1024 /** - * UTF-8 bytes of intermediate preview content one stream loop may send across - * all of its edits, well under the stream's 32 MiB replay budget. Past it the - * previews hold still; each edit's final snapshot is still sent on completion. + * UTF-8 bytes of preview content one turn may stream across all of its edits, + * final snapshots included, well under the stream's 32 MiB replay budget. */ -export const PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES = 8 * 1024 * 1024 +export const PREVIEW_TURN_CONTENT_BYTES = 8 * 1024 * 1024 + +/** Room left in a replay write for the envelope around a preview payload. */ +const PREVIEW_FRAME_ENVELOPE_HEADROOM_BYTES = 16 * 1024 + +/** The largest serialized preview payload the replay buffer can persist in one write. */ +export const PREVIEW_FRAME_MAX_BYTES = + getRedisBudgetLimits('copilot_stream').maxSingleWriteBytes - PREVIEW_FRAME_ENVELOPE_HEADROOM_BYTES /** The minimum gap between full snapshots of a preview this long. */ function snapshotIntervalMs(baseMs: number, previewText: string): number { @@ -364,28 +371,49 @@ export function buildPreviewContentUpdate( export interface FilePreviewAdapterState { editContentState: Map filePreviewState: Map - /** UTF-8 bytes of intermediate preview content sent by this stream loop. */ - intermediateBytes: number } export function createFilePreviewAdapterState(): FilePreviewAdapterState { return { editContentState: new Map(), filePreviewState: new Map(), - intermediateBytes: 0, } } +/** + * Sends one preview frame, and reports whether it did. A content frame that would + * not fit one replay write, or past the turn's preview budget, is skipped: the + * preview holds its last content, and on completion the client loads the stored + * file. A completion too large to send whole drops its tool output instead. + */ async function emitPreviewEvent( streamEvent: StreamEvent, options: Pick, + context: StreamingContext, payload: SyntheticFilePreviewPayload -): Promise { +): Promise { + let frame = payload + const frameBytes = Buffer.byteLength(JSON.stringify(frame), 'utf8') + if (frame.previewPhase === 'file_preview_content') { + const budget = context.filePreviewBudget + const contentBytes = budget.contentBytes + Buffer.byteLength(frame.content, 'utf8') + if (frameBytes > PREVIEW_FRAME_MAX_BYTES || contentBytes > PREVIEW_TURN_CONTENT_BYTES) { + return false + } + budget.contentBytes = contentBytes + } else if ( + frame.previewPhase === 'file_preview_complete' && + frameBytes > PREVIEW_FRAME_MAX_BYTES + ) { + const { output: _output, ...withoutOutput } = frame + frame = withoutOutput + } await options.onEvent?.({ type: MothershipStreamV1EventType.tool, - payload, + payload: frame, ...(streamEvent.scope ? { scope: streamEvent.scope } : {}), }) + return true } export async function processFilePreviewStreamEvent(input: { @@ -468,12 +496,12 @@ export async function processFilePreviewStreamEvent(input: { }) await persistFilePreviewSession(session) - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_start', }) - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_target', @@ -486,7 +514,7 @@ export async function processFilePreviewStreamEvent(input: { ...(title ? { title } : {}), }) if (edit) { - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_edit_meta', @@ -543,12 +571,12 @@ export async function processFilePreviewStreamEvent(input: { }) await persistFilePreviewSession(session) - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId: intent.toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_start', }) - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId: intent.toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_target', @@ -561,7 +589,7 @@ export async function processFilePreviewStreamEvent(input: { ...(intent.title ? { title: intent.title } : {}), }) if (intent.edit) { - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId: intent.toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_edit_meta', @@ -605,13 +633,8 @@ export async function processFilePreviewStreamEvent(input: { previewVersion: (currentPreview?.session.previewVersion ?? 0) + 1, updatedAt: new Date().toISOString(), } - filePreviewState.set(patchDeleteIntent.toolCallId, { - session: nextSession, - lastEmittedPreviewText: previewText, - lastSnapshotAt: Date.now(), - }) await persistFilePreviewSession(nextSession) - await emitPreviewEvent(streamEvent, options, { + const sent = await emitPreviewEvent(streamEvent, options, context, { toolCallId: nextSession.toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_content', @@ -624,6 +647,20 @@ export async function processFilePreviewStreamEvent(input: { ...(nextSession.operation ? { operation: nextSession.operation } : {}), ...(nextSession.edit ? { edit: nextSession.edit } : {}), }) + filePreviewState.set( + patchDeleteIntent.toolCallId, + sent + ? { + session: nextSession, + lastEmittedPreviewText: previewText, + lastSnapshotAt: Date.now(), + } + : { + session: nextSession, + lastEmittedPreviewText: currentPreview?.lastEmittedPreviewText ?? '', + lastSnapshotAt: currentPreview?.lastSnapshotAt ?? 0, + } + ) } } @@ -727,35 +764,30 @@ export async function processFilePreviewStreamEvent(input: { now, nextSession.operation ) - const intermediateBytes = - state.intermediateBytes + Buffer.byteLength(previewUpdate.content, 'utf8') - if (intermediateBytes > PREVIEW_INTERMEDIATE_SNAPSHOT_BYTES) { - filePreviewState.set(editIntent.toolCallId, { - ...currentPreview, - session: nextSession, - }) - } else { - filePreviewState.set(editIntent.toolCallId, { - session: nextSession, - lastEmittedPreviewText: nextSession.previewText, - lastSnapshotAt: previewUpdate.lastSnapshotAt, - }) - state.intermediateBytes = intermediateBytes - - await emitPreviewEvent(streamEvent, options, { - toolCallId: nextSession.toolCallId, - toolName: 'prepare_file_edit', - previewPhase: 'file_preview_content', - content: previewUpdate.content, - contentMode: previewUpdate.contentMode, - previewVersion: nextSession.previewVersion, - fileName: nextSession.fileName, - ...(nextSession.fileId ? { fileId: nextSession.fileId } : {}), - ...(nextSession.targetKind ? { targetKind: nextSession.targetKind } : {}), - ...(nextSession.operation ? { operation: nextSession.operation } : {}), - ...(nextSession.edit ? { edit: nextSession.edit } : {}), - }) - } + const sent = await emitPreviewEvent(streamEvent, options, context, { + toolCallId: nextSession.toolCallId, + toolName: 'prepare_file_edit', + previewPhase: 'file_preview_content', + content: previewUpdate.content, + contentMode: previewUpdate.contentMode, + previewVersion: nextSession.previewVersion, + fileName: nextSession.fileName, + ...(nextSession.fileId ? { fileId: nextSession.fileId } : {}), + ...(nextSession.targetKind ? { targetKind: nextSession.targetKind } : {}), + ...(nextSession.operation ? { operation: nextSession.operation } : {}), + ...(nextSession.edit ? { edit: nextSession.edit } : {}), + }) + // A skipped frame leaves the client on the last content it did receive. + filePreviewState.set( + editIntent.toolCallId, + sent + ? { + session: nextSession, + lastEmittedPreviewText: nextSession.previewText, + lastSnapshotAt: previewUpdate.lastSnapshotAt, + } + : { ...currentPreview, session: nextSession } + ) } } else { filePreviewState.set(editIntent.toolCallId, currentPreview) @@ -792,7 +824,7 @@ export async function processFilePreviewStreamEvent(input: { lastEmittedPreviewText: currentPreview.session.previewText, lastSnapshotAt: Date.now(), }) - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId: currentPreview.session.toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_content', @@ -826,7 +858,7 @@ export async function processFilePreviewStreamEvent(input: { await persistFilePreviewSession(completedSession) } - await emitPreviewEvent(streamEvent, options, { + await emitPreviewEvent(streamEvent, options, context, { toolCallId: editResultIntent.toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_complete', diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index 2a0719bf5a6..75b2013d361 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -149,6 +149,7 @@ function createStreamingContext(): StreamingContext { wasAborted: false, errors: [], activeFileIntents: new Map(), + filePreviewBudget: { contentBytes: 0 }, trace: new TraceCollector(), toolPermissions: { enabled: false, @@ -784,8 +785,9 @@ describe('copilot go stream helpers', () => { expect(fetch).toHaveBeenCalledTimes(1) }) - it('does not retry network errors because Go may already be executing the request', async () => { - vi.mocked(fetch).mockRejectedValueOnce(new TypeError('fetch failed')) + it('reports a worker it could not reach without the raw network error', async () => { + const networkError = new TypeError('fetch failed') + vi.mocked(fetch).mockRejectedValueOnce(networkError) const context = createStreamingContext() const execContext: ExecutionContext = { @@ -797,7 +799,11 @@ describe('copilot go stream helpers', () => { runStreamLoop('https://example.com/mothership/stream', {}, context, execContext, { timeout: 1000, }) - ).rejects.toThrow('fetch failed') + ).rejects.toMatchObject({ + name: 'WorkerUnreachableError', + message: 'The agent service is temporarily unavailable. Please try again.', + cause: networkError, + }) expect(fetch).toHaveBeenCalledTimes(1) }) diff --git a/apps/sim/lib/mothership/request/go/stream.ts b/apps/sim/lib/mothership/request/go/stream.ts index 769306034a7..4369839993c 100644 --- a/apps/sim/lib/mothership/request/go/stream.ts +++ b/apps/sim/lib/mothership/request/go/stream.ts @@ -70,6 +70,17 @@ const BACKEND_UNAVAILABLE_MESSAGE = 'The agent service is temporarily unavailable. Please try again.' const BACKEND_REJECTED_MESSAGE = 'The agent service could not process this request.' +/** + * The request never reached a worker: the connection failed before any response + * headers arrived. The network error stays on `cause` for logs. + */ +export class WorkerUnreachableError extends Error { + constructor(cause: unknown) { + super(BACKEND_UNAVAILABLE_MESSAGE, { cause }) + this.name = 'WorkerUnreachableError' + } +} + /** * A worker rejection message the user can act on: short, one line, plain text, * and free of identifiers (`userId`, `protocol_version_mismatch`) that only mean @@ -195,7 +206,8 @@ export async function runStreamLoop( headersMs: Math.round(performance.now() - fetchStart), } context.trace.endSpan(fetchSpan, abortSignal?.aborted ? 'cancelled' : 'error') - throw error + if (requestSignal.aborted) throw error + throw new WorkerUnreachableError(error) } const headersElapsedMs = Math.round(performance.now() - fetchStart) fetchSpan.attributes = { diff --git a/apps/sim/lib/mothership/request/handlers/handlers.test.ts b/apps/sim/lib/mothership/request/handlers/handlers.test.ts index 88a75b6e976..d458500329e 100644 --- a/apps/sim/lib/mothership/request/handlers/handlers.test.ts +++ b/apps/sim/lib/mothership/request/handlers/handlers.test.ts @@ -120,6 +120,7 @@ describe('sse-handlers tool lifecycle', () => { toolCalls: new Map(), pendingToolPromises: new Map(), activeFileIntents: new Map(), + filePreviewBudget: { contentBytes: 0 }, seenToolCalls: new Set(), seenToolResults: new Set(), currentThinkingBlock: null, diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 011bd49fc86..fa35ccdc6ee 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -121,9 +121,17 @@ vi.mock('@/lib/mothership/request/go/stream', () => { } } + class WorkerUnreachableError extends Error { + constructor(cause: unknown) { + super('The agent service is temporarily unavailable. Please try again.', { cause }) + this.name = 'WorkerUnreachableError' + } + } + return { BillingLimitError, CopilotBackendError, + WorkerUnreachableError, STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE, StreamEndedWithoutTerminalError, runStreamLoop: mockRunStreamLoop, @@ -2913,6 +2921,50 @@ describe('runCopilotLifecycle', () => { expect(JSON.stringify(persisted)).not.toMatch(/ { + vi.useFakeTimers() + try { + let attempts = 0 + mockRunStreamLoop.mockImplementation( + async ( + _url: string, + _init: RequestInit, + _context: StreamingContext, + _exec: ExecutionContext, + options: { onEvent?: (event: unknown) => Promise } + ): Promise => { + attempts++ + await options.onEvent?.({ type: 'session', payload: { kind: 'start' } }) + throw new TypeError("Cannot read properties of undefined (reading 'payload')") + } + ) + + const pending = runCopilotLifecycle( + { message: 'hello', messageId: 'stream-repeated-type-error' }, + { + userId: 'user-1', + workspaceId: 'ws-1', + chatId: 'chat-1', + executionId: 'exec-1', + runId: 'run-1', + executionContext: { + userId: 'user-1', + workflowId: '', + workspaceId: 'ws-1', + chatId: 'chat-1', + }, + } + ) + await vi.advanceTimersByTimeAsync(60_000) + + expect(attempts).toBe(4) + expect(await pending).toEqual(expect.objectContaining({ success: false })) + } finally { + mockRunStreamLoop.mockReset() + vi.useRealTimers() + } + }) + it('stops a failure that repeats after every reattach at three retries', async () => { vi.useFakeTimers() try { diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index 550c21a1e8e..59a326dd13a 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -875,7 +875,7 @@ async function runResumeLegWithRetry( if (backoff !== null) { leg.errors.length = errorsBeforeAttempt logger.warn('Child resume leg failed, retrying', { - attempt: retry.attempt + 1, + attempt: retry.attempts + 1, backoffMs: backoff, error: toError(error).message, }) @@ -1171,7 +1171,7 @@ async function runCheckpointLoop( { route, isResume, - ...(isResume ? { attempt: retry.attempt } : {}), + ...(isResume ? { attempt: retry.attempts } : {}), } ) context.trace.setActiveSpan(streamSpan) @@ -1179,7 +1179,7 @@ async function runCheckpointLoop( logger.info('Starting stream loop', { route, isResume, - resumeAttempt: retry.attempt, + resumeAttempt: retry.attempts, pendingToolPromises: context.pendingToolPromises.size, toolCallCount: context.toolCalls.size, hasCheckpoint: !!context.awaitingAsyncContinuation, @@ -1237,7 +1237,7 @@ async function runCheckpointLoop( logger.warn( isResume ? 'Resume stream failed, retrying' : 'Initial stream failed, retrying', { - attempt: (retry?.attempt ?? 0) + 1, + attempt: (retry?.attempts ?? 0) + 1, backoffMs: backoff, error: toError(streamError).message, } diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index bd4e4e1ac55..0e77a2c9f29 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -2,6 +2,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { CopilotBackendError, StreamEndedWithoutTerminalError, + WorkerUnreachableError, } from '@/lib/mothership/request/go/stream' import { StreamRetryWindow } from '@/lib/mothership/request/lifecycle/stream-retry' @@ -81,7 +82,7 @@ describe('stream recovery budget', () => { it('restarts the unreachable window once the worker answers again', () => { vi.useFakeTimers() const retry = new StreamRetryWindow() - const unreachable = new TypeError('fetch failed') + const unreachable = new WorkerUnreachableError(new TypeError('fetch failed')) for (let outage = 0; outage < 2; outage++) { const start = Date.now() while (Date.now() - start < 100_000) { @@ -93,8 +94,22 @@ describe('stream recovery budget', () => { } }) + it('keeps a TypeError from handling a delivered stream to the reachable budget', () => { + vi.useFakeTimers() + const error = new TypeError("Cannot read properties of undefined (reading 'payload')") + const retry = new StreamRetryWindow() + for (let index = 0; index < 3; index++) { + retry.recovered() + const delay = retry.nextDelay(error) + expect(delay).not.toBeNull() + vi.advanceTimersByTime(delay ?? 0) + } + retry.recovered() + expect(retry.nextDelay(error)).toBeNull() + }) + it.each([ - new TypeError('fetch failed'), + new WorkerUnreachableError(new TypeError('fetch failed')), new CopilotBackendError('Unavailable', { status: 502 }), new CopilotBackendError('Unavailable', { status: 504, diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts index d3987af37a2..d99d4fb26c4 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -4,6 +4,7 @@ import { StreamContinuityError } from '@/lib/mothership/request/go/parser' import { CopilotBackendError, StreamEndedWithoutTerminalError, + WorkerUnreachableError, } from '@/lib/mothership/request/go/stream' const MAX_STREAM_RETRIES = 3 @@ -36,6 +37,11 @@ export class StreamRetryWindow { this.deadline = Date.now() + timeoutMs } + /** Retries taken across both budgets, for logs and spans. */ + get attempts(): number { + return this.attempt + this.unreachableAttempt + } + remainingMs(): number { const remaining = this.deadline - Date.now() if (remaining <= 0) @@ -78,15 +84,15 @@ function backoff(attempt: number): number { } /** - * No worker answered: a gateway error page from the load balancer, or no - * connection at all. A JSON 5xx comes from a reachable worker, so it gets the - * short recovery budget instead. + * No worker answered: a gateway error page from the load balancer, or a request + * that failed before any response headers. A JSON 5xx, or a failure after the + * response began, comes from a reachable worker and gets the short budget. */ function isWorkerUnreachable(error: unknown): boolean { if (error instanceof CopilotBackendError) { return error.status !== undefined && GATEWAY_STATUSES.has(error.status) && !isJson(error.body) } - return error instanceof TypeError + return error instanceof WorkerUnreachableError } function isJson(body: string | undefined): boolean { @@ -108,5 +114,5 @@ function isRetryableStreamError(error: unknown): boolean { if (error instanceof CopilotBackendError) { return error.status !== undefined && error.status >= 500 } - return error instanceof TypeError + return error instanceof WorkerUnreachableError || error instanceof TypeError } 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 eb128d122c5..a55706e2468 100644 --- a/apps/sim/lib/mothership/request/session/replay-budget.integration.ts +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -63,7 +63,7 @@ vi.mock('@/lib/mothership/request/lifecycle/run', () => ({ worker.runs.push(run) for (const event of worker.script) { if (typeof event === 'function') { - await event() + await event(options) continue } try { @@ -98,6 +98,10 @@ import { TraceSpan } from '@/lib/mothership/generated/trace-spans-v1' import { readChatStream } from '@/lib/mothership/request/application/recover-stream' import { createStreamingContext } from '@/lib/mothership/request/context/request-context' import { restoreStreamingContext } from '@/lib/mothership/request/context/restore' +import { + createFilePreviewAdapterState, + processFilePreviewStreamEvent, +} from '@/lib/mothership/request/go/file-preview-adapter' import { finalizeStream } from '@/lib/mothership/request/lifecycle/finalize' import { createSSEStream } from '@/lib/mothership/request/lifecycle/start' import { acquirePendingChatStream } from '@/lib/mothership/request/session/abort' @@ -116,6 +120,10 @@ import { STREAM_STRING_PREVIEW_UNITS } from '@/lib/mothership/request/session/re 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' +import { + type PendingFileIntent, + storeFileIntent, +} from '@/lib/mothership/tools/server/files/file-intent-store' import { GET as streamGET } from '@/app/api/copilot/chat/stream/route' const MB = 1024 * 1024 @@ -593,6 +601,72 @@ describe.runIf(Boolean(redisUrl))('a turn whose stream exhausts its replay budge expect(await redis().ttl(`mothership_stream:${streamId}:events`)).toBeLessThanOrEqual(300) }) + it('completes a turn that previews a file larger than one replay write', async () => { + const editToolCallId = generateId() + const previewToolCallId = generateId() + const toolFrame = (payload: Record): StreamEvent => + ({ + type: 'tool', + payload: { + toolCallId: editToolCallId, + toolName: 'apply_file_edit', + executor: 'sim', + mode: 'async', + ...payload, + }, + }) as StreamEvent + const { runId, frames } = await runTurn([ + async (options: { onEvent?: (event: unknown) => Promise }) => { + // The real preview adapter, fed the worker's patch of a 1.2 MB stored file. + const fileId = generateId() + const anchor = { strategy: 'anchored', mode: 'insert_after', anchor: 'ANCHOR' } + await storeFileIntent(workspaceId, fileId, { + operation: 'patch', + fileId, + workspaceId, + userId, + fileRecord: {} as PendingFileIntent['fileRecord'], + existingContent: `ANCHOR\n${'existing line of the stored file\n'.repeat(40_000)}`, + edit: anchor, + createdAt: Date.now(), + }) + const context = createStreamingContext() + context.activeFileIntents.set('', { + toolCallId: previewToolCallId, + operation: 'patch', + target: { kind: 'file_id', fileId, fileName: 'large.md' }, + edit: anchor, + }) + const state = createFilePreviewAdapterState() + const preview = (streamEvent: StreamEvent) => + processFilePreviewStreamEvent({ + streamId: generateId(), + streamEvent, + context, + execContext: { userId, workflowId: '', workspaceId }, + options: { onEvent: (event) => options.onEvent?.(event) }, + state, + }) + await preview(toolFrame({ phase: 'args_delta', argumentsDelta: '{"content":"' })) + for (let chunk = 0; chunk < 3; chunk++) { + await preview( + toolFrame({ phase: 'args_delta', argumentsDelta: `inserted line ${chunk}` }) + ) + } + await preview(toolFrame({ phase: 'result', success: true, status: 'success' })) + }, + text('The file is updated.'), + ]) + + expect(frames.map((frame) => frame.type)).not.toContain('error') + expect(frames.at(-1)).toMatchObject({ type: 'complete', payload: { status: 'complete' } }) + expect(frames.some((frame) => frame.payload.previewPhase === 'file_preview_complete')).toBe( + true + ) + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('complete') + }) + it("leaves its successor's stream untouched when the lease is lost while ending a refused turn", async () => { const successorToken = `successor\n${generateId()}` worker.hooks.onAbort = async () => { diff --git a/apps/sim/lib/mothership/request/types.ts b/apps/sim/lib/mothership/request/types.ts index 1c11e6b7d4c..ce53d57fa8e 100644 --- a/apps/sim/lib/mothership/request/types.ts +++ b/apps/sim/lib/mothership/request/types.ts @@ -218,6 +218,8 @@ export interface StreamingContext { * main-lane / no-scope intent (file writes there are always sequential). */ activeFileIntents: Map + /** File preview content this turn has streamed; one object shared by every leg. */ + filePreviewBudget: { contentBytes: number } trace: TraceCollector subAgentTraceSpans?: Map /** From 77e7e4ae8ad59c0e17a75aff44bda61ef9b2253a Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 22:48:42 -0700 Subject: [PATCH 13/18] fix(mothership): omit the largest field of an event no cut can bound, and persist closed lanes - An event whose many short values survive every string and array cut, such as an object with thousands of keys, now has its largest bulk field replaced by a size note when it would still exceed one replay write, instead of ending the turn. Identity fields and client-executed arguments are never replaced; such an event is still refused cleanly. - finalizeResidualToolCalls reports closing an open subagent lane as a change, so the error finalizer persists it instead of leaving a stale lane. --- .../home/hooks/stream/stream-helpers.test.ts | 7 +++ .../home/hooks/stream/stream-helpers.ts | 3 +- .../request/session/replay-compaction.test.ts | 46 +++++++++++++++++++ .../request/session/replay-compaction.ts | 34 +++++++++++++- 4 files changed, 88 insertions(+), 2 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts index 2d0e905bb31..501e1fd2de0 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts @@ -27,4 +27,11 @@ describe('finalizeResidualToolCalls', () => { expect(finalizeResidualToolCalls(open, 'error')).toBe(true) expect(finalizeResidualToolCalls(settled, 'error')).toBe(false) }) + + it('reports closing an open subagent lane as a change to persist', () => { + const blocks: ContentBlock[] = [{ type: 'subagent', content: 'research' }] + + expect(finalizeResidualToolCalls(blocks, 'error')).toBe(true) + expect(blocks[0].endedAt).toEqual(expect.any(Number)) + }) }) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts index 3067ba815ba..100cbda798a 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.ts @@ -72,7 +72,7 @@ export function asPayloadRecord(value: unknown): StreamPayload | undefined { * Settles every unfinished tool row (running, pending, or awaiting approval) at * a turn terminal by propagating the turn's outcome: a clean `complete` settles * a straggler `success`, a stop `cancelled`, an error `error`. Also closes any - * open subagent lane. Returns whether any tool row was settled. + * open subagent lane. Returns whether it settled a row or closed a lane. */ export function finalizeResidualToolCalls( blocks: ContentBlock[], @@ -94,6 +94,7 @@ export function finalizeResidualToolCalls( // transport-based gating. if (block.type === 'subagent' && block.endedAt === undefined) { block.endedAt = endedAt + settled = true continue } const tc = block.toolCall diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index e208e561c6e..f205efe0927 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -14,6 +14,7 @@ import { describe, expect, it } from 'vitest' import { compactStreamEvent, STREAM_EVENT_COMPACTION_THRESHOLD_BYTES, + STREAM_EVENT_MAX_PAYLOAD_BYTES, STREAM_STRING_PREVIEW_UNITS, } from '@/lib/mothership/request/session/replay-compaction' import type { StreamEvent } from '@/lib/mothership/request/session/types' @@ -237,4 +238,49 @@ describe('compactStreamEvent', () => { expect(output.data).toEqual({ results: citations }) expect(toRecord(payloadOf(event).output).rows).toHaveLength(40_000) }) + + it('omits the largest remaining field when cuts alone leave the event over one replay write', () => { + const blocks = Object.fromEntries( + Array.from({ length: 5_000 }, (_, index) => [`block-${index}`, 'b'.repeat(300)]) + ) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_workflows_state_get', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + status: 'success', + output: { blocks }, + }, + } + + const compacted = payloadOf(compactStreamEvent(event)) + + expect(Buffer.byteLength(JSON.stringify(compacted))).toBeLessThanOrEqual( + STREAM_EVENT_MAX_PAYLOAD_BYTES + ) + expect(compacted.output).toMatch(/^…\[omitted, [\d.]+ MB total\]$/) + expect(compacted).toMatchObject({ toolCallId: 'c', success: true, status: 'success' }) + }) + + it('never omits identity to make room for client-executed arguments it must keep whole', () => { + const args = { workflowId: 'wf-1', input: 'w'.repeat(1.5 * MB) } + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'run_workflow', + executor: 'client', + mode: 'async', + phase: 'call', + ui: { clientExecutable: true }, + arguments: args, + }, + } + + expect(compactStreamEvent(event)).toBe(event) + }) }) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index 99f76f979c5..6e4c861b6b1 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -1,4 +1,5 @@ import { isRecordLike, toRecordOrNull } from '@sim/utils/object' +import { getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' import { MothershipStreamV1EventType, MothershipStreamV1ToolPhase, @@ -13,6 +14,13 @@ import type { StreamEvent } from './types' */ export const STREAM_EVENT_COMPACTION_THRESHOLD_BYTES = 256 * 1024 +/** Room left in a replay write for the envelope around an event's payload. */ +const ENVELOPE_HEADROOM_BYTES = 16 * 1024 + +/** The largest serialized payload the replay buffer can persist in one write. */ +export const STREAM_EVENT_MAX_PAYLOAD_BYTES = + getRedisBudgetLimits('copilot_stream').maxSingleWriteBytes - ENVELOPE_HEADROOM_BYTES + /** UTF-16 units kept at the head of a long string. */ export const STREAM_STRING_PREVIEW_UNITS = 8 * 1024 @@ -100,11 +108,32 @@ function mapLeaves( return copy ?? value } +/** A field must be at least this large to be omitted as a last resort. */ +const OMITTABLE_FIELD_MIN_BYTES = 64 * 1024 + +/** + * Replaces the payload's largest field with a note of its size: the last resort + * for an event whose many short values no cut could bound, such as an object + * with thousands of keys. + */ +function omitLargestField(payload: Record, skipKeys: ReadonlySet) { + let largest: { key: string; bytes: number } | undefined + for (const [key, field] of Object.entries(payload)) { + if (skipKeys.has(key)) continue + const bytes = Buffer.byteLength(JSON.stringify(field) ?? '', 'utf8') + if (!largest || bytes > largest.bytes) largest = { key, bytes } + } + // Identity fields are never this large, so only bulk is ever replaced. + if (!largest || largest.bytes <= OMITTABLE_FIELD_MIN_BYTES) return payload + return { ...payload, [largest.key]: `…[omitted, ${formatFileSize(largest.bytes)} total]` } +} + /** * Bounds an outgoing stream event so the replay buffer can persist it. Applied * only to the copy the writer delivers and persists; the caller keeps the full * event for dispatch. Long strings are cut to their head in place, so every - * object keeps its shape; if that is not enough, long arrays keep their head. + * object keeps its shape; if that is not enough, long arrays keep their head, + * and past one replay write the largest field is replaced by a size note. * Assistant text, file previews, and the arguments of calls the browser * executes are never cut; an event * still too large is refused by the buffer, which ends the turn with an error. @@ -126,5 +155,8 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { if (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { compacted = trimArrays(compacted, skipKeys) } + if (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_MAX_PAYLOAD_BYTES) { + compacted = omitLargestField(toRecordOrNull(compacted) ?? {}, skipKeys) + } return compacted === payload ? event : ({ ...event, payload: compacted } as StreamEvent) } From 7989cdf12023e3523f1fcc479dc7f78a05c28a39 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 23:05:37 -0700 Subject: [PATCH 14/18] fix(mothership): keep sibling fields when omitting bulk, and never cache a skipped preview as file content - The last-resort omission for an event no cut can bound now walks down from the largest field while one child holds most of its parent's bytes and replaces only that node, so the ids, resources, and status fields beside the bulk keep driving the UI's resource updates. - On completion the client seeds the file content cache from preview text only when it received the completed version. When the server skipped the final content frame, the text is an earlier draft, so the stored file loads instead. --- .../preview/apply-file-preview-phase.test.ts | 27 +++++++- .../hooks/preview/apply-file-preview-phase.ts | 14 +++++ .../preview/use-file-preview-controller.ts | 7 ++- .../request/session/replay-compaction.test.ts | 32 +++++++++- .../request/session/replay-compaction.ts | 63 ++++++++++++++----- 5 files changed, 125 insertions(+), 18 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts index 8ed6eb19032..dfe8bbb1046 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from 'vitest' import type { FilePreviewSession } from '@/lib/mothership/request/session' -import { deriveFilePreviewSession } from './apply-file-preview-phase' +import { deriveFilePreviewSession, previewHoldsFinalContent } from './apply-file-preview-phase' const NOW = '2026-06-08T00:00:00.000Z' @@ -94,3 +94,28 @@ describe('deriveFilePreviewSession', () => { expect(third?.previewText).toBe('ABC') }) }) + +describe('previewHoldsFinalContent', () => { + const complete = (previewVersion?: number) => ({ + previewPhase: 'file_preview_complete' as const, + toolCallId: 'tool-1', + toolName: 'prepare_file_edit' as const, + ...(previewVersion !== undefined ? { previewVersion } : {}), + }) + + it('holds the final content when the last content received is the completed version', () => { + const prev = session({ previewText: 'final text', previewVersion: 7 }) + + expect(previewHoldsFinalContent(prev, complete(7))).toBe(true) + }) + + it('does not hold it when later versions were never received, so the stored file must load', () => { + const prev = session({ previewText: 'an earlier draft', previewVersion: 5 }) + + expect(previewHoldsFinalContent(prev, complete(7))).toBe(false) + }) + + it('does not hold it when no content was received at all', () => { + expect(previewHoldsFinalContent(undefined, complete(7))).toBe(false) + }) +}) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.ts index f9d4c083070..a2d678e9856 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.ts @@ -99,3 +99,17 @@ export function deriveFilePreviewSession( } } } + +/** + * Whether the preview text the client holds is the edit's final content. The + * server skips a content frame too large for the stream, so a completion whose + * version is newer than the last content received means the text is an earlier + * draft, and only the stored file is the saved result. + */ +export function previewHoldsFinalContent( + prev: FilePreviewSession | undefined, + completion: Extract +): boolean { + if (!prev || prev.previewText.length === 0) return false + return completion.previewVersion === undefined || prev.previewVersion >= completion.previewVersion +} diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/use-file-preview-controller.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/use-file-preview-controller.ts index f4916b7c51a..15f6ea479c3 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/use-file-preview-controller.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/use-file-preview-controller.ts @@ -12,7 +12,10 @@ import { useQueryClient } from '@tanstack/react-query' import type { SyntheticFilePreviewPayload } from '@/lib/mothership/request/session' import type { FilePreviewSession } from '@/lib/mothership/request/session/file-preview-session-contract' import { invalidateResourceQueries } from '@/app/workspace/[workspaceId]/home/components/mothership-view/components/resource-registry' -import { deriveFilePreviewSession } from '@/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase' +import { + deriveFilePreviewSession, + previewHoldsFinalContent, +} from '@/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase' import { buildCompletedPreviewSessions, type FilePreviewSessionsState, @@ -381,7 +384,7 @@ export function useFilePreviewController({ sessionId: nextSession.id, suppressActivation: !shouldActivateOnComplete, }) - if (hasRenderableFilePreviewContent(nextSession)) { + if (previewHoldsFinalContent(prevSession, payload)) { seedCompletedPreviewContentCache(fileId, nextSession.previewText) } if (workspaceId) invalidateResourceQueries(queryClient, workspaceId, 'file', fileId) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index f205efe0927..4b811b3ba32 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -262,10 +262,40 @@ describe('compactStreamEvent', () => { expect(Buffer.byteLength(JSON.stringify(compacted))).toBeLessThanOrEqual( STREAM_EVENT_MAX_PAYLOAD_BYTES ) - expect(compacted.output).toMatch(/^…\[omitted, [\d.]+ MB total\]$/) + expect(toRecord(compacted.output).blocks).toMatch(/^…\[omitted, [\d.]+ MB total\]$/) expect(compacted).toMatchObject({ toolCallId: 'c', success: true, status: 'success' }) }) + it('omits only the bulk inside an output, keeping the ids and resources the UI reads', () => { + const blocks = Object.fromEntries( + Array.from({ length: 5_000 }, (_, index) => [`block-${index}`, 'b'.repeat(300)]) + ) + const resources = [{ type: 'workflow', id: 'wf-1', title: 'Pipeline' }] + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_workflows_state_get', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + status: 'success', + output: { workflowId: 'wf-1', resources, data: { id: 'wf-1', name: 'Pipeline', blocks } }, + }, + } + + const output = toRecord(payloadOf(compactStreamEvent(event)).output) + + expect(output.workflowId).toBe('wf-1') + expect(output.resources).toEqual(resources) + expect(output.data).toEqual({ + id: 'wf-1', + name: 'Pipeline', + blocks: expect.stringMatching(/^…\[omitted, [\d.]+ MB total\]$/), + }) + }) + it('never omits identity to make room for client-executed arguments it must keep whole', () => { const args = { workflowId: 'wf-1', input: 'w'.repeat(1.5 * MB) } const event: StreamEvent = { diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index 6e4c861b6b1..ad175ade0e7 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -111,21 +111,56 @@ function mapLeaves( /** A field must be at least this large to be omitted as a last resort. */ const OMITTABLE_FIELD_MIN_BYTES = 64 * 1024 +type Child = { key: string | number; value: unknown; bytes: number } + +/** The largest child of an object or array, by serialized size. */ +function largestChild(value: unknown, skipKeys: ReadonlySet): Child | undefined { + const entries: Array<[string | number, unknown]> = Array.isArray(value) + ? value.map((item, index) => [index, item]) + : isRecordLike(value) + ? Object.entries(value).filter(([key]) => !skipKeys.has(key)) + : [] + let largest: Child | undefined + for (const [key, child] of entries) { + const bytes = Buffer.byteLength(JSON.stringify(child) ?? '', 'utf8') + if (!largest || bytes > largest.bytes) largest = { key, value: child, bytes } + } + return largest +} + +/** Copies `value` along `path`, replacing the value at its end. */ +function replaceAt(value: unknown, path: Array, replacement: unknown): unknown { + if (path.length === 0) return replacement + const [key, ...rest] = path + if (Array.isArray(value)) { + const copy = [...value] + copy[Number(key)] = replaceAt(copy[Number(key)], rest, replacement) + return copy + } + const copy = { ...toRecordOrNull(value) } + copy[String(key)] = replaceAt(copy[String(key)], rest, replacement) + return copy +} + /** - * Replaces the payload's largest field with a note of its size: the last resort - * for an event whose many short values no cut could bound, such as an object - * with thousands of keys. + * The last resort for an event whose many short values no cut could bound, + * such as an object with thousands of keys. Walks down from the largest field + * while one child holds most of its parent's bytes, and replaces only that + * dominating node with a note of its size, so the small siblings the UI reads + * (ids, resources, status) stay intact. */ -function omitLargestField(payload: Record, skipKeys: ReadonlySet) { - let largest: { key: string; bytes: number } | undefined - for (const [key, field] of Object.entries(payload)) { - if (skipKeys.has(key)) continue - const bytes = Buffer.byteLength(JSON.stringify(field) ?? '', 'utf8') - if (!largest || bytes > largest.bytes) largest = { key, bytes } - } +function omitDominantBulk(payload: Record, skipKeys: ReadonlySet) { + let node = largestChild(payload, skipKeys) // Identity fields are never this large, so only bulk is ever replaced. - if (!largest || largest.bytes <= OMITTABLE_FIELD_MIN_BYTES) return payload - return { ...payload, [largest.key]: `…[omitted, ${formatFileSize(largest.bytes)} total]` } + if (!node || node.bytes <= OMITTABLE_FIELD_MIN_BYTES) return payload + const path = [node.key] + for (;;) { + const child = largestChild(node.value, NO_KEYS) + if (!child || child.bytes * 2 <= node.bytes || child.bytes <= OMITTABLE_FIELD_MIN_BYTES) break + path.push(child.key) + node = child + } + return replaceAt(payload, path, `…[omitted, ${formatFileSize(node.bytes)} total]`) } /** @@ -133,7 +168,7 @@ function omitLargestField(payload: Record, skipKeys: ReadonlySe * only to the copy the writer delivers and persists; the caller keeps the full * event for dispatch. Long strings are cut to their head in place, so every * object keeps its shape; if that is not enough, long arrays keep their head, - * and past one replay write the largest field is replaced by a size note. + * and past one replay write the dominating bulk is replaced by a size note. * Assistant text, file previews, and the arguments of calls the browser * executes are never cut; an event * still too large is refused by the buffer, which ends the turn with an error. @@ -156,7 +191,7 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { compacted = trimArrays(compacted, skipKeys) } if (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_MAX_PAYLOAD_BYTES) { - compacted = omitLargestField(toRecordOrNull(compacted) ?? {}, skipKeys) + compacted = omitDominantBulk(toRecordOrNull(compacted) ?? {}, skipKeys) } return compacted === payload ? event : ({ ...event, payload: compacted } as StreamEvent) } From fc2149c56b0a639e5bf8a11cb47f1d37a1395239 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Tue, 29 Sep 2026 23:39:35 -0700 Subject: [PATCH 15/18] fix(mothership): keep omitting bulk until an event fits one replay write A result with more than one large object of short fields left the second one over the write limit after the first was omitted, so the event was refused and the turn ended. Omission now repeats until the event fits or nothing large enough is left; each pass removes at least 64 KiB, so it always terminates. --- .../request/session/replay-compaction.test.ts | 30 +++++++++++++++++++ .../request/session/replay-compaction.ts | 11 +++++-- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index 4b811b3ba32..bc3134aac5e 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -296,6 +296,36 @@ describe('compactStreamEvent', () => { }) }) + it('keeps omitting bulk until the event fits when more than one large object remains', () => { + const manyKeys = (prefix: string, count: number) => + Object.fromEntries( + Array.from({ length: count }, (_, index) => [`${prefix}-${index}`, 'v'.repeat(300)]) + ) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_workflows_state_get', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + status: 'success', + output: { workflowId: 'wf-1', a: manyKeys('a', 4_300), b: manyKeys('b', 3_600) }, + }, + } + + const compacted = payloadOf(compactStreamEvent(event)) + const output = toRecord(compacted.output) + + expect(Buffer.byteLength(JSON.stringify(compacted))).toBeLessThanOrEqual( + STREAM_EVENT_MAX_PAYLOAD_BYTES + ) + expect(output.workflowId).toBe('wf-1') + expect(output.a).toMatch(/^…\[omitted, [\d.]+ MB total\]$/) + expect(output.b).toMatch(/^…\[omitted, [\d.]+ MB total\]$/) + }) + it('never omits identity to make room for client-executed arguments it must keep whole', () => { const args = { workflowId: 'wf-1', input: 'w'.repeat(1.5 * MB) } const event: StreamEvent = { diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index ad175ade0e7..5f372656935 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -168,7 +168,8 @@ function omitDominantBulk(payload: Record, skipKeys: ReadonlySe * only to the copy the writer delivers and persists; the caller keeps the full * event for dispatch. Long strings are cut to their head in place, so every * object keeps its shape; if that is not enough, long arrays keep their head, - * and past one replay write the dominating bulk is replaced by a size note. + * and past one replay write each dominating bulk is replaced by a size note + * until the event fits. * Assistant text, file previews, and the arguments of calls the browser * executes are never cut; an event * still too large is refused by the buffer, which ends the turn with an error. @@ -190,8 +191,12 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { if (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { compacted = trimArrays(compacted, skipKeys) } - if (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_MAX_PAYLOAD_BYTES) { - compacted = omitDominantBulk(toRecordOrNull(compacted) ?? {}, skipKeys) + // Each pass replaces at least OMITTABLE_FIELD_MIN_BYTES, and stops when nothing is left to omit. + while (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_MAX_PAYLOAD_BYTES) { + const record = toRecordOrNull(compacted) ?? {} + const omitted = omitDominantBulk(record, skipKeys) + if (omitted === record) break + compacted = omitted } return compacted === payload ? event : ({ ...event, payload: compacted } as StreamEvent) } From 7aa6002332708b3d4b664bb249d67be13660c7e9 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 00:02:27 -0700 Subject: [PATCH 16/18] fix(mothership): shed oversized replay events in one pass and retry mid-body stream cuts - Replace the replay-compaction omit loop with one bounded pass over memoized exact serialized sizes: recurse into the smallest child that covers the overage, otherwise drop the largest children whole, keeping small siblings - Map a non-abort body-read failure from the worker stream to a retryable WorkerStreamInterruptedError on the reachable budget, with the generic message and the original error on cause - Log the cause message on retry warnings and orchestration failures - Count replayed file preview content toward a restored turn's preview budget --- .../lib/mothership/request/context/restore.ts | 6 + .../lib/mothership/request/go/stream.test.ts | 47 ++++++ apps/sim/lib/mothership/request/go/stream.ts | 20 ++- .../request/handlers/handlers.test.ts | 23 +++ .../mothership/request/lifecycle/run.test.ts | 8 + .../lib/mothership/request/lifecycle/run.ts | 9 ++ .../request/lifecycle/stream-retry.test.ts | 13 ++ .../request/lifecycle/stream-retry.ts | 7 +- .../request/session/replay-compaction.test.ts | 126 +++++++++++++++ .../request/session/replay-compaction.ts | 147 ++++++++++++------ 10 files changed, 357 insertions(+), 49 deletions(-) diff --git a/apps/sim/lib/mothership/request/context/restore.ts b/apps/sim/lib/mothership/request/context/restore.ts index bb02d7239a4..6a6cdcea684 100644 --- a/apps/sim/lib/mothership/request/context/restore.ts +++ b/apps/sim/lib/mothership/request/context/restore.ts @@ -13,6 +13,12 @@ export async function restoreStreamingContext( execContext: ExecutionContext ): Promise { for (const saved of events) { + // Preview content already in the replay counts toward this turn's preview budget. + if (saved.type === 'tool' && 'previewPhase' in saved.payload) { + if (saved.payload.previewPhase === 'file_preview_content') { + context.filePreviewBudget.contentBytes += Buffer.byteLength(saved.payload.content, 'utf8') + } + } const event = reconcileTextEvent(saved, context.accumulatedContent) if (!event) continue const replay: StreamEvent = diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index 75b2013d361..fc7db309247 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -785,6 +785,53 @@ describe('copilot go stream helpers', () => { expect(fetch).toHaveBeenCalledTimes(1) }) + it('reports a stream cut mid-body without the raw socket error', async () => { + const socketError = Object.assign( + new Error( + 'The socket connection was closed unexpectedly. For more information, pass `verbose: true`' + ), + { code: 'ECONNRESET' } + ) + const first = createEvent({ + streamId: 'cut-stream', + cursor: '1', + seq: 1, + requestId: 'req-cut', + type: 'text', + payload: { channel: 'assistant', text: 'partial' }, + }) + vi.mocked(fetch).mockResolvedValueOnce( + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(first)}\n\n`)) + }, + pull(controller) { + controller.error(socketError) + }, + }), + { status: 200, headers: { 'Content-Type': 'text/event-stream' } } + ) + ) + + await expect( + runStreamLoop( + 'https://example.com/mothership/stream', + {}, + createStreamingContext(), + turnScopedExecContext(), + { + timeout: 1000, + flushAfterEvent: false, + } + ) + ).rejects.toMatchObject({ + name: 'WorkerStreamInterruptedError', + message: 'The agent service is temporarily unavailable. Please try again.', + cause: socketError, + }) + }) + it('reports a worker it could not reach without the raw network error', async () => { const networkError = new TypeError('fetch failed') vi.mocked(fetch).mockRejectedValueOnce(networkError) diff --git a/apps/sim/lib/mothership/request/go/stream.ts b/apps/sim/lib/mothership/request/go/stream.ts index 4369839993c..b55b623458d 100644 --- a/apps/sim/lib/mothership/request/go/stream.ts +++ b/apps/sim/lib/mothership/request/go/stream.ts @@ -81,6 +81,18 @@ export class WorkerUnreachableError extends Error { } } +/** + * The worker's response body failed mid-stream (the connection was reset or + * closed). The worker answered, so a retry reattaches under the short budget. + * The read error stays on `cause` for logs. + */ +export class WorkerStreamInterruptedError extends Error { + constructor(cause: unknown) { + super(BACKEND_UNAVAILABLE_MESSAGE, { cause }) + this.name = 'WorkerStreamInterruptedError' + } +} + /** * A worker rejection message the user can act on: short, one line, plain text, * and free of identifiers (`userId`, `protocol_version_mismatch`) that only mean @@ -293,7 +305,13 @@ export async function runStreamLoop( const rawReader = response.body.getReader() const reader: ReadableStreamDefaultReader = { async read() { - const result = await rawReader.read() + let result: ReadableStreamReadResult + try { + result = await rawReader.read() + } catch (error) { + if (requestSignal.aborted) throw error + throw new WorkerStreamInterruptedError(error) + } if (!result.done && result.value) { const now = performance.now() const gap = now - counters.lastChunkMs diff --git a/apps/sim/lib/mothership/request/handlers/handlers.test.ts b/apps/sim/lib/mothership/request/handlers/handlers.test.ts index d458500329e..2b4d8f04274 100644 --- a/apps/sim/lib/mothership/request/handlers/handlers.test.ts +++ b/apps/sim/lib/mothership/request/handlers/handlers.test.ts @@ -215,6 +215,29 @@ describe('sse-handlers tool lifecycle', () => { expect(context.subAgentTraceSpans?.size).toBe(0) }) + it('counts the preview content already in the replay toward the restored turn budget', async () => { + const preview = (content: string): StreamEvent => ({ + type: 'tool', + payload: { + toolCallId: 'file-edit', + toolName: 'prepare_file_edit', + previewPhase: 'file_preview_content', + content, + contentMode: 'delta', + previewVersion: 1, + fileName: 'notes.md', + }, + }) + + await restoreStreamingContext( + [preview('é'.repeat(1_000)), preview('x'.repeat(500))], + context, + execContext + ) + + expect(context.filePreviewBudget.contentBytes).toBe(2_500) + }) + it('restores a delivered prefix without executing its tools or consuming the next handoff', async () => { const call: StreamEvent = { type: 'tool', diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index fa35ccdc6ee..8fd87610b48 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -128,9 +128,17 @@ vi.mock('@/lib/mothership/request/go/stream', () => { } } + class WorkerStreamInterruptedError extends Error { + constructor(cause: unknown) { + super('The agent service is temporarily unavailable. Please try again.', { cause }) + this.name = 'WorkerStreamInterruptedError' + } + } + return { BillingLimitError, CopilotBackendError, + WorkerStreamInterruptedError, WorkerUnreachableError, STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE, StreamEndedWithoutTerminalError, diff --git a/apps/sim/lib/mothership/request/lifecycle/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index 59a326dd13a..9a9ac2cd667 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -593,6 +593,7 @@ export async function runCopilotLifecycle( // explained, not just reduced to a message string. logger.error('Copilot orchestration failed', { error: err.message, + ...causeForLog(err), name: err.name, ...(error instanceof CopilotBackendError ? { backendStatus: error.status, backendBody: error.body?.slice(0, 2000) } @@ -878,6 +879,7 @@ async function runResumeLegWithRetry( attempt: retry.attempts + 1, backoffMs: backoff, error: toError(error).message, + ...causeForLog(error), }) await interruptibleSleep(backoff, options.abortSignal) continue @@ -1240,6 +1242,7 @@ async function runCheckpointLoop( attempt: (retry?.attempts ?? 0) + 1, backoffMs: backoff, error: toError(streamError).message, + ...causeForLog(streamError), } ) await interruptibleSleep(backoff, options.abortSignal) @@ -1678,6 +1681,12 @@ async function withEnterpriseByokKey( return byokApiKey ? { ...refreshed, byokApiKey } : refreshed } +/** The underlying failure behind a generic user-facing error, for logs. */ +function causeForLog(error: unknown): { cause?: string } { + const cause = error instanceof Error ? error.cause : undefined + return cause === undefined ? {} : { cause: getErrorMessage(cause) } +} + function isAborted(options: CopilotLifecycleOptions, context: StreamingContext): boolean { return !!(options.abortSignal?.aborted || context.wasAborted) } diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index 0e77a2c9f29..10afce65508 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -2,6 +2,7 @@ import { afterEach, describe, expect, it, vi } from 'vitest' import { CopilotBackendError, StreamEndedWithoutTerminalError, + WorkerStreamInterruptedError, WorkerUnreachableError, } from '@/lib/mothership/request/go/stream' import { StreamRetryWindow } from '@/lib/mothership/request/lifecycle/stream-retry' @@ -150,4 +151,16 @@ describe('stream recovery budget', () => { expect(retry.nextDelay(new Error('Invalid operation'))).toBeNull() expect(retry.attempt).toBe(0) }) + + it('gives a stream cut mid-body the reachable budget, not the unreachable window', () => { + vi.useFakeTimers() + const error = new WorkerStreamInterruptedError(new Error('socket closed')) + const retry = new StreamRetryWindow() + for (let index = 0; index < 3; index++) { + const delay = retry.nextDelay(error) + expect(delay).not.toBeNull() + vi.advanceTimersByTime(delay ?? 0) + } + expect(retry.nextDelay(error)).toBeNull() + }) }) diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts index d99d4fb26c4..677d8bb8922 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -4,6 +4,7 @@ import { StreamContinuityError } from '@/lib/mothership/request/go/parser' import { CopilotBackendError, StreamEndedWithoutTerminalError, + WorkerStreamInterruptedError, WorkerUnreachableError, } from '@/lib/mothership/request/go/stream' @@ -108,7 +109,11 @@ function isJson(body: string | undefined): boolean { /** Initial sends and resumes both replay one durable identity after an ambiguous response. */ function isRetryableStreamError(error: unknown): boolean { if (error instanceof Error && error.name === 'AbortError') return false - if (error instanceof StreamEndedWithoutTerminalError || error instanceof StreamContinuityError) { + if ( + error instanceof StreamEndedWithoutTerminalError || + error instanceof StreamContinuityError || + error instanceof WorkerStreamInterruptedError + ) { return true } if (error instanceof CopilotBackendError) { diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index bc3134aac5e..06c384e4513 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -16,6 +16,7 @@ import { STREAM_EVENT_COMPACTION_THRESHOLD_BYTES, STREAM_EVENT_MAX_PAYLOAD_BYTES, STREAM_STRING_PREVIEW_UNITS, + serializedBytes, } from '@/lib/mothership/request/session/replay-compaction' import type { StreamEvent } from '@/lib/mothership/request/session/types' @@ -343,4 +344,129 @@ describe('compactStreamEvent', () => { expect(compactStreamEvent(event)).toBe(event) }) + + it('omits only the smallest sufficient bulk when no single child dominates, keeping its siblings', () => { + const manyKeys = (prefix: string) => + Object.fromEntries( + Array.from({ length: 2_300 }, (_, index) => [`${prefix}-${index}`, 'v'.repeat(300)]) + ) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_tables_rows_query', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + output: { fileId: 'file-1', rows: manyKeys('r'), cols: manyKeys('c') }, + }, + } + + const compacted = payloadOf(compactStreamEvent(event)) + const output = toRecord(compacted.output) + + expect(Buffer.byteLength(JSON.stringify(compacted))).toBeLessThanOrEqual( + STREAM_EVENT_MAX_PAYLOAD_BYTES + ) + expect(output.fileId).toBe('file-1') + expect([output.rows, output.cols].filter((value) => typeof value === 'string')).toHaveLength(1) + }) + + it('omits the smaller of two bulks when either alone would make the event fit', () => { + const manyKeys = (prefix: string, count: number) => + Object.fromEntries( + Array.from({ length: count }, (_, index) => [`${prefix}-${index}`, 'v'.repeat(300)]) + ) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_tables_rows_query', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + output: { fileId: 'file-1', rows: manyKeys('r', 3_000), cols: manyKeys('c', 1_900) }, + }, + } + + const output = toRecord(payloadOf(compactStreamEvent(event)).output) + + expect(output.fileId).toBe('file-1') + expect(Object.keys(toRecord(output.rows))).toHaveLength(3_000) + expect(output.cols).toMatch(/^…\[omitted, [\d.]+ KB total\]$/) + }) + + it('omits a sufficient bulk whole when its own large parts cannot cover the overage', () => { + const manyKeys = (prefix: string, count: number) => + Object.fromEntries( + Array.from({ length: count }, (_, index) => [`${prefix}-${index}`, 'v'.repeat(3_000)]) + ) + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_workflows_state_get', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + output: { workflowId: 'wf-1', state: { nested: manyKeys('n', 35), ...manyKeys('k', 400) } }, + }, + } + + const compacted = payloadOf(compactStreamEvent(event)) + const output = toRecord(compacted.output) + + expect(Buffer.byteLength(JSON.stringify(compacted))).toBeLessThanOrEqual( + STREAM_EVENT_MAX_PAYLOAD_BYTES + ) + expect(output.workflowId).toBe('wf-1') + expect(output.state).toMatch(/^…\[omitted, [\d.]+ MB total\]$/) + }) + + it('bounds a balanced 16 MB tree in one pass', () => { + const tree = (depth: number, leaf: number): unknown => + depth === 0 + ? 'x'.repeat(leaf) + : { l: tree(depth - 1, leaf + 100), r: tree(depth - 1, Math.max(leaf - 100, 1)) } + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_blocks_get', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + output: tree(12, 3_900), + }, + } + + const started = performance.now() + const compacted = compactStreamEvent(event) + const elapsedMs = performance.now() - started + + expect(Buffer.byteLength(JSON.stringify(compacted.payload))).toBeLessThanOrEqual( + STREAM_EVENT_MAX_PAYLOAD_BYTES + ) + expect(elapsedMs).toBeLessThan(2_000) + }) + + it('measures serialized size exactly without serializing', () => { + const values: unknown[] = [ + { a: 'é😀"\\\n', b: [1, null, true, { c: 'd' }], e: undefined, f: -1.5e-7 }, + [undefined, 'x', { 'kéy "q"': 0 }], + 'plain', + 42, + null, + {}, + [], + ] + + for (const value of values) { + expect(serializedBytes(value)).toBe(Buffer.byteLength(JSON.stringify(value))) + } + }) }) diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index 5f372656935..dcfb8e3883d 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -111,65 +111,120 @@ function mapLeaves( /** A field must be at least this large to be omitted as a last resort. */ const OMITTABLE_FIELD_MIN_BYTES = 64 * 1024 -type Child = { key: string | number; value: unknown; bytes: number } +/** + * The exact UTF-8 size of `JSON.stringify(value)`, computed bottom-up without + * building the string. `sizes` memoizes containers so a caller can ask about + * every node of one tree in a single linear pass. + */ +export function serializedBytes(value: unknown, sizes = new WeakMap()): number { + if ( + value === null || + typeof value !== 'object' || + typeof (value as { toJSON?: unknown }).toJSON === 'function' + ) { + return Buffer.byteLength(JSON.stringify(value) ?? 'null', 'utf8') + } + const cached = sizes.get(value) + if (cached !== undefined) return cached + let bytes = 2 + let members = 0 + if (Array.isArray(value)) { + for (const item of value) { + bytes += isSkippedByJson(item) ? 4 : serializedBytes(item, sizes) + members++ + } + } else { + for (const [key, field] of Object.entries(value)) { + if (isSkippedByJson(field)) continue + bytes += Buffer.byteLength(JSON.stringify(key), 'utf8') + 1 + serializedBytes(field, sizes) + members++ + } + } + bytes += Math.max(members - 1, 0) + sizes.set(value, bytes) + return bytes +} -/** The largest child of an object or array, by serialized size. */ -function largestChild(value: unknown, skipKeys: ReadonlySet): Child | undefined { +/** Values JSON leaves out of an object, or writes as `null` in an array. */ +function isSkippedByJson(value: unknown): boolean { + return value === undefined || typeof value === 'function' || typeof value === 'symbol' +} + +function omissionNote(bytes: number): string { + return `…[omitted, ${formatFileSize(bytes)} total]` +} + +type Sized = { key: string | number; value: unknown; bytes: number } + +/** + * Removes at least `need` serialized bytes from `value`, replacing as little as + * possible with size notes. Among the children larger than the omission floor, + * it recurses into the smallest one that alone covers what is still needed, so + * that subtree's siblings survive; failing that it replaces the largest whole + * and looks again. A nested node whose large parts cannot cover the need, or + * that has none, is replaced whole; the payload itself (`top`) never is. One + * pass over sizes that are each computed once. + */ +function shed( + value: unknown, + need: number, + sizes: WeakMap, + skipKeys: ReadonlySet, + top = false +): unknown { + const bytes = serializedBytes(value, sizes) + const isContainer = Array.isArray(value) || isRecordLike(value) + if (!top && (!isContainer || bytes <= OMITTABLE_FIELD_MIN_BYTES)) return omissionNote(bytes) const entries: Array<[string | number, unknown]> = Array.isArray(value) ? value.map((item, index) => [index, item]) : isRecordLike(value) ? Object.entries(value).filter(([key]) => !skipKeys.has(key)) : [] - let largest: Child | undefined - for (const [key, child] of entries) { - const bytes = Buffer.byteLength(JSON.stringify(child) ?? '', 'utf8') - if (!largest || bytes > largest.bytes) largest = { key, value: child, bytes } + // Identity fields are never this large, so only bulk is ever replaced. + const children: Sized[] = entries + .map(([key, child]) => ({ key, value: child, bytes: serializedBytes(child, sizes) })) + .filter((child) => child.bytes > OMITTABLE_FIELD_MIN_BYTES) + .sort((left, right) => right.bytes - left.bytes) + if (children.length === 0) return top ? value : omissionNote(bytes) + + const gain = (child: Sized) => + child.bytes - Buffer.byteLength(JSON.stringify(omissionNote(child.bytes)), 'utf8') + const replacements = new Map() + let remaining = need + for (let index = 0; index < children.length && remaining > 0; index++) { + let sufficient: Sized | undefined + for (let candidate = children.length - 1; candidate >= index; candidate--) { + if (gain(children[candidate]) >= remaining) { + sufficient = children[candidate] + break + } + } + if (sufficient) { + replacements.set(sufficient.key, shed(sufficient.value, remaining, sizes, NO_KEYS)) + remaining = 0 + } else { + replacements.set(children[index].key, omissionNote(children[index].bytes)) + remaining -= gain(children[index]) + } } - return largest -} + // A caller picks a nested node because replacing it whole covers the need. + if (remaining > 0 && !top) return omissionNote(bytes) -/** Copies `value` along `path`, replacing the value at its end. */ -function replaceAt(value: unknown, path: Array, replacement: unknown): unknown { - if (path.length === 0) return replacement - const [key, ...rest] = path if (Array.isArray(value)) { - const copy = [...value] - copy[Number(key)] = replaceAt(copy[Number(key)], rest, replacement) - return copy + return value.map((item, index) => (replacements.has(index) ? replacements.get(index) : item)) } const copy = { ...toRecordOrNull(value) } - copy[String(key)] = replaceAt(copy[String(key)], rest, replacement) + for (const [key, replacement] of replacements) copy[String(key)] = replacement return copy } -/** - * The last resort for an event whose many short values no cut could bound, - * such as an object with thousands of keys. Walks down from the largest field - * while one child holds most of its parent's bytes, and replaces only that - * dominating node with a note of its size, so the small siblings the UI reads - * (ids, resources, status) stay intact. - */ -function omitDominantBulk(payload: Record, skipKeys: ReadonlySet) { - let node = largestChild(payload, skipKeys) - // Identity fields are never this large, so only bulk is ever replaced. - if (!node || node.bytes <= OMITTABLE_FIELD_MIN_BYTES) return payload - const path = [node.key] - for (;;) { - const child = largestChild(node.value, NO_KEYS) - if (!child || child.bytes * 2 <= node.bytes || child.bytes <= OMITTABLE_FIELD_MIN_BYTES) break - path.push(child.key) - node = child - } - return replaceAt(payload, path, `…[omitted, ${formatFileSize(node.bytes)} total]`) -} - /** * Bounds an outgoing stream event so the replay buffer can persist it. Applied * only to the copy the writer delivers and persists; the caller keeps the full * event for dispatch. Long strings are cut to their head in place, so every * object keeps its shape; if that is not enough, long arrays keep their head, - * and past one replay write each dominating bulk is replaced by a size note - * until the event fits. + * and past one replay write the smallest sufficient bulk is replaced by a size + * note, keeping the fields beside it. * Assistant text, file previews, and the arguments of calls the browser * executes are never cut; an event * still too large is refused by the buffer, which ends the turn with an error. @@ -188,15 +243,13 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { ? ARGUMENTS_KEY : NO_KEYS let compacted = truncateStrings(payload, skipKeys) - if (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { + if (serializedBytes(compacted) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { compacted = trimArrays(compacted, skipKeys) } - // Each pass replaces at least OMITTABLE_FIELD_MIN_BYTES, and stops when nothing is left to omit. - while (Buffer.byteLength(JSON.stringify(compacted)) > STREAM_EVENT_MAX_PAYLOAD_BYTES) { - const record = toRecordOrNull(compacted) ?? {} - const omitted = omitDominantBulk(record, skipKeys) - if (omitted === record) break - compacted = omitted + const sizes = new WeakMap() + const bytes = serializedBytes(compacted, sizes) + if (bytes > STREAM_EVENT_MAX_PAYLOAD_BYTES) { + compacted = shed(compacted, bytes - STREAM_EVENT_MAX_PAYLOAD_BYTES, sizes, skipKeys, true) } return compacted === payload ? event : ({ ...event, payload: compacted } as StreamEvent) } From aea7ed0d9a4e671b98ffd887443b717b25e431b6 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 00:28:44 -0700 Subject: [PATCH 17/18] refactor(mothership): tidy replay compaction and stream retry after review - Share the replay write ceiling between compaction and file previews - Scan for the smallest sufficient child only once one can cover the need, and reuse one size memo across every compaction step - Stop retrying a raw TypeError; fetch and body-read failures are already wrapped as WorkerUnreachableError and WorkerStreamInterruptedError - Use absolute imports, a private retry counter, a typed unsettled-state set, and hasUnsettledTool naming - Cover a timed-out body read and the preview completion edge cases --- .../home/hooks/message-reconcile.ts | 4 +- .../preview/apply-file-preview-phase.test.ts | 17 ++++++- .../[workspaceId]/home/hooks/use-chat.ts | 4 +- .../lib/mothership/chat/display-message.ts | 8 +-- .../lib/mothership/chat/persisted-message.ts | 11 ++-- .../lib/mothership/request/context/restore.ts | 10 ++-- .../request/go/file-preview-adapter.test.ts | 4 +- .../request/go/file-preview-adapter.ts | 13 ++--- .../lib/mothership/request/go/stream.test.ts | 33 ++++++++++++ .../mothership/request/lifecycle/run.test.ts | 51 ++----------------- .../request/lifecycle/stream-retry.test.ts | 23 ++++----- .../request/lifecycle/stream-retry.ts | 4 +- .../request/session/replay-compaction.ts | 38 +++++++------- .../lib/mothership/request/session/writer.ts | 24 +++++---- 14 files changed, 124 insertions(+), 120 deletions(-) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts index 8686d925660..47d5373c2df 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/message-reconcile.ts @@ -146,11 +146,11 @@ export function buildAssistantSnapshotMessage(params: { } export function markMessageStopped(message: PersistedMessage): PersistedMessage { - const hasExecutingTool = message.contentBlocks?.some((block) => + const hasUnsettledTool = message.contentBlocks?.some((block) => isUnsettledToolState(block.toolCall?.state) ) const hasOpenBlock = message.contentBlocks?.some((block) => block.endedAt === undefined) - if (!hasExecutingTool && !hasOpenBlock) { + if (!hasUnsettledTool && !hasOpenBlock) { return message } diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts index dfe8bbb1046..2b83cc3d72d 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts @@ -1,6 +1,9 @@ import { describe, expect, it } from 'vitest' import type { FilePreviewSession } from '@/lib/mothership/request/session' -import { deriveFilePreviewSession, previewHoldsFinalContent } from './apply-file-preview-phase' +import { + deriveFilePreviewSession, + previewHoldsFinalContent, +} from '@/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase' const NOW = '2026-06-08T00:00:00.000Z' @@ -118,4 +121,16 @@ describe('previewHoldsFinalContent', () => { it('does not hold it when no content was received at all', () => { expect(previewHoldsFinalContent(undefined, complete(7))).toBe(false) }) + + it('does not hold it when the session exists but received no text', () => { + const prev = session({ previewText: '', previewVersion: 7 }) + + expect(previewHoldsFinalContent(prev, complete(7))).toBe(false) + }) + + it('holds the received text when the completion carries no version to compare', () => { + const prev = session({ previewText: 'final text', previewVersion: 3 }) + + expect(previewHoldsFinalContent(prev, complete())).toBe(true) + }) }) diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index d6b514fc733..92da00740fe 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -4370,11 +4370,11 @@ export function useChat( } else { setPendingMessages((prev) => prev.map((msg) => { - const hasExecutingTool = msg.contentBlocks?.some((block) => + const hasUnsettledTool = msg.contentBlocks?.some((block) => isUnsettledToolState(block.toolCall?.status) ) const hasOpenBlock = msg.contentBlocks?.some((block) => block.endedAt === undefined) - if (!hasExecutingTool && !hasOpenBlock) { + if (!hasUnsettledTool && !hasOpenBlock) { return msg } const updatedBlocks: ContentBlock[] = (msg.contentBlocks ?? []).map((block) => ({ diff --git a/apps/sim/lib/mothership/chat/display-message.ts b/apps/sim/lib/mothership/chat/display-message.ts index 56607544997..c2710d16cdc 100644 --- a/apps/sim/lib/mothership/chat/display-message.ts +++ b/apps/sim/lib/mothership/chat/display-message.ts @@ -1,4 +1,8 @@ import type { PersistedContentBlock } from '@/lib/api/contracts/copilot-messages' +import { getMothershipAttachmentPreviewUrl } from '@/lib/mothership/chat/attachment-preview' +import { isLiveAssistantMessageId } from '@/lib/mothership/chat/live-message-id' +import type { PersistedMessage } from '@/lib/mothership/chat/persisted-message' +import { isUnsettledToolState, withBlockTiming } from '@/lib/mothership/chat/persisted-message' import { MothershipStreamV1CompletionStatus, MothershipStreamV1EventType, @@ -17,10 +21,6 @@ import { type ToolCallInfo, ToolCallStatus, } from '@/app/workspace/[workspaceId]/home/types' -import { getMothershipAttachmentPreviewUrl } from './attachment-preview' -import { isLiveAssistantMessageId } from './live-message-id' -import type { PersistedMessage } from './persisted-message' -import { isUnsettledToolState, withBlockTiming } from './persisted-message' const STATE_TO_STATUS: Record = { [MothershipStreamV1ToolOutcome.success]: ToolCallStatus.success, diff --git a/apps/sim/lib/mothership/chat/persisted-message.ts b/apps/sim/lib/mothership/chat/persisted-message.ts index 2d0481cda1e..06394ab5267 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.ts @@ -21,7 +21,11 @@ import { MothershipStreamV1ToolOutcome, MothershipStreamV1ToolPhase, } from '@/lib/mothership/generated/mothership-stream-v1' -import type { ContentBlock, OrchestratorResult } from '@/lib/mothership/request/types' +import type { + ContentBlock, + LocalToolCallStatus, + OrchestratorResult, +} from '@/lib/mothership/request/types' import { RETIRED_BROWSER_REQUEST_TAKEOVER_ID } from '@/lib/mothership/tools/retired-tools' import { normalizeToolActivityDescription } from '@/lib/mothership/tools/tool-display' import type { BrowserTextSelection, TerminalTextSelection } from '@/stores/panel/types' @@ -362,7 +366,7 @@ export function buildPersistedAssistantMessage( return message } -const UNSETTLED_TOOL_STATES: ReadonlySet = new Set([ +const UNSETTLED_TOOL_STATES: ReadonlySet = new Set([ 'pending', 'executing', 'awaiting_approval', @@ -370,7 +374,8 @@ const UNSETTLED_TOOL_STATES: ReadonlySet = new Set([ /** A tool row that has not finished: waiting to run, running, or awaiting a decision. */ export function isUnsettledToolState(state: string | undefined): boolean { - return state !== undefined && UNSETTLED_TOOL_STATES.has(state) + const unsettled: ReadonlySet = UNSETTLED_TOOL_STATES + return unsettled.has(state) } /** Settles every unfinished tool row at a turn terminal so none reloads as a spinner. */ diff --git a/apps/sim/lib/mothership/request/context/restore.ts b/apps/sim/lib/mothership/request/context/restore.ts index 6a6cdcea684..36bc8630b78 100644 --- a/apps/sim/lib/mothership/request/context/restore.ts +++ b/apps/sim/lib/mothership/request/context/restore.ts @@ -14,10 +14,12 @@ export async function restoreStreamingContext( ): Promise { for (const saved of events) { // Preview content already in the replay counts toward this turn's preview budget. - if (saved.type === 'tool' && 'previewPhase' in saved.payload) { - if (saved.payload.previewPhase === 'file_preview_content') { - context.filePreviewBudget.contentBytes += Buffer.byteLength(saved.payload.content, 'utf8') - } + if ( + saved.type === 'tool' && + 'previewPhase' in saved.payload && + saved.payload.previewPhase === 'file_preview_content' + ) { + context.filePreviewBudget.contentBytes += Buffer.byteLength(saved.payload.content, 'utf8') } const event = reconcileTextEvent(saved, context.accumulatedContent) if (!event) continue diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts index dc1f0e28832..f83781dcee3 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts @@ -25,11 +25,11 @@ import { createStreamingContext } from '@/lib/mothership/request/context/request import { createFilePreviewAdapterState, type FilePreviewAdapterState, - PREVIEW_FRAME_MAX_BYTES, PREVIEW_TURN_CONTENT_BYTES, processFilePreviewStreamEvent, } from '@/lib/mothership/request/go/file-preview-adapter' import { createEvent, eventToStreamEvent } from '@/lib/mothership/request/session' +import { STREAM_EVENT_MAX_PAYLOAD_BYTES } from '@/lib/mothership/request/session/replay-compaction' import type { ActiveFileIntent, ExecutionContext, @@ -262,7 +262,7 @@ describe('processFilePreviewStreamEvent — preview byte rate', () => { for (const payload of payloads) { expect(Buffer.byteLength(JSON.stringify(payload))).toBeLessThanOrEqual( - PREVIEW_FRAME_MAX_BYTES + STREAM_EVENT_MAX_PAYLOAD_BYTES ) } expect(completed).toBe(true) diff --git a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts index 46efed8bd09..0dedab3b707 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts @@ -1,7 +1,6 @@ import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { isRecordLike } from '@sim/utils/object' -import { getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' import { executeCopilotFileUseCase } from '@/lib/mothership/application/execute-file-use-case' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' import { @@ -15,6 +14,7 @@ import { type SyntheticFilePreviewPayload, upsertFilePreviewSession, } from '@/lib/mothership/request/session' +import { STREAM_EVENT_MAX_PAYLOAD_BYTES } from '@/lib/mothership/request/session/replay-compaction' import type { ActiveFileIntent, ExecutionContext, @@ -70,13 +70,6 @@ const PREVIEW_SNAPSHOT_CHARS_PER_SECOND = 256 * 1024 */ export const PREVIEW_TURN_CONTENT_BYTES = 8 * 1024 * 1024 -/** Room left in a replay write for the envelope around a preview payload. */ -const PREVIEW_FRAME_ENVELOPE_HEADROOM_BYTES = 16 * 1024 - -/** The largest serialized preview payload the replay buffer can persist in one write. */ -export const PREVIEW_FRAME_MAX_BYTES = - getRedisBudgetLimits('copilot_stream').maxSingleWriteBytes - PREVIEW_FRAME_ENVELOPE_HEADROOM_BYTES - /** The minimum gap between full snapshots of a preview this long. */ function snapshotIntervalMs(baseMs: number, previewText: string): number { return Math.max(baseMs, (previewText.length / PREVIEW_SNAPSHOT_CHARS_PER_SECOND) * 1000) @@ -397,13 +390,13 @@ async function emitPreviewEvent( if (frame.previewPhase === 'file_preview_content') { const budget = context.filePreviewBudget const contentBytes = budget.contentBytes + Buffer.byteLength(frame.content, 'utf8') - if (frameBytes > PREVIEW_FRAME_MAX_BYTES || contentBytes > PREVIEW_TURN_CONTENT_BYTES) { + if (frameBytes > STREAM_EVENT_MAX_PAYLOAD_BYTES || contentBytes > PREVIEW_TURN_CONTENT_BYTES) { return false } budget.contentBytes = contentBytes } else if ( frame.previewPhase === 'file_preview_complete' && - frameBytes > PREVIEW_FRAME_MAX_BYTES + frameBytes > STREAM_EVENT_MAX_PAYLOAD_BYTES ) { const { output: _output, ...withoutOutput } = frame frame = withoutOutput diff --git a/apps/sim/lib/mothership/request/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index fc7db309247..51b07fb5ef7 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -832,6 +832,39 @@ describe('copilot go stream helpers', () => { }) }) + it('keeps the timeout error when the body read fails after the request timed out', async () => { + const first = createEvent({ + streamId: 'timed-out-stream', + cursor: '1', + seq: 1, + requestId: 'req-timed-out', + type: 'text', + payload: { channel: 'assistant', text: 'partial' }, + }) + vi.mocked(fetch).mockImplementationOnce(async (_url, init) => { + const signal = init?.signal + return new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(first)}\n\n`)) + signal?.addEventListener('abort', () => controller.error(signal.reason)) + }, + }), + { status: 200, headers: { 'Content-Type': 'text/event-stream' } } + ) + }) + + await expect( + runStreamLoop( + 'https://example.com/mothership/stream', + {}, + createStreamingContext(), + turnScopedExecContext(), + { timeout: 20, flushAfterEvent: false } + ) + ).rejects.toMatchObject({ name: 'TimeoutError' }) + }) + it('reports a worker it could not reach without the raw network error', async () => { const networkError = new TypeError('fetch failed') vi.mocked(fetch).mockRejectedValueOnce(networkError) diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 8fd87610b48..963aeb78ee0 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.test.ts @@ -221,6 +221,7 @@ import { CopilotBackendError, STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE, StreamEndedWithoutTerminalError, + WorkerUnreachableError, } from '@/lib/mothership/request/go/stream' import { runCopilotLifecycle } from '@/lib/mothership/request/lifecycle/run' import { @@ -2718,7 +2719,7 @@ describe('runCopilotLifecycle', () => { headers.push(new Headers(request.headers)) context.accumulatedContent = 'Saved partial answer' context.errors.push('connection interrupted') - throw new TypeError('fetch failed') + throw new WorkerUnreachableError(new TypeError('fetch failed')) } ) } @@ -2811,7 +2812,9 @@ describe('runCopilotLifecycle', () => { vi.useFakeTimers() try { const controller = new AbortController() - mockRunStreamLoop.mockRejectedValueOnce(new TypeError('fetch failed')) + mockRunStreamLoop.mockRejectedValueOnce( + new WorkerUnreachableError(new TypeError('fetch failed')) + ) const pending = runCopilotLifecycle( { message: 'hello', messageId: 'stopped-outage' }, { @@ -2929,50 +2932,6 @@ describe('runCopilotLifecycle', () => { expect(JSON.stringify(persisted)).not.toMatch(/ { - vi.useFakeTimers() - try { - let attempts = 0 - mockRunStreamLoop.mockImplementation( - async ( - _url: string, - _init: RequestInit, - _context: StreamingContext, - _exec: ExecutionContext, - options: { onEvent?: (event: unknown) => Promise } - ): Promise => { - attempts++ - await options.onEvent?.({ type: 'session', payload: { kind: 'start' } }) - throw new TypeError("Cannot read properties of undefined (reading 'payload')") - } - ) - - const pending = runCopilotLifecycle( - { message: 'hello', messageId: 'stream-repeated-type-error' }, - { - userId: 'user-1', - workspaceId: 'ws-1', - chatId: 'chat-1', - executionId: 'exec-1', - runId: 'run-1', - executionContext: { - userId: 'user-1', - workflowId: '', - workspaceId: 'ws-1', - chatId: 'chat-1', - }, - } - ) - await vi.advanceTimersByTimeAsync(60_000) - - expect(attempts).toBe(4) - expect(await pending).toEqual(expect.objectContaining({ success: false })) - } finally { - mockRunStreamLoop.mockReset() - vi.useRealTimers() - } - }) - it('stops a failure that repeats after every reattach at three retries', async () => { vi.useFakeTimers() try { diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts index 10afce65508..7ef1e679e49 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -95,18 +95,11 @@ describe('stream recovery budget', () => { } }) - it('keeps a TypeError from handling a delivered stream to the reachable budget', () => { - vi.useFakeTimers() - const error = new TypeError("Cannot read properties of undefined (reading 'payload')") + it('never retries a TypeError thrown by our own stream handling', () => { const retry = new StreamRetryWindow() - for (let index = 0; index < 3; index++) { - retry.recovered() - const delay = retry.nextDelay(error) - expect(delay).not.toBeNull() - vi.advanceTimersByTime(delay ?? 0) - } - retry.recovered() - expect(retry.nextDelay(error)).toBeNull() + expect( + retry.nextDelay(new TypeError("Cannot read properties of undefined (reading 'payload')")) + ).toBeNull() }) it.each([ @@ -135,9 +128,9 @@ describe('stream recovery budget', () => { vi.useFakeTimers() const retry = new StreamRetryWindow(120_000) vi.advanceTimersByTime(119_999) - expect(retry.nextDelay(new TypeError('fetch failed'))).toBeNull() + expect(retry.nextDelay(new WorkerUnreachableError(new TypeError('fetch failed')))).toBeNull() vi.advanceTimersByTime(120_000) - expect(retry.nextDelay(new TypeError('fetch failed'))).toBeNull() + expect(retry.nextDelay(new WorkerUnreachableError(new TypeError('fetch failed')))).toBeNull() expect(() => retry.remainingMs()).toThrow('could not be restored') }) @@ -145,7 +138,9 @@ describe('stream recovery budget', () => { const retry = new StreamRetryWindow() const controller = new AbortController() controller.abort() - expect(retry.nextDelay(new TypeError('fetch failed'), controller.signal)).toBeNull() + expect( + retry.nextDelay(new WorkerUnreachableError(new TypeError('fetch failed')), controller.signal) + ).toBeNull() expect(retry.nextDelay(new DOMException('Stopped', 'AbortError'))).toBeNull() expect(retry.nextDelay(new CopilotBackendError('Forbidden', { status: 403 }))).toBeNull() expect(retry.nextDelay(new Error('Invalid operation'))).toBeNull() diff --git a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts index 677d8bb8922..08da02415cc 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -32,7 +32,7 @@ export class StreamRetryWindow { private firstFailureAt?: number private firstUnreachableAt?: number private unreachableAttempt = 0 - attempt = 0 + private attempt = 0 constructor(timeoutMs = ORCHESTRATION_TIMEOUT_MS) { this.deadline = Date.now() + timeoutMs @@ -119,5 +119,5 @@ function isRetryableStreamError(error: unknown): boolean { if (error instanceof CopilotBackendError) { return error.status !== undefined && error.status >= 500 } - return error instanceof WorkerUnreachableError || error instanceof TypeError + return error instanceof WorkerUnreachableError } diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index dcfb8e3883d..8e5c10ff96f 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -4,9 +4,9 @@ import { MothershipStreamV1EventType, MothershipStreamV1ToolPhase, } from '@/lib/mothership/generated/mothership-stream-v1' +import type { StreamEvent } from '@/lib/mothership/request/session/types' import { isClientExecutedToolCall } from '@/lib/mothership/tools/client-executed-tools' import { formatFileSize } from '@/lib/uploads/utils/file-utils' -import type { StreamEvent } from './types' /** * Payloads whose strings could serialize past this are compacted before they are @@ -28,7 +28,7 @@ const NO_KEYS: ReadonlySet = new Set() const ARGUMENTS_KEY: ReadonlySet = new Set(['arguments']) /** Items kept at the head of an array that string cuts alone could not bound. */ -export const STREAM_ARRAY_HEAD_ITEMS = 100 +const STREAM_ARRAY_HEAD_ITEMS = 100 /** * A cheap estimate of a value's serialized size, without serializing it. It can @@ -52,7 +52,7 @@ function estimateBytes(value: unknown): number { * far shorter than the cut, and text fields keep their head, so nothing but the * named top-level keys is exempt. */ -function truncateStrings(value: unknown, skipKeys: ReadonlySet = NO_KEYS): unknown { +function truncateStrings(value: unknown, skipKeys: ReadonlySet): unknown { return mapLeaves(value, skipKeys, (leaf) => { if (typeof leaf !== 'string' || leaf.length <= STREAM_STRING_PREVIEW_UNITS) return leaf const end = STREAM_STRING_PREVIEW_UNITS @@ -180,26 +180,27 @@ function shed( : isRecordLike(value) ? Object.entries(value).filter(([key]) => !skipKeys.has(key)) : [] - // Identity fields are never this large, so only bulk is ever replaced. + // Only children over the floor are candidates; if they cannot cover the need, + // a nested node is replaced whole, its own identity fields included. const children: Sized[] = entries .map(([key, child]) => ({ key, value: child, bytes: serializedBytes(child, sizes) })) .filter((child) => child.bytes > OMITTABLE_FIELD_MIN_BYTES) .sort((left, right) => right.bytes - left.bytes) if (children.length === 0) return top ? value : omissionNote(bytes) - const gain = (child: Sized) => - child.bytes - Buffer.byteLength(JSON.stringify(omissionNote(child.bytes)), 'utf8') + const gain = (child: Sized) => child.bytes - serializedBytes(omissionNote(child.bytes)) const replacements = new Map() let remaining = need for (let index = 0; index < children.length && remaining > 0; index++) { - let sufficient: Sized | undefined - for (let candidate = children.length - 1; candidate >= index; candidate--) { - if (gain(children[candidate]) >= remaining) { - sufficient = children[candidate] - break + // Children are sorted largest first, so if this one cannot cover the rest, none can. + if (gain(children[index]) >= remaining) { + let sufficient = children[index] + for (let candidate = children.length - 1; candidate > index; candidate--) { + if (gain(children[candidate]) >= remaining) { + sufficient = children[candidate] + break + } } - } - if (sufficient) { replacements.set(sufficient.key, shed(sufficient.value, remaining, sizes, NO_KEYS)) remaining = 0 } else { @@ -224,10 +225,9 @@ function shed( * event for dispatch. Long strings are cut to their head in place, so every * object keeps its shape; if that is not enough, long arrays keep their head, * and past one replay write the smallest sufficient bulk is replaced by a size - * note, keeping the fields beside it. - * Assistant text, file previews, and the arguments of calls the browser - * executes are never cut; an event - * still too large is refused by the buffer, which ends the turn with an error. + * note, keeping the fields beside it. Assistant text, file previews, and the + * arguments of calls the browser executes are never cut; an event still too + * large is refused by the buffer, which ends the turn with an error. */ export function compactStreamEvent(event: StreamEvent): StreamEvent { const payload = toRecordOrNull(event.payload) @@ -242,11 +242,11 @@ export function compactStreamEvent(event: StreamEvent): StreamEvent { payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutedToolCall(toolName, args) ? ARGUMENTS_KEY : NO_KEYS + const sizes = new WeakMap() let compacted = truncateStrings(payload, skipKeys) - if (serializedBytes(compacted) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { + if (serializedBytes(compacted, sizes) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { compacted = trimArrays(compacted, skipKeys) } - const sizes = new WeakMap() const bytes = serializedBytes(compacted, sizes) if (bytes > STREAM_EVENT_MAX_PAYLOAD_BYTES) { compacted = shed(compacted, bytes - STREAM_EVENT_MAX_PAYLOAD_BYTES, sizes, skipKeys, true) diff --git a/apps/sim/lib/mothership/request/session/writer.ts b/apps/sim/lib/mothership/request/session/writer.ts index 400d3f0dfce..38889bd84a5 100644 --- a/apps/sim/lib/mothership/request/session/writer.ts +++ b/apps/sim/lib/mothership/request/session/writer.ts @@ -2,14 +2,14 @@ import { createLogger } from '@sim/logger' import { toError } from '@sim/utils/errors' import { encodeSSEComment } from '@/lib/core/utils/sse' import { MothershipStreamV1EventType } from '@/lib/mothership/generated/mothership-stream-v1' -import { appendEvents } from './buffer' -import type { PersistedStreamEventEnvelope } from './contract' -import type { ChatStreamLease } from './controller-lease' -import { createEvent } from './event' -import { StreamReplayBudgetExhaustedError } from './replay-budget' -import { compactStreamEvent } from './replay-compaction' -import { encodeSSEEnvelope } from './sse' -import type { StreamEvent } from './types' +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 { 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 type { StreamEvent } from '@/lib/mothership/request/session/types' const logger = createLogger('StreamWriter') @@ -130,9 +130,11 @@ export class StreamWriter { } /** - * Delivers an event and records it for replay. Oversized strings are cut first - * ({@link compactStreamEvent}); the client and the replay buffer receive the - * same compacted copy, while the caller keeps the full event for dispatch. + * Delivers an event and records it for replay. An oversized event is compacted + * first ({@link compactStreamEvent}): long strings are cut to their head, then + * long arrays to their head, then any bulk still past one replay write is + * replaced by a size note. The client and the replay buffer receive the same + * 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 From 77dedc2c112e33b913da51e3fb7640a7d3bccc16 Mon Sep 17 00:00:00 2001 From: Waleed Latif Date: Wed, 30 Sep 2026 00:44:36 -0700 Subject: [PATCH 18/18] fix(mothership): bound preview metadata frames and settle stopped rows on every save - Compact every preview phase except content and completion, which the preview adapter already bounds; a model-written patch search string could otherwise exceed one replay write in an edit_meta frame - Settle unfinished tool rows as stopped in buildPersistedAssistantMessage for a cancelled turn, so background and API callers that persist the result directly never save pending, executing, or awaiting-approval rows --- .../mothership/chat/persisted-message.test.ts | 22 ++++++++++++++++++ .../lib/mothership/chat/persisted-message.ts | 19 ++++++++++----- .../request/session/replay-compaction.test.ts | 23 ++++++++++++++++++- .../request/session/replay-compaction.ts | 15 ++++++++---- 4 files changed, 68 insertions(+), 11 deletions(-) diff --git a/apps/sim/lib/mothership/chat/persisted-message.test.ts b/apps/sim/lib/mothership/chat/persisted-message.test.ts index 9d08415e68a..e8cc7dd0451 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.test.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.test.ts @@ -802,3 +802,25 @@ describe('buildPersistedAssistantMessage on a completed turn', () => { } ) }) + +describe('buildPersistedAssistantMessage on a cancelled turn', () => { + it.each(['pending', 'executing', 'awaiting_approval'] as const)( + 'settles a %s tool row as stopped for a caller that persists the result directly', + (status) => { + const persisted = buildPersistedAssistantMessage({ + success: false, + cancelled: true, + content: 'Partial answer', + toolCalls: [], + contentBlocks: [ + { type: 'tool_call', timestamp: 1, toolCall: { id: 'call-1', name: 'read', status } }, + ], + }) + + expect(persisted.contentBlocks?.[0].toolCall).toMatchObject({ + state: 'cancelled', + display: { title: 'Stopped by user' }, + }) + } + ) +}) diff --git a/apps/sim/lib/mothership/chat/persisted-message.ts b/apps/sim/lib/mothership/chat/persisted-message.ts index 06394ab5267..3f09e52e716 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.ts @@ -359,13 +359,18 @@ export function buildPersistedAssistantMessage( return normalized } - // A completed turn settles its stragglers as the live view did at `complete`. - if (result.success && message.contentBlocks) { - message.contentBlocks = settleUnfinishedToolCalls(message.contentBlocks, 'success') + // A finished turn settles its stragglers as the live view did at its terminal; + // background and API callers persist a stopped turn without withStoppedContentBlock. + if (message.contentBlocks) { + message.contentBlocks = result.success + ? settleUnfinishedToolCalls(message.contentBlocks, 'success') + : settleUnfinishedToolCalls(message.contentBlocks, 'cancelled', STOPPED_TOOL_DISPLAY) } return message } +const STOPPED_TOOL_DISPLAY = { title: 'Stopped by user' } as const + const UNSETTLED_TOOL_STATES: ReadonlySet = new Set([ 'pending', 'executing', @@ -392,9 +397,11 @@ function settleUnfinishedToolCalls( } export function withStoppedContentBlock(message: PersistedMessage): PersistedMessage { - const contentBlocks = settleUnfinishedToolCalls(message.contentBlocks ?? [], 'cancelled', { - title: 'Stopped by user', - }) + const contentBlocks = settleUnfinishedToolCalls( + message.contentBlocks ?? [], + 'cancelled', + STOPPED_TOOL_DISPLAY + ) const hasAssistantText = contentBlocks.some( (block) => block.type === MothershipStreamV1EventType.text && diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts index 06c384e4513..b43ffe17de6 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.test.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -147,7 +147,7 @@ describe('compactStreamEvent', () => { expect(args.stdin).toBe(`${'x'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 1 MB total]`) }) - it('never compacts a file preview', () => { + it('never compacts preview content, which the client applies as exact deltas', () => { const preview: StreamEvent = { type: 'tool', payload: { @@ -164,6 +164,27 @@ describe('compactStreamEvent', () => { expect(compactStreamEvent(preview)).toBe(preview) }) + it('bounds a preview edit whose model-written search text is past one replay write', () => { + const meta: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'prepare_file_edit', + previewPhase: 'file_preview_edit_meta', + edit: { strategy: 'search_replace', search: 's'.repeat(2 * MB), replaceAll: false }, + }, + } + + const payload = payloadOf(compactStreamEvent(meta)) + const edit = toRecord(payload.edit) + + expect(serializedBytes(payload)).toBeLessThanOrEqual(STREAM_EVENT_MAX_PAYLOAD_BYTES) + expect(payload.previewPhase).toBe('file_preview_edit_meta') + expect(edit.strategy).toBe('search_replace') + expect(edit.replaceAll).toBe(false) + expect(edit.search).toBe(`${'s'.repeat(STREAM_STRING_PREVIEW_UNITS)}…[truncated, 2 MB total]`) + }) + it('compacts a copy and leaves the caller’s event whole for dispatch', () => { const stdout = 's'.repeat(MB) const event: StreamEvent = { diff --git a/apps/sim/lib/mothership/request/session/replay-compaction.ts b/apps/sim/lib/mothership/request/session/replay-compaction.ts index 8e5c10ff96f..feef020c9da 100644 --- a/apps/sim/lib/mothership/request/session/replay-compaction.ts +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -225,14 +225,21 @@ function shed( * event for dispatch. Long strings are cut to their head in place, so every * object keeps its shape; if that is not enough, long arrays keep their head, * and past one replay write the smallest sufficient bulk is replaced by a size - * note, keeping the fields beside it. Assistant text, file previews, and the - * arguments of calls the browser executes are never cut; an event still too - * large is refused by the buffer, which ends the turn with an error. + * note, keeping the fields beside it. Assistant text, the arguments of calls the + * browser executes, and preview content and completions are never cut: the + * client applies preview content as exact deltas, and the preview adapter + * bounds both itself. An event still too large is refused by the buffer, which + * ends the turn with an error. */ export function compactStreamEvent(event: StreamEvent): StreamEvent { const payload = toRecordOrNull(event.payload) // Text length is part of the receipt the worker and a replacement check. - if (!payload || event.type === MothershipStreamV1EventType.text || 'previewPhase' in payload) { + if ( + !payload || + event.type === MothershipStreamV1EventType.text || + payload.previewPhase === 'file_preview_content' || + payload.previewPhase === 'file_preview_complete' + ) { return event } if (estimateBytes(payload) <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) return event