diff --git a/apps/desktop/src/app/chat/session-tile-actions.ts b/apps/desktop/src/app/chat/session-tile-actions.ts index 8a35d14124..48faa43cf2 100644 --- a/apps/desktop/src/app/chat/session-tile-actions.ts +++ b/apps/desktop/src/app/chat/session-tile-actions.ts @@ -48,7 +48,11 @@ import { truncateSubmitParams } from '../session/hooks/use-prompt-actions/rewind' import { useSubmitPrompt } from '../session/hooks/use-prompt-actions/submit' -import { type SubmitTextOptions } from '../session/hooks/use-prompt-actions/utils' +import { + markSessionRecentlyInterrupted, + shouldInterruptBeforeRewind, + type SubmitTextOptions +} from '../session/hooks/use-prompt-actions/utils' import { upsertOptimisticSession } from '../session/hooks/use-session-actions/utils' import type { ComposerScope } from './composer/scope' @@ -258,6 +262,9 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses const cancelRun = useCallback(async () => { const sessionId = runtimeIdRef.current + // Frontend busy clears immediately; gateway wind-down can lag (#83855). + markSessionRecentlyInterrupted(sessionId) + update(state => ({ ...state, messages: finalizeInterruptedMessages(state.messages, state.streamId), @@ -456,13 +463,16 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses resetSessionBackground(sessionId) clearPreviewArtifacts(sessionId) - const wasBusy = readState()?.busy ?? false + const interruptFirst = shouldInterruptBeforeRewind({ + busy: readState()?.busy ?? false, + sessionId + }) update(state => applyRewindOptimistic(state, plan.sourceIndex)) try { applySurvivorRowIds( - await submitRewind(plan.text, plan.truncateOrdinal, wasBusy, plan.truncateMessageId, plan.truncateRowId) + await submitRewind(plan.text, plan.truncateOrdinal, interruptFirst, plan.truncateMessageId, plan.truncateRowId) ) } catch (err) { update(state => ({ ...state, busy: false, awaitingResponse: false, messages })) @@ -487,13 +497,16 @@ export function useSessionTileActions({ runtimeId, scope, storedSessionId }: Ses resetSessionBackground(sessionId) clearPreviewArtifacts(sessionId) - const wasBusy = readState()?.busy ?? false + const interruptFirst = shouldInterruptBeforeRewind({ + busy: readState()?.busy ?? false, + sessionId + }) update(state => applyRewindOptimistic(state, plan.sourceIndex, plan.editedMessage)) try { applySurvivorRowIds( - await submitRewind(plan.text, plan.truncateOrdinal, wasBusy, plan.truncateMessageId, plan.truncateRowId) + await submitRewind(plan.text, plan.truncateOrdinal, interruptFirst, plan.truncateMessageId, plan.truncateRowId) ) } catch (err) { update(state => ({ ...state, busy: false, awaitingResponse: false, messages })) 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 fdd32fdf74..709ea7026d 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 @@ -70,8 +70,10 @@ import { friendlyRemoteAttachError, type GatewayRequest, inlineErrorMessage, + markSessionRecentlyInterrupted, readFileDataUrlForAttach, readImageForRemoteAttach, + shouldInterruptBeforeRewind, type SubmitTextOptions, withSessionNotFoundResume } from './utils' @@ -649,6 +651,10 @@ export function usePromptActions({ return } + // Frontend busy clears immediately; gateway wind-down can lag. Mark so a + // fast edit/resend still interrupt-first instead of racing 4009 (#83855). + markSessionRecentlyInterrupted(sessionId) + updateSessionState(sessionId, state => { const streamId = state.streamId const messages = finalizeInterruptedMessages(state.messages, streamId) @@ -908,6 +914,13 @@ export function usePromptActions({ resetSessionBackground(sessionId) clearPreviewArtifacts(sessionId) + // Capture before optimistic busy=true — otherwise interruptFirst is always + // true and idle restores wrongly interrupt (and Stop→edit misses cooldown). + const interruptFirst = shouldInterruptBeforeRewind({ + busy: busyRef.current || $busy.get(), + sessionId + }) + clearNotifications() setMutableRef(busyRef, true) setBusy(true) @@ -920,7 +933,7 @@ export function usePromptActions({ plan.text, plan.truncateOrdinal, plan.truncateMessageId, - busyRef.current || $busy.get(), + interruptFirst, plan.truncateRowId ) @@ -965,6 +978,12 @@ export function usePromptActions({ resetSessionBackground(sessionId) clearPreviewArtifacts(sessionId) + // Before optimistic busy=true — see restoreToMessage (#83855). + const interruptFirst = shouldInterruptBeforeRewind({ + busy: busyRef.current || $busy.get(), + sessionId + }) + clearNotifications() setMutableRef(busyRef, true) setBusy(true) @@ -977,7 +996,7 @@ export function usePromptActions({ plan.text, plan.truncateOrdinal, plan.truncateMessageId, - busyRef.current || $busy.get(), + interruptFirst, plan.truncateRowId ) 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 be6264748e..422d79ef64 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 @@ -38,12 +38,13 @@ import { resolveSessionProfile } from '../use-session-actions/utils' import { finalizeInterruptedMessages } from './rewind' import { - _submitInFlight, + acquireSubmitInFlight, type GatewayRequest, inlineErrorMessage, isProviderSetupError, isSessionBusyError, isTargetSessionBusy, + releaseSubmitInFlight, SessionRecoveryAborted, type SubmitTextOptions, withSessionBusyRetry, @@ -296,17 +297,16 @@ export function useSubmitPrompt(deps: SubmitPromptDeps) { // session switch; this per-session lock makes that safe. const submitLockKey = targetStoredSessionId || sessionId || startingActiveSessionId || '__pending_new__' - if (_submitInFlight.has(submitLockKey)) { + if (!acquireSubmitInFlight(submitLockKey)) { return false } - _submitInFlight.add(submitLockKey) let submitLockReleased = false const releaseSubmitLock = () => { if (!submitLockReleased) { submitLockReleased = true - _submitInFlight.delete(submitLockKey) + releaseSubmitInFlight(submitLockKey) } } diff --git a/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.test.ts b/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.test.ts index c2b1ac231c..e00efecf66 100644 --- a/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.test.ts +++ b/apps/desktop/src/app/session/hooks/use-prompt-actions/utils.test.ts @@ -1,11 +1,14 @@ import type { AppendMessage } from '@assistant-ui/react' -import { describe, expect, it, vi } from 'vitest' +import { afterEach, describe, expect, it, vi } from 'vitest' import type { ChatMessage } from '@/lib/chat-messages' import { + acquireSubmitInFlight, appendText, base64FromDataUrl, + clearSessionRecentlyInterrupted, + clearSubmitInFlight, friendlyRemoteAttachError, type GatewayRequest, imageFilenameFromPath, @@ -13,15 +16,89 @@ import { isSessionBusyError, isSessionIdCandidate, isSessionNotFoundError, + isSessionRecentlyInterrupted, + isSubmitInFlight, + markSessionRecentlyInterrupted, + RECENT_INTERRUPT_COOLDOWN_MS, readFileDataUrlForAttach, + releaseSubmitInFlight, renderRpcResult, SessionRecoveryAborted, + shouldInterruptBeforeRewind, slashStatusText, + SUBMIT_IN_FLIGHT_TTL_MS, visibleUserIndexAtOrdinal, visibleUserOrdinal, withSessionNotFoundResume } from './utils' +afterEach(() => { + clearSessionRecentlyInterrupted() + clearSubmitInFlight() +}) + +describe('recent interrupt cooldown', () => { + it('is true within the cooldown and false after expiry', () => { + const sessionId = 'sess-cooldown' + const t0 = 1_000_000 + + markSessionRecentlyInterrupted(sessionId, t0) + + expect(isSessionRecentlyInterrupted(sessionId, t0)).toBe(true) + expect(isSessionRecentlyInterrupted(sessionId, t0 + RECENT_INTERRUPT_COOLDOWN_MS - 1)).toBe(true) + expect(isSessionRecentlyInterrupted(sessionId, t0 + RECENT_INTERRUPT_COOLDOWN_MS)).toBe(false) + }) + + it('returns false after mark + elapsed past cooldown', () => { + const sessionId = 'sess-elapsed' + const t0 = 5_000_000 + + markSessionRecentlyInterrupted(sessionId, t0) + expect(isSessionRecentlyInterrupted(sessionId, t0 + RECENT_INTERRUPT_COOLDOWN_MS + 1)).toBe(false) + }) + + it('shouldInterruptBeforeRewind is true when recently interrupted even if not busy', () => { + const sessionId = 'sess-edit-after-stop' + const t0 = 9_000_000 + + markSessionRecentlyInterrupted(sessionId, t0) + + expect(shouldInterruptBeforeRewind({ busy: false, sessionId, now: t0 + 500 })).toBe(true) + expect(shouldInterruptBeforeRewind({ busy: false, sessionId, now: t0 + RECENT_INTERRUPT_COOLDOWN_MS + 1 })).toBe( + false + ) + }) + + it('shouldInterruptBeforeRewind stays false for idle sessions with no recent interrupt', () => { + expect(shouldInterruptBeforeRewind({ busy: false, sessionId: 'idle-sess' })).toBe(false) + expect(shouldInterruptBeforeRewind({ busy: true, sessionId: 'busy-sess' })).toBe(true) + }) +}) + +describe('submit in-flight TTL', () => { + it('blocks a second acquire while fresh and frees after TTL without explicit release', () => { + const key = 'lock-ttl' + const t0 = 2_000_000 + + expect(acquireSubmitInFlight(key, t0)).toBe(true) + expect(isSubmitInFlight(key, t0 + 1)).toBe(true) + expect(acquireSubmitInFlight(key, t0 + 1)).toBe(false) + + expect(isSubmitInFlight(key, t0 + SUBMIT_IN_FLIGHT_TTL_MS)).toBe(false) + expect(acquireSubmitInFlight(key, t0 + SUBMIT_IN_FLIGHT_TTL_MS)).toBe(true) + }) + + it('release clears the lock immediately', () => { + const key = 'lock-release' + const t0 = 3_000_000 + + expect(acquireSubmitInFlight(key, t0)).toBe(true) + releaseSubmitInFlight(key) + expect(isSubmitInFlight(key, t0 + 1)).toBe(false) + expect(acquireSubmitInFlight(key, t0 + 1)).toBe(true) + }) +}) + describe('isSessionIdCandidate', () => { it('accepts the timestamped and hex id forms', () => { expect(isSessionIdCandidate('20260101_120000_abc123')).toBe(true) 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 aba343355b..77227229e5 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 @@ -259,12 +259,101 @@ export async function withSessionBusyRetry(call: () => Promise): Promise() + +export function markSessionRecentlyInterrupted(sessionId: string, now = Date.now()): void { + if (!sessionId) { + return + } + + _recentlyInterruptedUntil.set(sessionId, now + RECENT_INTERRUPT_COOLDOWN_MS) +} + +export function isSessionRecentlyInterrupted(sessionId: string, now = Date.now()): boolean { + const until = _recentlyInterruptedUntil.get(sessionId) + + if (until === undefined) { + return false + } + + if (now >= until) { + _recentlyInterruptedUntil.delete(sessionId) + + return false + } + + return true +} + +export function clearSessionRecentlyInterrupted(sessionId?: string): void { + if (sessionId) { + _recentlyInterruptedUntil.delete(sessionId) + + return + } + + _recentlyInterruptedUntil.clear() +} + +/** Whether a rewind/edit should interrupt before submit — busy OR recent Stop. */ +export function shouldInterruptBeforeRewind(opts: { + busy: boolean + sessionId: string + now?: number +}): boolean { + return opts.busy || isSessionRecentlyInterrupted(opts.sessionId, opts.now) +} + // Hard guard: at most one prompt.submit in flight per session. Every submit // path — user Enter, queue drain, busy-retry, slash fallthrough — funnels // through submitPromptText. Without this, a stalled turn (e.g. a context-bloated // session whose first call hangs) let the SAME prompt launch several real turns // at once (the "message stacked 5×" bug). Keyed by stored/active session id. -export const _submitInFlight = new Set() +// Entries expire so a hung submit cannot permanently block the session (#83855). +export const SUBMIT_IN_FLIGHT_TTL_MS = 30_000 + +const _submitInFlightAt = new Map() + +export function isSubmitInFlight(key: string, now = Date.now()): boolean { + const acquiredAt = _submitInFlightAt.get(key) + + if (acquiredAt === undefined) { + return false + } + + if (now - acquiredAt >= SUBMIT_IN_FLIGHT_TTL_MS) { + _submitInFlightAt.delete(key) + + return false + } + + return true +} + +/** Returns true when the lock was acquired; false when another fresh hold blocks. */ +export function acquireSubmitInFlight(key: string, now = Date.now()): boolean { + if (isSubmitInFlight(key, now)) { + return false + } + + _submitInFlightAt.set(key, now) + + return true +} + +export function releaseSubmitInFlight(key: string): void { + _submitInFlightAt.delete(key) +} + +export function clearSubmitInFlight(): void { + _submitInFlightAt.clear() +} export function base64FromDataUrl(dataUrl: string): string { const comma = dataUrl.indexOf(',')