From c7ee632516c1c1d8168729d2f0112f27798eec4d Mon Sep 17 00:00:00 2001 From: Siddharth Ganesan Date: Thu, 24 Sep 2026 12:23:42 -0700 Subject: [PATCH] Fix staging Stop admission and oversized response uploads --- .../app/api/copilot/chat/stop/route.test.ts | 84 ++++++++++----- apps/sim/app/api/copilot/chat/stop/route.ts | 32 +++--- .../home/hooks/use-chat.mount-send.test.tsx | 100 ++++++++++++++---- .../[workspaceId]/home/hooks/use-chat.ts | 42 ++------ .../sim/lib/mothership/chat/terminal-state.ts | 34 ++++-- .../request/session/explicit-abort.test.ts | 7 +- .../request/session/explicit-abort.ts | 20 +--- 7 files changed, 195 insertions(+), 124 deletions(-) diff --git a/apps/sim/app/api/copilot/chat/stop/route.test.ts b/apps/sim/app/api/copilot/chat/stop/route.test.ts index f2334667761..87488ef95f1 100644 --- a/apps/sim/app/api/copilot/chat/stop/route.test.ts +++ b/apps/sim/app/api/copilot/chat/stop/route.test.ts @@ -148,42 +148,70 @@ describe('copilot chat stop route', () => { expect(mockAppendCopilotChatMessages).not.toHaveBeenCalled() }) - it('appends a stopped assistant message even with no content', async () => { + it('persists a response larger than the HTTP limit from an identifiers-only Stop', async () => { mockReads({ chat: { workspaceId: 'ws-1', conversationId: 'stream-1', model: null }, last: { messageId: 'stream-1', role: 'user' }, }) - - const response = await stopRequest( - createRequest({ chatId: 'chat-1', streamId: 'stream-1', content: '' }) - ) - - expect(mockReadEvents).toHaveBeenCalledWith('stream-1', '0') - expect(response.status).toBe(200) - expect(await response.json()).toEqual({ success: true }) - - const setArg = dbChainMockFns.set.mock.calls[0]?.[0] as Record - expect(setArg.conversationId).toBeNull() - expect(Object.hasOwn(setArg, 'messages')).toBe(false) - - expect(mockAppendCopilotChatMessages).toHaveBeenCalledTimes(1) - const [, appended] = mockAppendCopilotChatMessages.mock.calls[0] - expect(appended[0]).toMatchObject({ - role: 'assistant', - content: '', - contentBlocks: [{ type: 'complete', status: 'cancelled' }], - }) - - expect(mockPublishStatusChanged).toHaveBeenCalledWith( - expect.objectContaining({ workspaceId: 'ws-1' }), + const content = 'x'.repeat(11 * 1024 * 1024) + const envelope = { v: 1, ts: '2026-09-24T19:00:00Z', stream: { streamId: 'stream-1' } } + mockReadEvents.mockResolvedValue([ + { ...envelope, seq: 1, type: 'text', payload: { channel: 'assistant', text: content } }, { - chatId: 'chat-1', - type: 'completed', - streamId: 'stream-1', - } + ...envelope, + seq: 2, + type: 'tool', + payload: { + phase: 'call', + toolCallId: 'call-1', + toolName: 'run_code', + arguments: { code: 'preserve me' }, + status: 'executing', + }, + }, + ]) + const request = createRequest({ chatId: 'chat-1', streamId: 'stream-1' }) + expect((await request.clone().text()).length).toBeLessThan(1024) + const response = await stopRequest(request) + expect(response.status).toBe(200) + expect(mockReadEvents).toHaveBeenCalledOnce() + expect(mockAppendCopilotChatMessages).toHaveBeenCalledOnce() + const saved = mockAppendCopilotChatMessages.mock.calls[0][1][0] + expect(saved.content).toBe(content) + expect(saved.contentBlocks).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + type: 'tool', + toolCall: expect.objectContaining({ id: 'call-1', params: { code: 'preserve me' } }), + }), + { type: 'complete', status: 'cancelled' }, + ]) ) + expect(dbChainMockFns.set.mock.calls[0][0].conversationId).toBeNull() + expect(mockPublishStatusChanged).toHaveBeenCalledOnce() }) + it.each([{ seqs: [] }, { seqs: [2, 3] }, { seqs: [1, 3] }])( + 'leaves incomplete replay $seqs to the run owner without erasing its response', + async ({ seqs }) => { + mockReadEvents.mockResolvedValue( + seqs.map((seq) => ({ + v: 1, + seq, + ts: '2026-09-24T19:00:00Z', + stream: { streamId: 'stream-1' }, + type: 'text', + payload: { channel: 'assistant', text: 'tail only' }, + })) + ) + const response = await stopRequest(createRequest({ chatId: 'chat-1', streamId: 'stream-1' })) + expect(response.status).toBe(200) + expect(mockAppendCopilotChatMessages).not.toHaveBeenCalled() + expect(dbChainMockFns.set).not.toHaveBeenCalled() + expect(mockPublishStatusChanged).not.toHaveBeenCalled() + } + ) + it('appends a stopped assistant message if the stream marker was already cleared', async () => { mockReads({ chat: { workspaceId: 'ws-1', conversationId: null, model: null }, diff --git a/apps/sim/app/api/copilot/chat/stop/route.ts b/apps/sim/app/api/copilot/chat/stop/route.ts index 2ed78fa7208..2473b8ddf91 100644 --- a/apps/sim/app/api/copilot/chat/stop/route.ts +++ b/apps/sim/app/api/copilot/chat/stop/route.ts @@ -11,7 +11,10 @@ import { type PersistedMessage, withStoppedContentBlock, } from '@/lib/mothership/chat/persisted-message' -import { finalizeAssistantTurn } from '@/lib/mothership/chat/terminal-state' +import { + finalizeAssistantTurn, + readStoppedAssistantMessage, +} from '@/lib/mothership/chat/terminal-state' import { publishChatStatusChanged } from '@/lib/mothership/chat-status' import { CopilotChatFinalizeOutcome, @@ -61,23 +64,28 @@ export const POST = withRouteHandler((req: NextRequest) => : hasContent ? [{ type: 'text', channel: 'assistant', content }] : [] - const assistantMessage: PersistedMessage = withStoppedContentBlock( - normalizeMessage({ - id: generateId(), - role: 'assistant', - content, - timestamp: new Date().toISOString(), - contentBlocks: assistantBlocks, - ...(requestId ? { requestId } : {}), - }) - ) + const assistantMessage: PersistedMessage | null = + hasContent || hasBlocks + ? withStoppedContentBlock( + normalizeMessage({ + id: generateId(), + role: 'assistant', + content, + timestamp: new Date().toISOString(), + contentBlocks: assistantBlocks, + ...(requestId ? { requestId } : {}), + }) + ) + : await readStoppedAssistantMessage(streamId) + /** The run owner retains the full response if replay was trimmed or has not started. */ + if (!assistantMessage) return NextResponse.json({ success: true }) const result = await finalizeAssistantTurn({ chatId, userId: session.user.id, userMessageId: streamId, assistantMessage, streamMarkerPolicy: 'active-or-cleared', - preferServerReplay: true, + preferServerReplay: hasContent || hasBlocks, }) span.setAttribute(TraceAttr.CopilotStopAppendedAssistant, result.appendedAssistant) const stopOutcome = !result.found diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx index 618e6e87624..8f50c51dc3e 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.mount-send.test.tsx @@ -81,7 +81,6 @@ import { seedDeploymentShape, } from '@/lib/core/config/deployment-shape' import { MothershipHandoffStorage } from '@/lib/core/utils/browser-storage' -import { normalizeMessage } from '@/lib/mothership/chat/persisted-message' import type { MothershipStreamV1EventEnvelope } from '@/lib/mothership/generated/mothership-stream-v1' import { createSearchResource } from '@/lib/mothership/resources/search' import { getChatResourceSelectionId } from '@/lib/mothership/resources/types' @@ -117,6 +116,7 @@ interface NetworkState { abortSettlements: boolean[] abortBodies: CopilotChatAbortBody[] stopBodies: CopilotChatStopBody[] + toolInputPadding?: string abortTraceparents: Array } @@ -221,11 +221,14 @@ async function fetchStub(input: RequestInfo | URL, init?: RequestInit): Promise< stream: { streamId }, payload: { phase: 'call', - executor: 'client', - mode: 'async', - toolName: 'run_workflow', + executor: state.toolInputPadding ? 'go' : 'client', + mode: state.toolInputPadding ? 'sync' : 'async', + toolName: state.toolInputPadding ? 'run_code' : 'run_workflow', toolCallId: 'this-chat-tool', - arguments: { workflowId: 'this-chat-workflow' }, + arguments: { + workflowId: 'this-chat-workflow', + ...(state.toolInputPadding ? { padding: state.toolInputPadding } : {}), + }, }, } return new Response( @@ -751,6 +754,7 @@ describe('useChat remount send recovery', () => { state.pendingAdmissions.clear() state.abortSettlements = [] state.abortBodies = [] + state.toolInputPadding = undefined state.stopBodies = [] state.abortTraceparents = [] mockRequestJson.mockResolvedValue({ chats: [] }) @@ -1450,6 +1454,41 @@ describe('useChat remount send recovery', () => { } ) + it.each([false, true])( + 'does not interrupt or send a queued edit before submission (explicit ID: %s)', + async (explicitId) => { + state.postBehavior = 'task' + const { getResult } = renderUseChatInChat('chat-a') + await act(async () => { + void getResult().sendMessage('Original request') + }) + await waitFor(() => state.postBodies.length === 1 && getResult().isSending) + const beforeRender = getResult() + let queuedId = '' + await act(async () => { + void beforeRender.sendMessage('Unfinished correction') + queuedId = allQueuedMessages()[0].id + beforeRender.editQueuedMessage(queuedId) + void beforeRender.sendNow(explicitId ? queuedId : undefined) + }) + expect(state.abortBodies).toHaveLength(0) + expect(state.postBodies).toHaveLength(1) + expect(allQueuedMessages()).toEqual([ + expect.objectContaining({ id: queuedId, content: 'Unfinished correction' }), + ]) + expect(useMothershipQueueStore.getState().editing['chat-a']).toBe(queuedId) + state.postBehavior = 'hang' + await act(async () => { + void getResult().sendMessage('Finished correction') + void getResult().sendNow() + }) + await waitFor(() => state.postBodies.length === 2) + expect(state.postBodies[1].message).toBe('Finished correction') + expect(state.abortBodies).toHaveLength(1) + expect(allQueuedMessages()).toHaveLength(0) + } + ) + it('sends the live queue head once without waiting for a render, after Stop settles', async () => { state.postBehavior = 'task' const { getResult } = renderUseChatInChat('chat-a') @@ -1713,7 +1752,7 @@ describe('useChat remount send recovery', () => { } }) - it('preserves a visible workflow watch when Stop persists the partial response', async () => { + it('preserves the visible workflow watch while Stop sends only identifiers', async () => { state.postBehavior = 'task' const { getResult } = renderUseChatInChat('chat-a') await act(async () => { @@ -1730,20 +1769,43 @@ describe('useChat remount send recovery', () => { await getResult().stopGeneration() }) expect(state.stopBodies).toHaveLength(1) - const saved = state.stopBodies[0] - const restored = normalizeMessage({ - id: 'saved-assistant', - role: 'assistant', - content: saved.content, - contentBlocks: saved.contentBlocks, - }) - expect(restored.contentBlocks?.find((block) => block.type === 'task')?.task).toEqual({ - taskId: 'watch-1', - kind: 'workflow_run', - status: 'pending', - target: { workflowId: 'workflow-1', executionId: 'watched-execution' }, - note: 'Check the completed invoice run', + expect(state.stopBodies[0]).toEqual({ + chatId: 'chat-a', + streamId: state.postBodies[0].userMessageId, + }) + const task = getResult() + .messages.flatMap((message) => message.contentBlocks ?? []) + .find((block) => block.type === 'task')?.task + expect(task?.taskId).toBe('watch-1') + }) + + it('sends a queued correction after stopping with more than 10 MiB of tool input', async () => { + state.postBehavior = 'tool' + state.toolInputPadding = 'x'.repeat(11 * 1024 * 1024) + const { getResult } = renderUseChatInChat('chat-a') + await act(async () => { + void getResult().sendMessage('Start working') }) + await waitFor(() => + expect( + getResult().messages.some((message) => + message.contentBlocks?.some((block) => block.toolCall?.id === 'this-chat-tool') + ) + ).toBe(true) + ) + state.postBehavior = 'hang' + await act(async () => { + await getResult().sendMessage('Use the correction') + void getResult().sendNow() + }) + await waitFor(() => expect(state.postBodies).toHaveLength(2)) + expect(state.postBodies[1].message).toBe('Use the correction') + expect(state.stopBodies).toHaveLength(1) + expect(state.stopBodies[0]).not.toHaveProperty('content') + expect(state.stopBodies[0]).not.toHaveProperty('contentBlocks') + expect(new TextEncoder().encode(JSON.stringify(state.stopBodies[0])).length).toBeLessThan(1024) + expect(allQueuedMessages()).toHaveLength(0) + expect(getResult().error).toBeNull() }) it.each([false, true])( diff --git a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts index ed353410510..7bad75a9d2e 100644 --- a/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts +++ b/apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts @@ -152,7 +152,6 @@ import { type ReconnectReplaySelection, reconcileLiveAssistantTurn, selectReconnectReplayState, - toRawPersistedContentBlock, } from './message-reconcile' import { clearQueuedSendHandoffClaim, @@ -3078,12 +3077,10 @@ export function useChat( } }, [recoverActiveStreamFromRedis]) - const persistPartialResponse = useCallback( + const persistStoppedResponse = useCallback( async (overrides?: { chatId?: string streamId?: string - content?: string - blocks?: ContentBlock[] // `stopGeneration` must snapshot these BEFORE clearActiveTurn() // nulls the refs, or the fetch sees undefined. requestId?: string @@ -3093,31 +3090,9 @@ export function useChat( const streamId = overrides?.streamId ?? streamIdRef.current if (!chatId || !streamId) return - const content = overrides?.content ?? streamingContentRef.current const requestId = overrides?.requestId ?? streamRequestIdRef.current const traceparent = overrides?.traceparent ?? streamTraceparentRef.current - const sourceBlocks = overrides?.blocks ?? streamingBlocksRef.current - const storedBlocks = sourceBlocks - .map((block) => { - const persisted = toRawPersistedContentBlock(block) - if ( - persisted?.toolCall && - (persisted.toolCall.state === 'executing' || persisted.toolCall.state === 'cancelled') - ) { - persisted.toolCall = { - ...persisted.toolCall, - state: 'cancelled', - display: { title: 'Stopped by user' }, - } - } - return persisted - }) - .filter((block) => block !== null) - if (storedBlocks.length > 0) { - storedBlocks.push({ type: 'complete', status: 'cancelled' }) - } - try { const res = await fetch(stopPathRef.current, { method: 'POST', @@ -3129,8 +3104,6 @@ export function useChat( body: JSON.stringify({ chatId, streamId, - content, - ...(storedBlocks.length > 0 && { contentBlocks: storedBlocks }), ...(requestId ? { requestId } : {}), }), }) @@ -4330,7 +4303,6 @@ export function useChat( throw err } - const stopContentSnapshot = streamingContentRef.current const stopNow = Date.now() const stopBlocksSnapshot = streamingBlocksRef.current.map((block) => ({ ...block, @@ -4489,11 +4461,9 @@ export function useChat( } if (wasSending && resolvedChatId) { - await persistPartialResponse({ + await persistStoppedResponse({ chatId: resolvedChatId, streamId: sid, - content: stopContentSnapshot, - blocks: stopBlocksSnapshot, requestId: stopRequestIdSnapshot, traceparent: stopTraceparentSnapshot, }) @@ -4560,7 +4530,7 @@ export function useChat( cancelActiveBrowserTools, invalidateChatQueries, notifyTurnEnded, - persistPartialResponse, + persistStoppedResponse, queryClient, resolveChatIdForStream, resetEphemeralPreviewState, @@ -4785,9 +4755,11 @@ export function useChat( const sendQueuedMessageImmediately = useCallback( async (id?: string) => { - const queue = useMothershipQueueStore.getState().queues[chatKeyRef.current] + const queueState = useMothershipQueueStore.getState() + const chatKey = chatKeyRef.current + const queue = queueState.queues[chatKey] const msg = id === undefined ? queue?.[0] : queue?.find((queued) => queued.id === id) - if (!msg) return + if (!msg || queueState.editing[chatKey] === msg.id) return if (queuedMessageDispatchIdsRef.current.has(msg.id)) return const admissionPending = hasPendingChatAdmission() diff --git a/apps/sim/lib/mothership/chat/terminal-state.ts b/apps/sim/lib/mothership/chat/terminal-state.ts index a23e9f388b5..f08f4157c51 100644 --- a/apps/sim/lib/mothership/chat/terminal-state.ts +++ b/apps/sim/lib/mothership/chat/terminal-state.ts @@ -39,6 +39,28 @@ export interface FinalizeAssistantTurnResult { outcome: (typeof CopilotChatFinalizeOutcome)[keyof typeof CopilotChatFinalizeOutcome] } +/** Rebuild a stopped response only when the server still has its complete event prefix. */ +export async function readStoppedAssistantMessage( + streamId: string +): Promise { + const events = await readEvents(streamId, '0') + /** StreamWriter starts at 1; Redis may trim oldest events or skip corrupt entries. */ + if (events.length === 0 || !events.every((event, index) => event.seq === index + 1)) return null + const replay = buildLiveAssistantMessage({ + streamId, + events: events.map(toStreamBatchEvent), + status: 'cancelled', + }) + if (!replay) return null + return withStoppedContentBlock({ + ...replay, + content: redactSensitiveContent(replay.content), + ...(replay.contentBlocks + ? { contentBlocks: mergeAndRedactPersistedBlocks(replay.contentBlocks) } + : {}), + }) +} + /** * Clear the active stream marker for a chat and optionally append the assistant * message once for its turn, including user steering accepted within that turn. @@ -154,17 +176,7 @@ export async function finalizeAssistantTurn({ if (assistantMessage && canAppendAssistant) { let response = assistantMessage if (preferServerReplay) { - const events = await readEvents(userMessageId, '0') - /** StreamWriter starts at 1; Redis trims oldest events and skips corrupt entries. */ - const replayIsComplete = - events.length > 0 && events.every((event, index) => event.seq === index + 1) - const replay = replayIsComplete - ? buildLiveAssistantMessage({ - streamId: userMessageId, - events: events.map(toStreamBatchEvent), - status: 'cancelled', - }) - : null + const replay = await readStoppedAssistantMessage(userMessageId) /** A stopped client's snapshot may be empty; preserve canonical output before the first finalizer commits. */ const replayHasContent = !!replay?.content.trim() || diff --git a/apps/sim/lib/mothership/request/session/explicit-abort.test.ts b/apps/sim/lib/mothership/request/session/explicit-abort.test.ts index d1283158600..b687a78721d 100644 --- a/apps/sim/lib/mothership/request/session/explicit-abort.test.ts +++ b/apps/sim/lib/mothership/request/session/explicit-abort.test.ts @@ -23,6 +23,7 @@ vi.mock('@/lib/mothership/server/agent-url', () => ({ getMothershipSourceEnvHeaders: vi.fn().mockReturnValue({ 'X-Sim-Source-Env': 'test' }), })) +import { BillingCallbackHeaders } from '@/lib/mothership/generated/billing' import { AbortRequest } from '@/lib/mothership/generated/protocol' import { requestExplicitStreamAbort } from '@/lib/mothership/request/session/explicit-abort' @@ -36,20 +37,22 @@ describe('requestExplicitStreamAbort', () => { }) }) - it('sends an explicit legacy protocol marker for strict Go admission', async () => { + it('sends a valid unbilled control request through worker admission', async () => { const result = await requestExplicitStreamAbort({ streamId: '11111111-1111-4111-8111-111111111111', userId: 'user-1', chatId: 'chat-1', }) expect(result).toEqual({ settled: false }) + const headers = mockFetchGo.mock.calls[0][1].headers + expect(BillingCallbackHeaders.safeParse(headers).success).toBe(true) + expect(headers).not.toHaveProperty('x-sim-billing-protocol') expect(mockFetchGo).toHaveBeenCalledWith( 'https://copilot.test/api/streams/explicit-abort', expect.objectContaining({ headers: expect.objectContaining({ 'x-api-key': 'sim-agent-key', - 'x-sim-billing-protocol': 'legacy-v0', }), }) ) diff --git a/apps/sim/lib/mothership/request/session/explicit-abort.ts b/apps/sim/lib/mothership/request/session/explicit-abort.ts index 2411e85349e..651a87d749a 100644 --- a/apps/sim/lib/mothership/request/session/explicit-abort.ts +++ b/apps/sim/lib/mothership/request/session/explicit-abort.ts @@ -1,19 +1,12 @@ import type { Context } from '@opentelemetry/api' import { sleep } from '@sim/utils/helpers' import { toRecordOrNull } from '@sim/utils/object' -import { - COPILOT_BILLING_PROTOCOL, - COPILOT_BILLING_PROTOCOL_HEADER, -} from '@/lib/billing/core/billing-attribution' -import { env } from '@/lib/core/config/env' import { AbortRequest, type AbortResponse } from '@/lib/mothership/generated/protocol' import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1' import { fetchGo } from '@/lib/mothership/request/go/fetch' +import { mothershipRequestHeaders } from '@/lib/mothership/request/headers' import { AbortReason } from '@/lib/mothership/request/session/abort' -import { - getMothershipBaseURL, - getMothershipSourceEnvHeaders, -} from '@/lib/mothership/server/agent-url' +import { getMothershipBaseURL } from '@/lib/mothership/server/agent-url' export const DEFAULT_EXPLICIT_ABORT_TIMEOUT_MS = 3000 @@ -32,14 +25,7 @@ export async function requestExplicitStreamAbort(params: { otelContext, } = params - const headers: Record = { - 'Content-Type': 'application/json', - [COPILOT_BILLING_PROTOCOL_HEADER]: COPILOT_BILLING_PROTOCOL.legacy, - } - if (env.COPILOT_API_KEY) { - headers['x-api-key'] = env.COPILOT_API_KEY - } - Object.assign(headers, getMothershipSourceEnvHeaders()) + const headers = mothershipRequestHeaders() const controller = new AbortController() const timeout = setTimeout(