diff --git a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx index ce72f306e1f96..cfa70c4468f48 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx +++ b/site/src/pages/AgentsPage/components/ChatConversation/chatStore.test.tsx @@ -236,6 +236,19 @@ const makeMessage = ( content: [{ type: "text", text }], }); +const makeMessageWithContent = ( + chatID: string, + id: number, + role: TypesGen.ChatMessageRole, + content: readonly TypesGen.ChatMessagePart[], +): TypesGen.ChatMessage => ({ + id, + chat_id: chatID, + created_at: "2025-01-01T00:00:00.000Z", + role, + content, +}); + const makeQueuedMessage = ( chatID: string, id: number, @@ -343,6 +356,445 @@ describe("useChatStore", () => { }); }); + it("keeps create_workspace durable call without result after preview_reset", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-create-workspace"; + const existingMessage = makeMessage(chatID, 1, "user", "create workspace"); + const assistantMessage = makeMessageWithContent(chatID, 2, "assistant", [ + { + type: "tool-call", + tool_call_id: "create-workspace-1", + tool_name: "create_workspace", + args: { name: "dev" }, + }, + ]); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + messagesByID: useChatSelector(store, selectMessagesByID), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + part: { + type: "tool-call", + tool_call_id: "create-workspace-1", + tool_name: "create_workspace", + args: { name: "dev" }, + }, + }, + }); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + await waitFor(() => { + expect( + result.current.streamState?.toolCalls["create-workspace-1"]?.name, + ).toBe("create_workspace"); + }); + + act(() => { + mockSocket.emitDataBatch([ + { type: "message", chat_id: chatID, message: assistantMessage }, + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await waitFor(() => { + expect(result.current.streamState).toBeNull(); + expect(result.current.orderedMessageIDs).toEqual([1, 2]); + expect(result.current.messagesByID.get(2)?.content).toEqual( + assistantMessage.content, + ); + }); + }); + + it("clears stream state when preview_reset arrives after durable tool result", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-tool-result"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const assistantMessage = makeMessageWithContent(chatID, 2, "assistant", [ + { + type: "tool-call", + tool_call_id: "tool-1", + tool_name: "read_template", + args: { template_id: "template-1" }, + }, + ]); + const toolMessage = makeMessageWithContent(chatID, 3, "tool", [ + { + type: "tool-result", + tool_call_id: "tool-1", + tool_name: "read_template", + result: { name: "Template" }, + }, + ]); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitData({ + type: "message_part", + chat_id: chatID, + message_part: { + part: { + type: "tool-call", + tool_call_id: "tool-1", + tool_name: "read_template", + args: { template_id: "template-1" }, + }, + }, + }); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + await waitFor(() => { + expect(result.current.streamState?.toolCalls["tool-1"]?.name).toBe( + "read_template", + ); + }); + + act(() => { + mockSocket.emitDataBatch([ + { type: "message", chat_id: chatID, message: assistantMessage }, + { type: "preview_reset", chat_id: chatID }, + { + type: "message_part", + chat_id: chatID, + message_part: { + part: { + type: "tool-result", + tool_call_id: "tool-1", + tool_name: "read_template", + result: { name: "Template" }, + }, + }, + }, + { type: "message", chat_id: chatID, message: toolMessage }, + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + expect(result.current.streamState).toBeNull(); + }); + }); + + it("preview_reset discards pending buffered parts", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-buffer"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "stale" }, + }, + }, + { type: "preview_reset", chat_id: chatID }, + ]); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + expect(result.current.streamState).toBeNull(); + }); + + it("keeps only post-reset parts after preview_reset in one batch", async () => { + vi.useFakeTimers({ shouldAdvanceTime: true }); + + const chatID = "chat-preview-reset-post-part"; + const existingMessage = makeMessage(chatID, 1, "user", "hello"); + const durableMessage = makeMessage(chatID, 2, "assistant", "done"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = createTestQueryClient(); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: [existingMessage], + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: [existingMessage], + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(watchChat).toHaveBeenCalledWith(chatID, 1); + }); + + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "stale" }, + }, + }, + { type: "message", chat_id: chatID, message: durableMessage }, + { type: "preview_reset", chat_id: chatID }, + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "fresh" }, + }, + }, + ]); + }); + + await act(async () => { + vi.advanceTimersByTime(1); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2]); + expect(result.current.streamState?.blocks).toEqual([ + { type: "response", text: "fresh" }, + ]); + }); + }); + + it("replaces messages after history_reset", async () => { + const chatID = "chat-history-reset"; + const initialMessages = [ + makeMessage(chatID, 1, "user", "old prompt"), + makeMessage(chatID, 2, "assistant", "old answer"), + makeMessage(chatID, 3, "user", "stale prompt"), + ]; + const replacementMessage = makeMessage(chatID, 1, "user", "new prompt"); + const mockSocket = createMockSocket(); + mockWatchChatReturn(mockSocket); + + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: Number.POSITIVE_INFINITY, + refetchOnWindowFocus: false, + networkMode: "offlineFirst", + }, + }, + }); + queryClient.setQueryData(chatMessagesKey(chatID), { + pages: [ + { + messages: [...initialMessages].reverse(), + queued_messages: [], + has_more: false, + }, + ], + pageParams: [undefined], + }); + const wrapper = ({ children }: PropsWithChildren) => ( + {children} + ); + const setChatErrorReason = vi.fn(); + const clearChatErrorReason = vi.fn(); + + const { result } = renderHook( + () => { + const { store } = useChatStore({ + chatID, + chatMessages: initialMessages, + chatRecord: makeChat(chatID), + chatMessagesData: { + messages: initialMessages, + queued_messages: [], + has_more: false, + }, + chatQueuedMessages: [], + setChatErrorReason, + clearChatErrorReason, + }); + return { + streamState: useChatSelector(store, selectStreamState), + messagesByID: useChatSelector(store, selectMessagesByID), + orderedMessageIDs: useChatSelector(store, selectOrderedMessageIDs), + }; + }, + { wrapper }, + ); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1, 2, 3]); + }); + + act(() => { + mockSocket.emitDataBatch([ + { + type: "message_part", + chat_id: chatID, + message_part: { + role: "assistant", + part: { type: "text", text: "stale" }, + }, + }, + { type: "history_reset", chat_id: chatID }, + { type: "message", chat_id: chatID, message: replacementMessage }, + ]); + }); + + await waitFor(() => { + expect(result.current.orderedMessageIDs).toEqual([1]); + expect(result.current.messagesByID.get(1)?.content).toEqual( + replacementMessage.content, + ); + expect(result.current.streamState).toBeNull(); + }); + + const cached = queryClient.getQueryData<{ + pages: TypesGen.ChatMessagesResponse[]; + pageParams: unknown[]; + }>(chatMessagesKey(chatID)); + expect(cached?.pages[0]?.messages.map((message) => message.id)).toEqual([ + 1, + ]); + }); + it("clears stream state when a new durable message arrives", async () => { immediateAnimationFrame(); diff --git a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts index 80e4a0254f8de..5e0a259fb47e8 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.test.ts @@ -4,6 +4,7 @@ import { getSubagentDescriptor } from "../ChatElements/tools/subagentDescriptor" import { buildSubagentMaps, getEditableUserMessagePayload, + getPendingToolCallIDs, mergeTools, parseMessageContent, parseMessagesWithMergedTools, @@ -519,6 +520,115 @@ describe("mergeTools", () => { expect(merged).toHaveLength(1); expect(merged[0].status).toBe("completed"); }); + + it("marks unresolved pending calls as running", () => { + const merged = mergeTools([{ id: "1", name: "bash" }], [], { + pendingToolCallIDs: new Set(["1"]), + }); + expect(merged).toHaveLength(1); + expect(merged[0].status).toBe("running"); + }); +}); + +describe("pending durable tool parsing", () => { + const msg = ( + id: number, + role: "assistant" | "tool" | "user", + parts: ChatMessagePart[], + ): ChatMessage => ({ + id, + chat_id: "chat-1", + created_at: new Date(2026, 0, id).toISOString(), + role, + content: parts, + }); + + const toolCall = ( + id: string, + name = "create_workspace", + ): ChatMessagePart => ({ + type: "tool-call", + tool_call_id: id, + tool_name: name, + args: { name: "dev" }, + }); + + const toolResult = ( + id: string, + name = "create_workspace", + isError = false, + ): ChatMessagePart => ({ + type: "tool-result", + tool_call_id: id, + tool_name: name, + result: { workspace_name: "dev", build_id: "build-1" }, + is_error: isError, + }); + + it("marks the latest unresolved assistant tool call as running in an active chat", () => { + const messages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-active")]), + ]; + + const parsed = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs: getPendingToolCallIDs(messages, "running"), + }); + + expect(parsed[1]?.parsed.tools[0]?.status).toBe("running"); + }); + + it("keeps unresolved historical tool calls completed in an inactive chat", () => { + const messages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-inactive")]), + ]; + + const parsed = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs: getPendingToolCallIDs(messages, "waiting"), + }); + + expect(parsed[1]?.parsed.tools[0]?.status).toBe("completed"); + }); + + it("uses completed or error status once a matching result exists", () => { + const successMessages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-success")]), + msg(3, "tool", [toolResult("call-success")]), + ]; + const errorMessages = [ + msg(1, "user", [{ type: "text", text: "create a workspace" }]), + msg(2, "assistant", [toolCall("call-error")]), + msg(3, "tool", [toolResult("call-error", "create_workspace", true)]), + ]; + + const successParsed = parseMessagesWithMergedTools(successMessages, { + pendingToolCallIDs: getPendingToolCallIDs(successMessages, "running"), + }); + const errorParsed = parseMessagesWithMergedTools(errorMessages, { + pendingToolCallIDs: getPendingToolCallIDs(errorMessages, "running"), + }); + + expect(successParsed[1]?.parsed.tools[0]?.status).toBe("completed"); + expect(errorParsed[1]?.parsed.tools[0]?.status).toBe("error"); + }); + + it("does not mark older unresolved rows running in an active chat", () => { + const messages = [ + msg(1, "user", [{ type: "text", text: "first" }]), + msg(2, "assistant", [toolCall("call-old", "read_file")]), + msg(3, "user", [{ type: "text", text: "second" }]), + msg(4, "assistant", [toolCall("call-latest", "create_workspace")]), + ]; + + const parsed = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs: getPendingToolCallIDs(messages, "running"), + }); + + expect(parsed[1]?.parsed.tools[0]?.status).toBe("completed"); + expect(parsed[3]?.parsed.tools[0]?.status).toBe("running"); + }); }); describe("parseMessagesWithMergedTools — killedBySignal annotation", () => { diff --git a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts index 9e118d0656671..3f48e6acfad6e 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/messageParsing.ts @@ -83,9 +83,65 @@ export const ensureToolBlock = ( return [...blocks, { type: "tool", id }]; }; +const isToolCallPart = ( + part: TypesGen.ChatMessagePart, +): part is TypesGen.ChatToolCallPart => part.type === "tool-call"; + +const isToolResultPart = ( + part: TypesGen.ChatMessagePart, +): part is TypesGen.ChatToolResultPart => part.type === "tool-result"; + +const chatHasActiveToolCalls = (status: TypesGen.ChatStatus | null): boolean => + status === "running" || status === "requires_action"; + +export const getPendingToolCallIDs = ( + messages: readonly TypesGen.ChatMessage[], + chatStatus: TypesGen.ChatStatus | null, +): ReadonlySet | undefined => { + if (!chatHasActiveToolCalls(chatStatus)) { + return undefined; + } + + const resultIDs = new Set(); + for (const message of messages) { + for (const part of message.content ?? []) { + if (isToolResultPart(part) && part.tool_call_id) { + resultIDs.add(part.tool_call_id); + } + } + } + + for (let index = messages.length - 1; index >= 0; index -= 1) { + const message = messages[index]; + if (!message) { + continue; + } + if (message.role === "user") { + return undefined; + } + if (message.role !== "assistant") { + continue; + } + const pendingToolCallIDs = (message.content ?? []) + .filter(isToolCallPart) + .map((part) => part.tool_call_id) + .filter((id): id is string => Boolean(id && !resultIDs.has(id))); + return pendingToolCallIDs.length > 0 + ? new Set(pendingToolCallIDs) + : undefined; + } + + return undefined; +}; + +type MergeToolsOptions = { + pendingToolCallIDs?: ReadonlySet; +}; + export const mergeTools = ( calls: ParsedToolCall[], results: ParsedToolResult[], + options: MergeToolsOptions = {}, ): MergedTool[] => { const resultById = new Map(results.map((r) => [r.id, r])); const seen = new Set(); @@ -100,13 +156,20 @@ export const mergeTools = ( typeof callArgs?.model_intent === "string" ? callArgs.model_intent : undefined; + const status = result + ? result.isError + ? "error" + : "completed" + : options.pendingToolCallIDs?.has(call.id) + ? "running" + : "completed"; merged.push({ id: call.id, name: call.name, args: call.args, result: result?.result, isError: result?.isError ?? false, - status: result ? (result.isError ? "error" : "completed") : "completed", + status, mcpServerConfigId: call.mcpServerConfigId || result?.mcpServerConfigId, modelIntent, parsedCommands: call.parsedCommands, @@ -272,8 +335,13 @@ export const getEditableUserMessagePayload = ( }; }; +type ParseMessagesWithMergedToolsOptions = { + pendingToolCallIDs?: ReadonlySet; +}; + export const parseMessagesWithMergedTools = ( messages: readonly TypesGen.ChatMessage[], + options: ParseMessagesWithMergedToolsOptions = {}, ): ParsedMessageEntry[] => { const rawParsed = messages.map((message) => ({ message, @@ -303,6 +371,7 @@ export const parseMessagesWithMergedTools = ( parsed.tools = mergeTools( parsed.toolCalls, Array.from(resultById.values()), + { pendingToolCallIDs: options.pendingToolCallIDs }, ); } diff --git a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts index d7fb56f0d3367..592a6d5145d1f 100644 --- a/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts +++ b/site/src/pages/AgentsPage/components/ChatConversation/useChatStore.ts @@ -184,6 +184,29 @@ export const useChatStore = ( [chatID, queryClient], ); + const replaceCacheMessages = useCallback( + (messages: readonly TypesGen.ChatMessage[]) => { + if (!chatID) { + return; + } + queryClient.setQueryData< + InfiniteData | undefined + >(chatMessagesKey(chatID), (currentData) => { + if (!currentData?.pages?.length) { + return currentData; + } + const firstPage = currentData.pages[0]; + const updatedMessages = [...messages].sort((a, b) => b.id - a.id); + return { + ...currentData, + pages: [{ ...firstPage, messages: updatedMessages, has_more: false }], + pageParams: currentData.pageParams.slice(0, 1), + }; + }); + }, + [chatID, queryClient], + ); + useEffect(() => { store.batch(() => { // When the active chat changes, clear stale messages @@ -393,9 +416,9 @@ export const useChatStore = ( }; // Discard buffered parts without applying them. Used when - // the stream is no longer active (pending, waiting, retry) + // the preview is reset or the stream is no longer active // so stale buffered parts are not applied after the - // status transition. + // boundary event. const discardBufferedParts = () => { partsBuf.length = 0; if (partsFlushTimer !== null) { @@ -427,6 +450,9 @@ export const useChatStore = ( // instead of N copies and N sorts. const pendingMessages: TypesGen.ChatMessage[] = []; let needsStreamReset = false; + let historyResetActive = false; + let historyResetOccurred = false; + const replacementMessages: TypesGen.ChatMessage[] = []; // Wrap all store mutations in a batch so subscribers // are notified exactly once at the end, not per event. @@ -447,12 +473,37 @@ export const useChatStore = ( continue; } + if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { + const nextStatus = streamEvent.status?.status; + if (streamEvent.type === "status" && nextStatus) { + store.setSubagentStatusOverride(streamEvent.chat_id, nextStatus); + } + continue; + } + + if (streamEvent.type === "history_reset") { + discardBufferedParts(); + store.clearStreamState(); + historyResetActive = true; + historyResetOccurred = true; + replacementMessages.length = 0; + pendingMessages.length = 0; + needsStreamReset = false; + continue; + } + + if (streamEvent.type === "preview_reset") { + discardBufferedParts(); + store.clearStreamState(); + continue; + } + // Only flush buffered parts before events that // need them applied first. `message` events // commit durable state that must include all // stream parts. `error` events should surface // partial output. Other events (status, retry, - // queue_update) must NOT flush — status changes + // queue_update) must not flush. Status changes // need to be visible before parts so the // "Thinking..." indicator can render, and retry // clears stream state which a flush would @@ -467,11 +518,12 @@ export const useChatStore = ( if (!message) { continue; } - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } store.clearRetryState(); - pendingMessages.push(message); + if (historyResetActive) { + replacementMessages.push(message); + } else { + pendingMessages.push(message); + } if ( message.id !== undefined && (lastMessageIdRef.current === undefined || @@ -485,9 +537,6 @@ export const useChatStore = ( continue; } case "queue_update": - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } wsQueueUpdateReceivedRef.current = true; store.applyAuthoritativeQueuedMessages( streamEvent.queued_messages, @@ -500,14 +549,6 @@ export const useChatStore = ( continue; } - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - store.setSubagentStatusOverride( - streamEvent.chat_id, - nextStatus, - ); - continue; - } - wsStatusReceivedRef.current = true; store.clearRetryState(); store.setChatStatus(nextStatus); @@ -529,9 +570,6 @@ export const useChatStore = ( continue; } case "error": { - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } const reason = normalizeChatErrorPayload(streamEvent.error) ?? { kind: "generic", message: "Chat processing failed.", @@ -546,9 +584,6 @@ export const useChatStore = ( continue; } case "retry": { - if (streamEvent.chat_id && streamEvent.chat_id !== chatID) { - continue; - } const retry = streamEvent.retry; if (retry) { discardBufferedParts(); @@ -567,6 +602,11 @@ export const useChatStore = ( // non-message_part event above, this is a no-op. schedulePartsFlush(); + if (historyResetOccurred) { + store.replaceMessages(replacementMessages); + replaceCacheMessages(replacementMessages); + } + // Bulk-upsert all collected durable messages in one // pass: one Map copy + one sort instead of N each. if (pendingMessages.length > 0) { @@ -644,7 +684,14 @@ export const useChatStore = ( } activeChatIDRef.current = null; }; - }, [chatID, initialDataLoaded, queryClient, store, upsertCacheMessages]); + }, [ + chatID, + initialDataLoaded, + queryClient, + replaceCacheMessages, + store, + upsertCacheMessages, + ]); return { store, clearStreamError: () => { diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx index 43d81820663ba..68d330249a84d 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.stories.tsx @@ -1,6 +1,7 @@ import type { Meta, StoryObj } from "@storybook/react-vite"; import { expect, within } from "storybook/test"; import type * as TypesGen from "#/api/typesGenerated"; +import { ChatWorkspaceContext } from "../context/ChatWorkspaceContext"; import { createChatStore } from "./ChatConversation/chatStore"; import { buildStreamRenderState, @@ -136,6 +137,36 @@ export const SpacerVisibleWhenNotStreaming: Story = { }, }; +export const DurableUnresolvedWorkspaceToolRuns: Story = { + render: () => { + const store = createChatStore(); + store.replaceMessages([ + buildMessage(1, "user", [{ type: "text", text: "Create a workspace" }]), + buildMessage(2, "assistant", [ + { + type: "tool-call", + tool_call_id: "create-workspace-call", + tool_name: "create_workspace", + args: { name: "dev" }, + }, + ]), + ]); + store.setChatStatus("running"); + + return ( + + + + ); + }, + play: async ({ canvasElement }) => { + const canvas = within(canvasElement); + expect(canvas.getByText("Creating workspace…")).toBeInTheDocument(); + expect(canvas.queryByText("Created workspace")).toBeNull(); + expect(canvas.getByText("Loading build logs…")).toBeInTheDocument(); + }, +}; + export const HiddenAssistantPlaceholderDoesNotRender: Story = { render: () => { const store = createChatStore(); diff --git a/site/src/pages/AgentsPage/components/ChatPageContent.tsx b/site/src/pages/AgentsPage/components/ChatPageContent.tsx index 9599b273d16f9..02da881a8f073 100644 --- a/site/src/pages/AgentsPage/components/ChatPageContent.tsx +++ b/site/src/pages/AgentsPage/components/ChatPageContent.tsx @@ -34,6 +34,7 @@ import { import { LiveStreamTail } from "./ChatConversation/LiveStreamTail"; import { buildSubagentMaps, + getPendingToolCallIDs, parseMessagesWithMergedTools, } from "./ChatConversation/messageParsing"; import { useOnRenderProfiler } from "./ChatConversation/useOnRenderProfiler"; @@ -94,7 +95,10 @@ export const ChatPageTimeline: FC = ({ return message; }) .filter(isChatMessage); - const parsedMessages = parseMessagesWithMergedTools(messages); + const pendingToolCallIDs = getPendingToolCallIDs(messages, chatStatus); + const parsedMessages = parseMessagesWithMergedTools(messages, { + pendingToolCallIDs, + }); const { titles: subagentTitles, variants: subagentVariants } = buildSubagentMaps(parsedMessages); const onRenderProfiler = useOnRenderProfiler();