From b274b44442fbff173353d30c9717a7c6b4eef6bf Mon Sep 17 00:00:00 2001 From: Sidharth Gehlot Date: Sun, 12 Jul 2026 13:10:38 +0530 Subject: [PATCH] fix(desktop): submit first message after session creation Allow the intentional new-chat route promotion without weakening session-switch isolation. Validate the created runtime and canonical stored-session route before continuing, and cover both successful promotion and a late user switch with regression tests. --- .../hooks/use-prompt-actions/index.test.tsx | 103 +++++++++++++++++- .../hooks/use-prompt-actions/submit.ts | 31 ++++-- 2 files changed, 126 insertions(+), 8 deletions(-) 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 5f028876a8e37..c229a222103b7 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 @@ -1294,7 +1294,14 @@ describe('usePromptActions sleep/wake session recovery', () => { }) it('still creates a new session for a genuine new-chat draft (no stored session selected)', async () => { - const createBackendSessionForSend = vi.fn(async () => RUNTIME_SESSION_ID) + const activeSessionIdRef: MutableRefObject = { current: null } + + const createBackendSessionForSend = vi.fn(async () => { + activeSessionIdRef.current = RUNTIME_SESSION_ID + + return RUNTIME_SESSION_ID + }) + const calls: string[] = [] const requestGateway = vi.fn(async (method: string) => { @@ -1307,6 +1314,7 @@ describe('usePromptActions sleep/wake session recovery', () => { render( (handle = h)} refreshSessions={async () => undefined} @@ -1321,6 +1329,99 @@ describe('usePromptActions sleep/wake session recovery', () => { expect(createBackendSessionForSend).toHaveBeenCalledTimes(1) expect(calls).not.toContain('session.resume') }) + + it('submits the first message after session creation routes to the newly titled chat', async () => { + const activeSessionIdRef: MutableRefObject = { current: null } + const selectedStoredSessionIdRef: MutableRefObject = { current: null } + let routeToken = '/' + const calls: { method: string; params?: Record }[] = [] + + const createBackendSessionForSend = vi.fn(async () => { + activeSessionIdRef.current = RUNTIME_SESSION_ID + selectedStoredSessionIdRef.current = STORED_SESSION_ID + routeToken = `/${STORED_SESSION_ID}` + + return RUNTIME_SESSION_ID + }) + + const requestGateway = vi.fn(async (method: string, params?: Record) => { + calls.push({ method, params }) + + return {} as never + }) + + let handle: HarnessHandle | null = null + render( + routeToken} + onReady={h => (handle = h)} + refreshSessions={async () => undefined} + requestGateway={requestGateway} + selectedStoredSessionIdRef={selectedStoredSessionIdRef} + storedSessionId={null} + /> + ) + + const ok = await handle!.submitText('first message becomes the title and must also send') + + expect(ok).toBe(true) + expect(calls).toContainEqual({ + method: 'prompt.submit', + params: { + session_id: RUNTIME_SESSION_ID, + text: 'first message becomes the title and must also send' + } + }) + }) + + it('aborts when the user switches chats during post-create async work', async () => { + const activeSessionIdRef: MutableRefObject = { current: null } + const selectedStoredSessionIdRef: MutableRefObject = { current: null } + let routeToken = '/' + const calls: string[] = [] + + const createBackendSessionForSend = vi.fn(async () => { + // session.create performs its intentional promotion first. + activeSessionIdRef.current = RUNTIME_SESSION_ID + selectedStoredSessionIdRef.current = STORED_SESSION_ID + routeToken = `/${STORED_SESSION_ID}` + await Promise.resolve() + + // Then the user selects another chat while post-create work is pending. + activeSessionIdRef.current = 'runtime-session-b' + selectedStoredSessionIdRef.current = 'stored-session-b' + routeToken = '/stored-session-b' + + return RUNTIME_SESSION_ID + }) + + const requestGateway = vi.fn(async (method: string) => { + calls.push(method) + + return {} as never + }) + + let handle: HarnessHandle | null = null + render( + routeToken} + onReady={h => (handle = h)} + refreshSessions={async () => undefined} + requestGateway={requestGateway} + selectedStoredSessionIdRef={selectedStoredSessionIdRef} + storedSessionId={null} + /> + ) + + expect(await handle!.submitText('must not follow the user into chat B')).toBe(false) + expect(calls).not.toContain('prompt.submit') + }) }) describe('usePromptActions submit session-context isolation (#54527)', () => { 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 5dd01fc6d6d87..f0530c06a7f55 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 @@ -15,6 +15,7 @@ import { clearNotifications, notify, notifyError } from '@/store/notifications' import { requestDesktopOnboarding } from '@/store/onboarding' import { setAwaitingResponse, setBusy, setMessages } from '@/store/session' +import { sessionRoute } from '../../../routes' import type { ClientSessionState } from '../../../types' import { @@ -120,11 +121,11 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { // redirect the user's text into a different chat (#54527). const startingActiveSessionId = activeSessionIdRef.current const startingStoredSessionId = selectedStoredSessionIdRef.current - const startingRouteToken = getRouteToken() + let expectedStoredSessionId = startingStoredSessionId + let expectedRouteToken = getRouteToken() const sessionContextDrifted = (): boolean => - selectedStoredSessionIdRef.current !== startingStoredSessionId || - getRouteToken() !== startingRouteToken + selectedStoredSessionIdRef.current !== expectedStoredSessionId || getRouteToken() !== expectedRouteToken // 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. @@ -281,10 +282,6 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { return false } - if (sessionContextDrifted()) { - return abortForSessionSwitch(sessionId) - } - if (!sessionId) { dropOptimistic(null) releaseBusy() @@ -293,6 +290,26 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { return false } + // session.create deliberately promotes the fresh draft onto its new + // stored-session route. Accept only that exact promotion: post-create + // work (for example applying an armed YOLO mode) can await before this + // function regains control, so the user may genuinely switch chats in + // that window. Never rebase the guard onto an unrelated destination. + const promotedStoredSessionId = selectedStoredSessionIdRef.current + const promotedRouteToken = getRouteToken() + + const promotionIsCurrent = + activeSessionIdRef.current === sessionId && + (promotedStoredSessionId + ? promotedRouteToken === sessionRoute(promotedStoredSessionId) + : promotedRouteToken === expectedRouteToken) + + if (!promotionIsCurrent) { + return abortForSessionSwitch(sessionId) + } + + expectedStoredSessionId = promotedStoredSessionId + expectedRouteToken = promotedRouteToken seedOptimistic(sessionId) }