From 09085053482a8492f6cedd4daa2c6bec86065195 Mon Sep 17 00:00:00 2001 From: Kyle Carberry Date: Mon, 2 Mar 2026 12:00:00 -0500 Subject: [PATCH] fix(chats): archive chat tree with single query instead of loop (#22496) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem When archiving an agent with subagents, the children briefly flash in the sidebar as root-level items before disappearing. Two issues: 1. **Backend:** Archive used N+1 queries — a recursive DFS (`archiveChatTree`, no transaction) or BFS loop (`chatd.ArchiveChat`, N+1 queries in a tx) to walk the tree and archive each chat individually. 2. **Frontend:** The SSE `deleted` event handler only filtered out the parent chat from the cache. Children remained briefly, got promoted to root-level by `buildChatTree`, then disappeared on the next re-fetch. ## Fix **Backend:** Replace both tree-walk implementations with a single SQL query: ```sql UPDATE chats SET archived = true, updated_at = NOW() WHERE id = @id OR root_chat_id = @id; ``` This leverages the existing `root_chat_id` column (already indexed) to archive the entire tree atomically. **Frontend:** When a `deleted` event arrives, also filter out any chats whose `root_chat_id` matches the deleted chat, so children vanish from the sidebar immediately with the parent. ## Changes - `coderd/database/queries/chats.sql` — Added `ArchiveChatTreeByID` query - `coderd/chats.go` — Use single query, delete `archiveChatTree` function - `coderd/chatd/chatd.go` — Simplify `ArchiveChat` to use single query - `coderd/database/dbauthz/dbauthz.go` — Auth wrapper for new query - `coderd/chats_test.go` — Added `TestArchiveChat/ArchivesChildren` subtest - `site/src/pages/AgentsPage/AgentsPage.tsx` — Filter children in SSE handler - Generated files updated via `make gen` --- coderd/chatd/chatd.go | 34 +------------ coderd/chats.go | 20 +------- coderd/chats_test.go | 61 ++++++++++++++++++++++++ coderd/database/queries.sql.go | 3 +- coderd/database/queries/chats.sql | 3 +- site/src/pages/AgentsPage/AgentsPage.tsx | 5 +- 6 files changed, 72 insertions(+), 54 deletions(-) diff --git a/coderd/chatd/chatd.go b/coderd/chatd/chatd.go index bb43a7a9df..31be73030d 100644 --- a/coderd/chatd/chatd.go +++ b/coderd/chatd/chatd.go @@ -517,38 +517,8 @@ func (p *Server) ArchiveChat(ctx context.Context, chatID uuid.UUID) error { return xerrors.Errorf("get chat: %w", err) } - err = p.db.InTx(func(tx database.Store) error { - // Collect descendants breadth-first, then archive from leaves upward. - descendantIDs := make([]uuid.UUID, 0) - queue := []uuid.UUID{chatID} - for len(queue) > 0 { - parentID := queue[0] - queue = queue[1:] - - children, err := tx.ListChildChatsByParentID(ctx, parentID) - if err != nil { - return xerrors.Errorf("list children of chat %s: %w", parentID, err) - } - for _, child := range children { - descendantIDs = append(descendantIDs, child.ID) - queue = append(queue, child.ID) - } - } - - for i := len(descendantIDs) - 1; i >= 0; i-- { - if err := tx.ArchiveChatByID(ctx, descendantIDs[i]); err != nil { - return xerrors.Errorf("archive descendant chat %s: %w", descendantIDs[i], err) - } - } - - if err := tx.ArchiveChatByID(ctx, chatID); err != nil { - return xerrors.Errorf("archive chat: %w", err) - } - - return nil - }, nil) - if err != nil { - return err + if err := p.db.ArchiveChatByID(ctx, chatID); err != nil { + return xerrors.Errorf("archive chat: %w", err) } p.publishChatPubsubEvent(chat, coderdpubsub.ChatEventKindDeleted) diff --git a/coderd/chats.go b/coderd/chats.go index 1d176fce1e..c706ae2db1 100644 --- a/coderd/chats.go +++ b/coderd/chats.go @@ -409,12 +409,7 @@ func (api *API) archiveChat(rw http.ResponseWriter, r *http.Request) { return } - var err error - if api.chatDaemon != nil { - err = api.chatDaemon.ArchiveChat(ctx, chat.ID) - } else { - err = archiveChatTree(ctx, api.Database, chat.ID) - } + err := api.Database.ArchiveChatByID(ctx, chat.ID) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to archive chat.", @@ -454,19 +449,6 @@ func (api *API) unarchiveChat(rw http.ResponseWriter, r *http.Request) { rw.WriteHeader(http.StatusNoContent) } -func archiveChatTree(ctx context.Context, store database.Store, chatID uuid.UUID) error { - children, err := store.ListChildChatsByParentID(ctx, chatID) - if err != nil { - return xerrors.Errorf("list child chats: %w", err) - } - for _, child := range children { - if err := archiveChatTree(ctx, store, child.ID); err != nil { - return err - } - } - return store.ArchiveChatByID(ctx, chatID) -} - // EXPERIMENTAL: this endpoint is experimental and is subject to change. func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() diff --git a/coderd/chats_test.go b/coderd/chats_test.go index ffac2eb42e..1bc342c704 100644 --- a/coderd/chats_test.go +++ b/coderd/chats_test.go @@ -1209,6 +1209,67 @@ func TestArchiveChat(t *testing.T) { err := client.ArchiveChat(ctx, uuid.New()) requireSDKError(t, err, http.StatusNotFound) }) + + t.Run("ArchivesChildren", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + user := coderdtest.CreateFirstUser(t, client) + modelConfig := createChatModelConfig(t, client) + + // Create a parent chat via the API. + parentChat, err := client.CreateChat(ctx, codersdk.CreateChatRequest{ + Content: []codersdk.ChatInputPart{ + { + Type: codersdk.ChatInputPartTypeText, + Text: "parent chat", + }, + }, + }) + require.NoError(t, err) + + // Insert child chats directly via the database. + child1, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OwnerID: user.UserID, + LastModelConfigID: modelConfig.ID, + Title: "child 1", + ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + }) + require.NoError(t, err) + + child2, err := db.InsertChat(dbauthz.AsSystemRestricted(ctx), database.InsertChatParams{ + OwnerID: user.UserID, + LastModelConfigID: modelConfig.ID, + Title: "child 2", + ParentChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + RootChatID: uuid.NullUUID{UUID: parentChat.ID, Valid: true}, + }) + require.NoError(t, err) + + // Archive the parent via the API. + err = client.ArchiveChat(ctx, parentChat.ID) + require.NoError(t, err) + + // List chats — none of the family should appear. + chats, err := client.ListChats(ctx) + require.NoError(t, err) + for _, c := range chats { + require.NotEqual(t, parentChat.ID, c.ID, "parent should not appear") + require.NotEqual(t, child1.ID, c.ID, "child1 should not appear") + require.NotEqual(t, child2.ID, c.ID, "child2 should not appear") + } + + // Verify children are archived directly in the DB. + dbChild1, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), child1.ID) + require.NoError(t, err) + require.True(t, dbChild1.Archived, "child1 should be archived") + + dbChild2, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), child2.ID) + require.NoError(t, err) + require.True(t, dbChild2.Archived, "child2 should be archived") + }) } func TestPostChatMessages(t *testing.T) { diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index f43d198ffb..3a3ef330a2 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2876,7 +2876,8 @@ func (q *sqlQuerier) AcquireChat(ctx context.Context, arg AcquireChatParams) (Ch } const archiveChatByID = `-- name: ArchiveChatByID :exec -UPDATE chats SET archived = true, updated_at = NOW() WHERE id = $1::uuid +UPDATE chats SET archived = true, updated_at = NOW() +WHERE id = $1 OR root_chat_id = $1 ` func (q *sqlQuerier) ArchiveChatByID(ctx context.Context, id uuid.UUID) error { diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 26a4b5aff5..bf75a12e62 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -1,5 +1,6 @@ -- name: ArchiveChatByID :exec -UPDATE chats SET archived = true, updated_at = NOW() WHERE id = @id::uuid; +UPDATE chats SET archived = true, updated_at = NOW() +WHERE id = @id OR root_chat_id = @id; -- name: UnarchiveChatByID :exec UPDATE chats SET archived = false, updated_at = NOW() WHERE id = @id::uuid; diff --git a/site/src/pages/AgentsPage/AgentsPage.tsx b/site/src/pages/AgentsPage/AgentsPage.tsx index a0826d8497..9945324817 100644 --- a/site/src/pages/AgentsPage/AgentsPage.tsx +++ b/site/src/pages/AgentsPage/AgentsPage.tsx @@ -322,7 +322,10 @@ const AgentsPage: FC = () => { queryClient.setQueryData( chatsKey, (prev: TypesGen.Chat[] | undefined) => - prev?.filter((c) => c.id !== updatedChat.id), + prev?.filter( + (c) => + c.id !== updatedChat.id && c.root_chat_id !== updatedChat.id, + ), ); queryClient.removeQueries({ queryKey: chatKey(updatedChat.id),