From fab99e828ac42cdfa64d4ee207f2647ec9e24eaf Mon Sep 17 00:00:00 2001 From: bb Date: Fri, 7 Aug 2026 12:49:31 -0500 Subject: [PATCH] refactor(desktop): one resolver for stale runtime-session recovery MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit main carried three hand-rolled copies of the same recovery policy — prompt.submit (submit.ts), session.interrupt (cancelRun) and session.redirect (steering) each open-coded isSessionNotFoundError -> resolveSessionProfile -> session.resume -> retry. Three copies of one policy is how call sites drift apart, and it is why the RPCs that were never given a copy (attach, /compress, checkpoint restore) still surface a raw "session not found" after sleep/wake while plain text silently recovers. Introduce withSessionNotFoundResume() as the single resolver and move all three existing copies onto it. Profile resolution is injected rather than imported so the helper is unit-testable without reaching through resolveStoredSession -> getSession(); a drift callback lets each caller keep its own abort semantics via SessionRecoveryAborted. Co-authored-by: xxxigm --- .../src/app/chat/session-tile-actions.ts | 18 +- .../session/hooks/use-prompt-actions/index.ts | 250 +++++++++--------- .../hooks/use-prompt-actions/submit.ts | 99 +++---- .../session/hooks/use-prompt-actions/utils.ts | 135 ++++++++++ 4 files changed, 324 insertions(+), 178 deletions(-) diff --git a/apps/desktop/src/app/chat/session-tile-actions.ts b/apps/desktop/src/app/chat/session-tile-actions.ts index 16b47b995207d..fdaf7126d5cd3 100644 --- a/apps/desktop/src/app/chat/session-tile-actions.ts +++ b/apps/desktop/src/app/chat/session-tile-actions.ts @@ -157,12 +157,20 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses sessionId: string, attachments: ComposerAttachment[], options: { updateComposerAttachments?: boolean } = {} - ): Promise => { + ): Promise<{ attachments: ComposerAttachment[]; sessionId: string }> => { const remote = $connection.get()?.mode === 'remote' + let liveSessionId = sessionId const synced: ComposerAttachment[] = [] + // A tile owns its own runtime binding, so a recovery here rebinds the + // tile's ref rather than the foreground session's. + const onSessionRecovered = (recoveredId: string) => { + liveSessionId = recoveredId + runtimeIdRef.current = recoveredId + } + for (const attachment of attachments) { - if (!attachment.path || attachment.attachedSessionId === sessionId) { + if (!attachment.path || attachment.attachedSessionId === liveSessionId) { synced.push(attachment) continue @@ -173,7 +181,9 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses backendCwd: readState()?.cwd, remote, requestGateway, - sessionId + sessionId: liveSessionId, + storedSessionId: storedIdRef.current, + onSessionRecovered }) if (options.updateComposerAttachments ?? true) { @@ -188,7 +198,7 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses synced.push(attachment) } - return synced + return { attachments: synced, sessionId: liveSessionId } }, [requestGateway, scope.attachments] ) 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 186773a874ffb..fcfe1d78dfb44 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 @@ -27,6 +27,7 @@ import { $connection, $currentCwd, $messages, + setActiveSessionId, setAwaitingResponse, setBusy, setMessages, @@ -45,7 +46,6 @@ import type { ImageAttachResponse, SessionRedirectResponse } from '../../../types' -import { resolveSessionProfile } from '../use-session-actions/utils' import { applyBranchVisibility, @@ -66,10 +66,10 @@ import { friendlyRemoteAttachError, type GatewayRequest, inlineErrorMessage, - isSessionNotFoundError, readFileDataUrlForAttach, readImageForRemoteAttach, - type SubmitTextOptions + type SubmitTextOptions, + withSessionNotFoundResume } from './utils' interface HandoffResult { @@ -98,88 +98,106 @@ function attachmentPathNeedsUpload(path: string, backendCwd?: null | string): bo */ export async function uploadComposerAttachment( attachment: ComposerAttachment, - opts: { backendCwd?: null | string; remote: boolean; requestGateway: GatewayRequest; sessionId: string } + opts: { + backendCwd?: null | string + remote: boolean + requestGateway: GatewayRequest + sessionId: string + /** Durable id used to re-register after sleep/wake or a backend restart. */ + storedSessionId?: null | string + /** Called when the attach recovered onto a fresh live id. */ + onSessionRecovered?: (sessionId: string) => void + } ): Promise { - const { backendCwd, remote, requestGateway, sessionId } = opts + const { backendCwd, remote, requestGateway, storedSessionId, onSessionRecovered } = opts const path = attachment.path ?? '' const label = attachment.label || pathLabel(path) const uploadBytes = remote || attachmentPathNeedsUpload(path, backendCwd) - if (attachment.kind === 'image') { - let result: ImageAttachResponse - - if (uploadBytes) { - let payload: Awaited> - - try { - payload = await readImageForRemoteAttach(path) - } catch (err) { - throw friendlyRemoteAttachError(err, label) - } - - if (!payload) { - throw new Error(`Could not read ${label}`) - } - - result = await requestGateway('image.attach_bytes', { - session_id: sessionId, - content_base64: payload.contentBase64, - filename: payload.filename - }) - } else { - result = await requestGateway('image.attach', { - path, - session_id: sessionId - }) - } - - if (!result.attached) { - throw new Error(result.message || `Could not attach ${label}`) - } - - const attachedPath = result.path || path - - return { - ...attachment, - attachedSessionId: sessionId, - label: attachedPath ? pathLabel(attachedPath) : attachment.label, - path: attachedPath, - uploadState: undefined - } - } - - // Non-image file. - let dataUrl: string | null = null + // Read bytes/paths ONCE, outside the retry. Only the session-scoped RPC is + // replayed on recovery — re-reading a multi-MB file to retry a dead session + // id would double the disk/IPC cost of every recovered attach. + let imagePayload: Awaited> | null = null + let fileDataUrl: null | string = null if (uploadBytes) { try { - dataUrl = await readFileDataUrlForAttach(path) + if (attachment.kind === 'image') { + imagePayload = await readImageForRemoteAttach(path) + } else { + fileDataUrl = await readFileDataUrlForAttach(path) + } } catch (err) { throw friendlyRemoteAttachError(err, label) } - if (!dataUrl) { + if (attachment.kind === 'image' ? !imagePayload : !fileDataUrl) { throw new Error(`Could not read ${label}`) } } - const result = await requestGateway('file.attach', { - name: label, - path, - session_id: sessionId, - ...(dataUrl ? { data_url: dataUrl } : {}) - }) + const stageForSession = async (liveSessionId: string): Promise => { + if (attachment.kind === 'image') { + const result = imagePayload + ? await requestGateway('image.attach_bytes', { + session_id: liveSessionId, + content_base64: imagePayload.contentBase64, + filename: imagePayload.filename + }) + : await requestGateway('image.attach', { + path, + session_id: liveSessionId + }) - if (!result.attached || !result.ref_text) { - throw new Error(result.message || `Could not attach ${label}`) + if (!result.attached) { + throw new Error(result.message || `Could not attach ${label}`) + } + + const attachedPath = result.path || path + + return { + ...attachment, + attachedSessionId: liveSessionId, + label: attachedPath ? pathLabel(attachedPath) : attachment.label, + path: attachedPath, + uploadState: undefined + } + } + + const result = await requestGateway('file.attach', { + name: label, + path, + session_id: liveSessionId, + ...(fileDataUrl ? { data_url: fileDataUrl } : {}) + }) + + if (!result.attached || !result.ref_text) { + throw new Error(result.message || `Could not attach ${label}`) + } + + return { + ...attachment, + attachedSessionId: liveSessionId, + refText: result.ref_text, + uploadState: undefined + } } - return { - ...attachment, - attachedSessionId: sessionId, - refText: result.ref_text, - uploadState: undefined + // Attach runs BEFORE prompt.submit, so submit's own recovery never gets a + // chance: a stale runtime id fails here first and the user sees "session not + // found" on an image while plain text works. + const { result, sessionId: usedSessionId } = await withSessionNotFoundResume( + opts.sessionId, + storedSessionId, + stageForSession, + { requestGateway } + ) + + if (usedSessionId !== opts.sessionId) { + onSessionRecovered?.(usedSessionId) } + + return result } interface PromptActionsOptions { @@ -300,11 +318,19 @@ export function usePromptActions({ sessionId: string, attachments: ComposerAttachment[], options: { updateComposerAttachments?: boolean } = {} - ): Promise => { + ): Promise<{ attachments: ComposerAttachment[]; sessionId: string }> => { const updateComposerAttachments = options.updateComposerAttachments ?? true const remote = $connection.get()?.mode === 'remote' + const storedSessionId = selectedStoredSessionIdRef.current + let liveSessionId = sessionId const synced: ComposerAttachment[] = [] + const onSessionRecovered = (recoveredId: string) => { + liveSessionId = recoveredId + activeSessionIdRef.current = recoveredId + setActiveSessionId(recoveredId) + } + for (const original of attachments) { let attachment = original @@ -322,8 +348,10 @@ export function usePromptActions({ // Already-synced or pathless refs (terminal, url, etc.) pass through. // A drop-time eager upload may already have staged this one (matching - // attachedSessionId) — don't re-upload it. - if (!attachment.path || attachment.attachedSessionId === sessionId) { + // attachedSessionId) — don't re-upload it. Compare against the LIVE id: + // after a mid-loop recovery an earlier chip's attachedSessionId points + // at the dead runtime and must be re-staged. + if (!attachment.path || attachment.attachedSessionId === liveSessionId) { synced.push(attachment) continue @@ -334,7 +362,9 @@ export function usePromptActions({ backendCwd: $currentCwd.get(), remote, requestGateway, - sessionId + sessionId: liveSessionId, + storedSessionId, + onSessionRecovered }) // Update-only: never resurrect a chip the user removed mid-upload. @@ -350,9 +380,9 @@ export function usePromptActions({ synced.push(attachment) } - return synced + return { attachments: synced, sessionId: liveSessionId } }, - [requestGateway] + [activeSessionIdRef, requestGateway, selectedStoredSessionIdRef] ) // Stage a freshly dropped file as soon as it lands (when a session already @@ -627,38 +657,22 @@ export function usePromptActions({ clearClarifyRequest(undefined, sessionId) try { - await requestGateway('session.interrupt', { session_id: sessionId }) + await withSessionNotFoundResume( + sessionId, + selectedStoredSessionIdRef.current, + liveId => requestGateway('session.interrupt', { session_id: liveId }), + { + requestGateway, + onRecovered: recoveredId => { + activeSessionIdRef.current = recoveredId + setActiveSessionId(recoveredId) + } + } + ) releaseBusy() } catch (err) { - let stopError = err - - if (isSessionNotFoundError(err) && selectedStoredSessionIdRef.current) { - try { - const resumeProfile = await resolveSessionProfile(selectedStoredSessionIdRef.current) - - const resumed = await requestGateway<{ session_id: string }>('session.resume', { - session_id: selectedStoredSessionIdRef.current, - source: 'desktop', - omit_messages: true, - ...(resumeProfile ? { profile: resumeProfile } : {}) - }) - - const recoveredId = resumed?.session_id - - if (recoveredId) { - activeSessionIdRef.current = recoveredId - await requestGateway('session.interrupt', { session_id: recoveredId }) - releaseBusy() - - return - } - } catch (resumeErr) { - stopError = resumeErr - } - } - releaseBusy() - notifyError(stopError, copy.stopFailed) + notifyError(err, copy.stopFailed) } }, [activeSessionIdRef, busyRef, copy.stopFailed, requestGateway, selectedStoredSessionIdRef, updateSessionState]) @@ -733,33 +747,19 @@ export function usePromptActions({ } try { - return await send(sessionId) - } catch (err) { - // A stale runtime id after reconnect 404s ("session not found"): resume - // the stored session and retry once, mirroring stopPrompt so a + // A stale runtime id after reconnect 404s ("session not found"): the + // shared resolver resumes the stored session and retries once, so a // correction right after a reconnect isn't lost to the race. - if (isSessionNotFoundError(err) && selectedStoredSessionIdRef.current) { - try { - const resumeProfile = await resolveSessionProfile(selectedStoredSessionIdRef.current) - - const resumed = await requestGateway<{ session_id: string }>('session.resume', { - session_id: selectedStoredSessionIdRef.current, - source: 'desktop', - omit_messages: true, - ...(resumeProfile ? { profile: resumeProfile } : {}) - }) - - const recoveredId = resumed?.session_id - - if (recoveredId) { - activeSessionIdRef.current = recoveredId - - return await send(recoveredId) - } - } catch { - // fall through — caller queues so nothing is lost + const { result } = await withSessionNotFoundResume(sessionId, selectedStoredSessionIdRef.current, send, { + requestGateway, + onRecovered: recoveredId => { + activeSessionIdRef.current = recoveredId + setActiveSessionId(recoveredId) } - } + }) + + return result + } catch { // Swallow — caller queues the text so nothing is lost. } 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 36c8bfd82d49f..5e48b28903044 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 @@ -23,6 +23,7 @@ import { requestDesktopOnboarding } from '@/store/onboarding' import { $sessions, resolveComposerSessionKey, + setActiveSessionId, setAwaitingResponse, setBusy, setMessages, @@ -39,13 +40,13 @@ import { _submitInFlight, type GatewayRequest, inlineErrorMessage, - isGatewayTimeoutError, isProviderSetupError, isSessionBusyError, - isSessionNotFoundError, isTargetSessionBusy, + SessionRecoveryAborted, type SubmitTextOptions, - withSessionBusyRetry + withSessionBusyRetry, + withSessionNotFoundResume } from './utils' interface SubmitPromptDeps { @@ -63,7 +64,7 @@ interface SubmitPromptDeps { sessionId: string, attachments: ComposerAttachment[], options?: { updateComposerAttachments?: boolean } - ) => Promise + ) => Promise<{ attachments: ComposerAttachment[]; sessionId: string }> updateSessionState: ( sessionId: string, updater: (state: ClientSessionState) => ClientSessionState, @@ -584,23 +585,34 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { } try { - const syncedAttachments = await syncAttachmentsForSubmit(sessionId, attachments, { + // Attach runs BEFORE prompt.submit, so a stale runtime id fails there + // first and submit's own recovery never runs — that asymmetry is why + // plain text survived sleep/wake but images reported "session not + // found". The attach path recovers and reports the live id back here. + const attachResult = await syncAttachmentsForSubmit(sessionId, attachments, { updateComposerAttachments: usingComposerAttachments }) + const syncedAttachments = attachResult.attachments + // Always a live string; pin it so TS narrows past the outer + // `string | null` sessionId binding for prompt.submit. + const liveSessionId = attachResult.sessionId + + sessionId = liveSessionId + const attachmentsDrift = sessionDriftReason() if (attachmentsDrift) { console.warn('[submit-drift-abort]', attachmentsDrift, { phase: 'post-attachments' }) - return abortForSessionSwitch(sessionId) + return abortForSessionSwitch(liveSessionId) } // 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.) attachmentRefs = syncedAttachments.map(optimisticAttachmentRef).filter((r): r is string => Boolean(r)) - rewriteOptimistic(sessionId) + rewriteOptimistic(liveSessionId) const text = buildContextText(syncedAttachments) const submitParams = (targetId: string) => ({ @@ -616,56 +628,45 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { }) // On sleep/wake the gateway's in-memory session may have been cleared - // while the desktop app still holds the old session ID. Detect this, - // resume the stored session to re-register it, and retry once. + // while the desktop app still holds the old session ID. The shared + // resolver re-registers the stored session and retries once; every + // other session-scoped RPC (attach, /compress, rewind, interrupt) goes + // through the same helper so one policy covers the whole bug class. let submitErr: unknown = null try { - await withSessionBusyRetry(() => - requestGateway('prompt.submit', submitParams(sessionId), PROMPT_SUBMIT_REQUEST_TIMEOUT_MS) - ) - } catch (firstErr) { const recoverStoredSessionId = targetStoredSessionId ?? selectedStoredSessionIdRef.current - if ((isSessionNotFoundError(firstErr) || isGatewayTimeoutError(firstErr)) && recoverStoredSessionId) { - // Re-register the session in the gateway and get a fresh live ID. - // Timeouts recover the same way as "session not found": a starved - // backend loop (#55578 symptom d) rejects the submit even though - // the stored session is fine — resume + retry instead of erroring - // out and losing the session binding. - const resumeProfile = await resolveSessionProfile(recoverStoredSessionId) - - const resumed = await requestGateway<{ session_id: string }>('session.resume', { - session_id: recoverStoredSessionId, - source: 'desktop', - omit_messages: true, - ...(resumeProfile ? { profile: resumeProfile } : {}) - }) - - const resumeRetryDrift = sessionDriftReason() - - if (resumeRetryDrift) { - console.warn('[submit-drift-abort]', resumeRetryDrift, { phase: 'post-resume-retry' }) - - return abortForSessionSwitch(sessionId) - } - - const recoveredId = resumed?.session_id - - if (recoveredId) { - if (targetIsCurrentView()) { - activeSessionIdRef.current = recoveredId + await withSessionNotFoundResume( + sessionId, + recoverStoredSessionId, + liveId => + withSessionBusyRetry(() => + requestGateway('prompt.submit', submitParams(liveId), PROMPT_SUBMIT_REQUEST_TIMEOUT_MS) + ), + { + requestGateway, + driftReason: sessionDriftReason, + onRecovered: recoveredId => { + if (targetIsCurrentView()) { + activeSessionIdRef.current = recoveredId + setActiveSessionId(recoveredId) + } } + }, + // A starved backend loop (#55578 symptom d) rejects the submit even + // though the stored session is fine — recover it like a dead id + // instead of erroring out and losing the session binding. + { alsoTimeout: true } + ) + } catch (firstErr) { + if (firstErr instanceof SessionRecoveryAborted) { + console.warn('[submit-drift-abort]', firstErr.reason, { phase: 'post-resume-retry' }) - await withSessionBusyRetry(() => - requestGateway('prompt.submit', submitParams(recoveredId), PROMPT_SUBMIT_REQUEST_TIMEOUT_MS) - ) - } else { - submitErr = firstErr - } - } else { - submitErr = firstErr + return abortForSessionSwitch(sessionId) } + + submitErr = firstErr } if (submitErr !== null) { diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.ts index cc9ba164ad750..52c2555669d41 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.ts @@ -52,6 +52,141 @@ export function isSessionNotFoundError(error: unknown): boolean { return /session not found/i.test(message) } +/** + * Thrown when a stale-session recovery resumed successfully but the caller's + * drift check says the user has since moved on (profile swap, route rebind, + * a different chat in the foreground). The retry is deliberately NOT attempted: + * landing it would run the prompt against a session the user is no longer + * looking at. Callers unwind through their own abort path (#66889). + */ +export class SessionRecoveryAborted extends Error { + constructor( + readonly reason: string, + readonly recoveredSessionId: string + ) { + super(`session recovery aborted: ${reason}`) + this.name = 'SessionRecoveryAborted' + } +} + +export interface SessionRecoveryDeps { + requestGateway: GatewayRequest + /** + * Owning profile for a stored session. A resume without it lands on + * whichever gateway is active and forks the conversation into the wrong + * profile's DB (#67603). + * + * Injected rather than imported so this module stays free of the session + * store and the REST layer: the default implementation reaches through + * `resolveStoredSession` → `getSession()`, a real fetch that makes any unit + * test of this helper depend on leftover `$sessions` / `$profiles` state. + */ + resolveProfile?: (storedSessionId: string) => Promise + /** + * Publish the fresh live id. Implementations must update BOTH the hot ref + * and the `$activeSessionId` atom — a ref-only write leaves the atom + * pointing at the dead runtime and every atom-reading surface desyncs + * (#62471). + */ + onRecovered?: (liveSessionId: string) => void + /** + * Non-null reason ⇒ abort instead of retrying. Evaluated AFTER the resume + * and BEFORE the retry, because the resume is the slow await during which a + * profile switch or route rebind can land. + */ + driftReason?: () => null | string +} + +async function defaultResolveProfile(storedSessionId: string): Promise { + // Lazy so utils.ts has no init-time cycle with use-session-actions. + const { resolveSessionProfile } = await import('../use-session-actions/utils') + + return resolveSessionProfile(storedSessionId) +} + +/** + * Re-register a durable stored session after the gateway dropped its + * in-memory runtime id (sleep/wake, remote backend restart, long idle). + * Returns the fresh live id, or null when the resume yields none. + */ +export async function resumeStoredRuntimeSession( + storedSessionId: string, + deps: SessionRecoveryDeps +): Promise { + const resolveProfile = deps.resolveProfile ?? defaultResolveProfile + const profile = await resolveProfile(storedSessionId) + + const resumed = await deps.requestGateway<{ session_id: string }>('session.resume', { + session_id: storedSessionId, + source: 'desktop', + omit_messages: true, + ...(profile ? { profile } : {}) + }) + + return resumed?.session_id ?? null +} + +/** + * Single resolver for "the runtime session id I hold is dead." + * + * Every session-scoped RPC needs this, not just `prompt.submit`. Attach, + * `/compress`, checkpoint restore, and interrupt all run against the same + * runtime id and all used to surface a raw "session not found" after sleep — + * while plain text silently recovered, which is why the bug reads as "text + * works, images don't." + * + * Runs `call(sessionId)`. On a stale-session error it resumes the stored + * session ONCE, republishes the fresh id, and retries. Bounded to a single + * retry: a second failure is a real error, not a stale binding. + * + * A resume that itself 404s (a never-persisted first-submit draft has no DB + * row until its first successful submit) rethrows the ORIGINAL error rather + * than the confusing secondary one (#67539). + */ +export async function withSessionNotFoundResume( + sessionId: string, + storedSessionId: null | string | undefined, + call: (liveSessionId: string) => Promise, + deps: SessionRecoveryDeps, + options?: { alsoTimeout?: boolean } +): Promise<{ recovered: boolean; result: T; sessionId: string }> { + try { + return { recovered: false, result: await call(sessionId), sessionId } + } catch (err) { + // A starved backend loop rejects with a timeout that is indistinguishable + // from a dead runtime on the client side (#55578). Opt-in per caller: + // submit recovers from it, a compress/attach retry should not mask a + // genuinely slow LLM-bound call. + const recoverable = isSessionNotFoundError(err) || (Boolean(options?.alsoTimeout) && isGatewayTimeoutError(err)) + + if (!recoverable || !storedSessionId) { + throw err + } + + let recoveredId: null | string + + try { + recoveredId = await resumeStoredRuntimeSession(storedSessionId, deps) + } catch { + throw err + } + + if (!recoveredId) { + throw err + } + + const drift = deps.driftReason?.() + + if (drift) { + throw new SessionRecoveryAborted(drift, recoveredId) + } + + deps.onRecovered?.(recoveredId) + + return { recovered: true, result: await call(recoveredId), sessionId: recoveredId } + } +} + /** * Is the session a prompt is about to run against currently mid-turn? *