diff --git a/coderd/chatd/chatd_test.go b/coderd/chatd/chatd_test.go index 028df0bef2..b76bfac772 100644 --- a/coderd/chatd/chatd_test.go +++ b/coderd/chatd/chatd_test.go @@ -176,7 +176,7 @@ func TestSubagentChatExcludesWorkspaceProvisioningTools(t *testing.T) { if getErr != nil { return false } - if got.Chat.Status != codersdk.ChatStatusWaiting && got.Chat.Status != codersdk.ChatStatusError { + if got.Status != codersdk.ChatStatusWaiting && got.Status != codersdk.ChatStatusError { return false } // Also ensure the subagent LLM call has been made. @@ -1055,32 +1055,35 @@ func TestCreateWorkspaceTool_EndToEnd(t *testing.T) { }) require.NoError(t, err) - var chatWithMessages codersdk.ChatWithMessages + var chatResult codersdk.Chat require.Eventually(t, func() bool { got, getErr := client.GetChat(ctx, chat.ID) if getErr != nil { return false } - chatWithMessages = got - return got.Chat.Status == codersdk.ChatStatusWaiting || got.Chat.Status == codersdk.ChatStatusError + chatResult = got + return got.Status == codersdk.ChatStatusWaiting || got.Status == codersdk.ChatStatusError }, testutil.WaitLong, testutil.IntervalFast) - if chatWithMessages.Chat.Status == codersdk.ChatStatusError { + if chatResult.Status == codersdk.ChatStatusError { lastError := "" - if chatWithMessages.Chat.LastError != nil { - lastError = *chatWithMessages.Chat.LastError + if chatResult.LastError != nil { + lastError = *chatResult.LastError } require.FailNowf(t, "chat run failed", "last_error=%q", lastError) } - require.NotNil(t, chatWithMessages.Chat.WorkspaceID) - workspaceID := *chatWithMessages.Chat.WorkspaceID + require.NotNil(t, chatResult.WorkspaceID) + workspaceID := *chatResult.WorkspaceID workspace, err := client.Workspace(ctx, workspaceID) require.NoError(t, err) require.Equal(t, workspaceName, workspace.Name) + chatMsgs, err := client.GetChatMessages(ctx, chat.ID) + require.NoError(t, err) + var foundCreateWorkspaceResult bool - for _, message := range chatWithMessages.Messages { + for _, message := range chatMsgs.Messages { if message.Role != "tool" { continue } @@ -1223,33 +1226,36 @@ func TestStartWorkspaceTool_EndToEnd(t *testing.T) { }) require.NoError(t, err) - var chatWithMessages codersdk.ChatWithMessages + var chatResult codersdk.Chat require.Eventually(t, func() bool { got, getErr := client.GetChat(ctx, chat.ID) if getErr != nil { return false } - chatWithMessages = got - return got.Chat.Status == codersdk.ChatStatusWaiting || got.Chat.Status == codersdk.ChatStatusError + chatResult = got + return got.Status == codersdk.ChatStatusWaiting || got.Status == codersdk.ChatStatusError }, testutil.WaitSuperLong, testutil.IntervalFast) - if chatWithMessages.Chat.Status == codersdk.ChatStatusError { + if chatResult.Status == codersdk.ChatStatusError { lastError := "" - if chatWithMessages.Chat.LastError != nil { - lastError = *chatWithMessages.Chat.LastError + if chatResult.LastError != nil { + lastError = *chatResult.LastError } require.FailNowf(t, "chat run failed", "last_error=%q", lastError) } // Verify the workspace was started. - require.NotNil(t, chatWithMessages.Chat.WorkspaceID) + require.NotNil(t, chatResult.WorkspaceID) updatedWorkspace, err := client.Workspace(ctx, workspace.ID) require.NoError(t, err) require.Equal(t, codersdk.WorkspaceTransitionStart, updatedWorkspace.LatestBuild.Transition) + chatMsgs, err := client.GetChatMessages(ctx, chat.ID) + require.NoError(t, err) + // Verify start_workspace tool result exists in the chat messages. var foundStartWorkspaceResult bool - for _, message := range chatWithMessages.Messages { + for _, message := range chatMsgs.Messages { if message.Role != "tool" { continue } diff --git a/coderd/chatd/integration_test.go b/coderd/chatd/integration_test.go index be65ded0f9..b70a30d718 100644 --- a/coderd/chatd/integration_test.go +++ b/coderd/chatd/integration_test.go @@ -92,17 +92,19 @@ func TestAnthropicWebSearchRoundTrip(t *testing.T) { // Verify the chat completed and messages were persisted. chatData, err := client.GetChat(ctx, chat.ID) require.NoError(t, err) + chatMsgs, err := client.GetChatMessages(ctx, chat.ID) + require.NoError(t, err) t.Logf("Chat status after step 1: %s, messages: %d", - chatData.Chat.Status, len(chatData.Messages)) - logMessages(t, chatData.Messages) + chatData.Status, len(chatMsgs.Messages)) + logMessages(t, chatMsgs.Messages) - require.Equal(t, codersdk.ChatStatusWaiting, chatData.Chat.Status, + require.Equal(t, codersdk.ChatStatusWaiting, chatData.Status, "chat should be in waiting status after step 1") // Find the first assistant message and verify it has the // content parts the UI needs to render web search results: // tool-call(PE), source, tool-result(PE), and text. - assistantMsg := findAssistantWithText(t, chatData.Messages) + assistantMsg := findAssistantWithText(t, chatMsgs.Messages) require.NotNil(t, assistantMsg, "expected an assistant message with text content after step 1") @@ -152,17 +154,19 @@ func TestAnthropicWebSearchRoundTrip(t *testing.T) { // Verify the follow-up completed and produced content. chatData2, err := client.GetChat(ctx, chat.ID) require.NoError(t, err) + chatMsgs2, err := client.GetChatMessages(ctx, chat.ID) + require.NoError(t, err) t.Logf("Chat status after step 2: %s, messages: %d", - chatData2.Chat.Status, len(chatData2.Messages)) - logMessages(t, chatData2.Messages) + chatData2.Status, len(chatMsgs2.Messages)) + logMessages(t, chatMsgs2.Messages) - require.Equal(t, codersdk.ChatStatusWaiting, chatData2.Chat.Status, + require.Equal(t, codersdk.ChatStatusWaiting, chatData2.Status, "chat should be in waiting status after step 2") - require.Greater(t, len(chatData2.Messages), len(chatData.Messages), + require.Greater(t, len(chatMsgs2.Messages), len(chatMsgs.Messages), "follow-up should have added more messages") // The last assistant message should have text. - lastAssistant := findLastAssistantWithText(t, chatData2.Messages) + lastAssistant := findLastAssistantWithText(t, chatMsgs2.Messages) require.NotNil(t, lastAssistant, "expected an assistant message with text in the follow-up") diff --git a/coderd/chats.go b/coderd/chats.go index 31c9ffa3c3..959cc764aa 100644 --- a/coderd/chats.go +++ b/coderd/chats.go @@ -360,6 +360,15 @@ func (api *API) listChatModels(rw http.ResponseWriter, r *http.Request) { // //nolint:revive // HTTP handler writes to ResponseWriter. func (api *API) getChat(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + chat := httpmw.ChatParam(r) + httpapi.Write(ctx, rw, http.StatusOK, convertChat(chat, nil)) +} + +// EXPERIMENTAL: this endpoint is experimental and is subject to change. +// +//nolint:revive // HTTP handler writes to ResponseWriter. +func (api *API) getChatMessages(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() chat := httpmw.ChatParam(r) chatID := chat.ID @@ -385,8 +394,7 @@ func (api *API) getChat(rw http.ResponseWriter, r *http.Request) { return } - httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatWithMessages{ - Chat: convertChat(chat, nil), + httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatMessagesResponse{ Messages: convertChatMessages(messages), QueuedMessages: convertChatQueuedMessages(queuedMessages), }) diff --git a/coderd/chats_test.go b/coderd/chats_test.go index d451be1f04..5b452f6e85 100644 --- a/coderd/chats_test.go +++ b/coderd/chats_test.go @@ -85,12 +85,14 @@ func TestPostChats(t *testing.T) { require.NotNil(t, chat.RootChatID) require.Equal(t, chat.ID, *chat.RootChatID) - chatWithMessages, err := client.GetChat(ctx, chat.ID) + chatResult, err := client.GetChat(ctx, chat.ID) require.NoError(t, err) - require.Equal(t, chat.ID, chatWithMessages.Chat.ID) + messagesResult, err := client.GetChatMessages(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, chat.ID, chatResult.ID) foundUserMessage := false - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { if message.Role != "user" { continue } @@ -123,9 +125,9 @@ func TestPostChats(t *testing.T) { }) require.NoError(t, err) - chatWithMessages, err := client.GetChat(ctx, chat.ID) + messagesResult, err := client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { require.NotEqual(t, "system", message.Role) } }) @@ -1322,19 +1324,21 @@ func TestGetChat(t *testing.T) { }) require.NoError(t, err) - chatWithMessages, err := client.GetChat(ctx, createdChat.ID) + chatResult, err := client.GetChat(ctx, createdChat.ID) require.NoError(t, err) - require.Equal(t, createdChat.ID, chatWithMessages.Chat.ID) - require.Equal(t, firstUser.UserID, chatWithMessages.Chat.OwnerID) - require.Equal(t, modelConfig.ID, chatWithMessages.Chat.LastModelConfigID) - require.Equal(t, "get chat route payload", chatWithMessages.Chat.Title) - require.NotZero(t, chatWithMessages.Chat.CreatedAt) - require.NotZero(t, chatWithMessages.Chat.UpdatedAt) - require.NotEmpty(t, chatWithMessages.Messages) - require.Empty(t, chatWithMessages.QueuedMessages) + messagesResult, err := client.GetChatMessages(ctx, createdChat.ID) + require.NoError(t, err) + require.Equal(t, createdChat.ID, chatResult.ID) + require.Equal(t, firstUser.UserID, chatResult.OwnerID) + require.Equal(t, modelConfig.ID, chatResult.LastModelConfigID) + require.Equal(t, "get chat route payload", chatResult.Title) + require.NotZero(t, chatResult.CreatedAt) + require.NotZero(t, chatResult.UpdatedAt) + require.NotEmpty(t, messagesResult.Messages) + require.Empty(t, messagesResult.QueuedMessages) foundUserMessage := false - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { require.Equal(t, createdChat.ID, message.ChatID) require.NotEqual(t, "system", message.Role) for _, part := range message.Content { @@ -1646,19 +1650,19 @@ func TestPostChatMessages(t *testing.T) { require.True(t, hasTextPart(created.QueuedMessage.Content, messageText)) require.Eventually(t, func() bool { - chatWithMessages, getErr := client.GetChat(ctx, chat.ID) + messagesResult, getErr := client.GetChatMessages(ctx, chat.ID) if getErr != nil { return false } - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { if queued.ID == created.QueuedMessage.ID && queued.ChatID == chat.ID && hasTextPart(queued.Content, messageText) { return true } } - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { if message.Role == "user" && hasTextPart(message.Content, messageText) { return true } @@ -1674,11 +1678,11 @@ func TestPostChatMessages(t *testing.T) { require.True(t, hasTextPart(created.Message.Content, messageText)) require.Eventually(t, func() bool { - chatWithMessages, getErr := client.GetChat(ctx, chat.ID) + messagesResult, getErr := client.GetChatMessages(ctx, chat.ID) if getErr != nil { return false } - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { if message.ID == created.Message.ID && message.Role == "user" && hasTextPart(message.Content, messageText) { @@ -1784,11 +1788,11 @@ func TestChatMessageWithFileReferences(t *testing.T) { var found bool require.Eventually(t, func() bool { - chatWithMessages, getErr := client.GetChat(ctx, chat.ID) + messagesResult, getErr := client.GetChatMessages(ctx, chat.ID) if getErr != nil { return false } - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { if message.Role != "user" { continue } @@ -1802,7 +1806,7 @@ func TestChatMessageWithFileReferences(t *testing.T) { } // The message may have been queued. if created.Queued && created.QueuedMessage != nil { - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { for _, part := range queued.Content { if part.Type == codersdk.ChatMessagePartTypeText && part.Text == wantText { @@ -1842,11 +1846,11 @@ func TestChatMessageWithFileReferences(t *testing.T) { "```lib/utils.ts\nconst x = 1;\n```" require.Eventually(t, func() bool { - chatWithMessages, getErr := client.GetChat(ctx, chat.ID) + messagesResult, getErr := client.GetChatMessages(ctx, chat.ID) if getErr != nil { return false } - for _, msg := range chatWithMessages.Messages { + for _, msg := range messagesResult.Messages { for _, part := range msg.Content { if part.Type == codersdk.ChatMessagePartTypeText && part.Text == wantText { return true @@ -1854,7 +1858,7 @@ func TestChatMessageWithFileReferences(t *testing.T) { } } if created.Queued && created.QueuedMessage != nil { - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { for _, part := range queued.Content { if part.Type == codersdk.ChatMessagePartTypeText && part.Text == wantText { return true @@ -1889,11 +1893,11 @@ func TestChatMessageWithFileReferences(t *testing.T) { // No fenced code block when content is empty. wantText := "[file-reference] README.md:1" require.Eventually(t, func() bool { - chatWithMessages, getErr := client.GetChat(ctx, chat.ID) + messagesResult, getErr := client.GetChatMessages(ctx, chat.ID) if getErr != nil { return false } - for _, msg := range chatWithMessages.Messages { + for _, msg := range messagesResult.Messages { for _, part := range msg.Content { if part.Type == codersdk.ChatMessagePartTypeText && part.Text == wantText { return true @@ -1901,7 +1905,7 @@ func TestChatMessageWithFileReferences(t *testing.T) { } } if created.Queued && created.QueuedMessage != nil { - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { for _, part := range queued.Content { if part.Type == codersdk.ChatMessagePartTypeText && part.Text == wantText { return true @@ -1937,11 +1941,11 @@ func TestChatMessageWithFileReferences(t *testing.T) { "```server.go\nfunc main() {\n\tfmt.Println()\n}\n```" require.Eventually(t, func() bool { - chatWithMessages, getErr := client.GetChat(ctx, chat.ID) + messagesResult, getErr := client.GetChatMessages(ctx, chat.ID) if getErr != nil { return false } - for _, msg := range chatWithMessages.Messages { + for _, msg := range messagesResult.Messages { for _, part := range msg.Content { if part.Type == codersdk.ChatMessagePartTypeText && part.Text == wantText { return true @@ -1949,7 +1953,7 @@ func TestChatMessageWithFileReferences(t *testing.T) { } } if created.Queued && created.QueuedMessage != nil { - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { for _, part := range queued.Content { if part.Type == codersdk.ChatMessagePartTypeText && part.Text == wantText { return true @@ -2017,7 +2021,7 @@ func TestChatMessageWithFileReferences(t *testing.T) { } require.Eventually(t, func() bool { - chatWithMessages, getErr := client.GetChat(ctx, chat.ID) + messagesResult, getErr := client.GetChatMessages(ctx, chat.ID) if getErr != nil { return false } @@ -2042,13 +2046,13 @@ func TestChatMessageWithFileReferences(t *testing.T) { return true } - for _, msg := range chatWithMessages.Messages { + for _, msg := range messagesResult.Messages { if msg.Role == "user" && checkParts(msg.Content) { return true } } if created.Queued && created.QueuedMessage != nil { - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { if checkParts(queued.Content) { return true } @@ -2201,9 +2205,9 @@ func TestChatMessageWithFiles(t *testing.T) { } // Verify file parts omit inline data in the API response. - chatWithMessages, err := client.GetChat(ctx, chat.ID) + messagesResult, err := client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) - for _, msg := range chatWithMessages.Messages { + for _, msg := range messagesResult.Messages { for _, part := range msg.Content { if part.Type == codersdk.ChatMessagePartTypeFile { require.True(t, part.FileID.Valid, "file part should have a valid file_id") @@ -2297,11 +2301,11 @@ func TestPatchChatMessage(t *testing.T) { }) require.NoError(t, err) - chatWithMessages, err := client.GetChat(ctx, chat.ID) + messagesResult, err := client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) var userMessageID int64 - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { if message.Role == "user" { userMessageID = message.ID break @@ -2329,11 +2333,11 @@ func TestPatchChatMessage(t *testing.T) { } require.True(t, foundEditedText) - updatedChat, err := client.GetChat(ctx, chat.ID) + messagesResult, err = client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) foundEditedInChat := false foundOriginalInChat := false - for _, message := range updatedChat.Messages { + for _, message := range messagesResult.Messages { if message.Role != "user" { continue } @@ -2382,11 +2386,11 @@ func TestPatchChatMessage(t *testing.T) { require.NoError(t, err) // Find the user message ID. - chatWithMessages, err := client.GetChat(ctx, chat.ID) + messagesResult, err := client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) var userMessageID int64 - for _, message := range chatWithMessages.Messages { + for _, message := range messagesResult.Messages { if message.Role == "user" { userMessageID = message.ID break @@ -2424,12 +2428,12 @@ func TestPatchChatMessage(t *testing.T) { require.True(t, foundText, "edited message should contain updated text") require.True(t, foundFile, "edited message should preserve file_id") - // GET the chat and verify the file_id persists. - updatedChat, err := client.GetChat(ctx, chat.ID) + // GET the chat messages and verify the file_id persists. + messagesResult, err = client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) var foundTextInChat, foundFileInChat bool - for _, message := range updatedChat.Messages { + for _, message := range messagesResult.Messages { if message.Role != "user" { continue } @@ -3037,9 +3041,9 @@ func TestDeleteChatQueuedMessage(t *testing.T) { res.Body.Close() require.Equal(t, http.StatusNoContent, res.StatusCode) - chatWithMessages, err := client.GetChat(ctx, chat.ID) + messagesResult, err := client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { require.NotEqual(t, queuedMessage.ID, queued.ID) } @@ -3136,9 +3140,9 @@ func TestPromoteChatQueuedMessage(t *testing.T) { } require.True(t, foundPromotedText) - chatWithMessages, err := client.GetChat(ctx, chat.ID) + messagesResult, err := client.GetChatMessages(ctx, chat.ID) require.NoError(t, err) - for _, queued := range chatWithMessages.QueuedMessages { + for _, queued := range messagesResult.QueuedMessages { require.NotEqual(t, queuedMessage.ID, queued.ID) } diff --git a/coderd/coderd.go b/coderd/coderd.go index f1c842639d..ced5b3968b 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -1173,6 +1173,7 @@ func New(options *Options) *API { r.Get("/git/watch", api.watchChatGit) r.Post("/archive", api.archiveChat) r.Post("/unarchive", api.unarchiveChat) + r.Get("/messages", api.getChatMessages) r.Post("/messages", api.postChatMessages) r.Patch("/messages/{message}", api.patchChatMessage) r.Get("/stream", api.streamChat) diff --git a/codersdk/chats.go b/codersdk/chats.go index 009055ec6a..9117714789 100644 --- a/codersdk/chats.go +++ b/codersdk/chats.go @@ -167,9 +167,8 @@ type UploadChatFileResponse struct { ID uuid.UUID `json:"id" format:"uuid"` } -// ChatWithMessages is a chat along with its messages. -type ChatWithMessages struct { - Chat Chat `json:"chat"` +// ChatMessagesResponse contains the messages and queued messages for a chat. +type ChatMessagesResponse struct { Messages []ChatMessage `json:"messages"` QueuedMessages []ChatQueuedMessage `json:"queued_messages"` } @@ -980,20 +979,34 @@ func (c *Client) StreamChat(ctx context.Context, chatID uuid.UUID, opts *StreamC }), nil } -// GetChat returns a chat by ID, including its messages. -func (c *Client) GetChat(ctx context.Context, chatID uuid.UUID) (ChatWithMessages, error) { +// GetChat returns a chat by ID. +func (c *Client) GetChat(ctx context.Context, chatID uuid.UUID) (Chat, error) { res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/experimental/chats/%s", chatID), nil) if err != nil { - return ChatWithMessages{}, err + return Chat{}, err } defer res.Body.Close() if res.StatusCode != http.StatusOK { - return ChatWithMessages{}, ReadBodyAsError(res) + return Chat{}, ReadBodyAsError(res) } - var chat ChatWithMessages + var chat Chat return chat, json.NewDecoder(res.Body).Decode(&chat) } +// GetChatMessages returns the messages and queued messages for a chat. +func (c *Client) GetChatMessages(ctx context.Context, chatID uuid.UUID) (ChatMessagesResponse, error) { + res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/experimental/chats/%s/messages", chatID), nil) + if err != nil { + return ChatMessagesResponse{}, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return ChatMessagesResponse{}, ReadBodyAsError(res) + } + var resp ChatMessagesResponse + return resp, json.NewDecoder(res.Body).Decode(&resp) +} + func (c *Client) ArchiveChat(ctx context.Context, chatID uuid.UUID) error { res, err := c.Request(ctx, http.MethodPost, fmt.Sprintf("/api/experimental/chats/%s/archive", chatID), nil) if err != nil { diff --git a/site/src/api/api.ts b/site/src/api/api.ts index 854518e9b4..26ee6570cc 100644 --- a/site/src/api/api.ts +++ b/site/src/api/api.ts @@ -2948,12 +2948,20 @@ class ApiMethods { ); return response.data; }; - getChat = async (chatId: string): Promise => { - const response = await this.axios.get( + getChat = async (chatId: string): Promise => { + const response = await this.axios.get( `/api/experimental/chats/${chatId}`, ); return response.data; }; + getChatMessages = async ( + chatId: string, + ): Promise => { + const response = await this.axios.get( + `/api/experimental/chats/${chatId}/messages`, + ); + return response.data; + }; createChat = async ( req: TypesGen.CreateChatRequest, diff --git a/site/src/api/queries/chats.test.ts b/site/src/api/queries/chats.test.ts index 8959617eb7..b9fb6f7ec8 100644 --- a/site/src/api/queries/chats.test.ts +++ b/site/src/api/queries/chats.test.ts @@ -55,15 +55,6 @@ const makeChat = ( ...overrides, }); -const makeChatWithMessages = ( - chatId: string, - overrides?: Partial, -): TypesGen.ChatWithMessages => ({ - chat: makeChat(chatId, overrides), - messages: [], - queued_messages: [], -}); - const createTestQueryClient = (): QueryClient => new QueryClient({ defaultOptions: { @@ -99,17 +90,15 @@ describe("archiveChat optimistic update", () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; seedInfiniteChats(queryClient, [makeChat(chatId)]); - queryClient.setQueryData(chatKey(chatId), makeChatWithMessages(chatId)); + queryClient.setQueryData(chatKey(chatId), makeChat(chatId)); vi.mocked(API.archiveChat).mockResolvedValue(); const mutation = archiveChat(queryClient); await mutation.onMutate(chatId); - const cachedChat = queryClient.getQueryData( - chatKey(chatId), - ); - expect(cachedChat?.chat.archived).toBe(true); + const cachedChat = queryClient.getQueryData(chatKey(chatId)); + expect(cachedChat?.archived).toBe(true); }); it("rolls back the chats list on error by invalidating", async () => { @@ -117,7 +106,7 @@ describe("archiveChat optimistic update", () => { const chatId = "chat-1"; const initialChats = [makeChat(chatId)]; seedInfiniteChats(queryClient, initialChats); - queryClient.setQueryData(chatKey(chatId), makeChatWithMessages(chatId)); + queryClient.setQueryData(chatKey(chatId), makeChat(chatId)); const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); const mutation = archiveChat(queryClient); @@ -139,22 +128,19 @@ describe("archiveChat optimistic update", () => { const queryClient = createTestQueryClient(); const chatId = "chat-1"; seedInfiniteChats(queryClient, [makeChat(chatId)]); - queryClient.setQueryData(chatKey(chatId), makeChatWithMessages(chatId)); + queryClient.setQueryData(chatKey(chatId), makeChat(chatId)); const mutation = archiveChat(queryClient); const context = await mutation.onMutate(chatId); expect( - queryClient.getQueryData(chatKey(chatId))?.chat - .archived, + queryClient.getQueryData(chatKey(chatId))?.archived, ).toBe(true); mutation.onError(new Error("server error"), chatId, context); - const rolledBack = queryClient.getQueryData( - chatKey(chatId), - ); - expect(rolledBack?.chat.archived).toBe(false); + const rolledBack = queryClient.getQueryData(chatKey(chatId)); + expect(rolledBack?.archived).toBe(false); }); it("handles error rollback gracefully when context is undefined", () => { @@ -226,15 +212,14 @@ describe("unarchiveChat optimistic update", () => { seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })]); queryClient.setQueryData( chatKey(chatId), - makeChatWithMessages(chatId, { archived: true }), + makeChat(chatId, { archived: true }), ); const mutation = unarchiveChat(queryClient); await mutation.onMutate(chatId); expect( - queryClient.getQueryData(chatKey(chatId))?.chat - .archived, + queryClient.getQueryData(chatKey(chatId))?.archived, ).toBe(false); }); @@ -244,7 +229,7 @@ describe("unarchiveChat optimistic update", () => { seedInfiniteChats(queryClient, [makeChat(chatId, { archived: true })]); queryClient.setQueryData( chatKey(chatId), - makeChatWithMessages(chatId, { archived: true }), + makeChat(chatId, { archived: true }), ); const invalidateSpy = vi.spyOn(queryClient, "invalidateQueries"); @@ -254,8 +239,7 @@ describe("unarchiveChat optimistic update", () => { // Verify optimistic update. expect(readInfiniteChats(queryClient)?.[0].archived).toBe(false); expect( - queryClient.getQueryData(chatKey(chatId))?.chat - .archived, + queryClient.getQueryData(chatKey(chatId))?.archived, ).toBe(false); // Roll back. @@ -267,8 +251,7 @@ describe("unarchiveChat optimistic update", () => { }); // The individual chat cache is restored directly. expect( - queryClient.getQueryData(chatKey(chatId))?.chat - .archived, + queryClient.getQueryData(chatKey(chatId))?.archived, ).toBe(true); }); diff --git a/site/src/api/queries/chats.ts b/site/src/api/queries/chats.ts index d6586a61f4..e058050d37 100644 --- a/site/src/api/queries/chats.ts +++ b/site/src/api/queries/chats.ts @@ -4,6 +4,8 @@ import type { QueryClient, UseInfiniteQueryOptions } from "react-query"; export const chatsKey = ["chats"] as const; export const chatKey = (chatId: string) => ["chats", chatId] as const; +export const chatMessagesKey = (chatId: string) => + ["chats", chatId, "messages"] as const; /** * Updates a single chat inside every page of the infinite chats query @@ -96,12 +98,17 @@ export const chat = (chatId: string) => ({ queryFn: () => API.getChat(chatId), }); +export const chatMessages = (chatId: string) => ({ + queryKey: chatMessagesKey(chatId), + queryFn: () => API.getChatMessages(chatId), +}); + export const archiveChat = (queryClient: QueryClient) => ({ mutationFn: (chatId: string) => API.archiveChat(chatId), onMutate: async (chatId: string) => { await queryClient.cancelQueries({ queryKey: chatsKey }); await queryClient.cancelQueries({ queryKey: chatKey(chatId) }); - const previousChat = queryClient.getQueryData( + const previousChat = queryClient.getQueryData( chatKey(chatId), ); updateInfiniteChatsCache(queryClient, (chats) => @@ -110,9 +117,9 @@ export const archiveChat = (queryClient: QueryClient) => ({ ), ); if (previousChat) { - queryClient.setQueryData(chatKey(chatId), { + queryClient.setQueryData(chatKey(chatId), { ...previousChat, - chat: { ...previousChat.chat, archived: true }, + archived: true, }); } return { previousChat }; @@ -122,14 +129,14 @@ export const archiveChat = (queryClient: QueryClient) => ({ chatId: string, context: | { - previousChat?: TypesGen.ChatWithMessages; + previousChat?: TypesGen.Chat; } | undefined, ) => { // Rollback: invalidate to re-fetch the correct state. void queryClient.invalidateQueries({ queryKey: chatsKey }); if (context?.previousChat) { - queryClient.setQueryData( + queryClient.setQueryData( chatKey(chatId), context.previousChat, ); @@ -146,7 +153,7 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ onMutate: async (chatId: string) => { await queryClient.cancelQueries({ queryKey: chatsKey }); await queryClient.cancelQueries({ queryKey: chatKey(chatId) }); - const previousChat = queryClient.getQueryData( + const previousChat = queryClient.getQueryData( chatKey(chatId), ); updateInfiniteChatsCache(queryClient, (chats) => @@ -155,9 +162,9 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ ), ); if (previousChat) { - queryClient.setQueryData(chatKey(chatId), { + queryClient.setQueryData(chatKey(chatId), { ...previousChat, - chat: { ...previousChat.chat, archived: false }, + archived: false, }); } return { previousChat }; @@ -167,14 +174,14 @@ export const unarchiveChat = (queryClient: QueryClient) => ({ chatId: string, context: | { - previousChat?: TypesGen.ChatWithMessages; + previousChat?: TypesGen.Chat; } | undefined, ) => { // Rollback: invalidate to re-fetch the correct state. void queryClient.invalidateQueries({ queryKey: chatsKey }); if (context?.previousChat) { - queryClient.setQueryData( + queryClient.setQueryData( chatKey(chatId), context.previousChat, ); @@ -215,6 +222,7 @@ export const editChatMessage = (queryClient: QueryClient, chatId: string) => ({ onSuccess: () => { void queryClient.invalidateQueries({ queryKey: chatsKey }); void queryClient.invalidateQueries({ queryKey: chatKey(chatId) }); + void queryClient.invalidateQueries({ queryKey: chatMessagesKey(chatId) }); }, }); @@ -233,6 +241,7 @@ export const deleteChatQueuedMessage = ( API.deleteChatQueuedMessage(chatId, queuedMessageId), onSuccess: async () => { await queryClient.invalidateQueries({ queryKey: chatKey(chatId) }); + await queryClient.invalidateQueries({ queryKey: chatMessagesKey(chatId) }); }, }); @@ -245,6 +254,7 @@ export const promoteChatQueuedMessage = ( onSuccess: () => { void queryClient.invalidateQueries({ queryKey: chatsKey }); void queryClient.invalidateQueries({ queryKey: chatKey(chatId) }); + void queryClient.invalidateQueries({ queryKey: chatMessagesKey(chatId) }); }, }); diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index e0c2f29e35..681d77e8e0 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -1229,6 +1229,15 @@ export interface ChatMessageUsage { readonly context_limit?: number; } +// From codersdk/chats.go +/** + * ChatMessagesResponse contains the messages and queued messages for a chat. + */ +export interface ChatMessagesResponse { + readonly messages: readonly ChatMessage[]; + readonly queued_messages: readonly ChatQueuedMessage[]; +} + // From codersdk/chats.go /** * ChatModel represents a model in the chat model catalog. @@ -1631,16 +1640,6 @@ export interface ChatSystemPromptResponse { readonly system_prompt: string; } -// From codersdk/chats.go -/** - * ChatWithMessages is a chat along with its messages. - */ -export interface ChatWithMessages { - readonly chat: Chat; - readonly messages: readonly ChatMessage[]; - readonly queued_messages: readonly ChatQueuedMessage[]; -} - // From codersdk/client.go /** * CoderDesktopTelemetryHeader contains a JSON-encoded representation of Desktop telemetry diff --git a/site/src/pages/AgentsPage/AgentDetail.stories.tsx b/site/src/pages/AgentsPage/AgentDetail.stories.tsx index 51f9c6e0b2..7fc3c0d746 100644 --- a/site/src/pages/AgentsPage/AgentDetail.stories.tsx +++ b/site/src/pages/AgentsPage/AgentDetail.stories.tsx @@ -15,6 +15,7 @@ import { chatDiffContentsKey, chatDiffStatusKey, chatKey, + chatMessagesKey, chatModelsKey, chatsKey, } from "api/queries/chats"; @@ -131,13 +132,15 @@ index abc1234..def5678 100644 } `; -/** Build `parameters.queries` entries for a given chat data object. */ +/** Build `parameters.queries` entries for a given chat and messages. */ const buildQueries = ( - chatData: TypesGen.ChatWithMessages, + chat: TypesGen.Chat, + messagesData: TypesGen.ChatMessagesResponse, opts?: { diffUrl?: string }, ) => [ - { key: chatKey(CHAT_ID), data: chatData }, - { key: chatsKey, data: [chatData.chat] }, + { key: chatKey(CHAT_ID), data: chat }, + { key: chatMessagesKey(CHAT_ID), data: messagesData }, + { key: chatsKey, data: [chat] }, { key: chatDiffStatusKey(CHAT_ID), data: { @@ -217,12 +220,12 @@ export const WithMessageHistory: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Markdown rendering showcase", - status: "completed", - }, + id: CHAT_ID, + ...baseChatFields, + title: "Markdown rendering showcase", + status: "completed", + }, + { messages: [ // -- Turn 1: user asks for a summary -- { @@ -533,8 +536,8 @@ export const WithMessageHistory: Story = { }, }; -/** Skeleton placeholder when no query data is available yet. */ -export const Loading: Story = {}; +/** Skeleton placeholder when no query data is available yet. */ export const Loading: Story = + {}; /** Full layout with actions menu and diff panel portaled to the right slot. */ export const CompletedWithDiffPanel: Story = { @@ -545,15 +548,12 @@ export const CompletedWithDiffPanel: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Build a feature", - status: "completed", - }, - messages: [], - queued_messages: [], + id: CHAT_ID, + ...baseChatFields, + title: "Build a feature", + status: "completed", }, + { messages: [], queued_messages: [] }, { diffUrl: "https://github.com/coder/coder/pull/123" }, ), }, @@ -583,15 +583,12 @@ export const NoDiffUrl: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "No diff yet", - status: "completed", - }, - messages: [], - queued_messages: [], + id: CHAT_ID, + ...baseChatFields, + title: "No diff yet", + status: "completed", }, + { messages: [], queued_messages: [] }, { diffUrl: undefined }, ), }, @@ -602,12 +599,12 @@ export const WithSubagentCards: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Parent agent", - status: "running", - }, + id: CHAT_ID, + ...baseChatFields, + title: "Parent agent", + status: "running", + }, + { messages: [ { id: 1, @@ -654,12 +651,12 @@ export const WithReasoningCollapsed: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Reasoning title", - status: "completed", - }, + id: CHAT_ID, + ...baseChatFields, + title: "Reasoning title", + status: "completed", + }, + { messages: [ { id: 1, @@ -706,15 +703,12 @@ export const StreamedSubagentTitle: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Streaming title", - status: "running", - }, - messages: [], - queued_messages: [], + id: CHAT_ID, + ...baseChatFields, + title: "Streaming title", + status: "running", }, + { messages: [], queued_messages: [] }, { diffUrl: undefined }, ), webSocket: { @@ -761,15 +755,12 @@ export const SidebarWithPRAndRepos: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Full sidebar demo", - status: "completed", - }, - messages: [], - queued_messages: [], + id: CHAT_ID, + ...baseChatFields, + title: "Full sidebar demo", + status: "completed", }, + { messages: [], queued_messages: [] }, { diffUrl: "https://github.com/coder/coder/pull/456" }, ), webSocket: { @@ -945,15 +936,12 @@ export const SidebarWithSingleRepo: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Single repo sidebar", - status: "completed", - }, - messages: [], - queued_messages: [], + id: CHAT_ID, + ...baseChatFields, + title: "Single repo sidebar", + status: "completed", }, + { messages: [], queued_messages: [] }, { diffUrl: undefined }, ), webSocket: { @@ -1010,15 +998,12 @@ export const StreamedReasoningCollapsed: Story = { parameters: { queries: buildQueries( { - chat: { - id: CHAT_ID, - ...baseChatFields, - title: "Streaming reasoning title", - status: "running", - }, - messages: [], - queued_messages: [], + id: CHAT_ID, + ...baseChatFields, + title: "Streaming reasoning title", + status: "running", }, + { messages: [], queued_messages: [] }, { diffUrl: undefined }, ), webSocket: { diff --git a/site/src/pages/AgentsPage/AgentDetail.tsx b/site/src/pages/AgentsPage/AgentDetail.tsx index 4d80d03e38..4b9b10a76f 100644 --- a/site/src/pages/AgentsPage/AgentDetail.tsx +++ b/site/src/pages/AgentsPage/AgentDetail.tsx @@ -2,6 +2,7 @@ import { API, watchWorkspace } from "api/api"; import { chat, chatDiffStatus, + chatMessages, chatModelConfigs, chatModels, chats, @@ -603,8 +604,12 @@ const AgentDetail: FC = () => { ...chat(agentId ?? ""), enabled: Boolean(agentId), }); + const chatMessagesQuery = useQuery({ + ...chatMessages(agentId ?? ""), + enabled: Boolean(agentId), + }); const chatsQuery = useQuery(chats()); - const workspaceId = chatQuery.data?.chat?.workspace_id; + const workspaceId = chatQuery.data?.workspace_id; const workspaceQuery = useQuery({ ...workspaceById(workspaceId ?? ""), enabled: Boolean(workspaceId), @@ -669,11 +674,11 @@ const AgentDetail: FC = () => { [proxy.preferredWildcardHostname, workspaceAgent, workspace], ); - const chatData = chatQuery.data; - const chatRecord = chatData?.chat; + const chatRecord = chatQuery.data; + const chatMessagesData = chatMessagesQuery.data; const isArchived = chatRecord?.archived ?? false; - const chatMessages = chatData?.messages; - const chatQueuedMessages = chatData?.queued_messages; + const chatMessagesList = chatMessagesData?.messages; + const chatQueuedMessages = chatMessagesData?.queued_messages; const chatLastModelConfigID = chatRecord?.last_model_config_id; const modelOptions = useMemo( @@ -729,9 +734,9 @@ const AgentDetail: FC = () => { const { store, clearStreamError } = useChatStore({ chatID: agentId, - chatMessages, + chatMessages: chatMessagesList, chatRecord, - chatData, + chatMessagesData, chatQueuedMessages, setChatErrorReason, clearChatErrorReason, @@ -980,7 +985,7 @@ const AgentDetail: FC = () => { inputValueRef, }); - const chatTitle = chatQuery.data?.chat?.title; + const chatTitle = chatQuery.data?.title; const titleElement = ( @@ -988,7 +993,7 @@ const AgentDetail: FC = () => { ); - const parentChatID = getParentChatID(chatQuery.data?.chat); + const parentChatID = getParentChatID(chatQuery.data); const parentChat = parentChatID ? chatsQuery.data?.find((chat) => chat.id === parentChatID) : undefined; @@ -1074,7 +1079,7 @@ const AgentDetail: FC = () => { requestUnarchiveAgent(agentId); }; - if (chatQuery.isLoading) { + if (chatQuery.isLoading || chatMessagesQuery.isLoading) { return ( { ); } - if (!chatQuery.data || !agentId) { + if (!chatQuery.data || !chatMessagesQuery.data || !agentId) { return ( { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -331,8 +330,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -406,8 +404,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -500,8 +497,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -574,8 +570,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -649,8 +644,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -737,8 +731,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [queuedMessage], }, @@ -781,11 +774,7 @@ describe("useChatStore", () => { rerender({ ...initialOptions, - chatData: { - chat: { - ...makeChat(chatID), - updated_at: "2025-01-01T00:00:01.000Z", - }, + chatMessagesData: { messages: [existingMessage], queued_messages: [queuedMessage], }, @@ -818,8 +807,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [queuedMessage], }, @@ -853,11 +841,7 @@ describe("useChatStore", () => { // data with an empty queue (no queue_update from WS yet). rerender({ ...staleOptions, - chatData: { - chat: { - ...makeChat(chatID), - updated_at: "2025-01-01T00:00:02.000Z", - }, + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -888,12 +872,11 @@ describe("useChatStore", () => { }, }, }); - const initialChatData: TypesGen.ChatWithMessages = { - chat: makeChat(chatID), + const initialChatMessagesData: TypesGen.ChatMessagesResponse = { messages: [existingMessage], queued_messages: [queuedMessage], }; - queryClient.setQueryData(chatKey(chatID), initialChatData); + queryClient.setQueryData(chatMessagesKey(chatID), initialChatMessagesData); const wrapper = ({ children }: PropsWithChildren) => ( {children} @@ -907,7 +890,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: initialChatData, + chatMessagesData: initialChatMessagesData, chatQueuedMessages: [queuedMessage], setChatErrorReason, clearChatErrorReason, @@ -935,8 +918,8 @@ describe("useChatStore", () => { expect(result.current.queuedMessages).toEqual([]); }); expect( - queryClient.getQueryData( - chatKey(chatID), + queryClient.getQueryData( + chatMessagesKey(chatID), )?.queued_messages, ).toEqual([]); }); @@ -970,8 +953,7 @@ describe("useChatStore", () => { chatID: chatID1, chatMessages: [msg1] as TypesGen.ChatMessage[], chatRecord: makeChat(chatID1), - chatData: { - chat: makeChat(chatID1), + chatMessagesData: { messages: [msg1], queued_messages: [] as TypesGen.ChatQueuedMessage[], }, @@ -1019,8 +1001,7 @@ describe("useChatStore", () => { chatID: chatID2, chatMessages: [msg2], chatRecord: makeChat(chatID2), - chatData: { - chat: makeChat(chatID2), + chatMessagesData: { messages: [msg2], queued_messages: [], }, @@ -1057,8 +1038,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [queuedMessage], }, @@ -1113,8 +1093,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -1215,8 +1194,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -1311,8 +1289,7 @@ describe("useChatStore", () => { chatID: chatID1, chatMessages: [msg1] as TypesGen.ChatMessage[], chatRecord: makeChat(chatID1), - chatData: { - chat: makeChat(chatID1), + chatMessagesData: { messages: [msg1], queued_messages: [] as TypesGen.ChatQueuedMessage[], }, @@ -1360,8 +1337,7 @@ describe("useChatStore", () => { chatID: chatID2, chatMessages: [msg2], chatRecord: makeChat(chatID2), - chatData: { - chat: makeChat(chatID2), + chatMessagesData: { messages: [msg2], queued_messages: [], }, @@ -1402,8 +1378,7 @@ describe("useChatStore", () => { chatID: chatID1, chatMessages: [msg1] as TypesGen.ChatMessage[], chatRecord: makeChat(chatID1), - chatData: { - chat: makeChat(chatID1), + chatMessagesData: { messages: [msg1], queued_messages: [queuedMsg], }, @@ -1438,8 +1413,7 @@ describe("useChatStore", () => { chatID: chatID2, chatMessages: [], chatRecord: makeChat(chatID2), - chatData: { - chat: makeChat(chatID2), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1474,8 +1448,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1544,8 +1517,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1605,8 +1577,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1658,8 +1629,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1719,8 +1689,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1794,8 +1763,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1856,8 +1824,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1928,8 +1895,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -1994,8 +1960,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2047,8 +2012,7 @@ describe("useChatStore", () => { chatID, chatMessages: [msg], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [msg], queued_messages: [], }, @@ -2112,8 +2076,7 @@ describe("useChatStore", () => { chatID, chatMessages: [existingMessage], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [existingMessage], queued_messages: [], }, @@ -2236,8 +2199,7 @@ describe("useChatStore", () => { chatID, chatMessages: [], chatRecord: makeChat(chatID), - chatData: { - chat: makeChat(chatID), + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2305,8 +2267,7 @@ describe("updateSidebarChat via stream events", () => { chatID, chatMessages: [], chatRecord: initialChat, - chatData: { - chat: initialChat, + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2369,8 +2330,7 @@ describe("updateSidebarChat via stream events", () => { chatID, chatMessages: [], chatRecord: initialChat, - chatData: { - chat: initialChat, + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2442,8 +2402,7 @@ describe("updateSidebarChat via stream events", () => { chatID, chatMessages: [], chatRecord: initialChat, - chatData: { - chat: initialChat, + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2508,8 +2467,7 @@ describe("updateSidebarChat via stream events", () => { chatID, chatMessages: [], chatRecord: activeChat, - chatData: { - chat: activeChat, + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2581,8 +2539,7 @@ describe("updateSidebarChat via stream events", () => { chatID, chatMessages: [], chatRecord: initialChat, - chatData: { - chat: initialChat, + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2652,8 +2609,7 @@ describe("updateSidebarChat via stream events", () => { chatID, chatMessages: [], chatRecord: initialChat, - chatData: { - chat: initialChat, + chatMessagesData: { messages: [], queued_messages: [], }, @@ -2718,8 +2674,7 @@ describe("updateSidebarChat via stream events", () => { chatID, chatMessages: [], chatRecord: initialChat, - chatData: { - chat: initialChat, + chatMessagesData: { messages: [], queued_messages: [], }, diff --git a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts index 6217cd76f4..d7b68af90e 100644 --- a/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts +++ b/site/src/pages/AgentsPage/AgentDetail/ChatContext.ts @@ -1,5 +1,5 @@ import { watchChat } from "api/api"; -import { chatKey, updateInfiniteChatsCache } from "api/queries/chats"; +import { chatMessagesKey, updateInfiniteChatsCache } from "api/queries/chats"; import type * as TypesGen from "api/typesGenerated"; import { asRecord, asString } from "components/ai-elements/runtimeTypeUtils"; import { @@ -417,7 +417,7 @@ interface UseChatStoreOptions { chatID: string | undefined; chatMessages: readonly TypesGen.ChatMessage[] | undefined; chatRecord: TypesGen.Chat | undefined; - chatData: TypesGen.ChatWithMessages | undefined; + chatMessagesData: TypesGen.ChatMessagesResponse | undefined; chatQueuedMessages: readonly TypesGen.ChatQueuedMessage[] | undefined; setChatErrorReason: (chatID: string, reason: string) => void; clearChatErrorReason: (chatID: string) => void; @@ -444,7 +444,7 @@ export const useChatStore = ( chatID, chatMessages, chatRecord, - chatData, + chatMessagesData, chatQueuedMessages, setChatErrorReason, clearChatErrorReason, @@ -519,22 +519,22 @@ export const useChatStore = ( return; } const nextQueuedMessages = queuedMessages ?? []; - queryClient.setQueryData( - chatKey(chatID), - (currentChat) => { - if (!currentChat) { - return currentChat; + queryClient.setQueryData( + chatMessagesKey(chatID), + (currentData) => { + if (!currentData) { + return currentData; } if ( chatQueuedMessagesEqualByID( - currentChat.queued_messages, + currentData.queued_messages, nextQueuedMessages, ) ) { - return currentChat; + return currentData; } return { - ...currentChat, + ...currentData, queued_messages: nextQueuedMessages, }; }, @@ -568,7 +568,7 @@ export const useChatStore = ( }, [chatID, store]); useEffect(() => { - if (!chatID || !chatData) { + if (!chatID || !chatMessagesData) { return; } // Allow re-hydration from REST as long as the WebSocket hasn't @@ -584,7 +584,7 @@ export const useChatStore = ( } queuedMessagesHydratedChatIDRef.current = chatID; store.setQueuedMessages(chatQueuedMessages); - }, [chatData, chatID, chatQueuedMessages, store]); + }, [chatMessagesData, chatID, chatQueuedMessages, store]); useEffect(() => { cancelScheduledStreamReset(); diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index 5a1424a899..d96d31dc0e 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -426,7 +426,7 @@ const AgentsPage: FC = () => { } return chats; }); - queryClient.setQueryData( + queryClient.setQueryData( chatKey(updatedChat.id), (previousChat) => { if (!previousChat) { @@ -434,15 +434,12 @@ const AgentsPage: FC = () => { } return { ...previousChat, - chat: { - ...previousChat.chat, - ...(isStatusEvent && { status: updatedChat.status }), - ...(isTitleEvent && { title: updatedChat.title }), - updated_at: - previousChat.chat.updated_at > updatedChat.updated_at - ? previousChat.chat.updated_at - : updatedChat.updated_at, - }, + ...(isStatusEvent && { status: updatedChat.status }), + ...(isTitleEvent && { title: updatedChat.title }), + updated_at: + previousChat.updated_at > updatedChat.updated_at + ? previousChat.updated_at + : updatedChat.updated_at, }; }, );