diff --git a/site/src/api/queries/chatMessageEdits.test.ts b/site/src/api/queries/chatMessageEdits.test.ts new file mode 100644 index 0000000000..0cf726ff85 --- /dev/null +++ b/site/src/api/queries/chatMessageEdits.test.ts @@ -0,0 +1,44 @@ +import { describe, expect, it } from "vitest"; +import type * as TypesGen from "#/api/typesGenerated"; +import { buildOptimisticEditedMessage } from "./chatMessageEdits"; + +const makeUserMessage = ( + content: readonly TypesGen.ChatMessagePart[] = [ + { type: "text", text: "original" }, + ], +): TypesGen.ChatMessage => ({ + id: 1, + chat_id: "chat-1", + created_at: "2025-01-01T00:00:00.000Z", + role: "user", + content, +}); + +describe("buildOptimisticEditedMessage", () => { + it("preserves image MIME types for newly attached files", () => { + const message = buildOptimisticEditedMessage({ + requestContent: [{ type: "file", file_id: "image-1" }], + originalMessage: makeUserMessage(), + attachmentMediaTypes: new Map([["image-1", "image/png"]]), + }); + + expect(message.content).toEqual([ + { type: "file", file_id: "image-1", media_type: "image/png" }, + ]); + }); + + it("reuses existing file parts before local attachment metadata", () => { + const existingFilePart: TypesGen.ChatFilePart = { + type: "file", + file_id: "existing-1", + media_type: "image/jpeg", + }; + const message = buildOptimisticEditedMessage({ + requestContent: [{ type: "file", file_id: "existing-1" }], + originalMessage: makeUserMessage([existingFilePart]), + attachmentMediaTypes: new Map([["existing-1", "text/plain"]]), + }); + + expect(message.content).toEqual([existingFilePart]); + }); +}); diff --git a/site/src/api/queries/chatMessageEdits.ts b/site/src/api/queries/chatMessageEdits.ts new file mode 100644 index 0000000000..2fbefa1274 --- /dev/null +++ b/site/src/api/queries/chatMessageEdits.ts @@ -0,0 +1,148 @@ +import type { InfiniteData } from "react-query"; +import type * as TypesGen from "#/api/typesGenerated"; + +const buildOptimisticEditedContent = ({ + requestContent, + originalMessage, + attachmentMediaTypes, +}: { + requestContent: readonly TypesGen.ChatInputPart[]; + originalMessage: TypesGen.ChatMessage; + attachmentMediaTypes?: ReadonlyMap; +}): readonly TypesGen.ChatMessagePart[] => { + const existingFilePartsByID = new Map(); + for (const part of originalMessage.content ?? []) { + if (part.type === "file" && part.file_id) { + existingFilePartsByID.set(part.file_id, part); + } + } + + return requestContent.map((part): TypesGen.ChatMessagePart => { + if (part.type === "text") { + return { type: "text", text: part.text ?? "" }; + } + if (part.type === "file-reference") { + return { + type: "file-reference", + file_name: part.file_name ?? "", + start_line: part.start_line ?? 1, + end_line: part.end_line ?? 1, + content: part.content ?? "", + }; + } + const fileId = part.file_id ?? ""; + return ( + existingFilePartsByID.get(fileId) ?? { + type: "file", + file_id: part.file_id, + media_type: + attachmentMediaTypes?.get(fileId) ?? "application/octet-stream", + } + ); + }); +}; + +export const buildOptimisticEditedMessage = ({ + requestContent, + originalMessage, + attachmentMediaTypes, +}: { + requestContent: readonly TypesGen.ChatInputPart[]; + originalMessage: TypesGen.ChatMessage; + attachmentMediaTypes?: ReadonlyMap; +}): TypesGen.ChatMessage => ({ + ...originalMessage, + content: buildOptimisticEditedContent({ + requestContent, + originalMessage, + attachmentMediaTypes, + }), +}); + +const sortMessagesDescending = ( + messages: readonly TypesGen.ChatMessage[], +): TypesGen.ChatMessage[] => [...messages].sort((a, b) => b.id - a.id); + +const upsertFirstPageMessage = ( + messages: readonly TypesGen.ChatMessage[], + message: TypesGen.ChatMessage, +): TypesGen.ChatMessage[] => { + const byID = new Map( + messages.map((existingMessage) => [existingMessage.id, existingMessage]), + ); + byID.set(message.id, message); + return sortMessagesDescending(Array.from(byID.values())); +}; + +export const projectEditedConversationIntoCache = ({ + currentData, + editedMessageId, + replacementMessage, + queuedMessages, +}: { + currentData: InfiniteData | undefined; + editedMessageId: number; + replacementMessage?: TypesGen.ChatMessage; + queuedMessages?: readonly TypesGen.ChatQueuedMessage[]; +}): InfiniteData | undefined => { + if (!currentData?.pages?.length) { + return currentData; + } + + const truncatedPages = currentData.pages.map((page, pageIndex) => { + const truncatedMessages = page.messages.filter( + (message) => message.id < editedMessageId, + ); + const nextPage = { + ...page, + ...(pageIndex === 0 && queuedMessages !== undefined + ? { queued_messages: queuedMessages } + : {}), + }; + if (pageIndex !== 0 || !replacementMessage) { + return { ...nextPage, messages: truncatedMessages }; + } + return { + ...nextPage, + messages: upsertFirstPageMessage(truncatedMessages, replacementMessage), + }; + }); + + return { + ...currentData, + pages: truncatedPages, + }; +}; + +export const reconcileEditedMessageInCache = ({ + currentData, + optimisticMessageId, + responseMessage, +}: { + currentData: InfiniteData | undefined; + optimisticMessageId: number; + responseMessage: TypesGen.ChatMessage; +}): InfiniteData | undefined => { + if (!currentData?.pages?.length) { + return currentData; + } + + const replacedPages = currentData.pages.map((page, pageIndex) => { + const preservedMessages = page.messages.filter( + (message) => + message.id !== optimisticMessageId && message.id !== responseMessage.id, + ); + if (pageIndex !== 0) { + return { ...page, messages: preservedMessages }; + } + return { + ...page, + messages: upsertFirstPageMessage(preservedMessages, responseMessage), + }; + }); + + return { + ...currentData, + pages: replacedPages, + }; +}; diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 1201fe9cea..148b23a9b3 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -2,6 +2,7 @@ import { QueryClient } from "react-query"; import { describe, expect, it, vi } from "vitest"; import { API } from "#/api/api"; import type * as TypesGen from "#/api/typesGenerated"; +import { buildOptimisticEditedMessage } from "./chatMessageEdits"; import { archiveChat, cancelChatListRefetches, @@ -795,14 +796,44 @@ describe("mutation invalidation scope", () => { content: [{ type: "text" as const, text: `msg ${id}` }], }); + const makeQueuedMessage = ( + chatId: string, + id: number, + ): TypesGen.ChatQueuedMessage => ({ + id, + chat_id: chatId, + created_at: `2025-01-01T00:10:${String(id).padStart(2, "0")}Z`, + content: [{ type: "text" as const, text: `queued ${id}` }], + }); + const editReq = { content: [{ type: "text" as const, text: "edited" }], }; - it("editChatMessage optimistically removes truncated messages from cache", async () => { + const requireMessage = ( + messages: readonly TypesGen.ChatMessage[], + messageId: number, + ): TypesGen.ChatMessage => { + const message = messages.find((candidate) => candidate.id === messageId); + if (!message) { + throw new Error(`missing message ${messageId}`); + } + return message; + }; + + const buildOptimisticMessage = (message: TypesGen.ChatMessage) => + buildOptimisticEditedMessage({ + originalMessage: message, + requestContent: editReq.content, + }); + + it("editChatMessage writes the optimistic replacement into cache", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id)); + const optimisticMessage = buildOptimisticMessage( + requireMessage(messages, 3), + ); queryClient.setQueryData(chatMessagesKey(chatId), { pages: [{ messages, queued_messages: [], has_more: false }], @@ -812,18 +843,58 @@ describe("mutation invalidation scope", () => { const mutation = editChatMessage(queryClient, chatId); const context = await mutation.onMutate({ messageId: 3, + optimisticMessage, req: editReq, }); const data = queryClient.getQueryData(chatMessagesKey(chatId)); - expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([2, 1]); + expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 3, 2, 1, + ]); + expect(data?.pages[0]?.messages[0]?.content).toEqual( + optimisticMessage.content, + ); expect(context?.previousData?.pages[0]?.messages).toHaveLength(5); }); + it("editChatMessage clears queued messages in cache during optimistic history edit", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id)); + const optimisticMessage = buildOptimisticMessage( + requireMessage(messages, 3), + ); + const queuedMessages = [makeQueuedMessage(chatId, 11)]; + + queryClient.setQueryData(chatMessagesKey(chatId), { + pages: [ + { + messages, + queued_messages: queuedMessages, + has_more: false, + }, + ], + pageParams: [undefined], + }); + + const mutation = editChatMessage(queryClient, chatId); + await mutation.onMutate({ + messageId: 3, + optimisticMessage, + req: editReq, + }); + + const data = queryClient.getQueryData(chatMessagesKey(chatId)); + expect(data?.pages[0]?.queued_messages).toEqual([]); + }); + it("editChatMessage restores cache on error", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id)); + const optimisticMessage = buildOptimisticMessage( + requireMessage(messages, 3), + ); queryClient.setQueryData(chatMessagesKey(chatId), { pages: [{ messages, queued_messages: [], has_more: false }], @@ -833,22 +904,85 @@ describe("mutation invalidation scope", () => { const mutation = editChatMessage(queryClient, chatId); const context = await mutation.onMutate({ messageId: 3, + optimisticMessage, req: editReq, }); expect( queryClient.getQueryData(chatMessagesKey(chatId))?.pages[0] ?.messages, - ).toHaveLength(2); + ).toHaveLength(3); mutation.onError( new Error("network failure"), - { messageId: 3, req: editReq }, + { messageId: 3, optimisticMessage, req: editReq }, context, ); const data = queryClient.getQueryData(chatMessagesKey(chatId)); - expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([5, 4, 3, 2, 1]); + expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 5, 4, 3, 2, 1, + ]); + }); + + it("editChatMessage preserves websocket-upserted newer messages on success", async () => { + const queryClient = createTestQueryClient(); + const chatId = "chat-1"; + const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id)); + const optimisticMessage = buildOptimisticMessage( + requireMessage(messages, 3), + ); + const responseMessage = { + ...makeMsg(chatId, 9), + content: [{ type: "text" as const, text: "edited authoritative" }], + }; + const websocketMessage = { + ...makeMsg(chatId, 10), + content: [{ type: "text" as const, text: "assistant follow-up" }], + role: "assistant" as const, + }; + + queryClient.setQueryData(chatMessagesKey(chatId), { + pages: [{ messages, queued_messages: [], has_more: false }], + pageParams: [undefined], + }); + + const mutation = editChatMessage(queryClient, chatId); + await mutation.onMutate({ + messageId: 3, + optimisticMessage, + req: editReq, + }); + queryClient.setQueryData( + chatMessagesKey(chatId), + (current) => { + if (!current) { + return current; + } + return { + ...current, + pages: [ + { + ...current.pages[0], + messages: [websocketMessage, ...current.pages[0].messages], + }, + ...current.pages.slice(1), + ], + }; + }, + ); + mutation.onSuccess( + { message: responseMessage }, + { messageId: 3, optimisticMessage, req: editReq }, + ); + + const data = queryClient.getQueryData(chatMessagesKey(chatId)); + expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 10, 9, 2, 1, + ]); + expect(data?.pages[0]?.messages[1]?.content).toEqual( + responseMessage.content, + ); }); it("editChatMessage onMutate is a no-op when cache is empty", async () => { @@ -890,13 +1024,14 @@ describe("mutation invalidation scope", () => { expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([3, 2, 1]); }); - it("editChatMessage onMutate filters across multiple pages", async () => { + it("editChatMessage onMutate updates the first page and preserves older pages", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; // Page 0 (newest): IDs 10–6. Page 1 (older): IDs 5–1. const page0 = [10, 9, 8, 7, 6].map((id) => makeMsg(chatId, id)); const page1 = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id)); + const optimisticMessage = buildOptimisticMessage(requireMessage(page0, 7)); queryClient.setQueryData(chatMessagesKey(chatId), { pages: [ @@ -907,19 +1042,28 @@ describe("mutation invalidation scope", () => { }); const mutation = editChatMessage(queryClient, chatId); - await mutation.onMutate({ messageId: 7, req: editReq }); + await mutation.onMutate({ + messageId: 7, + optimisticMessage, + req: editReq, + }); const data = queryClient.getQueryData(chatMessagesKey(chatId)); - // Page 0: only ID 6 survives (< 7). - expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([6]); - // Page 1: all survive (all < 7). - expect(data?.pages[1]?.messages.map((m) => m.id)).toEqual([5, 4, 3, 2, 1]); + expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 7, 6, + ]); + expect(data?.pages[1]?.messages.map((message) => message.id)).toEqual([ + 5, 4, 3, 2, 1, + ]); }); - it("editChatMessage onMutate editing the first message empties all pages", async () => { + it("editChatMessage onMutate keeps the optimistic replacement when editing the first message", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id)); + const optimisticMessage = buildOptimisticMessage( + requireMessage(messages, 1), + ); queryClient.setQueryData(chatMessagesKey(chatId), { pages: [{ messages, queued_messages: [], has_more: false }], @@ -927,20 +1071,25 @@ describe("mutation invalidation scope", () => { }); const mutation = editChatMessage(queryClient, chatId); - await mutation.onMutate({ messageId: 1, req: editReq }); + await mutation.onMutate({ + messageId: 1, + optimisticMessage, + req: editReq, + }); const data = queryClient.getQueryData(chatMessagesKey(chatId)); - // All messages have id >= 1, so the page is empty. - expect(data?.pages[0]?.messages).toHaveLength(0); - // Sibling fields survive the spread. + expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([1]); expect(data?.pages[0]?.queued_messages).toEqual([]); expect(data?.pages[0]?.has_more).toBe(false); }); - it("editChatMessage onMutate editing the latest message keeps earlier ones", async () => { + it("editChatMessage onMutate keeps earlier messages when editing the latest message", async () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; const messages = [5, 4, 3, 2, 1].map((id) => makeMsg(chatId, id)); + const optimisticMessage = buildOptimisticMessage( + requireMessage(messages, 5), + ); queryClient.setQueryData(chatMessagesKey(chatId), { pages: [{ messages, queued_messages: [], has_more: false }], @@ -948,10 +1097,19 @@ describe("mutation invalidation scope", () => { }); const mutation = editChatMessage(queryClient, chatId); - await mutation.onMutate({ messageId: 5, req: editReq }); + await mutation.onMutate({ + messageId: 5, + optimisticMessage, + req: editReq, + }); const data = queryClient.getQueryData(chatMessagesKey(chatId)); - expect(data?.pages[0]?.messages.map((m) => m.id)).toEqual([4, 3, 2, 1]); + expect(data?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 5, 4, 3, 2, 1, + ]); + expect(data?.pages[0]?.messages[0]?.content).toEqual( + optimisticMessage.content, + ); }); it("interruptChat does not invalidate unrelated queries", async () => { diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index 7b4076a0e3..6e9677aa7e 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -6,6 +6,10 @@ import type { import { API } from "#/api/api"; import type * as TypesGen from "#/api/typesGenerated"; import type { UsePaginatedQueryOptions } from "#/hooks/usePaginatedQuery"; +import { + projectEditedConversationIntoCache, + reconcileEditedMessageInCache, +} from "./chatMessageEdits"; export const chatsKey = ["chats"] as const; export const chatKey = (chatId: string) => ["chats", chatId] as const; @@ -601,13 +605,21 @@ export const createChatMessage = ( type EditChatMessageMutationArgs = { messageId: number; + optimisticMessage?: TypesGen.ChatMessage; req: TypesGen.EditChatMessageRequest; }; +type EditChatMessageMutationContext = { + previousData?: InfiniteData | undefined; +}; + export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ mutationFn: ({ messageId, req }: EditChatMessageMutationArgs) => API.experimental.editChatMessage(chatId, messageId, req), - onMutate: async ({ messageId }: EditChatMessageMutationArgs) => { + onMutate: async ({ + messageId, + optimisticMessage, + }: EditChatMessageMutationArgs): Promise => { // Cancel in-flight refetches so they don't overwrite the // optimistic update before the mutation completes. await queryClient.cancelQueries({ @@ -619,40 +631,23 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ InfiniteData >(chatMessagesKey(chatId)); - // Optimistically remove the edited message and everything - // after it. The server soft-deletes these and inserts a - // replacement with a new ID. Without this, the WebSocket - // handler's upsertCacheMessages adds new messages to the - // React Query cache without removing the soft-deleted ones, - // causing deleted messages to flash back into view until - // the full REST refetch resolves. queryClient.setQueryData< InfiniteData | undefined - >(chatMessagesKey(chatId), (current) => { - if (!current?.pages?.length) { - return current; - } - return { - ...current, - pages: current.pages.map((page) => ({ - ...page, - messages: page.messages.filter((m) => m.id < messageId), - })), - }; - }); + >(chatMessagesKey(chatId), (current) => + projectEditedConversationIntoCache({ + currentData: current, + editedMessageId: messageId, + replacementMessage: optimisticMessage, + queuedMessages: [], + }), + ); return { previousData }; }, onError: ( _error: unknown, _variables: EditChatMessageMutationArgs, - context: - | { - previousData?: - | InfiniteData - | undefined; - } - | undefined, + context: EditChatMessageMutationContext | undefined, ) => { // Restore the cache on failure so the user sees the // original messages again. @@ -660,6 +655,20 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ queryClient.setQueryData(chatMessagesKey(chatId), context.previousData); } }, + onSuccess: ( + response: TypesGen.EditChatMessageResponse, + variables: EditChatMessageMutationArgs, + ) => { + queryClient.setQueryData< + InfiniteData | undefined + >(chatMessagesKey(chatId), (current) => + reconcileEditedMessageInCache({ + currentData: current, + optimisticMessageId: variables.messageId, + responseMessage: response.message, + }), + ); + }, onSettled: () => { // Always reconcile with the server regardless of whether // the mutation succeeded or failed. On success this picks diff --git a/site/src/pages/AgentsPage/AgentChatPage.test.ts b/site/src/pages/AgentsPage/AgentChatPage.test.ts index 44535d8041..2c1287e36a 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.test.ts +++ b/site/src/pages/AgentsPage/AgentChatPage.test.ts @@ -4,9 +4,12 @@ import { beforeEach, describe, expect, it, vi } from "vitest"; import { draftInputStorageKeyPrefix, getPersistedDraftInputValue, + restoreOptimisticRequestSnapshot, useConversationEditingState, } from "./AgentChatPage"; import type { ChatMessageInputRef } from "./components/AgentChatInput"; +import { createChatStore } from "./components/ChatConversation/chatStore"; +import type { PendingAttachment } from "./components/ChatPageContent"; type MockChatInputHandle = { handle: ChatMessageInputRef; @@ -84,6 +87,41 @@ describe("getPersistedDraftInputValue", () => { }); }); +describe("restoreOptimisticRequestSnapshot", () => { + it("restores queued messages, stream output, status, and stream error", () => { + const store = createChatStore(); + store.setQueuedMessages([ + { + id: 9, + chat_id: "chat-abc-123", + created_at: "2025-01-01T00:00:00.000Z", + content: [{ type: "text" as const, text: "queued" }], + }, + ]); + store.setChatStatus("running"); + store.applyMessagePart({ type: "text", text: "partial response" }); + store.setStreamError({ kind: "generic", message: "old error" }); + const previousSnapshot = store.getSnapshot(); + + store.batch(() => { + store.setQueuedMessages([]); + store.setChatStatus("pending"); + store.clearStreamState(); + store.clearStreamError(); + }); + + restoreOptimisticRequestSnapshot(store, previousSnapshot); + + const restoredSnapshot = store.getSnapshot(); + expect(restoredSnapshot.queuedMessages).toEqual( + previousSnapshot.queuedMessages, + ); + expect(restoredSnapshot.chatStatus).toBe(previousSnapshot.chatStatus); + expect(restoredSnapshot.streamState).toBe(previousSnapshot.streamState); + expect(restoredSnapshot.streamError).toEqual(previousSnapshot.streamError); + }); +}); + describe("useConversationEditingState", () => { const chatID = "chat-abc-123"; const expectedKey = `${draftInputStorageKeyPrefix}${chatID}`; @@ -327,6 +365,64 @@ describe("useConversationEditingState", () => { unmount(); }); + it("forwards pending attachments through history-edit send", async () => { + const { result, onSend, unmount } = renderEditing(); + const attachments: PendingAttachment[] = [ + { fileId: "file-1", mediaType: "image/png" }, + ]; + + act(() => { + result.current.handleEditUserMessage(7, "hello"); + }); + + await act(async () => { + await result.current.handleSendFromInput("hello", attachments); + }); + + expect(onSend).toHaveBeenCalledWith("hello", attachments, 7); + unmount(); + }); + + it("restores the edit draft and file-block seed when an edit submission fails", async () => { + const { result, onSend, unmount } = renderEditing(); + const mockInput = createMockChatInputHandle("edited message"); + const fileBlocks = [ + { type: "file", file_id: "file-1", media_type: "image/png" }, + ] as const; + result.current.chatInputRef.current = mockInput.handle; + onSend.mockRejectedValueOnce(new Error("boom")); + const editorState = JSON.stringify({ + root: { + children: [ + { + children: [{ text: "edited message" }], + type: "paragraph", + }, + ], + type: "root", + }, + }); + + act(() => { + result.current.handleEditUserMessage(7, "edited message", fileBlocks); + result.current.handleContentChange("edited message", editorState, false); + }); + + await act(async () => { + await expect( + result.current.handleSendFromInput("edited message"), + ).rejects.toThrow("boom"); + }); + + expect(mockInput.clear).toHaveBeenCalled(); + expect(result.current.inputValueRef.current).toBe("edited message"); + expect(result.current.editingMessageId).toBe(7); + expect(result.current.editingFileBlocks).toEqual(fileBlocks); + expect(result.current.editorInitialValue).toBe("edited message"); + expect(result.current.initialEditorState).toBe(editorState); + unmount(); + }); + it("clears the composer and persisted draft after a successful send", async () => { localStorage.setItem(expectedKey, "draft to clear"); const { result, onSend, unmount } = renderEditing(); diff --git a/site/src/pages/AgentsPage/AgentChatPage.tsx b/site/src/pages/AgentsPage/AgentChatPage.tsx index 43580f86fa..920878a582 100644 --- a/site/src/pages/AgentsPage/AgentChatPage.tsx +++ b/site/src/pages/AgentsPage/AgentChatPage.tsx @@ -11,6 +11,7 @@ import { toast } from "sonner"; import type { UrlTransform } from "streamdown"; import { API, watchWorkspace } from "#/api/api"; import { isApiError } from "#/api/errors"; +import { buildOptimisticEditedMessage } from "#/api/queries/chatMessageEdits"; import { chat, chatDesktopEnabled, @@ -51,11 +52,14 @@ import { getWorkspaceAgent, } from "./components/ChatConversation/chatHelpers"; import { + type ChatStore, + type ChatStoreState, selectChatStatus, useChatSelector, useChatStore, } from "./components/ChatConversation/chatStore"; import { useWorkspaceCreationWatcher } from "./components/ChatConversation/useWorkspaceCreationWatcher"; +import type { PendingAttachment } from "./components/ChatPageContent"; import { getDefaultMCPSelection, getSavedMCPSelection, @@ -101,12 +105,47 @@ export function getPersistedDraftInputValue( ).text; } +/** @internal Exported for testing. */ +export const restoreOptimisticRequestSnapshot = ( + store: Pick< + ChatStore, + | "batch" + | "setChatStatus" + | "setQueuedMessages" + | "setStreamError" + | "setStreamState" + >, + snapshot: Pick< + ChatStoreState, + "chatStatus" | "queuedMessages" | "streamError" | "streamState" + >, +): void => { + store.batch(() => { + store.setQueuedMessages(snapshot.queuedMessages); + store.setChatStatus(snapshot.chatStatus); + store.setStreamState(snapshot.streamState); + store.setStreamError(snapshot.streamError); + }); +}; + +const buildAttachmentMediaTypes = ( + attachments?: readonly PendingAttachment[], +): ReadonlyMap | undefined => { + if (!attachments?.length) { + return undefined; + } + + return new Map( + attachments.map(({ fileId, mediaType }) => [fileId, mediaType]), + ); +}; + /** @internal Exported for testing. */ export function useConversationEditingState(deps: { chatID: string | undefined; onSend: ( message: string, - fileIds?: string[], + attachments?: readonly PendingAttachment[], editedMessageID?: number, ) => Promise; onDeleteQueuedMessage: (id: number) => Promise; @@ -130,6 +169,9 @@ export function useConversationEditingState(deps: { }; }, ); + const serializedEditorStateRef = useRef( + initialEditorState, + ); // Monotonic counter to force LexicalComposer remount. const [remountKey, setRemountKey] = useState(0); @@ -176,6 +218,7 @@ export function useConversationEditingState(deps: { editorInitialValue: text, initialEditorState: undefined, }); + serializedEditorStateRef.current = undefined; setRemountKey((k) => k + 1); inputValueRef.current = text; setEditingFileBlocks(fileBlocks ?? []); @@ -188,6 +231,7 @@ export function useConversationEditingState(deps: { editorInitialValue: savedText, initialEditorState: savedState, }); + serializedEditorStateRef.current = savedState; setRemountKey((k) => k + 1); inputValueRef.current = savedText; setEditingMessageId(null); @@ -221,6 +265,7 @@ export function useConversationEditingState(deps: { editorInitialValue: text, initialEditorState: undefined, }); + serializedEditorStateRef.current = undefined; setRemountKey((k) => k + 1); inputValueRef.current = text; setEditingFileBlocks(fileBlocks); @@ -233,6 +278,7 @@ export function useConversationEditingState(deps: { editorInitialValue: savedText, initialEditorState: savedState, }); + serializedEditorStateRef.current = savedState; setRemountKey((k) => k + 1); inputValueRef.current = savedText; setEditingQueuedMessageID(null); @@ -240,25 +286,48 @@ export function useConversationEditingState(deps: { setEditingFileBlocks([]); }; - // Wraps the parent onSend to clear local input/editing state - // and handle queue-edit deletion. - const handleSendFromInput = async (message: string, fileIds?: string[]) => { - const editedMessageID = - editingMessageId !== null ? editingMessageId : undefined; - const queueEditID = editingQueuedMessageID; + // Clears the composer for an in-flight history edit and + // returns a rollback function that restores the editing draft + // if the send fails. + const clearInputForHistoryEdit = (message: string) => { + const snapshot = { + editorState: serializedEditorStateRef.current, + fileBlocks: editingFileBlocks, + messageId: editingMessageId, + }; - await onSend(message, fileIds, editedMessageID); - // Clear input and editing state on success. + chatInputRef.current?.clear(); + inputValueRef.current = ""; + setEditingMessageId(null); + + return () => { + setDraftState({ + editorInitialValue: message, + initialEditorState: snapshot.editorState, + }); + serializedEditorStateRef.current = snapshot.editorState; + setRemountKey((k) => k + 1); + inputValueRef.current = message; + setEditingMessageId(snapshot.messageId); + setEditingFileBlocks(snapshot.fileBlocks); + }; + }; + + // Clears all input and editing state after a successful send. + const finalizeSuccessfulSend = ( + editedMessageID: number | undefined, + queueEditID: number | null, + ) => { chatInputRef.current?.clear(); if (!isMobileViewport()) { chatInputRef.current?.focus(); } inputValueRef.current = ""; + serializedEditorStateRef.current = undefined; if (draftStorageKey) { localStorage.removeItem(draftStorageKey); } - if (editingMessageId !== null) { - setEditingMessageId(null); + if (editedMessageID !== undefined) { setDraftBeforeHistoryEdit(null); setEditingFileBlocks([]); } @@ -270,12 +339,41 @@ export function useConversationEditingState(deps: { } }; + // Wraps the parent onSend to clear local input/editing state + // and handle queue-edit deletion. + const handleSendFromInput = async ( + message: string, + attachments?: readonly PendingAttachment[], + ) => { + const editedMessageID = + editingMessageId !== null ? editingMessageId : undefined; + const queueEditID = editingQueuedMessageID; + const sendPromise = onSend(message, attachments, editedMessageID); + + // For history edits, clear input immediately and prepare + // a rollback in case the send fails. + const rollback = + editedMessageID !== undefined + ? clearInputForHistoryEdit(message) + : undefined; + + try { + await sendPromise; + } catch (error) { + rollback?.(); + throw error; + } + + finalizeSuccessfulSend(editedMessageID, queueEditID); + }; + const handleContentChange = ( content: string, serializedEditorState: string, hasFileReferences: boolean, ) => { inputValueRef.current = content; + serializedEditorStateRef.current = serializedEditorState; // Don't overwrite the persisted draft while editing a // history or queued message — the original draft (possibly @@ -430,9 +528,6 @@ const AgentChatPage: FC = () => { } = useOutletContext(); const queryClient = useQueryClient(); const [selectedModel, setSelectedModel] = useState(""); - const [pendingEditMessageId, setPendingEditMessageId] = useState< - number | null - >(null); const scrollToBottomRef = useRef<(() => void) | null>(null); const chatInputRef = useRef(null); const inputValueRef = useRef( @@ -775,7 +870,7 @@ const AgentChatPage: FC = () => { const handleSend = async ( message: string, - fileIds?: string[], + attachments?: readonly PendingAttachment[], editedMessageID?: number, ) => { const chatInputHandle = ( @@ -790,7 +885,9 @@ const AgentChatPage: FC = () => { (p) => p.type === "file-reference", ); const hasContent = - message.trim() || (fileIds && fileIds.length > 0) || hasFileReferences; + message.trim() || + (attachments && attachments.length > 0) || + hasFileReferences; if (!hasContent || isSubmissionPending || !agentId || !hasModelOptions) { return; } @@ -818,28 +915,41 @@ const AgentChatPage: FC = () => { } } - // Add pre-uploaded file references. - if (fileIds && fileIds.length > 0) { - for (const fileId of fileIds) { + // Add pre-uploaded file attachments. + if (attachments && attachments.length > 0) { + for (const { fileId } of attachments) { content.push({ type: "file", file_id: fileId }); } } if (editedMessageID !== undefined) { const request: TypesGen.EditChatMessageRequest = { content }; + const originalEditedMessage = chatMessagesList?.find( + (existingMessage) => existingMessage.id === editedMessageID, + ); + const optimisticMessage = originalEditedMessage + ? buildOptimisticEditedMessage({ + requestContent: request.content, + originalMessage: originalEditedMessage, + attachmentMediaTypes: buildAttachmentMediaTypes(attachments), + }) + : undefined; + const previousSnapshot = store.getSnapshot(); clearChatErrorReason(agentId); clearStreamError(); - setPendingEditMessageId(editedMessageID); + store.batch(() => { + store.setQueuedMessages([]); + store.setChatStatus("running"); + store.clearStreamState(); + }); scrollToBottomRef.current?.(); try { await editMessage({ messageId: editedMessageID, + optimisticMessage, req: request, }); - store.clearStreamState(); - store.setChatStatus("running"); - setPendingEditMessageId(null); } catch (error) { - setPendingEditMessageId(null); + restoreOptimisticRequestSnapshot(store, previousSnapshot); handleUsageLimitError(error); throw error; } @@ -918,10 +1028,8 @@ const AgentChatPage: FC = () => { const handlePromoteQueuedMessage = async (id: number) => { const previousSnapshot = store.getSnapshot(); - const previousQueuedMessages = previousSnapshot.queuedMessages; - const previousChatStatus = previousSnapshot.chatStatus; store.setQueuedMessages( - previousQueuedMessages.filter((message) => message.id !== id), + previousSnapshot.queuedMessages.filter((message) => message.id !== id), ); store.clearStreamState(); if (agentId) { @@ -937,8 +1045,7 @@ const AgentChatPage: FC = () => { store.upsertDurableMessage(promotedMessage); upsertCacheMessages([promotedMessage]); } catch (error) { - store.setQueuedMessages(previousQueuedMessages); - store.setChatStatus(previousChatStatus); + restoreOptimisticRequestSnapshot(store, previousSnapshot); handleUsageLimitError(error); throw error; } @@ -1133,7 +1240,6 @@ const AgentChatPage: FC = () => { workspaceAgent={workspaceAgent} store={store} editing={editing} - pendingEditMessageId={pendingEditMessageId} effectiveSelectedModel={effectiveSelectedModel} setSelectedModel={setSelectedModel} modelOptions={modelOptions} diff --git a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx index 2b1837bfd3..bd0a5c90e8 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.stories.tsx @@ -113,7 +113,6 @@ const StoryAgentChatPageView: FC = ({ editing, ...overrides }) => { parentChat: undefined as TypesGen.Chat | undefined, isArchived: false, store: createChatStore(), - pendingEditMessageId: null as number | null, effectiveSelectedModel: defaultModelConfigID, setSelectedModel: fn(), modelOptions: defaultModelOptions, @@ -474,16 +473,6 @@ const buildStoreWithMessages = ( return store; }; -const gapTestStore = createChatStore(); -gapTestStore.replaceMessages([ - buildMessage(1, "user", "Explain the layout."), - buildMessage(2, "assistant", "Here is the explanation."), - buildMessage(3, "user", "Can you elaborate?"), -]); -gapTestStore.applyMessageParts([ - { type: "text", text: "Certainly, here are more details..." }, -]); - // --------------------------------------------------------------------------- // Editing flow stories // --------------------------------------------------------------------------- @@ -515,42 +504,6 @@ export const EditingMessage: Story = { ), }; -/** The saving state while an edit is in progress — shows the pending - * indicator on the message being saved. */ -export const EditingSaving: Story = { - render: () => ( - - ), -}; - -export const ConsistentGapBetweenTimelineAndStream: Story = { - render: () => , - play: async ({ canvasElement }) => { - const wrapper = canvasElement.querySelector( - '[data-testid="chat-timeline-wrapper"]', - ); - expect(wrapper).not.toBeNull(); - - const outerGap = window.getComputedStyle(wrapper!).rowGap; - expect(outerGap).toBe("8px"); - - const timeline = wrapper!.querySelector( - '[data-testid="conversation-timeline"]', - ); - expect(timeline).not.toBeNull(); - const innerGap = window.getComputedStyle(timeline!).rowGap; - expect(innerGap).toBe("8px"); - }, -}; - // --------------------------------------------------------------------------- // AgentChatPageNotFoundView stories // --------------------------------------------------------------------------- diff --git a/site/src/pages/AgentsPage/AgentChatPageView.tsx b/site/src/pages/AgentsPage/AgentChatPageView.tsx index 4079d79749..f8baa75568 100644 --- a/site/src/pages/AgentsPage/AgentChatPageView.tsx +++ b/site/src/pages/AgentsPage/AgentChatPageView.tsx @@ -36,6 +36,7 @@ import { import type { useChatStore } from "./components/ChatConversation/chatStore"; import type { ModelSelectorOption } from "./components/ChatElements"; import { DesktopPanelContext } from "./components/ChatElements/tools/DesktopPanelContext"; +import type { PendingAttachment } from "./components/ChatPageContent"; import { ChatPageInput, ChatPageTimeline } from "./components/ChatPageContent"; import { ChatScrollContainer } from "./components/ChatScrollContainer"; import { ChatTopBar } from "./components/ChatTopBar"; @@ -69,7 +70,10 @@ interface EditingState { fileBlocks: readonly ChatMessagePart[], ) => void; handleCancelQueueEdit: () => void; - handleSendFromInput: (message: string, fileIds?: string[]) => void; + handleSendFromInput: ( + message: string, + attachments?: readonly PendingAttachment[], + ) => void; handleContentChange: ( content: string, serializedEditorState: string, @@ -92,7 +96,6 @@ interface AgentChatPageViewProps { // Editing state. editing: EditingState; - pendingEditMessageId: number | null; // Model/input configuration. effectiveSelectedModel: string; @@ -179,7 +182,6 @@ export const AgentChatPageView: FC = ({ workspace, store, editing, - pendingEditMessageId, effectiveSelectedModel, setSelectedModel, modelOptions, @@ -387,7 +389,6 @@ export const AgentChatPageView: FC = ({ persistedError={persistedError} onEditUserMessage={editing.handleEditUserMessage} editingMessageId={editing.editingMessageId} - savingMessageId={pendingEditMessageId} urlTransform={urlTransform} mcpServers={mcpServers} /> diff --git a/site/src/pages/AgentsPage/components/AgentChatInput.tsx b/site/src/pages/AgentsPage/components/AgentChatInput.tsx index 681dab3785..f5530d71ff 100644 --- a/site/src/pages/AgentsPage/components/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/components/AgentChatInput.tsx @@ -668,9 +668,8 @@ export const AgentChatInput: FC = ({
- {isLoading - ? "Saving edit..." - : "Editing will delete all subsequent messages and restart the conversation here."} + Editing will delete all subsequent messages and restart the + conversation here.
)} {hasFileBlocks && ( @@ -704,7 +692,6 @@ const StickyUserMessage = memo<{ fileBlocks?: readonly TypesGen.ChatMessagePart[], ) => void; editingMessageId?: number | null; - savingMessageId?: number | null; isAfterEditingMessage?: boolean; }>( ({ @@ -712,7 +699,6 @@ const StickyUserMessage = memo<{ parsed, onEditUserMessage, editingMessageId, - savingMessageId, isAfterEditingMessage = false, }) => { const [isStuck, setIsStuck] = useState(false); @@ -937,7 +923,6 @@ const StickyUserMessage = memo<{ parsed={parsed} onEditUserMessage={handleEditUserMessage} editingMessageId={editingMessageId} - savingMessageId={savingMessageId} isAfterEditingMessage={isAfterEditingMessage} /> @@ -980,7 +965,6 @@ const StickyUserMessage = memo<{ parsed={parsed} onEditUserMessage={handleEditUserMessage} editingMessageId={editingMessageId} - savingMessageId={savingMessageId} isAfterEditingMessage={isAfterEditingMessage} fadeFromBottom /> @@ -1002,7 +986,6 @@ interface ConversationTimelineProps { fileBlocks?: readonly TypesGen.ChatMessagePart[], ) => void; editingMessageId?: number | null; - savingMessageId?: number | null; urlTransform?: UrlTransform; mcpServers?: readonly TypesGen.MCPServerConfig[]; computerUseSubagentIds?: Set; @@ -1016,7 +999,6 @@ export const ConversationTimeline = memo( subagentTitles, onEditUserMessage, editingMessageId, - savingMessageId, urlTransform, mcpServers, computerUseSubagentIds, @@ -1053,7 +1035,6 @@ export const ConversationTimeline = memo( parsed={parsed} onEditUserMessage={onEditUserMessage} editingMessageId={editingMessageId} - savingMessageId={savingMessageId} isAfterEditingMessage={afterEditingMessageIds.has(message.id)} /> ); @@ -1067,7 +1048,6 @@ export const ConversationTimeline = memo( key={message.id} message={message} parsed={parsed} - savingMessageId={savingMessageId} urlTransform={urlTransform} isAfterEditingMessage={afterEditingMessageIds.has(message.id)} hideActions={!isLastInChain} diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts index cd4504178e..45dd3d4930 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.createStore.test.ts @@ -190,6 +190,27 @@ describe("setChatStatus", () => { }); }); +// --------------------------------------------------------------------------- +// setStreamState +// --------------------------------------------------------------------------- + +describe("setStreamState", () => { + it("does not notify when setting the same stream state reference", () => { + const store = createChatStore(); + store.applyMessagePart({ type: "text", text: "hello" }); + const streamState = store.getSnapshot().streamState; + expect(streamState).not.toBeNull(); + + let notified = false; + store.subscribe(() => { + notified = true; + }); + + store.setStreamState(streamState); + expect(notified).toBe(false); + }); +}); + // --------------------------------------------------------------------------- // setStreamError / clearStreamError // --------------------------------------------------------------------------- diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index 01105e1478..50c4d74a18 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -4213,6 +4213,96 @@ describe("store/cache desync protection", () => { expect(result.current.orderedMessageIDs).toEqual([1]); }); }); + + it("reflects optimistic and authoritative history-edit cache updates through the normal sync effect", async () => { + immediateAnimationFrame(); + + const chatID = "chat-local-edit-sync"; + const msg1 = makeMessage(chatID, 1, "user", "first"); + const msg2 = makeMessage(chatID, 2, "assistant", "second"); + const msg3 = makeMessage(chatID, 3, "user", "third"); + const optimisticReplacement = { + ...msg3, + content: [{ type: "text" as const, text: "edited draft" }], + }; + const authoritativeReplacement = makeMessage(chatID, 9, "user", "edited"); + + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper: FC = ({ children }) => ( + {children} + ); + const initialOptions = { + chatID, + chatMessages: [msg1, msg2, msg3], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [msg1, msg2, msg3], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [] as TypesGen.ChatQueuedMessage[], + setChatErrorReason: vi.fn(), + clearChatErrorReason: vi.fn(), + }; + + const { result, rerender } = renderHook( + (options: Parameters[0]) => { + const { store } = useChatStore(options); + return { + store, + messagesByID: useChatSelector(store, selectMessagesByID), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { initialProps: initialOptions, wrapper }, + ); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + + act(() => { + mockSocket.emitOpen(); + }); + + rerender({ + ...initialOptions, + chatMessages: [msg1, msg2, optimisticReplacement], + chatMessagesData: { + messages: [msg1, msg2, optimisticReplacement], + queued_messages: [], + has_more: false, + }, + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + expect(result.current.messagesByID.get(3)?.content).toEqual( + optimisticReplacement.content, + ); + }); + + rerender({ + ...initialOptions, + chatMessages: [msg1, msg2, authoritativeReplacement], + chatMessagesData: { + messages: [msg1, msg2, authoritativeReplacement], + queued_messages: [], + has_more: false, + }, + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 9]); + expect(result.current.messagesByID.has(3)).toBe(false); + expect(result.current.messagesByID.get(9)?.content).toEqual( + authoritativeReplacement.content, + ); + }); + }); }); describe("parse errors", () => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts index 7c71ddf8dd..972da499dd 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.ts @@ -174,6 +174,7 @@ export type ChatStore = { queuedMessages: readonly TypesGen.ChatQueuedMessage[] | undefined, ) => void; setChatStatus: (status: TypesGen.ChatStatus | null) => void; + setStreamState: (streamState: StreamState | null) => void; setStreamError: (reason: ChatDetailError | null) => void; clearStreamError: () => void; setRetryState: (state: RetryState | null) => void; @@ -412,6 +413,20 @@ export const createChatStore = (): ChatStore => { chatStatus: status, })); }, + setStreamState: (streamState) => { + if (state.streamState === streamState) { + return; + } + setState((current) => { + if (current.streamState === streamState) { + return current; + } + return { + ...current, + streamState, + }; + }); + }, setStreamError: (reason) => { setState((current) => { if (chatDetailErrorsEqual(current.streamError, reason)) { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index 1dad72051e..9d20cda79c 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -206,10 +206,9 @@ export const useChatStore = ( const fetchedIDs = new Set(chatMessages.map((m) => m.id)); // Only classify a store-held ID as stale if it was // present in the PREVIOUS sync's fetched data. IDs - // added to the store after the last sync (by the WS - // handler or handleSend) are new, not stale, and - // must not trigger the destructive replaceMessages - // path. + // added to the store after the last sync (for example + // by the WS handler) are new, not stale, and must not + // trigger the destructive replaceMessages path. const prevIDs = new Set(prev.map((m) => m.id)); const hasStaleEntries = contentChanged && diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 609961eec9..8c5c071c31 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -48,7 +48,6 @@ interface ChatPageTimelineProps { fileBlocks?: readonly TypesGen.ChatMessagePart[], ) => void; editingMessageId?: number | null; - savingMessageId?: number | null; urlTransform?: UrlTransform; mcpServers?: readonly TypesGen.MCPServerConfig[]; } @@ -59,7 +58,6 @@ export const ChatPageTimeline: FC = ({ persistedError, onEditUserMessage, editingMessageId, - savingMessageId, urlTransform, mcpServers, }) => { @@ -100,7 +98,6 @@ export const ChatPageTimeline: FC = ({ subagentTitles={subagentTitles} onEditUserMessage={onEditUserMessage} editingMessageId={editingMessageId} - savingMessageId={savingMessageId} urlTransform={urlTransform} mcpServers={mcpServers} computerUseSubagentIds={computerUseSubagentIds} @@ -121,10 +118,18 @@ export const ChatPageTimeline: FC = ({ ); }; +export type PendingAttachment = { + fileId: string; + mediaType: string; +}; + interface ChatPageInputProps { store: ChatStoreHandle; compressionThreshold: number | undefined; - onSend: (message: string, fileIds?: string[]) => void; + onSend: ( + message: string, + attachments?: readonly PendingAttachment[], + ) => Promise | void; onDeleteQueuedMessage: (id: number) => Promise; onPromoteQueuedMessage: (id: number) => Promise; onInterrupt: () => void; @@ -315,9 +320,10 @@ export const ChatPageInput: FC = ({ { void (async () => { - // Collect file IDs from already-uploaded attachments. - // Skip files in error state (e.g. too large). - const fileIds: string[] = []; + // Collect uploaded attachment metadata for the optimistic + // transcript builder while keeping the server payload + // shape unchanged downstream. + const pendingAttachments: PendingAttachment[] = []; let skippedErrors = 0; for (const file of attachments) { const state = uploadStates.get(file); @@ -326,7 +332,10 @@ export const ChatPageInput: FC = ({ continue; } if (state?.status === "uploaded" && state.fileId) { - fileIds.push(state.fileId); + pendingAttachments.push({ + fileId: state.fileId, + mediaType: file.type || "application/octet-stream", + }); } } if (skippedErrors > 0) { @@ -334,9 +343,10 @@ export const ChatPageInput: FC = ({ `${skippedErrors} attachment${skippedErrors > 1 ? "s" : ""} could not be sent (upload failed)`, ); } - const fileArg = fileIds.length > 0 ? fileIds : undefined; + const attachmentArg = + pendingAttachments.length > 0 ? pendingAttachments : undefined; try { - await onSend(message, fileArg); + await onSend(message, attachmentArg); } catch { // Attachments preserved for retry on failure. return;