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
Original file line number Diff line number Diff line change
Expand Up @@ -1696,6 +1696,139 @@ describe('useChat remount send recovery', () => {
expect(state.postBodies[0]).toHaveProperty('effort')
})

it.each(['initial', 'tail'] as const)(
'recovers a silent %s connection after a tool group without refresh or resending',
async (connection) => {
vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
let unmount: (() => void) | undefined
try {
const history: MothershipChatHistory = {
id: 'chat-silent-stream',
title: 'Silent stream',
messages: [],
activeStreamId: null,
resources: [],
}
const cancelled = vi.fn()
const cursors: string[] = []
let recovered = false
let tailReads = 0
let streamId = ''
const textEvent = (): MothershipStreamV1EventEnvelope => ({
v: 1,
seq: 3,
ts: new Date().toISOString(),
type: 'text',
stream: { streamId },
payload: { channel: 'assistant', text: 'The work continued.' },
})
vi.stubGlobal('fetch', async (input: RequestInfo | URL, init?: RequestInit) => {
const url = String(input)
if (url === '/api/mothership/chat' && init?.method === 'POST') {
const sent = JSON.parse(String(init.body))
state.postBodies.push(sent)
streamId = sent.userMessageId
const events: MothershipStreamV1EventEnvelope[] = [
{
v: 1,
seq: 1,
ts: new Date().toISOString(),
type: 'tool',
stream: { streamId },
payload: {
phase: 'call',
executor: 'go',
mode: 'sync',
toolName: 'run_code',
toolCallId: 'finished-tool',
arguments: { code: 'return 1' },
},
},
{
v: 1,
seq: 2,
ts: new Date().toISOString(),
type: 'tool',
stream: { streamId },
payload: {
phase: 'result',
toolName: 'run_code',
toolCallId: 'finished-tool',
success: true,
output: { value: 1 },
},
},
]
return new Response(
new ReadableStream<Uint8Array>({
start(controller) {
for (const event of events)
controller.enqueue(
new TextEncoder().encode(`data: ${JSON.stringify(event)}\n\n`)
)
if (connection === 'tail') controller.close()
},
cancel: cancelled,
}),
{
headers: {
'Content-Type': 'text/event-stream',
'x-mothership-chat-id': history.id,
},
}
)
}
if (url.includes('/api/mothership/chat/stream')) {
const params = new URL(url, 'https://sim.test').searchParams
if (params.get('batch') === 'true') {
cursors.push(params.get('after') ?? '')
recovered = connection === 'initial' || tailReads > 0
return Response.json({
success: true,
status: 'streaming',
events: recovered ? [{ eventId: 3, streamId, event: textEvent() }] : [],
})
}
tailReads++
return new Response(new ReadableStream<Uint8Array>({ cancel: cancelled }), {
headers: { 'Content-Type': 'text/event-stream' },
})
}
return fetchStub(input, init)
})
const mounted = renderUseChatInChat(history.id, history)
unmount = mounted.unmount
const { getResult } = mounted
await act(async () => {
void getResult().sendMessage('Keep working')
})
await act(async () => vi.advanceTimersByTimeAsync(0))
expect(
getResult()
.messages.flatMap((message) => message.contentBlocks ?? [])
.find((block) => block.toolCall?.id === 'finished-tool')?.toolCall?.status
).toBe('success')
expect(recovered).toBe(false)
await act(async () => vi.advanceTimersByTimeAsync(45_000))
expect(recovered).toBe(true)
expect(cursors.every((cursor) => cursor === '2')).toBe(true)
expect(cancelled).toHaveBeenCalledTimes(1)
expect(
getResult()
.messages.filter((message) => message.role === 'assistant')
.map((message) => message.content)
).toEqual(['The work continued.'])
expect(getResult().isSending).toBe(true)
expect(getResult().error).toBeNull()
expect(state.postBodies).toHaveLength(1)
expect(state.abortBodies).toHaveLength(0)
} finally {
unmount?.()
vi.useRealTimers()
}
}
)

it('recovers a running turn after reconnect exhaustion without reloading or resending', async () => {
vi.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] })
try {
Expand Down
3 changes: 3 additions & 0 deletions apps/sim/app/workspace/[workspaceId]/home/hooks/use-chat.ts
Original file line number Diff line number Diff line change
Expand Up @@ -287,6 +287,8 @@ const RECONNECT_BASE_DELAY_MS = 1000
const RECONNECT_MAX_DELAY_MS = 30_000
const RECONNECT_EXHAUSTED_RECHECK_MS = 30_000
const STREAM_BATCH_FETCH_TIMEOUT_MS = 10_000
/** Both live transports heartbeat every 15s; three missed heartbeats trigger cursor recovery. */
const STREAM_IDLE_TIMEOUT_MS = 45_000
const STREAM_CHAT_ID_RESOLVE_TIMEOUT_MS = 10_000
const CHAT_HISTORY_RECOVERY_TIMEOUT_MS = 10_000
const STOP_REQUEST_TIMEOUT_MS = 15_000
Expand Down Expand Up @@ -2219,6 +2221,7 @@ export function useChat(

try {
await readSSELines(reader, {
idleTimeoutMs: STREAM_IDLE_TIMEOUT_MS,
onData: (raw) => {
if (state.sawCompleteEvent) return true
if (ops.isStale()) return
Expand Down
49 changes: 49 additions & 0 deletions apps/sim/lib/core/utils/sse.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -551,6 +551,55 @@ describe('readSSEEvents', () => {
})

describe('readSSELines', () => {
it('rejects a silent open connection and cancels its reader without waiting for cancellation', async () => {
vi.useFakeTimers()
const cancel = vi.fn(() => new Promise<void>(() => {}))
const reader = new ReadableStream<Uint8Array>({ cancel }).getReader()
const rejected = vi.fn()
const reading = readSSELines(reader, { onData: vi.fn(), idleTimeoutMs: 45_000 }).catch(rejected)
try {
await vi.advanceTimersByTimeAsync(45_000)
expect(rejected).toHaveBeenCalledWith(
expect.objectContaining({ name: 'SSEIdleTimeoutError' })
)
expect(cancel).toHaveBeenCalledTimes(1)
expect(vi.getTimerCount()).toBe(0)
} finally {
void reader.cancel()
await reading
vi.useRealTimers()
}
})

it('counts keepalive comments as activity while a long tool is running', async () => {
vi.useFakeTimers()
let controller!: ReadableStreamDefaultController<Uint8Array>
const stream = new ReadableStream<Uint8Array>({
start: (value) => {
controller = value
},
})
const onData = vi.fn()
const rejected = vi.fn()
const reading = readSSELines(stream, { onData, idleTimeoutMs: 45_000 }).catch(rejected)
try {
for (let heartbeat = 0; heartbeat < 6; heartbeat++) {
await vi.advanceTimersByTimeAsync(15_000)
controller.enqueue(new TextEncoder().encode(': keepalive\n\n'))
await vi.advanceTimersByTimeAsync(0)
}
expect(rejected).not.toHaveBeenCalled()
expect(onData).not.toHaveBeenCalled()
controller.enqueue(new TextEncoder().encode('data: tool finished\n\n'))
controller.close()
await reading
expect(onData).toHaveBeenCalledWith('tool finished')
expect(vi.getTimerCount()).toBe(0)
} finally {
vi.useRealTimers()
}
})

it('delivers raw (un-parsed) data payloads', async () => {
const stream = streamFromStringChunks(['data: raw-one\n\n', 'data: {"keep":"asString"}\n\n'])
const lines: string[] = []
Expand Down
25 changes: 23 additions & 2 deletions apps/sim/lib/core/utils/sse.ts
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,8 @@ export interface ReadSSELinesOptions {
onData: (rawData: string) => SSEStopSignal
/** Aborts the read; checked before each chunk and between events. */
signal?: AbortSignal
/** Reconnect-capable consumers can bound a silent read; comments also keep it alive. */
idleTimeoutMs?: number
}

/**
Expand Down Expand Up @@ -133,16 +135,35 @@ function stripCarriageReturn(line: string): string {
* @param options - The `onData` callback plus an optional `signal`.
*/
export async function readSSELines(source: SSESource, options: ReadSSELinesOptions): Promise<void> {
const { onData, signal } = options
const { onData, signal, idleTimeoutMs } = options
const { reader, ownsLock } = toReader(source)
const decoder = new TextDecoder()
let buffer = ''

const readChunk = async () => {
if (idleTimeoutMs === undefined) return reader.read()
let timer: ReturnType<typeof setTimeout> | undefined
try {
return await new Promise<ReadableStreamReadResult<Uint8Array>>((resolve, reject) => {
timer = setTimeout(() => {
const error = new Error(`SSE connection was silent for ${idleTimeoutMs}ms`)
error.name = 'SSEIdleTimeoutError'
reject(error)
/** Release the transport without waiting for an unresponsive underlying source. */
void reader.cancel(error).catch(() => {})
}, idleTimeoutMs)
reader.read().then(resolve, reject)
})
} finally {
clearTimeout(timer)
}
}

try {
while (true) {
if (signal?.aborted) break

const { done, value } = await reader.read()
const { done, value } = await readChunk()

buffer += done ? decoder.decode() : decoder.decode(value, { stream: true })
const lines = buffer.split('\n')
Expand Down
Loading