Skip to content

Commit ef42424

Browse files
authored
fix(sse): bound workspace SSE connection lifetime (#7058)
* fix(sse): bound workspace SSE connection lifetime Teardown ran only from the request abort listener and the stream cancel callback, both of which fire only when the runtime reports a client disconnect. Nothing else bounded the connection, so a missed report left the pub/sub handler, the heartbeat timer, and the stream's undrained queue held for the life of the process. Add a jittered lifetime ceiling checked on the existing heartbeat tick, tighten reclaim for a vanished consumer via desiredSize, remove the abort listener on every teardown path, and run full teardown when a heartbeat enqueue fails. Log the close reason so opens minus closes is observable. * fix(mothership): resync chat caches after an SSE reconnect gap task_status events are transient and never replayed, so any window with no open connection can drop a create, rename, delete, or completion. The chat hook reconnected silently and reconciled nothing, leaving list and detail caches stale until an unrelated action refreshed them. Resync on reconnect, on the first open of a re-subscription, and on a first open that only succeeded after an error, matching the pattern useMcpToolsEvents already uses for the same gap. * fix(mothership): keep reconnect resync off locally streaming chats The resync invalidated every chat detail, including one whose stream this client is rendering optimistically. Refetching there replaces the local transcript with a server copy that does not yet hold the in-flight message, which is exactly what status events avoid via shouldSkipDetailInvalidationForStreamEvent. Filter the detail invalidation with the same isLocalOptimisticActiveStream check. Those chats reconcile when their own stream finishes. * fix(mothership): only skip resync for a stream still running Optimistic markers alone were the skip condition, but a finished turn can leave activeStreamId and its live-assistant message cached when finalization skips detail invalidation for a queued follow-up. That chat would then be excluded from every future resync — permanently, since only a refetch clears the markers, and the resync was the refetch. Gate the skip on a non-terminal streamSnapshot status so it covers turns that are genuinely still streaming. Exports isTerminalStreamStatus, which was already the private check for this in effective-transcript. * fix(sse): raise the ceiling and narrow reconnect resync to the lists Deciding from cache whether a chat is still streaming is not reliable — the optimistic markers outlive the turn, and each refinement of that predicate exposed another state where it answers wrongly. Drop it: the resync now invalidates only the workspace lists, which is always safe, and chat detail reconciliation stays as it is today rather than being half-solved here. Raise the ceiling to 4h, matching lib/realtime/event-stream-route.ts. A healthy client is drained and so is never unread; the unread check is what reclaims a vanished consumer, and it does so within minutes. A short ceiling would therefore only force reconnects on the connections that are working, and every reconnect is a window where a transient event can be missed. Retention stays bounded by the ceiling instead of by process uptime.
1 parent 49a3399 commit ef42424

4 files changed

Lines changed: 263 additions & 26 deletions

File tree

apps/sim/hooks/use-mothership-chat-events.test.ts

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,10 @@ vi.mock('@/lib/browser-agent/transport', () => ({ suspendBrowserScope }))
1414
vi.mock('@/lib/terminal/transport', () => ({ suspendTerminalScope }))
1515

1616
import { mothershipChatKeys } from '@/hooks/queries/mothership-chats'
17-
import { handleMothershipChatStatusEvent } from '@/hooks/use-mothership-chat-events'
17+
import {
18+
handleMothershipChatStatusEvent,
19+
resyncMothershipChatCaches,
20+
} from '@/hooks/use-mothership-chat-events'
1821

1922
describe('handleMothershipChatStatusEvent', () => {
2023
const queryClient = {
@@ -419,3 +422,30 @@ describe('handleMothershipChatStatusEvent', () => {
419422
expect(queryClient.removeQueries).not.toHaveBeenCalled()
420423
})
421424
})
425+
426+
describe('resyncMothershipChatCaches', () => {
427+
const queryClient = {
428+
invalidateQueries: vi.fn().mockResolvedValue(undefined),
429+
} satisfies Pick<QueryClient, 'invalidateQueries'>
430+
431+
beforeEach(() => {
432+
vi.clearAllMocks()
433+
})
434+
435+
it('invalidates the workspace lists', () => {
436+
resyncMothershipChatCaches(queryClient, 'ws-1')
437+
438+
expect(queryClient.invalidateQueries).toHaveBeenCalledTimes(1)
439+
expect(queryClient.invalidateQueries).toHaveBeenCalledWith({
440+
queryKey: mothershipChatKeys.workspaceLists('ws-1'),
441+
})
442+
})
443+
444+
it('leaves chat details untouched so a mounted stream cannot be refetched mid-turn', () => {
445+
resyncMothershipChatCaches(queryClient, 'ws-1')
446+
447+
expect(queryClient.invalidateQueries).not.toHaveBeenCalledWith(
448+
expect.objectContaining({ queryKey: mothershipChatKeys.details() })
449+
)
450+
})
451+
})

apps/sim/hooks/use-mothership-chat-events.ts

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@ import { type MothershipChatHistory, mothershipChatKeys } from '@/hooks/queries/
99

1010
const logger = createLogger('MothershipChatEvents')
1111

12+
/** Workspaces this process has subscribed to before, so a re-subscribe can be told from a first one. */
13+
const everSubscribed = new Set<string>()
14+
1215
const CHAT_STATUS_TYPES = ['started', 'completed', 'created', 'deleted', 'renamed'] as const
1316
type ChatStatusEventType = (typeof CHAT_STATUS_TYPES)[number]
1417
const CHAT_STATUS_TYPE_SET = new Set<string>(CHAT_STATUS_TYPES)
@@ -128,6 +131,28 @@ export function handleMothershipChatStatusEvent(
128131
}
129132
}
130133

134+
/**
135+
* Re-syncs the workspace chat lists after a gap in the event stream.
136+
*
137+
* `task_status` events are transient — nothing replays what was published while
138+
* no connection was open — so a reconnect may have missed a create, rename, or
139+
* delete. The lists carry that workspace-level state, and refetching them is
140+
* always safe.
141+
*
142+
* Chat details are deliberately left alone. The only one that would refetch is
143+
* the mounted chat, which may be rendering an in-flight stream, and refetching
144+
* there replaces the optimistic transcript with a server copy that does not yet
145+
* hold the streaming message. Cached state cannot reliably say whether a turn is
146+
* still running — the optimistic markers outlive it — so detail reconciliation
147+
* stays as it is today and belongs with the streaming state that can answer it.
148+
*/
149+
export function resyncMothershipChatCaches(
150+
queryClient: Pick<QueryClient, 'invalidateQueries'>,
151+
workspaceId: string
152+
): void {
153+
queryClient.invalidateQueries({ queryKey: mothershipChatKeys.workspaceLists(workspaceId) })
154+
}
155+
131156
/**
132157
* Subscribes to chat status SSE events and invalidates chat caches on changes.
133158
* The SSE event name remains `task_status` for wire compatibility.
@@ -154,7 +179,28 @@ export function useMothershipChatEvents(workspaceId: string | undefined) {
154179
)
155180
})
156181

182+
// `onopen` fires on the initial connect and on every auto-reconnect. Re-sync
183+
// whenever a gap could have swallowed an event: on any reconnect, on the
184+
// first open of a RE-subscription (switching workspace away and back tears
185+
// the connection down, and the list/detail queries remount inside their
186+
// stale times so they do not refetch on their own), and on a first open that
187+
// only succeeded after an error (the initial queries may have failed during
188+
// that gap and will not retry themselves). Skip only a clean first
189+
// subscription — those queries fetch fresh on their own initial mount.
190+
const isResubscribe = everSubscribed.has(workspaceId)
191+
everSubscribed.add(workspaceId)
192+
let opened = false
193+
let erroredBeforeOpen = false
194+
195+
eventSource.onopen = () => {
196+
if (opened || isResubscribe || erroredBeforeOpen) {
197+
resyncMothershipChatCaches(queryClient, workspaceId)
198+
}
199+
opened = true
200+
}
201+
157202
eventSource.onerror = () => {
203+
if (!opened) erroredBeforeOpen = true
158204
logger.warn(`SSE connection error for workspace ${workspaceId}`)
159205
}
160206

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
5+
import { authMockFns, permissionsMock, permissionsMockFns } from '@sim/testing'
6+
import { NextRequest } from 'next/server'
7+
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
8+
import {
9+
createWorkspaceSSE,
10+
HEARTBEAT_INTERVAL_MS,
11+
MAX_CONNECTION_JITTER_MS,
12+
MAX_CONNECTION_MS,
13+
MAX_UNDRAINED_CHUNKS,
14+
} from '@/lib/events/sse-endpoint'
15+
16+
vi.mock('@/lib/workspaces/permissions/utils', () => permissionsMock)
17+
18+
const PAST_CEILING_MS = MAX_CONNECTION_MS + MAX_CONNECTION_JITTER_MS + HEARTBEAT_INTERVAL_MS
19+
20+
/** Enough undrained heartbeats to trip the unread check, and no more. */
21+
const PAST_UNREAD_MS = (MAX_UNDRAINED_CHUNKS + 2) * HEARTBEAT_INTERVAL_MS
22+
23+
async function openConnection(signal: AbortSignal = new AbortController().signal) {
24+
const unsubscribe = vi.fn()
25+
const handler = createWorkspaceSSE({
26+
label: 'test',
27+
subscriptions: [{ subscribe: () => unsubscribe }],
28+
})
29+
const request = new NextRequest(new URL('https://sim.test/api/test/events?workspaceId=ws-1'), {
30+
signal,
31+
})
32+
const response = await handler(request)
33+
34+
return { body: response.body as ReadableStream<Uint8Array>, unsubscribe }
35+
}
36+
37+
/** Resolves once the stream closes. */
38+
async function drain(body: ReadableStream<Uint8Array>): Promise<void> {
39+
const reader = body.getReader()
40+
while (true) {
41+
const { done } = await reader.read()
42+
if (done) return
43+
}
44+
}
45+
46+
describe('createWorkspaceSSE', () => {
47+
beforeEach(() => {
48+
vi.clearAllMocks()
49+
vi.useFakeTimers()
50+
authMockFns.mockGetSession.mockResolvedValue({ user: { id: 'user-1' } })
51+
permissionsMockFns.mockGetUserEntityPermissions.mockResolvedValue('admin')
52+
})
53+
54+
afterEach(() => {
55+
vi.useRealTimers()
56+
})
57+
58+
it('releases subscriptions and closes the stream when the connection reaches its ceiling', async () => {
59+
const { body, unsubscribe } = await openConnection()
60+
const drained = drain(body)
61+
62+
await vi.advanceTimersByTimeAsync(PAST_CEILING_MS)
63+
64+
await drained
65+
expect(unsubscribe).toHaveBeenCalledTimes(1)
66+
})
67+
68+
it('releases subscriptions when the consumer stops draining the stream', async () => {
69+
const { unsubscribe } = await openConnection()
70+
71+
await vi.advanceTimersByTimeAsync(PAST_UNREAD_MS)
72+
73+
expect(unsubscribe).toHaveBeenCalledTimes(1)
74+
})
75+
76+
it('keeps a drained connection alive past the unread threshold', async () => {
77+
const { body, unsubscribe } = await openConnection()
78+
void drain(body)
79+
80+
await vi.advanceTimersByTimeAsync(PAST_UNREAD_MS)
81+
82+
expect(unsubscribe).not.toHaveBeenCalled()
83+
})
84+
85+
it('releases subscriptions when the request aborts', async () => {
86+
const controller = new AbortController()
87+
const { unsubscribe } = await openConnection(controller.signal)
88+
89+
controller.abort()
90+
91+
expect(unsubscribe).toHaveBeenCalledTimes(1)
92+
})
93+
94+
it('releases subscriptions when the consumer cancels the stream', async () => {
95+
const { body, unsubscribe } = await openConnection()
96+
97+
await body.cancel()
98+
99+
expect(unsubscribe).toHaveBeenCalledTimes(1)
100+
})
101+
102+
it('releases subscriptions once when abort and the ceiling both elapse', async () => {
103+
const controller = new AbortController()
104+
const { unsubscribe } = await openConnection(controller.signal)
105+
106+
controller.abort()
107+
await vi.advanceTimersByTimeAsync(PAST_CEILING_MS)
108+
109+
expect(unsubscribe).toHaveBeenCalledTimes(1)
110+
})
111+
})

apps/sim/lib/events/sse-endpoint.ts

Lines changed: 75 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
*/
77

88
import { createLogger } from '@sim/logger'
9+
import { randomFloat } from '@sim/utils/random'
910
import type { NextRequest } from 'next/server'
1011
import { getSession } from '@/lib/auth'
1112
import { SSE_HEADERS } from '@/lib/core/utils/sse'
@@ -23,7 +24,40 @@ interface WorkspaceSSEConfig {
2324
subscriptions: SSESubscription[]
2425
}
2526

26-
const HEARTBEAT_INTERVAL_MS = 30_000
27+
const encoder = new TextEncoder()
28+
29+
export const HEARTBEAT_INTERVAL_MS = 30_000
30+
31+
/**
32+
* Defensive ceiling on one connection's lifetime; `EventSource` reconnects past
33+
* this, so delivery continues across the boundary.
34+
*
35+
* `request.signal` abort and stream `cancel()` are the primary teardown paths,
36+
* but both fire only when the runtime reports the client disconnect, and the
37+
* unread check below only catches a consumer that has stopped draining. This
38+
* releases whatever both miss, so retention is bounded by the ceiling instead
39+
* of by process uptime.
40+
*
41+
* Matches the ceiling `lib/realtime/event-stream-route.ts` already uses for the
42+
* same purpose. It is deliberately far longer than the unread window: a healthy
43+
* client is drained and therefore never unread, so a short ceiling would only
44+
* force reconnects on the connections that are working, and every reconnect is
45+
* a window in which a transient event can be missed.
46+
*/
47+
export const MAX_CONNECTION_MS = 4 * 60 * 60 * 1000
48+
49+
/** Spreads reconnects so connections opened together do not expire together. */
50+
export const MAX_CONNECTION_JITTER_MS = 60_000
51+
52+
/**
53+
* Undrained chunk count that marks a connection unread, which shortens the
54+
* reclaim window for a vanished consumer that the ceiling above would
55+
* otherwise hold for its full duration. The default queuing strategy reports
56+
* `desiredSize` as `1 - queued`, so this trips only once the consumer has
57+
* pulled nothing for several minutes — well beyond the transient backpressure
58+
* of a slow but live client.
59+
*/
60+
export const MAX_UNDRAINED_CHUNKS = 16
2761

2862
export function createWorkspaceSSE(config: WorkspaceSSEConfig) {
2963
const logger = createLogger(`${config.label}-SSE`)
@@ -45,21 +79,30 @@ export function createWorkspaceSSE(config: WorkspaceSSEConfig) {
4579
return new Response('Access denied to workspace', { status: 403 })
4680
}
4781

48-
const encoder = new TextEncoder()
49-
const unsubscribers: Array<() => void> = []
82+
const teardowns: Array<() => void> = []
5083
let cleaned = false
5184

52-
const cleanup = () => {
85+
const cleanup = (reason: string) => {
5386
if (cleaned) return
5487
cleaned = true
55-
for (const unsub of unsubscribers) {
56-
unsub()
88+
for (const teardown of teardowns) {
89+
teardown()
5790
}
58-
logger.info(`SSE connection closed for workspace ${workspaceId}`)
91+
teardowns.length = 0
92+
logger.info(`SSE connection closed for workspace ${workspaceId}`, { reason })
5993
}
6094

6195
const stream = new ReadableStream({
6296
start(controller) {
97+
const close = (reason: string) => {
98+
cleanup(reason)
99+
try {
100+
controller.close()
101+
} catch {
102+
// Already closed
103+
}
104+
}
105+
63106
const send = (eventName: string, data: Record<string, unknown>) => {
64107
if (cleaned) return
65108
try {
@@ -72,40 +115,47 @@ export function createWorkspaceSSE(config: WorkspaceSSEConfig) {
72115
}
73116

74117
for (const subscription of config.subscriptions) {
75-
const unsub = subscription.subscribe(workspaceId, send)
76-
unsubscribers.push(unsub)
118+
teardowns.push(subscription.subscribe(workspaceId, send))
77119
}
78120

121+
const deadline = Date.now() + MAX_CONNECTION_MS + randomFloat() * MAX_CONNECTION_JITTER_MS
122+
79123
const heartbeat = setInterval(() => {
80124
if (cleaned) {
81125
clearInterval(heartbeat)
82126
return
83127
}
128+
if (Date.now() >= deadline) {
129+
close('expired')
130+
return
131+
}
132+
const desiredSize = controller.desiredSize
133+
if (desiredSize !== null && desiredSize <= -MAX_UNDRAINED_CHUNKS) {
134+
close('unread')
135+
return
136+
}
84137
try {
85138
controller.enqueue(encoder.encode(': heartbeat\n\n'))
86139
} catch {
87-
clearInterval(heartbeat)
140+
close('errored')
88141
}
89142
}, HEARTBEAT_INTERVAL_MS)
90-
unsubscribers.push(() => clearInterval(heartbeat))
91-
92-
request.signal.addEventListener(
93-
'abort',
94-
() => {
95-
cleanup()
96-
try {
97-
controller.close()
98-
} catch {
99-
// Already closed
100-
}
101-
},
102-
{ once: true }
103-
)
143+
teardowns.push(() => clearInterval(heartbeat))
144+
145+
// `once` only self-removes if abort fires; the expiry and unread paths
146+
// close the connection while the signal is still live, so the listener
147+
// needs its own removal or it retains this whole scope.
148+
const listenerScope = new AbortController()
149+
request.signal.addEventListener('abort', () => close('aborted'), {
150+
once: true,
151+
signal: listenerScope.signal,
152+
})
153+
teardowns.push(() => listenerScope.abort())
104154

105155
logger.info(`SSE connection opened for workspace ${workspaceId}`)
106156
},
107157
cancel() {
108-
cleanup()
158+
cleanup('cancelled')
109159
},
110160
})
111161

0 commit comments

Comments
 (0)