diff --git a/site/src/pages/AgentsPage/AgentChatInput.tsx b/site/src/pages/AgentsPage/AgentChatInput.tsx index 225f27e27e..d7decfcde2 100644 --- a/site/src/pages/AgentsPage/AgentChatInput.tsx +++ b/site/src/pages/AgentsPage/AgentChatInput.tsx @@ -21,6 +21,7 @@ import { CheckIcon, ImageIcon, MicIcon, + PencilIcon, Square, XIcon, } from "lucide-react"; @@ -520,6 +521,7 @@ export const AgentChatInput = memo( ) { return; } + onSend(text); if (!isMobileViewport()) { internalRef.current?.focus(); @@ -572,10 +574,20 @@ export const AgentChatInput = memo( } }; - const sendButtonLabel = editingQueuedMessageID !== null ? "Save" : "Send"; + const sendButtonLabel = + editingQueuedMessageID !== null + ? "Save" + : isEditingHistoryMessage + ? "Save Edit" + : "Send"; const content = ( -
+
{queuedMessages.length > 0 && ( ( className={cn( "rounded-2xl border border-border-default/80 bg-surface-secondary/45 p-1 shadow-sm has-[textarea:focus]:ring-2 has-[textarea:focus]:ring-content-link/40", isDragging && "ring-2 ring-content-link/40", + isEditingHistoryMessage && + "shadow-[0_0_0_2px_hsla(var(--border-warning),0.6)]", )} onKeyDown={handleKeyDown} onDragOver={onAttach ? handleDragOver : undefined} @@ -623,10 +637,12 @@ export const AgentChatInput = memo(
)} {isEditingHistoryMessage && editingQueuedMessageID === null && ( -
- - {isLoading && } - {isLoading ? "Saving edit..." : "Editing message"} +
+ + + {isLoading + ? "Saving edit..." + : "Editing message \u2014 all subsequent messages will be deleted"} @@ -662,8 +678,10 @@ export const AgentChatInput = memo( disabled={isDisabled || isLoading} autoFocus /> +
+ {" "} { expect(clearChatErrorReason).toHaveBeenCalledWith(chatID); }); }); + + it("removes stale messages when refetched set is smaller (edit truncation)", async () => { + immediateAnimationFrame(); + + const chatID = "chat-edit-truncation"; + const msg1 = makeMessage(chatID, 1, "user", "first"); + const msg2 = makeMessage(chatID, 2, "assistant", "second"); + const msg3 = makeMessage(chatID, 3, "user", "third"); + + const mockSocket = createMockSocket(); + vi.mocked(watchChat).mockReturnValue(mockSocket as never); + + const queryClient = createTestQueryClient(); + const wrapper: FC = ({ children }) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const noQueued: TypesGen.ChatQueuedMessage[] = []; + const initialMessages = [msg1, msg2, msg3]; + + const initialOptions = { + chatID, + chatMessages: initialMessages, + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: initialMessages, + queued_messages: noQueued, + has_more: false, + }, + chatQueuedMessages: noQueued, + setChatErrorReason, + clearChatErrorReason, + }; + + const { result, rerender } = renderHook( + (options: Parameters[0]) => { + const { store } = useChatStore(options); + return { + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { initialProps: initialOptions, wrapper }, + ); + + // All three messages should be in the store. + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + + // Simulate a post-edit refetch that only returns the first + // message (server truncated messages 2 and 3). + rerender({ + ...initialOptions, + chatMessages: [msg1], + chatMessagesData: { + messages: [msg1], + queued_messages: [], + has_more: false, + }, + }); + + // Messages 2 and 3 should be removed — replaceMessages should + // have been used instead of upsert because the store contained + // IDs not present in the fetched set. + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1]); + }); + }); }); describe("updateSidebarChat via stream events", () => { diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index e03d5512c5..806f176675 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -559,9 +559,22 @@ export const useChatStore = ( // of replacing the entire map. This preserves any messages the // WebSocket delivered via upsertDurableMessage that haven't // appeared in a REST page yet. + // + // However, if the fetched set is missing message IDs the store + // already has (e.g. after an edit truncation), a full replace + // is needed because upsert can only add/update, not remove. if (chatMessages) { - for (const message of chatMessages) { - store.upsertDurableMessage(message); + const fetchedIDs = new Set(chatMessages.map((m) => m.id)); + const storeSnap = store.getSnapshot(); + const hasStaleEntries = storeSnap.orderedMessageIDs.some( + (id) => !fetchedIDs.has(id), + ); + if (hasStaleEntries) { + store.replaceMessages(chatMessages); + } else { + for (const message of chatMessages) { + store.upsertDurableMessage(message); + } } } }, [chatID, chatMessages, store]); diff --git a/site/src/pages/AgentsPage/AgentDetail/ConversationTimeline.tsx b/site/src/pages/AgentsPage/AgentDetail/ConversationTimeline.tsx index 64f03651e6..311c6b2dc4 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ConversationTimeline.tsx +++ b/site/src/pages/AgentsPage/AgentDetail/ConversationTimeline.tsx @@ -12,7 +12,12 @@ import { WebSearchSources } from "components/ai-elements/tool"; import { Button } from "components/Button/Button"; import { FileReferenceChip } from "components/ChatMessageInput/FileReferenceNode"; import { Spinner } from "components/Spinner/Spinner"; -import { ChevronDownIcon } from "lucide-react"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "components/Tooltip/Tooltip"; +import { ChevronDownIcon, PencilIcon } from "lucide-react"; import { type FC, Fragment, @@ -296,6 +301,7 @@ const ChatMessageItem = memo<{ ) => void; editingMessageId?: number | null; savingMessageId?: number | null; + isAfterEditingMessage?: boolean; // When true, renders a gradient overlay inside the bubble // that fades text out toward the bottom. Used by the sticky // overlay to indicate truncated content. @@ -308,6 +314,7 @@ const ChatMessageItem = memo<{ onEditUserMessage, editingMessageId, savingMessageId, + isAfterEditingMessage = false, fadeFromBottom = false, urlTransform, }) => { @@ -369,18 +376,20 @@ const ChatMessageItem = memo<{ ); return ( - <> +
{isUser ? ( { - const fileBlocks = parsed.blocks.filter( - (b): b is Extract => - b.type === "file" && - b.media_type.startsWith("image/"), - ); - onEditUserMessage( - message.id, - parsed.markdown || "", - fileBlocks.length > 0 ? fileBlocks : undefined, - ); - } - : undefined - } >
@@ -431,6 +424,37 @@ const ChatMessageItem = memo<{ loading /> )} + {onEditUserMessage && !isSavingMessage && ( + + + + + Edit message + + )}
{(() => { const imageBlocks = parsed.blocks.filter( @@ -508,7 +532,7 @@ const ChatMessageItem = memo<{ onClose={() => setPreviewImage(null)} /> )} - +
); }, ); @@ -605,12 +629,14 @@ const StickyUserMessage: FC<{ ) => void; editingMessageId?: number | null; savingMessageId?: number | null; + isAfterEditingMessage?: boolean; }> = ({ message, parsed, onEditUserMessage, editingMessageId, savingMessageId, + isAfterEditingMessage = false, }) => { const [isStuck, setIsStuck] = useState(false); const [isReady, setIsReady] = useState(false); @@ -787,6 +813,7 @@ const StickyUserMessage: FC<{ onEditUserMessage={handleEditUserMessage} editingMessageId={editingMessageId} savingMessageId={savingMessageId} + isAfterEditingMessage={isAfterEditingMessage} />
@@ -831,6 +858,7 @@ const StickyUserMessage: FC<{ onEditUserMessage={handleEditUserMessage} editingMessageId={editingMessageId} savingMessageId={savingMessageId} + isAfterEditingMessage={isAfterEditingMessage} fadeFromBottom />
@@ -889,6 +917,24 @@ export const ConversationTimeline: FC = ({ const isUsageLimitError = detailError?.kind === "usage-limit"; const showUsageAction = onOpenAnalytics !== undefined && isUsageLimitError; + // Build a set of message IDs that appear after the message + // currently being edited so they can be visually faded. + const afterEditingMessageIds = new Set(); + if (editingMessageId != null) { + let found = false; + for (const section of parsedSections) { + for (const entry of section.entries) { + if (entry.message.id === editingMessageId) { + found = true; + continue; + } + if (found) { + afterEditingMessageIds.add(entry.message.id); + } + } + } + } + return (
{isEmpty && !hasStreamOutput ? ( @@ -924,6 +970,9 @@ export const ConversationTimeline: FC = ({ onEditUserMessage={onEditUserMessage} editingMessageId={editingMessageId} savingMessageId={savingMessageId} + isAfterEditingMessage={afterEditingMessageIds.has( + message.id, + )} /> ) : ( = ({ parsed={parsed} savingMessageId={savingMessageId} urlTransform={urlTransform} + isAfterEditingMessage={afterEditingMessageIds.has( + message.id, + )} /> ), - )} + )}{" "} {shouldRenderStreamInLastSection && sectionIdx === parsedSections.length - 1 && ( ({ + id, + chat_id: AGENT_ID, + created_at: new Date(Date.now() - (10 - id) * 60_000).toISOString(), + role, + content: [{ type: "text", text }], +}); + +const buildStoreWithMessages = ( + msgs: TypesGen.ChatMessage[], + status: TypesGen.ChatStatus = "completed", +) => { + const store = createChatStore(); + store.replaceMessages(msgs); + store.setChatStatus(status); + return store; +}; + +// --------------------------------------------------------------------------- +// Editing flow stories +// --------------------------------------------------------------------------- + +const editingMessages = [ + buildMessage(1, "user", "Say hi back"), + buildMessage(2, "assistant", "Hi!"), + buildMessage(3, "user", "Now tell me a joke"), + buildMessage( + 4, + "assistant", + "Why did the developer quit? Because they didn't get arrays.", + ), + buildMessage(5, "user", "That was terrible, try again"), +]; + +/** Editing a message in the middle of the conversation — shows the warning + * border on the edited message, faded subsequent messages, and the editing + * banner + outline on the chat input. */ +export const EditingMessage: Story = { + args: { + store: buildStoreWithMessages(editingMessages), + editing: { + ...defaultEditing, + editingMessageId: 3, + editorInitialValue: "Now tell me a joke", + }, + }, +}; + +/** The saving state while an edit is in progress — shows the pending + * indicator on the message being saved. */ +export const EditingSaving: Story = { + args: { + store: buildStoreWithMessages(editingMessages), + editing: { + ...defaultEditing, + editingMessageId: 3, + editorInitialValue: "Now tell me a better joke", + }, + pendingEditMessageId: 3, + isSubmissionPending: true, + }, +}; + // --------------------------------------------------------------------------- // AgentDetailNotFoundView stories // ---------------------------------------------------------------------------