diff --git a/src/renderer/src/lib/notes-carried-by-chat.test.ts b/src/renderer/src/lib/notes-carried-by-chat.test.ts index f02530ea683..f817479ab72 100644 --- a/src/renderer/src/lib/notes-carried-by-chat.test.ts +++ b/src/renderer/src/lib/notes-carried-by-chat.test.ts @@ -75,7 +75,10 @@ import { setLocalRuntimeCapabilitiesForTests } from '@/runtime/local-runtime-cap import { readOutbox } from '@/components/native-chat/structured-agent-session-outbox-storage' import { resetStructuredAgentSessionCarriedNotesForTests } from '@/components/native-chat/structured-agent-session-outbox-carried-notes' import { useStructuredAgentSessionOutbox } from '@/components/native-chat/use-structured-agent-session-outbox' -import { structuredAgentSessionDraftScopeKey } from '@/components/native-chat/native-chat-composer-draft-store' +import { + nativeChatComposerDraftWritesSettled, + structuredAgentSessionDraftScopeKey +} from '@/components/native-chat/native-chat-composer-draft-store' import { clearNativeChatDraftCacheForTests, readNativeChatDraftCache @@ -427,6 +430,12 @@ describe('notes sent to a chat already open', () => { expect(readNativeChatDraftCache(structuredAgentSessionDraftScopeKey(target.sessionId))).toBe( NOTES ) + // Until storage confirms the draft, a crash could lose the text: the notes stay, held. + expect(isNoteInFlight(KEY_A)).toBe(true) + expect(mocks.clearDeliveredDiffComments).not.toHaveBeenCalled() + await nativeChatComposerDraftWritesSettled() + await Promise.resolve() + await Promise.resolve() expect(isNoteInFlight(KEY_A)).toBe(false) expect(mocks.clearDeliveredDiffComments).toHaveBeenCalledExactlyOnceWith(WORKTREE_ID, [NOTE_A]) expect(cleared).toEqual([NOTES]) diff --git a/src/renderer/src/lib/notes-sent-by-chat.ts b/src/renderer/src/lib/notes-sent-by-chat.ts index c5ed6c60234..284c4ba4fa3 100644 --- a/src/renderer/src/lib/notes-sent-by-chat.ts +++ b/src/renderer/src/lib/notes-sent-by-chat.ts @@ -1,8 +1,16 @@ import { useAppStore } from '@/store' +import type { StructuredAgentSessionOutboxEntry } from '../../../shared/structured-agent-session-outbox' import { subscribeToStructuredAgentSessionEntryEndings } from '@/components/native-chat/structured-agent-session-entry-endings' +import { + isNativeChatComposerDraftUnsaved, + nativeChatComposerDraftWritesSettled, + structuredAgentSessionDraftScopeKey, + subscribeToNativeChatComposerDraft +} from '@/components/native-chat/native-chat-composer-draft-store' import { browserAnnotationSendKey, diffCommentSendKey, + holdNotesForSend, noteSendKeyOwner } from './notes-send-in-flight' @@ -100,35 +108,71 @@ function clearWhenLoaded(): void { } } +/** Resolves once the draft a returned message went back to is saved: storage writes are async, and + * until then a crash could lose the text, so its notes must still be there. A draft storage keeps + * refusing holds them until a later save lands. */ +async function returnedDraftSaved(sessionId: string): Promise { + const scopeKey = structuredAgentSessionDraftScopeKey(sessionId) + await nativeChatComposerDraftWritesSettled() + while (isNativeChatComposerDraftUnsaved(scopeKey)) { + await new Promise((resolve) => { + const unsubscribe = subscribeToNativeChatComposerDraft(scopeKey, () => { + unsubscribe() + resolve() + }) + }) + await nativeChatComposerDraftWritesSettled() + } +} + /** Clears the notes a chat message carried once it is used: the host has it, or its text went - * back to the composer. By whichever send, resend or Stop ended it, reload included. */ + * back to the composer and that draft is saved. By whichever send, resend or Stop ended it, + * reload included. */ export function installNotesSentByChat(): () => void { + let installed = true const unsubscribeEndings = subscribeToStructuredAgentSessionEntryEndings((entry, ending) => { if (ending === 'discarded') { return } - for (const key of entry.carriedNoteKeys ?? []) { - const parsed = noteSendKeyOwner(key) - if (parsed) { - const ownerId = ownerKey(parsed.kind, parsed.owner) - const pending = pendingByOwner.get(ownerId) ?? { - ...parsed, - keys: new Set(), - seen: NOT_SEEN + const keys = entry.carriedNoteKeys ?? [] + if (ending === 'returned' && keys.length > 0) { + // Held off the shelf meanwhile: the message that carried them has already left the outbox. + const saved = returnedDraftSaved(entry.sessionId) + holdNotesForSend(keys, saved) + void saved.then(() => { + if (installed) { + markNotesUsed(entry) } - pendingByOwner.set(ownerId, pending) - pending.keys.add(key) - pending.seen = NOT_SEEN - } + }) + return } - // An ending can fire inside a store update (a removed workspace's chats settle there), and a - // clear written into it could be overwritten by that update's own result. - queueMicrotask(clearWhenLoaded) + markNotesUsed(entry) }) return () => { + installed = false unsubscribeEndings() unsubscribeStore?.() unsubscribeStore = null pendingByOwner.clear() } } + +function markNotesUsed(entry: Pick): void { + for (const key of entry.carriedNoteKeys ?? []) { + const parsed = noteSendKeyOwner(key) + if (parsed) { + const ownerId = ownerKey(parsed.kind, parsed.owner) + const pending = pendingByOwner.get(ownerId) ?? { + ...parsed, + keys: new Set(), + seen: NOT_SEEN + } + pendingByOwner.set(ownerId, pending) + pending.keys.add(key) + pending.seen = NOT_SEEN + } + } + // An ending can fire inside a store update (a removed workspace's chats settle there), and a + // clear written into it could be overwritten by that update's own result. + queueMicrotask(clearWhenLoaded) +}