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/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 dccc95ad1a9..47d5373c2df 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, @@ -18,7 +22,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' @@ -142,18 +146,18 @@ export function buildAssistantSnapshotMessage(params: { } export function markMessageStopped(message: PersistedMessage): PersistedMessage { - const hasExecutingTool = message.contentBlocks?.some( - (block) => block.toolCall?.state === 'executing' + 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 } 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/preview/apply-file-preview-phase.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/preview/apply-file-preview-phase.test.ts index 8ed6eb19032..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 } 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' @@ -94,3 +97,40 @@ 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) + }) + + 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/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/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/stream-helpers.test.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts new file mode 100644 index 00000000000..501e1fd2de0 --- /dev/null +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/stream/stream-helpers.test.ts @@ -0,0 +1,37 @@ +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') + } + ) + + 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) + }) + + 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 99d455dd7b1..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 @@ -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, @@ -68,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 it settled a row or closed a lane. */ export function finalizeResidualToolCalls( blocks: ContentBlock[], turnTerminal: 'complete' | 'cancelled' | 'error' -): void { +): boolean { const endedAt = Date.now() + let settled = false const propagated = turnTerminal === 'cancelled' ? ToolCallStatus.cancelled @@ -93,10 +94,12 @@ export function finalizeResidualToolCalls( // transport-based gating. if (block.type === 'subagent' && block.endedAt === undefined) { block.endedAt = endedAt + settled = true continue } const tc = block.toolCall - if (!tc || tc.status !== ToolCallStatus.executing) continue + if (!tc || !isUnsettledToolState(tc.status)) continue + settled = true tc.status = propagated if (propagated === ToolCallStatus.cancelled) { tc.displayTitle = 'Stopped by user' @@ -105,6 +108,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 2885e0fd8f1..92da00740fe 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, @@ -82,8 +83,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' @@ -3227,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) @@ -4370,11 +4370,11 @@ export function useChat( } else { setPendingMessages((prev) => prev.map((msg) => { - const hasExecutingTool = msg.contentBlocks?.some( - (block) => block.toolCall?.status === 'executing' + 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.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..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,9 +21,6 @@ import { type ToolCallInfo, ToolCallStatus, } from '@/app/workspace/[workspaceId]/home/types' -import { getMothershipAttachmentPreviewUrl } from './attachment-preview' -import type { PersistedMessage } from './persisted-message' -import { 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, so 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 78861c7169c..e8cc7dd0451 100644 --- a/apps/sim/lib/mothership/chat/persisted-message.test.ts +++ b/apps/sim/lib/mothership/chat/persisted-message.test.ts @@ -761,3 +761,66 @@ 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') + } + ) +}) + +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') + } + ) +}) + +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 2e15cabb9da..3f09e52e716 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' @@ -336,8 +340,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 +350,7 @@ export function buildPersistedAssistantMessage( content: message.content, }, ] - : [])), + : []), { type: MothershipStreamV1EventType.error, content: buildMothershipErrorTag({ message: error }), @@ -354,25 +359,48 @@ export function buildPersistedAssistantMessage( return normalized } + // 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', + 'awaiting_approval', +]) + +/** A tool row that has not finished: waiting to run, running, or awaiting a decision. */ +export function isUnsettledToolState(state: string | undefined): boolean { + 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. */ +function settleUnfinishedToolCalls( + blocks: PersistedContentBlock[], + state: 'success' | '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 = (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 contentBlocks = settleUnfinishedToolCalls( + message.contentBlocks ?? [], + 'cancelled', + STOPPED_TOOL_DISPLAY ) const hasAssistantText = contentBlocks.some( (block) => 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/context/restore.ts b/apps/sim/lib/mothership/request/context/restore.ts index bb02d7239a4..36bc8630b78 100644 --- a/apps/sim/lib/mothership/request/context/restore.ts +++ b/apps/sim/lib/mothership/request/context/restore.ts @@ -13,6 +13,14 @@ 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 && + 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/file-preview-adapter.test.ts b/apps/sim/lib/mothership/request/go/file-preview-adapter.test.ts index af13587cad4..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 @@ -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, @@ -25,9 +25,11 @@ import { createStreamingContext } from '@/lib/mothership/request/context/request import { createFilePreviewAdapterState, type FilePreviewAdapterState, + 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, @@ -163,6 +165,122 @@ 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()) + + /** 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, + { turn = newTurn(), edit = 0, baseLines = 10_000 } = {} + ) { + vi.useFakeTimers() + 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}` + const intent = { + ...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) => { + turn.context.activeFileIntents.set('', intent) + await processFilePreviewStreamEvent({ + streamId: STREAM_ID, + streamEvent, + context: turn.context, + execContext, + options: { + onEvent: (event) => { + payloads.push((event as { payload: Record }).payload) + }, + }, + state: turn.state, + }) + } + const delta = (argumentsDelta: string) => + toolEvent({ + toolCallId: editToolCallId, + toolName: 'apply_file_edit', + phase: MothershipStreamV1ToolPhase.args_delta, + argumentsDelta, + }) + + let streamed = '' + await drive(delta('{"content":"')) + for (let elapsed = 0; elapsed < durationMs; elapsed += tickMs) { + vi.advanceTimersByTime(tickMs) + streamed += `word${elapsed} ` + await drive(delta(`word${elapsed} `)) + } + await drive( + toolEvent({ + toolCallId: editToolCallId, + toolName: 'apply_file_edit', + phase: MothershipStreamV1ToolPhase.result, + success: true, + }) + ) + + const contents = payloads.filter((payload) => payload.previewPhase === 'file_preview_content') + return { + payloads, + finalContent: base.replace('ANCHOR\n', `ANCHOR\n${streamed}\n`), + lastContent: contents.at(-1)?.content, + completed: payloads.some((payload) => payload.previewPhase === 'file_preview_complete'), + contentBytes: 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, contentBytes } = await streamPatch(10_000, 16) + + expect(contentBytes).toBeLessThan(8 * 1024 * 1024) + expect(lastContent).toBe(finalContent) + }) + + 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( + STREAM_EVENT_MAX_PAYLOAD_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, { turn, edit }) + expect(result.completed).toBe(true) + contentBytes += result.contentBytes + } + + expect(contentBytes).toBeLessThanOrEqual(PREVIEW_TURN_CONTENT_BYTES) + }) +}) + /** * 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..0dedab3b707 100644 --- a/apps/sim/lib/mothership/request/go/file-preview-adapter.ts +++ b/apps/sim/lib/mothership/request/go/file-preview-adapter.ts @@ -14,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, @@ -56,6 +57,23 @@ 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 + +/** + * 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_TURN_CONTENT_BYTES = 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) +} function asJsonRecord(value: unknown): JsonRecord | undefined { return isRecordLike(value) ? (value as JsonRecord) : undefined @@ -326,7 +344,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 { @@ -355,16 +373,40 @@ export function createFilePreviewAdapterState(): FilePreviewAdapterState { } } +/** + * 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 > STREAM_EVENT_MAX_PAYLOAD_BYTES || contentBytes > PREVIEW_TURN_CONTENT_BYTES) { + return false + } + budget.contentBytes = contentBytes + } else if ( + frame.previewPhase === 'file_preview_complete' && + frameBytes > STREAM_EVENT_MAX_PAYLOAD_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: { @@ -447,12 +489,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', @@ -465,7 +507,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', @@ -522,12 +564,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', @@ -540,7 +582,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', @@ -584,13 +626,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', @@ -603,6 +640,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, + } + ) } } @@ -691,12 +742,12 @@ 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, { + ...currentPreview, session: nextSession, - lastEmittedPreviewText: currentPreview.lastEmittedPreviewText, - lastSnapshotAt: currentPreview.lastSnapshotAt, }) } else { const previewUpdate = buildPreviewContentUpdate( @@ -706,14 +757,7 @@ export async function processFilePreviewStreamEvent(input: { now, nextSession.operation ) - - filePreviewState.set(editIntent.toolCallId, { - session: nextSession, - lastEmittedPreviewText: nextSession.previewText, - lastSnapshotAt: previewUpdate.lastSnapshotAt, - }) - - await emitPreviewEvent(streamEvent, options, { + const sent = await emitPreviewEvent(streamEvent, options, context, { toolCallId: nextSession.toolCallId, toolName: 'prepare_file_edit', previewPhase: 'file_preview_content', @@ -726,13 +770,20 @@ export async function processFilePreviewStreamEvent(input: { ...(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, { - session: currentPreview.session, - lastEmittedPreviewText: currentPreview.lastEmittedPreviewText, - lastSnapshotAt: currentPreview.lastSnapshotAt, - }) + filePreviewState.set(editIntent.toolCallId, currentPreview) } } } @@ -766,7 +817,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', @@ -800,7 +851,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/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/go/stream.test.ts b/apps/sim/lib/mothership/request/go/stream.test.ts index f412a194479..51b07fb5ef7 100644 --- a/apps/sim/lib/mothership/request/go/stream.test.ts +++ b/apps/sim/lib/mothership/request/go/stream.test.ts @@ -80,6 +80,7 @@ vi.mock('@/lib/mothership/tools/server/files/file-preview', async () => { import { buildPreviewContentUpdate, + CopilotBackendError, decodeJsonStringPrefix, extractEditContent, runStreamLoop, @@ -148,6 +149,7 @@ function createStreamingContext(): StreamingContext { wasAborted: false, errors: [], activeFileIntents: new Map(), + filePreviewBudget: { contentBytes: 0 }, trace: new TraceCollector(), toolPermissions: { enabled: false, @@ -250,6 +252,77 @@ 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.', + ], + [ + "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"}', + 'The agent service could not process this request.', + ], + [ + 'a plain-text request rejection', + 400, + 'text/plain', + 'Invalid request body', + 'The agent service could not process this request.', + ], + ])( + '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 } }) + ) + + 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 }) + } + ) + it('terminates the stream on an exhausted identity budget before forwarding later events', async () => { const identity = createProviderToolCallIdentity('exhausted-identity-run') identity.retainedBytes = PROVIDER_TOOL_CALL_IDENTITY_LIMITS.maxRetainedBytes @@ -712,8 +785,89 @@ 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 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('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) const context = createStreamingContext() const execContext: ExecutionContext = { @@ -725,7 +879,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 91dba0de262..b55b623458d 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,65 @@ 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.' + +/** + * 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' + } +} + +/** + * 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 + * something to the code that raised them. + */ +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 { constructor(public readonly userId: string) { super('Usage limit reached') @@ -158,7 +218,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 = { @@ -175,10 +236,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) { @@ -244,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 d8c93d82444..2b4d8f04274 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 { @@ -115,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, @@ -209,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', @@ -402,6 +431,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 1e81d1bb07c..cb65ad989ef 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, @@ -216,10 +214,13 @@ export async function prePersistClientExecutableToolCall( if (!isToolCallStreamEvent(event)) return const data = event.payload - if (data.replay) return const isGenerating = data.status === TOOL_CALL_STATUS.generating const isPartial = data.partial === true || isGenerating - if (isPartial) return + // 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 + } const ui = getToolCallUI(data) const catalogEntry = getToolEntry(data.toolName) diff --git a/apps/sim/lib/mothership/request/lifecycle/finalize.ts b/apps/sim/lib/mothership/request/lifecycle/finalize.ts index 2aaaa04c3ff..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') @@ -86,27 +87,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,36 +143,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, - 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 + ) ) } @@ -178,24 +183,46 @@ 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, since nothing else settles a run. A controller superseded + * while publishing leaves the run to its successor instead. + */ +async function publishThenSettle( + publisher: StreamWriter, + publish: () => Promise, + settle: () => Promise +): Promise { + try { + await publish() + await publisher.flush() + } catch (error) { + if (!(error instanceof StreamControllerSupersededError)) await settle() + throw error + } + await settle() +} + async function loggedRunStatusUpdate( runId: string, status: Parameters[1], diff --git a/apps/sim/lib/mothership/request/lifecycle/run.test.ts b/apps/sim/lib/mothership/request/lifecycle/run.test.ts index 38cd3f1faf6..963aeb78ee0 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' @@ -121,9 +121,25 @@ 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' + } + } + + 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, runStreamLoop: mockRunStreamLoop, @@ -196,6 +212,7 @@ vi.mock('@/lib/mothership/request/enterprise-byok', () => ({ resolveEnterpriseByokKey: mockResolveEnterpriseByokKey, })) +import { buildPersistedAssistantMessage } from '@/lib/mothership/chat/persisted-message' import { MothershipStreamV1CompletionStatus, MothershipStreamV1ToolOutcome, @@ -204,8 +221,13 @@ import { CopilotBackendError, STREAM_ENDED_WITHOUT_TERMINAL_MESSAGE, StreamEndedWithoutTerminalError, + WorkerUnreachableError, } from '@/lib/mothership/request/go/stream' import { runCopilotLifecycle } from '@/lib/mothership/request/lifecycle/run' +import { + REPLAY_BUDGET_EXHAUSTED_CODE, + StreamReplayBudgetExhaustedError, +} from '@/lib/mothership/request/session/replay-budget' import { executeToolAndReport } from '@/lib/mothership/request/tools/executor' const mockExecuteAppTool = toolsMockFns.mockExecuteTool @@ -1947,6 +1969,81 @@ 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], + ])( + '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 (_url: string, _init: RequestInit, context: StreamingContext): Promise => { + context.accumulatedContent = 'partial answer' + abortController.abort(refusal) + context.wasAborted = true + if (throws) throw refusal + } + ) + + const result = await runWithStreamAbort(abortController) + + expect(result).toEqual( + expect.objectContaining({ + success: false, + cancelled: false, + content: 'partial answer', + 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')) @@ -2622,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')) } ) } @@ -2715,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' }, { @@ -2798,6 +2897,180 @@ 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++ + 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-failure' }, + { + 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() + } + }) + + 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/run.ts b/apps/sim/lib/mothership/request/lifecycle/run.ts index 110a4bbf45f..9a9ac2cd667 100644 --- a/apps/sim/lib/mothership/request/lifecycle/run.ts +++ b/apps/sim/lib/mothership/request/lifecycle/run.ts @@ -53,6 +53,7 @@ import { StreamRetryWindow } from '@/lib/mothership/request/lifecycle/stream-ret import { recordDegraded } from '@/lib/mothership/request/metrics' import { AbortReason } from '@/lib/mothership/request/session/abort-reason' import { StreamControllerSupersededError } from '@/lib/mothership/request/session/controller-lease' +import { replayRefusal } from '@/lib/mothership/request/session/replay-budget' import { getToolCallTerminalData, requireToolCallStateResult, @@ -534,16 +535,21 @@ export async function runCopilotLifecycle( // the work the user watched succeed. const backendFinishedTurn = context.completionStatus === MothershipStreamV1CompletionStatus.complete + // A refused replay write aborts the turn to stop it, but the turn failed; it + // was not stopped by the user. + const refusal = replayRefusal(lifecycleOptions.abortSignal?.reason) // 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) + !refusal && + (context.completionStatus === MothershipStreamV1CompletionStatus.cancelled || + context.wasAborted || + (lifecycleOptions.abortSignal?.aborted ?? false)) const succeeded = + !refusal && !turnWasAborted && (backendFinishedTurn || (!context.completionStatus && context.errors.length === 0)) @@ -564,6 +570,7 @@ export async function runCopilotLifecycle( toolCalls: buildToolCallSummaries(context), chatId: context.chatId, requestId: context.requestId, + ...(refusal ? { error: refusal.userMessage, errorCode: refusal.code } : {}), errors: !succeeded && context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, @@ -586,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) } @@ -601,7 +609,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 wasCancelled = 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 @@ -616,7 +625,8 @@ export async function runCopilotLifecycle( toolCalls: buildToolCallSummaries(context), chatId: context.chatId, requestId: context.requestId, - error: err.message, + error: refusal?.userMessage ?? err.message, + ...(refusal ? { errorCode: refusal.code } : {}), errors: context.errors.length ? context.errors : undefined, usage: context.usage, cost: context.cost, @@ -836,6 +846,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 @@ -850,7 +868,7 @@ async function runResumeLegWithRetry( }, leg, execContext, - { ...options, timeout: retry.remainingMs() } + { ...legOptions, timeout: retry.remainingMs() } ) return } catch (error) { @@ -858,9 +876,10 @@ 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, + ...causeForLog(error), }) await interruptibleSleep(backoff, options.abortSignal) continue @@ -1128,6 +1147,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 && @@ -1153,7 +1173,7 @@ async function runCheckpointLoop( { route, isResume, - ...(isResume ? { attempt: retry.attempt } : {}), + ...(isResume ? { attempt: retry.attempts } : {}), } ) context.trace.setActiveSpan(streamSpan) @@ -1161,7 +1181,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, @@ -1219,9 +1239,10 @@ 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, + ...causeForLog(streamError), } ) await interruptibleSleep(backoff, options.abortSignal) @@ -1660,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/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 9efcbc019ed..4be603e2be8 100644 --- a/apps/sim/lib/mothership/request/lifecycle/start.ts +++ b/apps/sim/lib/mothership/request/lifecycle/start.ts @@ -57,10 +57,16 @@ 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 { + 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' +import type { OrchestratorResult } from '@/lib/mothership/request/types' import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url' export { SSE_RESPONSE_HEADERS } @@ -204,6 +210,75 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS const collector = new TraceCollector() + /** + * The replay refusal that ends this turn. Once the controller is aborted its + * reason is authoritative, so a refusal that follows a Stop stays a Stop. + */ + const refusalOf = (thrown?: unknown) => + replayRefusal(abortController.signal.aborted ? abortController.signal.reason : thrown) + + /** + * 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) => { + try { + await finalizeStream(...args) + turnEnded = true + } catch (error) { + turnEnded = !(error instanceof StreamControllerSupersededError) + throw error + } + } + + /** + * 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 endTurn( + { + content: '', + contentBlocks: [], + toolCalls: [], + ...result, + success: false, + cancelled: false, + error: refusal.userMessage, + errorCode: refusal.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) @@ -271,7 +346,7 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS outcome = RequestTraceV1Outcome.cancelled abortController.abort(AbortReason.UserStop) cancelReason = recordCancelled() - await finalizeStream( + await endTurn( { success: false, cancelled: true, @@ -348,7 +423,17 @@ 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( + replayRefusal(error) ?? new StreamControllerSupersededError() + ) + } throw error } }, @@ -360,17 +445,24 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS }) lifecycleResult = result + const refusal = refusalOf() // 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 + : refusal + ? 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 (refusal && !result.success) { + await finalizeAfterReplayRefusal(refusal, result) + } else { + await endTurn(result, publisher, runId, outcome, requestId) + } } catch (error) { if ( error instanceof StreamControllerSupersededError || @@ -380,6 +472,12 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS return } await assertControllerOwnership() + const refusal = refusalOf(error) + if (refusal) { + outcome = RequestTraceV1Outcome.error + await finalizeAfterReplayRefusal(refusal) + return + } const wasCancelled = abortController.signal.aborted outcome = wasCancelled ? RequestTraceV1Outcome.cancelled : RequestTraceV1Outcome.error if (outcome === RequestTraceV1Outcome.cancelled) { @@ -403,7 +501,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, @@ -423,15 +521,21 @@ export function createSSEStream(params: StreamingOrchestrationParams): ReadableS }) } unregisterActiveStream(streamId, abortController) - if (chatId) { - await releasePendingChatStream(chatId, streamId, lease) - } - processResourcesReleased = true - if (!(abortController.signal.reason instanceof StreamControllerSupersededError)) { + /* + The stream's buffer and abort marker belong to whoever holds the chat + 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 (turnEnded && (!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 ed75d7106e1..7ef1e679e49 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.test.ts @@ -2,18 +2,17 @@ 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' 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,23 +24,113 @@ describe('stream recovery budget', () => { expect(retry.remainingMs()).toBeGreaterThan(3_500_000) }) - it('bounds the recovery period from the first failure without shortening healthy work', () => { + 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('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 WorkerUnreachableError(new TypeError('fetch failed')) + for (let outage = 0; outage < 2; outage++) { + 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() + } + }) + + it('never retries a TypeError thrown by our own stream handling', () => { + const retry = new StreamRetryWindow() + expect( + retry.nextDelay(new TypeError("Cannot read properties of undefined (reading 'payload')")) + ).toBeNull() + }) + + it.each([ + new WorkerUnreachableError(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) - 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) + 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() 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') }) @@ -49,10 +138,24 @@ 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() 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 98e382edb41..08da02415cc 100644 --- a/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts +++ b/apps/sim/lib/mothership/request/lifecycle/stream-retry.ts @@ -4,21 +4,45 @@ import { StreamContinuityError } from '@/lib/mothership/request/go/parser' import { CopilotBackendError, StreamEndedWithoutTerminalError, + WorkerStreamInterruptedError, + WorkerUnreachableError, } from '@/lib/mothership/request/go/stream' 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 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 + * reattach, never a second run. + */ +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 recoveryDeadline?: number - attempt = 0 + private firstFailureAt?: number + private firstUnreachableAt?: number + private unreachableAttempt = 0 + private attempt = 0 constructor(timeoutMs = ORCHESTRATION_TIMEOUT_MS) { 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) @@ -26,25 +50,74 @@ export class StreamRetryWindow { return remaining } + /** The worker answered, so a later loss of it starts a fresh unreachable window. */ + recovered(): void { + this.firstUnreachableAt = undefined + this.unreachableAttempt = 0 + } + nextDelay(error: unknown, signal?: AbortSignal): number | null { if (signal?.aborted || !isRetryableStreamError(error)) return null - this.recoveryDeadline ??= Date.now() + STREAM_RECOVERY_WINDOW_MS + 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() if (this.attempt >= MAX_STREAM_RETRIES) return null - const delay = backoffWithJitter(this.attempt + 1, null, { baseMs: 250, maxMs: 5_000 }) - if (Date.now() + delay >= Math.min(this.deadline, this.recoveryDeadline)) 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 }) +} + +/** + * 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 WorkerUnreachableError +} + +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) { + if ( + error instanceof StreamEndedWithoutTerminalError || + error instanceof StreamContinuityError || + error instanceof WorkerStreamInterruptedError + ) { return true } if (error instanceof CopilotBackendError) { return error.status !== undefined && error.status >= 500 } - return error instanceof TypeError + return error instanceof WorkerUnreachableError } 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 new file mode 100644 index 00000000000..a55706e2468 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-budget.integration.ts @@ -0,0 +1,787 @@ +/** + * 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 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)) + await hooks.onAbort?.() + } + 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, + hooks, + /** Events, or steps to run between them, that the scripted worker streams in order. */ + script: [] as unknown[], + /** Controller lifecycles started, and what each sink call threw. */ + runs: [] as Array<{ dispatched: unknown[]; sinkErrors: unknown[] }>, + }, + } +}) + +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) { + if (typeof event === 'function') { + await event(options) + continue + } + 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 { 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 { 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' +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' +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 { + REPLAY_BUDGET_EXHAUSTED_CODE, + 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 { + 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 + +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 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 } } +} + +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)}…[truncated, 2 MB total]`, + }, + }) + 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)}…[truncated, 1.6 MB total]`, + }, + }) + }) + + 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 ownerMessage = new StreamReplayBudgetExhaustedError({ + resource: 'owner_redis_bytes', + currentBytes: 0, + limitBytes: 0, + attemptedBytes: 0, + }).userMessage + + 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' }) + authMockFns.mockGetSession.mockResolvedValue({ + user: { id: userId }, + session: { id: generateId() }, + }) + }) + + 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)) + }) + + /** 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)))! + const request = { + message: 'Summarize the 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, + recovery: { + kind: 'interactive_stream', + request, + goRoute: '/api/mothership', + clientToolPickupExpected: false, + }, + }, + }) + .returning() + await prepare?.(streamId) + worker.runs.length = 0 + worker.abortRequests.length = 0 + worker.script = script + 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 }, + }) + /** 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 } + } + + 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) + } + ) + + 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(ownerMessage) + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored).toMatchObject({ status: 'error', error: ownerMessage }) + expect(worker.abortRequests).toEqual([expect.objectContaining({ messageId: streamId })]) + expect(await redis().get(chatStreamLockKey(chatId))).toBeNull() + + const reconnect = await streamGET( + new NextRequest(`http://localhost:3000/api/copilot/chat/stream?streamId=${streamId}&after=0`), + { params: Promise.resolve({}) } + ) + const replayed = dataFrames(await reconnect.text()) + expect(replayed.map((frame) => frame.type)).toEqual(['session', 'error', 'complete']) + expect(replayed[1].payload.message).toBe(ownerMessage) + 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() + + const [controllerRun] = worker.runs + expect(controllerRun.dispatched).toEqual([]) + expect(controllerRun.sinkErrors[0]).toBeInstanceOf(StreamReplayBudgetExhaustedError) + }) + + 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']) + 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('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('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('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 () => { + 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() + await db.insert(copilotRuns).values({ + id: runId, + executionId: generateId(), + chatId, + userId, + workspaceId, + streamId, + status: 'paused_waiting_for_tool', + requestContext: { requestId: generateId(), controllerToken }, + }) + 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 }) + + 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 () => { + 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) + + const [stored] = await db.select().from(copilotRuns).where(eq(copilotRuns.id, runId)) + expect(stored.status).toBe('paused_waiting_for_tool') + }) +}) 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..c64a6203abf --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-budget.ts @@ -0,0 +1,36 @@ +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' + +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 + * 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' + } + + /** 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 new file mode 100644 index 00000000000..b43ffe17de6 --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-compaction.test.ts @@ -0,0 +1,493 @@ +/** + * 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 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; + * - 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 { + compactStreamEvent, + 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' + +const MB = 1024 * 1024 + +function payloadOf(event: StreamEvent): Record { + return toRecord(event.payload) +} + +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 resources = [{ type: 'file', id: 'file-1', title: 'out.csv' }] + 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, + citations, + reason: 'user_cancelled', + }, + }, + } + + 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: 'user_cancelled' }) + expect(Object.keys(output)).toEqual(Object.keys(toRecord(payloadOf(event).output))) + expect(payloadOf(compacted)).toMatchObject({ + toolCallId: 'call-1', + toolName: 'run_code', + phase: 'result', + success: true, + status: 'success', + }) + }) + + 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' }], + ['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: 'c', + toolName, + executor: 'go', + mode: 'async', + phase: 'call', + arguments: args, + }, + } as StreamEvent + + expect(compactStreamEvent(event)).toBe(event) + }) + + it('cuts the arguments of a call the browser does not execute', () => { + const event: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'sim_cli', + executor: 'go', + mode: 'async', + phase: 'call', + arguments: { activity: { id: 'a', title: 'Reading logs' }, stdin: 'x'.repeat(MB) }, + }, + } + + const args = toRecord(payloadOf(compactStreamEvent(event)).arguments) + + 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 preview content, which the client applies as exact deltas', () => { + const preview: StreamEvent = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'prepare_file_edit', + previewPhase: 'file_preview_content', + content: 'p'.repeat(400 * 1024), + contentMode: 'snapshot', + previewVersion: 3, + fileName: 'notes.md', + }, + } + + 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 = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'cli_logs_get', + executor: 'sim', + mode: 'async', + phase: 'call', + arguments: { stdout }, + }, + } + + expect(compactStreamEvent(event)).not.toBe(event) + expect(toRecord(payloadOf(event).arguments).stdout).toBe(stdout) + }) + + 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: 'c', + toolName: 'cli_logs_get', + executor: 'sim', + mode: 'async', + phase: 'result', + success: true, + output: { text }, + }, + } + + const cut = toRecord(payloadOf(compactStreamEvent(event)).output).text as string + + 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) + }) + + 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(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('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 = { + type: 'tool', + payload: { + toolCallId: 'c', + toolName: 'run_workflow', + executor: 'client', + mode: 'async', + phase: 'call', + ui: { clientExecutable: true }, + arguments: args, + }, + } + + 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 new file mode 100644 index 00000000000..feef020c9da --- /dev/null +++ b/apps/sim/lib/mothership/request/session/replay-compaction.ts @@ -0,0 +1,262 @@ +import { isRecordLike, toRecordOrNull } from '@sim/utils/object' +import { getRedisBudgetLimits } from '@/lib/core/redis/byte-budget.server' +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' + +/** + * 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 + +/** 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 + +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. */ +const STREAM_ARRAY_HEAD_ITEMS = 100 + +/** + * 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)) + 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 +} + +/** + * 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): unknown { + 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 = leaf.charCodeAt(end - 1) + const cut = lastUnit >= 0xd800 && lastUnit <= 0xdbff ? end - 1 : end + 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)) { + 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 leaf(value) + let copy: Record | undefined + for (const [key, field] of Object.entries(value)) { + if (skipKeys.has(key)) continue + const next = mapLeaves(field, NO_KEYS, leaf, array) + if (next !== field) (copy ??= { ...value })[key] = next + } + 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 + +/** + * 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 +} + +/** 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)) + : [] + // 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 - serializedBytes(omissionNote(child.bytes)) + const replacements = new Map() + let remaining = need + for (let index = 0; index < children.length && remaining > 0; index++) { + // 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 + } + } + 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]) + } + } + // A caller picks a nested node because replacing it whole covers the need. + if (remaining > 0 && !top) return omissionNote(bytes) + + if (Array.isArray(value)) { + return value.map((item, index) => (replacements.has(index) ? replacements.get(index) : item)) + } + const copy = { ...toRecordOrNull(value) } + for (const [key, replacement] of replacements) copy[String(key)] = replacement + return copy +} + +/** + * 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 the smallest sufficient bulk is replaced by a size + * 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 || + payload.previewPhase === 'file_preview_content' || + payload.previewPhase === 'file_preview_complete' + ) { + return event + } + if (estimateBytes(payload) <= STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) return event + const toolName = typeof payload.toolName === 'string' ? payload.toolName : '' + const args = isRecordLike(payload.arguments) ? payload.arguments : undefined + const skipKeys = + payload.phase === MothershipStreamV1ToolPhase.call && isClientExecutedToolCall(toolName, args) + ? ARGUMENTS_KEY + : NO_KEYS + const sizes = new WeakMap() + let compacted = truncateStrings(payload, skipKeys) + if (serializedBytes(compacted, sizes) > STREAM_EVENT_COMPACTION_THRESHOLD_BYTES) { + compacted = trimArrays(compacted, skipKeys) + } + 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) +} diff --git a/apps/sim/lib/mothership/request/session/writer.test.ts b/apps/sim/lib/mothership/request/session/writer.test.ts index 6fa5c4361a4..1eab0e0a636 100644 --- a/apps/sim/lib/mothership/request/session/writer.test.ts +++ b/apps/sim/lib/mothership/request/session/writer.test.ts @@ -13,6 +13,8 @@ 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 { StreamWriter } from '@/lib/mothership/request/session/writer' function decodeChunk(value: Uint8Array): string { @@ -77,7 +79,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 +95,36 @@ 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() - expect(writer.persistenceStopped).toBe(true) - await expect(writer.close()).rejects.toThrow('Stream replay byte budget exhausted') + writer.publish({ type: 'text', payload: { channel: 'assistant', text: 'later' } }) + ).rejects.toBeInstanceOf(StreamReplayBudgetExhaustedError) + expect(delivered).toEqual([]) + await expect(writer.flush()).resolves.toBeUndefined() + + await writer.publish({ + type: 'error', + payload: { message: 'stopped', code: 'replay_budget_exhausted' }, + }) + await writer.publish({ type: 'complete', payload: { status: 'error' } }) + await expect(writer.close()).resolves.toBeUndefined() + expect(delivered.map((frame) => JSON.parse(frame.replace(/^data: /, '')).type)).toEqual([ + 'error', + 'complete', + ]) + expect(writer.sawComplete).toBe(true) }) it('enqueues before persistence completes and flushes pending writes on close', async () => { diff --git a/apps/sim/lib/mothership/request/session/writer.ts b/apps/sim/lib/mothership/request/session/writer.ts index 05eb280b809..38889bd84a5 100644 --- a/apps/sim/lib/mothership/request/session/writer.ts +++ b/apps/sim/lib/mothership/request/session/writer.ts @@ -2,12 +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 { 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') @@ -26,6 +28,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 +62,7 @@ export class StreamWriter { private pendingEnvelopes: PersistedStreamEventEnvelope[] = [] private persistenceTail: Promise = Promise.resolve() private lastPersistenceError: Error | null = null + private replayBudgetError: StreamReplayBudgetExhaustedError | null = null private readonly lease?: ChatStreamLease constructor(options: StreamWriterOptions) { @@ -114,23 +129,48 @@ export class StreamWriter { } } + /** + * 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 + * 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 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 delivery = this.persistenceTail.then(async () => { - const result = await appendEvents( - [envelope], - { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, - this.lease - ) - if (!result.persisted) { - this._persistenceStopped = true - throw new Error('Stream replay byte budget exhausted') + const persistThenDeliver = async () => { + if (!this.replayBudgetError) { + const result = await appendEvents( + [envelope], + { streamId: this.streamId, ...(this.userId ? { userId: this.userId } : {}) }, + lease + ) + if (!result.persisted) { + this._persistenceStopped = true + this.replayBudgetError = new StreamReplayBudgetExhaustedError(result.refusal) + } } + 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() }) this.persistenceTail = delivery return delivery @@ -148,7 +188,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 diff --git a/apps/sim/lib/mothership/request/types.ts b/apps/sim/lib/mothership/request/types.ts index 527b111b13c..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 /** @@ -298,6 +300,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 } 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,