fix(web): queued-message edits keep the context you added when their run starts

apps/web/src/components/chat/queuedMessageEdit.test.ts

@@ -1,12 +1,18 @@
 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,
   useComposerDraftStore,
   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");

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);
   }

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<ComposerDraftStoreState>()(
             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<ComposerDraftStoreState>()(
               });
             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 = <Value>(
+              values: ReadonlyArray<Value>,
+              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<ComposerDraftStoreState>()(
             // 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),
               ),

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>,
+): 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";
 }