diff --git a/apps/web/src/components/chat/queuedMessageEdit.test.ts b/apps/web/src/components/chat/queuedMessageEdit.test.ts index f7aed151f2fcc..c22b3983c8a6a 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, @@ -86,6 +87,18 @@ describe("queued message file edits", () => { expect(store.getComposerDraft(editTarget)).toBeNull(); }); + it("removes a kept edit draft that held a thread mention", () => { + const store = useComposerDraftStore.getState(); + store.setPrompt(editTarget, "Edited message"); + store.setThreadContexts(editTarget, [ + threadContextRecord(scopeThreadRef(environmentId, ThreadId.make("thread:other")), "Other"), + ]); + expect( + recoverQueuedMessageEdit({ editTarget, threadTarget, originalText: "Original message" }), + ).toBe("kept"); + 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..ee7d954d3b4a4 100644 --- a/apps/web/src/components/chat/queuedMessageEdit.ts +++ b/apps/web/src/components/chat/queuedMessageEdit.ts @@ -28,9 +28,8 @@ export function recoverQueuedMessageEdit(input: { }); if (recovery === "kept") { store.moveComposerPromptAndImages(input.editTarget, input.threadTarget); - } else { - store.clearComposerContent(input.editTarget); } + store.clearComposerContent(input.editTarget); return recovery; } diff --git a/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx b/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx index 4a58d25e48757..0320dd6d6bb18 100644 --- a/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx +++ b/apps/web/src/components/chat/useQueuedRunEdit.browser.test.tsx @@ -9,11 +9,20 @@ import { } from "@t3tools/contracts"; import { useImperativeHandle, type RefObject } from "react"; import { describe, expect, it, vi } from "vite-plus/test"; -import { ATTACHMENT_ONLY_BOOTSTRAP_PROMPT } from "@t3tools/shared/composerContextReferences"; +import { + ATTACHMENT_ONLY_BOOTSTRAP_PROMPT, + threadContextRecord, +} from "@t3tools/shared/composerContextReferences"; import { toastManager } from "../ui/toast"; import { act, render } from "../../../test/render"; import { useComposerDraftStore, useComposerThreadDraft } from "../../composerDraftStore"; +import { formatInlineContextReference } from "../../lib/composerContextReferences"; +import { + environmentContextRecord, + environmentContextReference, + threadContextReference, +} from "../../lib/composerContextRecords"; import { deriveComposerSendState } from "../ChatView.logic"; import { ComposerPrimaryActions } from "./ComposerPrimaryActions"; import { AppAtomRegistryProvider, appAtomRegistry } from "../../rpc/atomRegistry"; @@ -74,8 +83,8 @@ const RUN = RunId.make("run-1"); type HookResult = ReturnType; -function projection(runs: ReadonlyArray) { - return { runs, messages: [] } as unknown as ThreadStructure; +function projection(runs: ReadonlyArray, messages: ReadonlyArray) { + return { runs, messages } as unknown as ThreadStructure; } function setRuns(ref: ScopedThreadRef, status: string) { @@ -86,13 +95,14 @@ function Harness(props: { readonly threadRef: ScopedThreadRef; readonly runStatus: string; readonly hookRef: RefObject; + readonly messages: ReadonlyArray; }) { const result = useQueuedRunEdit({ environmentId: ENV, routeThreadRef: props.threadRef, threadDraftTarget: props.threadRef, activeThreadId: props.threadRef.threadId, - serverProjection: projection([{ id: RUN, status: props.runStatus }]), + serverProjection: projection([{ id: RUN, status: props.runStatus }], props.messages), promptRef: { current: "" }, composerRef: { current: null }, }); @@ -124,10 +134,15 @@ function Harness(props: { ) : null; } -function view(threadRef: ScopedThreadRef, runStatus: string, hook: { current: HookResult | null }) { +function view( + threadRef: ScopedThreadRef, + runStatus: string, + hook: { current: HookResult | null }, + messages: ReadonlyArray = [], +) { return ( - + ); } @@ -262,6 +277,56 @@ describe("useQueuedRunEdit", () => { ); }); + it("opens an edit with only the thread and machine records its text still mentions", async () => { + const { hook } = setup(); + const thread = threadContextRecord(THREAD_B, "Support Any File Uploads"); + const machine = environmentContextRecord(ENV, "Build box"); + const unreferenced = threadContextRecord(THREAD_A, "Removed earlier"); + const messageId = MessageId.make("message-1"); + const text = `wait for ${formatInlineContextReference(threadContextReference(thread))} on ${formatInlineContextReference(environmentContextReference(machine))}`; + const messages = [ + { id: messageId, text, context: { version: 1, records: [thread, machine, unreferenced] } }, + ]; + await using _screen = await render(view(THREAD_A, "queued", hook, messages)); + + await act(() => requireHook(hook).beginEdit({ runId: RUN, messageId, text, attachments: [] })); + + const draft = useComposerDraftStore.getState().getComposerDraft(requireHook(hook).draftTarget); + expect(draft?.prompt).toBe(text); + expect(draft?.threadContexts).toEqual([thread]); + expect(draft?.environmentContexts).toEqual([machine]); + }); + + it("saves only the thread mentions still in the edited message", async () => { + const { hook } = setup(); + mocks.editQueuedRun.mockResolvedValue({ _tag: "Success" }); + const kept = threadContextRecord(THREAD_B, "Kept"); + const removed = threadContextRecord(THREAD_A, "Removed"); + const messageId = MessageId.make("message-1"); + const text = [kept, removed] + .map((record) => formatInlineContextReference(threadContextReference(record))) + .join(" and "); + const messages = [{ id: messageId, text, context: { version: 1, records: [kept, removed] } }]; + await using _screen = await render(view(THREAD_A, "queued", hook, messages)); + await act(() => requireHook(hook).beginEdit({ runId: RUN, messageId, text, attachments: [] })); + const target = requireHook(hook).draftTarget; + await act(() => useComposerDraftStore.getState().setThreadContexts(target, [kept])); + + const draft = useComposerDraftStore.getState().getComposerDraft(target); + await requireHook(hook).saveEdit( + { ...savedContent(draft?.prompt ?? ""), threadContexts: draft?.threadContexts ?? [] }, + () => {}, + ); + + expect(mocks.editQueuedRun).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + edit: expect.objectContaining({ context: { version: 1, records: [kept] } }), + }), + }), + ); + }); + it("keeps an open edit when the user leaves the thread and comes back", async () => { const { hook } = setup(); await using screen = await render(view(THREAD_A, "queued", hook)); diff --git a/apps/web/src/components/chat/useQueuedRunEdit.ts b/apps/web/src/components/chat/useQueuedRunEdit.ts index 8feb302dc6aae..f8f968487fc75 100644 --- a/apps/web/src/components/chat/useQueuedRunEdit.ts +++ b/apps/web/src/components/chat/useQueuedRunEdit.ts @@ -7,6 +7,7 @@ import type { ThreadStructure } from "@t3tools/client-runtime/state/shell"; import { PROVIDER_SEND_TURN_MAX_ATTACHMENTS, type ChatAttachment, + type ComposerContextRecord, type EnvironmentContextRecord, type EnvironmentId, type MessageId, @@ -40,6 +41,7 @@ import { startAttachmentUpload, } from "../../lib/attachmentUploadQueue"; import { buildMessageContext } from "../../lib/composerContextRecords"; +import { collectInlineContextIds } from "../../lib/composerContextReferences"; import { readFileAsDataUrl } from "../../lib/imageCompression"; import type { TerminalContextDraft } from "../../lib/terminalContext"; import type { ReviewCommentContext } from "../../reviewCommentContext"; @@ -85,6 +87,9 @@ interface QueuedRunEditContent { readonly environmentContexts: ReadonlyArray; } +const isDraftOwnedRecord = (record: ComposerContextRecord) => + record.kind === "thread" || record.kind === "environment"; + function queuedEditDraftTarget(routeThreadRef: ScopedThreadRef, runId: RunId) { return DraftId.make(`queued-edit:${scopedThreadKey(routeThreadRef)}:${runId}`); } @@ -127,10 +132,12 @@ async function submitQueuedRunEdit(input: { readImage: readFileAsDataUrl, uploadFiles, }); + // The edit draft owns thread and machine mentions, so a chip removed while editing is dropped. const keptRecords = (editing.context?.records ?? []).filter( (record) => - !("attachmentId" in record) || - editing.existingAttachments.some((attachment) => attachment.id === record.attachmentId), + !isDraftOwnedRecord(record) && + (!("attachmentId" in record) || + editing.existingAttachments.some((attachment) => attachment.id === record.attachmentId)), ); const newRecords = buildMessageContext({ @@ -233,9 +240,22 @@ export function useQueuedRunEdit(input: { if (editing !== null && editing.runId !== request.runId) clearEditDraft(editing.runId); const target = queuedEditDraftTarget(routeThreadRef, request.runId); const text = stripAttachmentOnlyBootstrapPrompt(request.text); + const context = serverProjection?.messages.find( + (message) => message.id === request.messageId, + )?.context; + const referenced = new Set(collectInlineContextIds(text)); + const records = (context?.records ?? []).filter((record) => referenced.has(record.contextId)); const store = useComposerDraftStore.getState(); store.clearComposerContent(target); store.setPrompt(target, text); + store.setThreadContexts( + target, + records.filter((record): record is ThreadContextRecord => record.kind === "thread"), + ); + store.setEnvironmentContexts( + target, + records.filter((record): record is EnvironmentContextRecord => record.kind === "environment"), + ); const threadRef = { environmentId, threadId: activeThreadId }; queuedRunEdits.begin( editKey, @@ -245,8 +265,7 @@ export function useQueuedRunEdit(input: { messageId: request.messageId, originalText: text, existingAttachments: request.attachments, - context: serverProjection?.messages.find((message) => message.id === request.messageId) - ?.context, + context, }, { runs: makeQueuedRunsAtom({