diff --git a/apps/web/src/components/chat/queuedMessageEdit.test.ts b/apps/web/src/components/chat/queuedMessageEdit.test.ts index f7aed151f2fcc..2f8e5e444081c 100644 --- a/apps/web/src/components/chat/queuedMessageEdit.test.ts +++ b/apps/web/src/components/chat/queuedMessageEdit.test.ts @@ -1,5 +1,6 @@ import { EnvironmentId, ThreadId } from "@t3tools/contracts"; import { scopeThreadRef } from "@t3tools/client-runtime/environment"; +import { threadContextRecord } from "@t3tools/shared/composerContextReferences"; import { beforeEach, describe, expect, it } from "vite-plus/test"; import { DraftId, @@ -7,6 +8,11 @@ import { type ComposerFileAttachment, type ComposerImageAttachment, } from "../../composerDraftStore"; +import { threadContextReference } from "../../lib/composerContextRecords"; +import { + collectInlineContextIds, + formatInlineContextReference, +} from "../../lib/composerContextReferences"; import { prepareQueuedEditAttachments, recoverQueuedMessageEdit } from "./queuedMessageEdit"; const environmentId = EnvironmentId.make("remote-environment"); @@ -86,6 +92,39 @@ describe("queued message file edits", () => { expect(store.getComposerDraft(editTarget)).toBeNull(); }); + it("keeps context added to an edit when another client starts the queued run", () => { + const store = useComposerDraftStore.getState(); + const mention = threadContextRecord( + scopeThreadRef(environmentId, ThreadId.make("thread:other")), + "Other thread", + ); + store.setPrompt(editTarget, "Original message"); + store.addThreadContexts(editTarget, [mention]); + expect( + recoverQueuedMessageEdit({ editTarget, threadTarget, originalText: "Original message" }), + ).toBe("kept"); + const kept = store.getComposerDraft(threadTarget); + expect(kept?.threadContexts).toEqual([mention]); + expect(collectInlineContextIds(kept?.prompt ?? "")).toEqual([mention.contextId]); + expect(store.getComposerDraft(editTarget)).toBeNull(); + }); + + it("leaves out a mention typed over and a chip with no record when keeping an edit", () => { + const store = useComposerDraftStore.getState(); + const threadRecord = (id: string) => + threadContextRecord(scopeThreadRef(environmentId, ThreadId.make(id)), id); + const deadChip = formatInlineContextReference(threadContextReference(threadRecord("dead"))); + store.setThreadContexts(editTarget, [threadRecord("typed-over")]); + store.setPrompt(editTarget, `Edited ${deadChip}`); + expect( + recoverQueuedMessageEdit({ editTarget, threadTarget, originalText: "Original message" }), + ).toBe("kept"); + const kept = store.getComposerDraft(threadTarget); + expect(kept?.prompt).toBe("Edited "); + expect(kept?.threadContexts).toEqual([]); + expect(store.getComposerDraft(editTarget)).toBeNull(); + }); + it("reports a removed queued attachment as an edit when the run leaves", () => { const store = useComposerDraftStore.getState(); store.setPrompt(editTarget, "Original message"); diff --git a/apps/web/src/components/chat/queuedMessageEdit.ts b/apps/web/src/components/chat/queuedMessageEdit.ts index 40a222a9f383a..468929d4cabaf 100644 --- a/apps/web/src/components/chat/queuedMessageEdit.ts +++ b/apps/web/src/components/chat/queuedMessageEdit.ts @@ -27,7 +27,7 @@ export function recoverQueuedMessageEdit(input: { !composerDraftHasUserContent(store.getComposerDraft(input.threadTarget)), }); if (recovery === "kept") { - store.moveComposerPromptAndImages(input.editTarget, input.threadTarget); + store.moveComposerContent(input.editTarget, input.threadTarget); } else { store.clearComposerContent(input.editTarget); } diff --git a/apps/web/src/composerDraftStore.ts b/apps/web/src/composerDraftStore.ts index 04715866949a0..2e0fc68050efb 100644 --- a/apps/web/src/composerDraftStore.ts +++ b/apps/web/src/composerDraftStore.ts @@ -1,4 +1,3 @@ -import { stripInlineContextReferences } from "./lib/composerContextReferences"; import { composerFileAttachmentMimeType } from "./components/chat/composerAttachmentFiles"; import { elementContextToPreviewAnnotation } from "./lib/elementContext"; import { @@ -74,12 +73,15 @@ import { import { appendInlineContextReference, type ComposerContextReference, + collectInlineContextIds, ensureInlineContextReferences, formatInlineContextReference, + keepInlineContextReferences, removeInlineContextReference, } from "./lib/composerContextReferences"; import { fileContextReference, + imageContextReference, previewAnnotationContextId, previewAnnotationContextReference, reviewCommentContextId, @@ -630,8 +632,8 @@ interface ComposerDraftStoreState { clearComposerContent: (threadRef: ComposerThreadTarget) => void; /** Clears text and attachments, preserving terminal/element contexts, preview annotations, and review comments. */ clearComposerPromptAndImages: (threadRef: ComposerThreadTarget) => void; - /** Keeps attachments over destination limits, other environments' uploads, terminal/element contexts, annotations, and review comments in source. */ - moveComposerPromptAndImages: (from: ComposerThreadTarget, to: ComposerThreadTarget) => void; + /** For a destination with no user content. Attachments over its limit or uploaded to another environment stay in source. */ + moveComposerContent: (from: ComposerThreadTarget, to: ComposerThreadTarget) => void; } export interface EffectiveComposerModelState { @@ -4111,7 +4113,7 @@ const composerDraftStore = create()( return { draftsByThreadKey: nextDraftsByThreadKey }; }); }, - moveComposerPromptAndImages: (from, to) => { + moveComposerContent: (from, to) => { const fromKey = resolveComposerDraftKey(get(), from) ?? ""; const toKey = resolveComposerDraftKey(get(), to) ?? ""; if (fromKey.length === 0 || toKey.length === 0 || fromKey === toKey) { @@ -4162,15 +4164,40 @@ const composerDraftStore = create()( }); const movedFileIds = new Set(movedFiles.map((file) => file.id)); const retainedFiles = source.files.filter((file) => !movedFileIds.has(file.id)); - const movedPrompt = ensureInlineContextReferences( - stripInlineContextReferences(source.prompt), - destination.terminalContexts.map(terminalContextReference), - ); + const referencedIds = new Set(collectInlineContextIds(source.prompt)); + const referenced = ( + values: ReadonlyArray, + reference: (value: Value) => ComposerContextReference, + ) => values.filter((value) => referencedIds.has(reference(value).contextId)); const nextDestination: ComposerThreadDraftState = { ...destination, - prompt: movedPrompt, + prompt: keepInlineContextReferences( + source.prompt, + new Set( + [ + ...source.terminalContexts.map(terminalContextReference), + ...source.previewAnnotations.map(previewAnnotationContextReference), + ...source.reviewComments.map(reviewCommentContextReference), + ...source.threadContexts.map(threadContextReference), + ...source.environmentContexts.map(environmentContextReference), + ...movedImages.map(imageContextReference), + ...movedFiles.map(fileContextReference), + ].map((reference) => reference.contextId), + ), + ), images: [...destination.images, ...movedImages], files: [...destination.files, ...movedFiles], + terminalContexts: referenced(source.terminalContexts, terminalContextReference), + previewAnnotations: referenced( + source.previewAnnotations, + previewAnnotationContextReference, + ), + reviewComments: referenced(source.reviewComments, reviewCommentContextReference), + threadContexts: referenced(source.threadContexts, threadContextReference), + environmentContexts: referenced( + source.environmentContexts, + environmentContextReference, + ), nonPersistedImageIds: [ ...destination.nonPersistedImageIds, ...source.nonPersistedImageIds.filter((imageId) => movedImageIds.has(imageId)), @@ -4185,12 +4212,14 @@ const composerDraftStore = create()( // Moved images still use these blob URLs at the destination; do not revoke them. const nextSource: ComposerThreadDraftState = { ...source, - prompt: ensureInlineContextReferences( - "", - source.terminalContexts.map(terminalContextReference), - ), + prompt: "", images: retainedImages, files: retainedFiles, + terminalContexts: [], + previewAnnotations: [], + reviewComments: [], + threadContexts: [], + environmentContexts: [], nonPersistedImageIds: source.nonPersistedImageIds.filter( (imageId) => !movedImageIds.has(imageId), ), diff --git a/apps/web/src/lib/composerContextReferences.ts b/apps/web/src/lib/composerContextReferences.ts index 51ceb029ca3f3..1587e8fdf7c0a 100644 --- a/apps/web/src/lib/composerContextReferences.ts +++ b/apps/web/src/lib/composerContextReferences.ts @@ -37,6 +37,16 @@ export function stripInlineContextReferences(prompt: string): string { return replaceComposerContextReferences(prompt, () => ""); } +/** Strips references with no record in `contextIds`, which would otherwise show as dead chips. */ +export function keepInlineContextReferences( + prompt: string, + contextIds: ReadonlySet, +): string { + return replaceComposerContextReferences(prompt, (occurrence) => + contextIds.has(occurrence.contextId) ? occurrence.source : "", + ); +} + function isBoundaryWhitespace(char: string | undefined): boolean { return char === undefined || char === " " || char === "\n" || char === "\t" || char === "\r"; }