diff --git a/apps/desktop/src/app/session/hooks/use-message-stream.ts b/apps/desktop/src/app/session/hooks/use-message-stream.ts index 3ee52ec8eb7d..99330d997d13 100644 --- a/apps/desktop/src/app/session/hooks/use-message-stream.ts +++ b/apps/desktop/src/app/session/hooks/use-message-stream.ts @@ -17,7 +17,7 @@ import { } from '@/lib/chat-messages' import { coerceGatewayText, coerceThinkingText, normalizePersonalityValue } from '@/lib/chat-runtime' import { playCompletionSound } from '@/lib/completion-sound' -import { gatewayEventRequiresSessionId } from '@/lib/gateway-events' +import { resolveGatewayEventSessionId } from '@/lib/gateway-events' import { dedupeGeneratedImageEchoesInParts, generatedImageEchoSources, @@ -262,6 +262,8 @@ export function useMessageStream({ sessionStateByRuntimeIdRef, updateSessionState }: MessageStreamOptions) { + const unscopedStreamSessionIdRef = useRef(null) + const sessionInterrupted = useCallback( (sessionId: string) => sessionStateByRuntimeIdRef.current.get(sessionId)?.interrupted ?? false, [sessionStateByRuntimeIdRef] @@ -715,11 +717,20 @@ export function useMessageStream({ const payload = event.payload as GatewayEventPayload | undefined const explicitSid = event.session_id || '' - if (!explicitSid && gatewayEventRequiresSessionId(event.type)) { + const route = resolveGatewayEventSessionId({ + activeSessionId: activeSessionIdRef.current, + eventType: event.type, + explicitSessionId: explicitSid, + unscopedStreamSessionId: unscopedStreamSessionIdRef.current + }) + + unscopedStreamSessionIdRef.current = route.nextUnscopedStreamSessionId + + if (route.drop) { return } - const sessionId = explicitSid || activeSessionIdRef.current + const sessionId = route.sessionId const isActiveEvent = !!sessionId && sessionId === activeSessionIdRef.current if (event.type === 'gateway.ready') { diff --git a/apps/desktop/src/lib/gateway-events.test.ts b/apps/desktop/src/lib/gateway-events.test.ts index d51a943611f0..7435d22d6ee4 100644 --- a/apps/desktop/src/lib/gateway-events.test.ts +++ b/apps/desktop/src/lib/gateway-events.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from 'vitest' -import { gatewayEventRequiresSessionId } from './gateway-events' +import { gatewayEventRequiresSessionId, resolveGatewayEventSessionId } from './gateway-events' describe('gateway event routing', () => { it('drops only unscoped subagent events (genuinely background work)', () => { @@ -24,4 +24,75 @@ describe('gateway event routing', () => { expect(gatewayEventRequiresSessionId('session.info')).toBe(false) expect(gatewayEventRequiresSessionId(undefined)).toBe(false) }) + + it('keeps unscoped stream events pinned to the session that started them', () => { + const started = resolveGatewayEventSessionId({ + activeSessionId: 'session-a', + eventType: 'message.start', + explicitSessionId: '', + unscopedStreamSessionId: null + }) + + expect(started).toEqual({ + drop: false, + nextUnscopedStreamSessionId: 'session-a', + sessionId: 'session-a' + }) + + const delta = resolveGatewayEventSessionId({ + activeSessionId: 'session-b', + eventType: 'message.delta', + explicitSessionId: '', + unscopedStreamSessionId: started.nextUnscopedStreamSessionId + }) + + expect(delta).toEqual({ + drop: false, + nextUnscopedStreamSessionId: 'session-a', + sessionId: 'session-a' + }) + + const completed = resolveGatewayEventSessionId({ + activeSessionId: 'session-b', + eventType: 'message.complete', + explicitSessionId: '', + unscopedStreamSessionId: delta.nextUnscopedStreamSessionId + }) + + expect(completed).toEqual({ + drop: false, + nextUnscopedStreamSessionId: null, + sessionId: 'session-a' + }) + }) + + it('routes a new unscoped stream start to the currently active session', () => { + const routed = resolveGatewayEventSessionId({ + activeSessionId: 'session-b', + eventType: 'message.start', + explicitSessionId: '', + unscopedStreamSessionId: 'session-a' + }) + + expect(routed).toEqual({ + drop: false, + nextUnscopedStreamSessionId: 'session-b', + sessionId: 'session-b' + }) + }) + + it('keeps explicit events scoped and clears a matching pinned stream on completion', () => { + const routed = resolveGatewayEventSessionId({ + activeSessionId: 'session-b', + eventType: 'message.complete', + explicitSessionId: 'session-a', + unscopedStreamSessionId: 'session-a' + }) + + expect(routed).toEqual({ + drop: false, + nextUnscopedStreamSessionId: null, + sessionId: 'session-a' + }) + }) }) diff --git a/apps/desktop/src/lib/gateway-events.ts b/apps/desktop/src/lib/gateway-events.ts index 673d1df8c6d6..d3f0e30999ee 100644 --- a/apps/desktop/src/lib/gateway-events.ts +++ b/apps/desktop/src/lib/gateway-events.ts @@ -11,6 +11,28 @@ function asRecord(payload: unknown): Record { return payload && typeof payload === 'object' ? (payload as Record) : {} } +const UNSCOPED_STREAM_EVENT_TYPES = new Set([ + 'approval.request', + 'browser.progress', + 'clarify.request', + 'error', + 'message.complete', + 'message.delta', + 'message.start', + 'reasoning.available', + 'reasoning.delta', + 'secret.request', + 'status.update', + 'sudo.request', + 'thinking.delta', + 'tool.complete', + 'tool.generating', + 'tool.progress', + 'tool.start' +]) + +const UNSCOPED_STREAM_END_EVENT_TYPES = new Set(['error', 'message.complete']) + /** * Whether an unscoped event (no `session_id`) must be dropped rather than * attributed to the focused chat. @@ -27,6 +49,63 @@ export function gatewayEventRequiresSessionId(eventType: string | undefined): bo return eventType?.startsWith('subagent.') ?? false } +export interface GatewayEventSessionRouteInput { + activeSessionId: null | string + eventType: string | undefined + explicitSessionId: string + unscopedStreamSessionId: null | string +} + +export interface GatewayEventSessionRoute { + drop: boolean + nextUnscopedStreamSessionId: null | string + sessionId: null | string +} + +export function resolveGatewayEventSessionId({ + activeSessionId, + eventType, + explicitSessionId, + unscopedStreamSessionId +}: GatewayEventSessionRouteInput): GatewayEventSessionRoute { + if (explicitSessionId) { + const nextUnscopedStreamSessionId = + eventType && UNSCOPED_STREAM_END_EVENT_TYPES.has(eventType) && explicitSessionId === unscopedStreamSessionId + ? null + : unscopedStreamSessionId + + return { + drop: false, + nextUnscopedStreamSessionId, + sessionId: explicitSessionId + } + } + + if (gatewayEventRequiresSessionId(eventType)) { + return { + drop: true, + nextUnscopedStreamSessionId: unscopedStreamSessionId, + sessionId: null + } + } + + const streamEvent = eventType ? UNSCOPED_STREAM_EVENT_TYPES.has(eventType) : false + const sessionId = eventType === 'message.start' ? activeSessionId : streamEvent ? unscopedStreamSessionId || activeSessionId : activeSessionId + let nextUnscopedStreamSessionId = unscopedStreamSessionId + + if (eventType === 'message.start' && activeSessionId) { + nextUnscopedStreamSessionId = activeSessionId + } else if (eventType && UNSCOPED_STREAM_END_EVENT_TYPES.has(eventType)) { + nextUnscopedStreamSessionId = null + } + + return { + drop: false, + nextUnscopedStreamSessionId, + sessionId + } +} + export function gatewayEventCompletedFileDiff(event: RpcEventLike): boolean { if (event.type !== 'tool.complete') { return false