diff --git a/apps/desktop/src/app/desktop-controller.tsx b/apps/desktop/src/app/desktop-controller.tsx index cb7b10f655b..2ff2ff176e7 100644 --- a/apps/desktop/src/app/desktop-controller.tsx +++ b/apps/desktop/src/app/desktop-controller.tsx @@ -812,6 +812,7 @@ export function DesktopController() { branchCurrentSession: branchInNewChat, busyRef, createBackendSessionForSend, + getRouteToken, handleSkinCommand, openMemoryGraph: openStarmap, refreshSessions, diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx index 045a4e7975d..5f028876a8e 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.test.tsx @@ -51,7 +51,9 @@ interface HarnessHandle { } function Harness({ + activeSessionIdRef: activeSessionIdRefProp, busyRef, + getRouteToken, onReady, onSeedState, openMemoryGraph, @@ -59,11 +61,14 @@ function Harness({ requestGateway, resumeStoredSession, seedMessages, + selectedStoredSessionIdRef: selectedStoredSessionIdRefProp, storedSessionId, activeSessionId, createBackendSessionForSend }: { + activeSessionIdRef?: MutableRefObject busyRef?: MutableRefObject + getRouteToken?: () => string onReady: (handle: HarnessHandle) => void onSeedState?: (state: Record) => void openMemoryGraph?: () => void @@ -71,15 +76,16 @@ function Harness({ requestGateway: (method: string, params?: Record) => Promise resumeStoredSession?: (storedSessionId: string) => Promise | void seedMessages?: unknown[] + selectedStoredSessionIdRef?: MutableRefObject storedSessionId?: null | string activeSessionId?: null | string createBackendSessionForSend?: () => Promise }) { - const activeSessionIdRef: MutableRefObject = { + const activeSessionIdRef: MutableRefObject = activeSessionIdRefProp ?? { current: activeSessionId === undefined ? RUNTIME_SESSION_ID : activeSessionId } - const selectedStoredSessionIdRef: MutableRefObject = { + const selectedStoredSessionIdRef: MutableRefObject = selectedStoredSessionIdRefProp ?? { current: storedSessionId === undefined ? RUNTIME_SESSION_ID : storedSessionId } @@ -98,6 +104,7 @@ function Harness({ branchCurrentSession: async () => true, busyRef: localBusyRef, createBackendSessionForSend: createBackendSessionForSend ?? (async () => RUNTIME_SESSION_ID), + getRouteToken: getRouteToken ?? (() => 'token'), handleSkinCommand: () => '', openMemoryGraph: openMemoryGraph ?? (() => undefined), refreshSessions, @@ -1239,7 +1246,7 @@ describe('usePromptActions sleep/wake session recovery', () => { expect(ok).toBe(true) expect(calls.map(c => c.method)).toEqual(['prompt.submit', 'session.resume', 'prompt.submit']) - expect(calls[1]?.params).toEqual({ session_id: STORED_SESSION_ID }) + expect(calls[1]?.params).toEqual({ session_id: STORED_SESSION_ID, source: 'desktop' }) expect(calls[2]?.params).toEqual({ session_id: RECOVERED_SESSION_ID, text: 'message during starved loop' @@ -1316,6 +1323,125 @@ describe('usePromptActions sleep/wake session recovery', () => { }) }) +describe('usePromptActions submit session-context isolation (#54527)', () => { + const STORED_SESSION_A = 'stored-project-a' + const STORED_SESSION_B = 'stored-project-b' + const RUNTIME_SESSION_B = 'rt-session-b-wrong' + + afterEach(() => { + cleanup() + vi.restoreAllMocks() + }) + + it('aborts submit when the user switches sessions during session.resume (no misroute)', async () => { + // Exact #54527 failure: user submits in Session A while its runtime binding + // is gone; before resume returns they switch to Session B. Without a pinned + // context the resumed runtime id belongs to B and A's text lands in the + // wrong chat — permanently lost from A. + let releaseResume: () => void = () => {} + const calls: { method: string; params?: Record }[] = [] + + const selectedStoredSessionIdRef: MutableRefObject = { current: STORED_SESSION_A } + const activeSessionIdRef: MutableRefObject = { current: null } + + const requestGateway = vi.fn(async (method: string, params?: Record) => { + calls.push({ method, params }) + + if (method === 'session.resume') { + await new Promise(resolve => { + releaseResume = resolve + }) + + // Simulate the user switching to Session B while resume is in flight. + selectedStoredSessionIdRef.current = STORED_SESSION_B + activeSessionIdRef.current = RUNTIME_SESSION_B + + return { session_id: RUNTIME_SESSION_B } as never + } + + return {} as never + }) + + let handle: HarnessHandle | null = null + render( + (handle = h)} + refreshSessions={async () => undefined} + requestGateway={requestGateway} + selectedStoredSessionIdRef={selectedStoredSessionIdRef} + storedSessionId={STORED_SESSION_A} + /> + ) + await waitFor(() => expect(handle).not.toBeNull()) + + const submitting = handle!.submitText('carefully composed prompt for project A') + await waitFor(() => expect(calls.some(c => c.method === 'session.resume')).toBe(true)) + releaseResume() + + expect(await submitting).toBe(false) + expect(calls.some(c => c.method === 'prompt.submit')).toBe(false) + expect(calls.find(c => c.method === 'session.resume')?.params).toEqual({ + session_id: STORED_SESSION_A + }) + }) + + it('aborts recovery submit when the user switches sessions during timeout resume', async () => { + const calls: { method: string; params?: Record }[] = [] + let submitAttempts = 0 + let releaseResume: () => void = () => {} + + const selectedStoredSessionIdRef: MutableRefObject = { current: STORED_SESSION_A } + + const requestGateway = vi.fn(async (method: string, params?: Record) => { + calls.push({ method, params }) + + if (method === 'prompt.submit') { + submitAttempts += 1 + + if (submitAttempts === 1) { + throw new Error('request timed out: prompt.submit') + } + } + + if (method === 'session.resume') { + await new Promise(resolve => { + releaseResume = resolve + }) + selectedStoredSessionIdRef.current = STORED_SESSION_B + + return { session_id: RUNTIME_SESSION_B } as never + } + + return {} as never + }) + + let handle: HarnessHandle | null = null + render( + (handle = h)} + refreshSessions={async () => undefined} + requestGateway={requestGateway} + selectedStoredSessionIdRef={selectedStoredSessionIdRef} + storedSessionId={STORED_SESSION_A} + /> + ) + await waitFor(() => expect(handle).not.toBeNull()) + + const submitting = handle!.submitText('message that must not land in session B') + await waitFor(() => expect(calls.some(c => c.method === 'session.resume')).toBe(true)) + releaseResume() + + expect(await submitting).toBe(false) + expect(submitAttempts).toBe(1) + expect(calls.filter(c => c.method === 'prompt.submit')).toHaveLength(1) + expect(calls.find(c => c.method === 'session.resume')?.params).toMatchObject({ + session_id: STORED_SESSION_A + }) + }) +}) + describe('usePromptActions eager attachment upload (drop-time)', () => { afterEach(() => { cleanup() diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts index 8ebe8889d16..51a3689ae4e 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/index.ts @@ -158,6 +158,7 @@ interface PromptActionsOptions { busyRef: MutableRefObject branchCurrentSession: () => Promise createBackendSessionForSend: (preview?: string | null) => Promise + getRouteToken: () => string handleSkinCommand: (arg: string) => string openMemoryGraph: () => void refreshSessions: () => Promise @@ -186,6 +187,7 @@ export function usePromptActions({ busyRef, branchCurrentSession, createBackendSessionForSend, + getRouteToken, handleSkinCommand, openMemoryGraph, refreshSessions, @@ -354,6 +356,7 @@ export function usePromptActions({ busyRef, copy, createBackendSessionForSend, + getRouteToken, requestGateway, selectedStoredSessionIdRef, syncAttachmentsForSubmit, diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/submit.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/submit.ts index 6b2049f892b..5dd01fc6d6d 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/submit.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/submit.ts @@ -35,6 +35,7 @@ interface SubmitPromptDeps { busyRef: MutableRefObject copy: Translations['desktop'] createBackendSessionForSend: (preview?: string | null) => Promise + getRouteToken: () => string requestGateway: GatewayRequest selectedStoredSessionIdRef: MutableRefObject syncAttachmentsForSubmit: ( @@ -57,6 +58,7 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { busyRef, copy, createBackendSessionForSend, + getRouteToken, requestGateway, selectedStoredSessionIdRef, syncAttachmentsForSubmit, @@ -113,9 +115,20 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { return false } + // Pin the session context for the whole async submit pipeline. Without + // this, a fast session switch during session.resume / file.attach can + // redirect the user's text into a different chat (#54527). + const startingActiveSessionId = activeSessionIdRef.current + const startingStoredSessionId = selectedStoredSessionIdRef.current + const startingRouteToken = getRouteToken() + + const sessionContextDrifted = (): boolean => + selectedStoredSessionIdRef.current !== startingStoredSessionId || + getRouteToken() !== startingRouteToken + // One submit in flight per session — drop any concurrent re-fire so a // stalled turn can't stack the same prompt into multiple real turns. - const submitLockKey = selectedStoredSessionIdRef.current || activeSessionId || '__pending_new__' + const submitLockKey = startingStoredSessionId || startingActiveSessionId || '__pending_new__' if (_submitInFlight.has(submitLockKey)) { return false @@ -166,7 +179,7 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { // (what made drained-after-interrupt sends go silent). interrupted: false }), - selectedStoredSessionIdRef.current + startingStoredSessionId ) // After sync rewrites refs, refresh the optimistic message in place so the @@ -178,7 +191,7 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { ...state, messages: state.messages.map(message => (message.id === optimisticId ? buildUserMessage() : message)) }), - selectedStoredSessionIdRef.current + startingStoredSessionId ) const dropOptimistic = (sid: null | string) => { @@ -197,10 +210,17 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { awaitingResponse: false, pendingBranchGroup: null }), - selectedStoredSessionIdRef.current + startingStoredSessionId ) } + const abortForSessionSwitch = (optimisticSessionId: null | string): false => { + dropOptimistic(optimisticSessionId) + releaseBusy() + + return false + } + setMutableRef(busyRef, true) setBusy(true) setAwaitingResponse(true) @@ -214,7 +234,7 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { setMessages(current => [...current, buildUserMessage()]) } - if (!sessionId && selectedStoredSessionIdRef.current) { + if (!sessionId && startingStoredSessionId) { // A stored session is SELECTED but its runtime binding is gone (the // live session was orphan-reaped, or a timeout/reconnect cleared // activeSessionId). Continuing the selected conversation must mean @@ -224,9 +244,13 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { // new-chat draft). try { const resumed = await requestGateway<{ session_id: string }>('session.resume', { - session_id: selectedStoredSessionIdRef.current + session_id: startingStoredSessionId }) + if (sessionContextDrifted()) { + return abortForSessionSwitch(sessionId) + } + if (resumed?.session_id) { sessionId = resumed.session_id activeSessionIdRef.current = sessionId @@ -237,6 +261,10 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { // the user's message. } + if (sessionContextDrifted()) { + return abortForSessionSwitch(sessionId) + } + if (sessionId) { seedOptimistic(sessionId) } @@ -253,6 +281,10 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { return false } + if (sessionContextDrifted()) { + return abortForSessionSwitch(sessionId) + } + if (!sessionId) { dropOptimistic(null) releaseBusy() @@ -269,6 +301,10 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { updateComposerAttachments: usingComposerAttachments }) + if (sessionContextDrifted()) { + return abortForSessionSwitch(sessionId) + } + // Rewrite the optimistic message + prompt text with the synced refs so // the gateway receives @file: paths that resolve in its workspace. // (Images keep their inline base64 preview — see optimisticAttachmentRef.) @@ -288,7 +324,7 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { } catch (firstErr) { if ( (isSessionNotFoundError(firstErr) || isGatewayTimeoutError(firstErr)) && - selectedStoredSessionIdRef.current + startingStoredSessionId ) { // Re-register the session in the gateway and get a fresh live ID. // Timeouts recover the same way as "session not found": a starved @@ -296,10 +332,14 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { // the stored session is fine — resume + retry instead of erroring // out and losing the session binding. const resumed = await requestGateway<{ session_id: string }>('session.resume', { - session_id: selectedStoredSessionIdRef.current, + session_id: startingStoredSessionId, source: 'desktop' }) + if (sessionContextDrifted()) { + return abortForSessionSwitch(sessionId) + } + const recoveredId = resumed?.session_id if (recoveredId) { @@ -375,6 +415,7 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { busyRef, copy, createBackendSessionForSend, + getRouteToken, requestGateway, selectedStoredSessionIdRef, syncAttachmentsForSubmit,