Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions apps/sim/lib/workflows/streaming/streaming.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1417,6 +1417,37 @@ describe('createStreamingResponse agent-events-v1', () => {
expect(events).toContainEqual({ event: 'error', error: 'Client cancelled request' })
})

it('stops writing frames once the client cancels the response stream', async () => {
let finishExecution!: (result: ExecutionResult) => void
let completeBlock!: (blockId: string, output: unknown) => Promise<void>
const stream = await createStreamingResponse({
requestId: 'request-client-cancel',
streamConfig: { selectedOutputs: ['block_answer'], includeFileBase64: false },
executeFn: async ({ onBlockComplete }) => {
completeBlock = onBlockComplete
return await new Promise<ExecutionResult>((resolve) => {
finishExecution = resolve
})
},
})
const reader = stream.getReader()
await reader.read()
await reader.cancel()

await expect(completeBlock('block', { answer: 'late answer' })).resolves.toBeUndefined()
finishExecution({ success: false, status: 'cancelled', output: {}, logs: [] })

await vi.waitFor(() =>
expect(workflowStreamingLogger.info).toHaveBeenCalledWith(
'[request-client-cancel] Streaming execution aborted by client disconnect'
)
)
expect(workflowStreamingLogger.error).not.toHaveBeenCalledWith(
'[request-client-cancel] Stream error',
expect.anything()
)
})

it('thinking never enters streamedChunks / log content rewrite', async () => {
const headers = new Headers({
'x-sim-stream-protocol': 'agent-events-v1',
Expand Down
45 changes: 26 additions & 19 deletions apps/sim/lib/workflows/streaming/streaming.ts
Original file line number Diff line number Diff line change
Expand Up @@ -523,6 +523,7 @@ export async function createStreamingResponse(
const maxThinkingChars = DEFAULT_MAX_THINKING_CHARS

let requestAborted = false
let streamCancelled = false
const onRequestAbort = () => {
requestAborted = true
timeoutController.abort()
Expand All @@ -549,8 +550,13 @@ export async function createStreamingResponse(

return new ReadableStream({
async start(controller) {
/** After `cancel()` the controller is closed and enqueue throws; the client is gone. */
const enqueue = (bytes: Uint8Array) => {
if (!streamCancelled) controller.enqueue(bytes)
}

/** Flush headers promptly and keep silent blocks alive through idle-limited proxies. */
controller.enqueue(encodeSSEComment('keepalive'))
enqueue(encodeSSEComment('keepalive'))
keepaliveId = setInterval(() => {
try {
controller.enqueue(encodeSSEComment('keepalive'))
Expand Down Expand Up @@ -586,7 +592,7 @@ export async function createStreamingResponse(
state.streamedSelectedOutputKeys.add(options.selectedOutputKey)
}
const frame: ChatStreamChunkFrame = { blockId, chunk }
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
state.processedOutputs.add(blockId)
}

Expand All @@ -601,7 +607,7 @@ export async function createStreamingResponse(
event: 'thinking',
data: forwarded,
}
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
}

const sendTool = (
Expand All @@ -619,7 +625,7 @@ export async function createStreamingResponse(
name,
...(phase === 'end' && status ? { status } : {}),
}
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
}

/**
Expand Down Expand Up @@ -667,7 +673,7 @@ export async function createStreamingResponse(
emittedSinceReset = true
} else {
const frame: ChatStreamChunkFrame = { blockId, chunk: text }
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
}
}

Expand All @@ -688,7 +694,7 @@ export async function createStreamingResponse(
} else if (sinkAnswerText && event.type === 'turn_end') {
if (event.turn === 'intermediate' && emittedSinceReset) {
const frame: ChatStreamChunkResetFrame = { blockId, event: 'chunk_reset' }
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
// Re-arm separator bookkeeping so re-streamed text starts clean.
emittedSinceReset = false
state.processedOutputs.delete(blockId)
Expand Down Expand Up @@ -729,7 +735,7 @@ export async function createStreamingResponse(
blockId,
error: getErrorMessage(error, 'Stream reading error'),
}
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
} finally {
unsubscribe?.()
}
Expand All @@ -746,7 +752,7 @@ export async function createStreamingResponse(
const selectedOutputBlockId = outputBlockId ?? blockId
state.completedBlockIds.add(selectedOutputBlockId)

if (!streamConfig.selectedOutputs?.length) {
if (streamCancelled || !streamConfig.selectedOutputs?.length) {
return
}

Expand Down Expand Up @@ -776,7 +782,7 @@ export async function createStreamingResponse(
: output

for (const descriptor of matchingOutputs) {
if (state.selectedOutputError) {
if (streamCancelled || state.selectedOutputError) {
break
}
try {
Expand Down Expand Up @@ -837,7 +843,7 @@ export async function createStreamingResponse(
event: 'output',
data: hydratedOutput,
}
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
state.selectedOutputBytes = nextSelectedOutputBytes
state.streamedSelectedOutputKeys.add(descriptor.key)
state.processedOutputs.add(selectedOutputBlockId)
Expand Down Expand Up @@ -869,7 +875,7 @@ export async function createStreamingResponse(
blockId: selectedOutputBlockId,
error: errorMessage,
}
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
break
}
}
Expand Down Expand Up @@ -908,7 +914,7 @@ export async function createStreamingResponse(
await result._streamingMetadata.loggingSession.markAsFailed(timeoutErrorMessage)
}
const frame: ChatStreamErrorFrame = { event: 'error', error: timeoutErrorMessage }
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
} else if (result.status === 'cancelled' && requestAborted) {
logger.info(`[${requestId}] Streaming execution aborted by client disconnect`)
if (result._streamingMetadata?.loggingSession) {
Expand All @@ -917,7 +923,7 @@ export async function createStreamingResponse(
}
// No `final` after abort; clients that already disconnected ignore these.
const frame: ChatStreamErrorFrame = { event: 'error', error: 'Client cancelled request' }
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
} else {
await completeLoggingSession(result)

Expand Down Expand Up @@ -953,18 +959,18 @@ export async function createStreamingResponse(
...(result.status === 'paused' && { status: 'paused' }),
},
}
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
}
}

// Terminal marker: always follows success `final` or a single terminal `error`.
controller.enqueue(encodeSSE('[DONE]'))
enqueue(encodeSSE('[DONE]'))

if (executionId) {
await cleanupExecutionBase64Cache(executionId)
}

controller.close()
if (!streamCancelled) controller.close()
} catch (error) {
logger.error(
`[${requestId}] Stream error`,
Expand All @@ -975,22 +981,23 @@ export async function createStreamingResponse(
? SELECTED_OUTPUT_TOO_LARGE_MESSAGE
: getErrorMessage(error, 'Stream processing error')
const frame: ChatStreamErrorFrame = { event: 'error', error: errorMessage }
controller.enqueue(encodeSSE(frame))
enqueue(encodeSSE(frame))
// Same terminal rule as timeout/abort: one error, then [DONE], never `final`.
controller.enqueue(encodeSSE('[DONE]'))
enqueue(encodeSSE('[DONE]'))

if (executionId) {
await cleanupExecutionBase64Cache(executionId)
}

controller.close()
if (!streamCancelled) controller.close()
} finally {
stopKeepalive()
cleanupRequestAbort()
timeoutController.cleanup()
}
},
async cancel(reason) {
streamCancelled = true
logger.info(
`[${requestId}] Streaming response cancelled`,
projectResolvedSecretDiagnosticError(reason, undefined)
Expand Down
Loading