diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 795c0fe262..f3d0c870ef 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -37,7 +37,6 @@ import ( "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/database/dbtime" "github.com/coder/coder/v2/coderd/externalauth" - coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/rbac/policy" "github.com/coder/coder/v2/coderd/util/ptr" @@ -1970,7 +1969,7 @@ func TestWatchChats(t *testing.T) { require.NotZero(t, got.UpdatedAt) }) - t.Run("DiffStatusChangeIncludesDiffStatus", func(t *testing.T) { + t.Run("DiffStatusChangeIncludesDiffStatusAndOmitsInjectedContext", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) @@ -1979,9 +1978,17 @@ func TestWatchChats(t *testing.T) { }) client := codersdk.NewExperimentalClient(rawClient) db := api.Database + chatDaemon := api.ChatDaemonForTest() user := coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) + lastInjectedContext, err := json.Marshal([]codersdk.ChatMessagePart{{ + Type: codersdk.ChatMessagePartTypeSkill, + SkillName: "large-skill", + SkillDescription: strings.Repeat("x", 9000), + }}) + require.NoError(t, err) + // Insert a chat and a diff status row. chat := dbgen.Chat(t, db, database.Chat{ OrganizationID: user.OrganizationID, @@ -1989,9 +1996,20 @@ func TestWatchChats(t *testing.T) { LastModelConfigID: modelConfig.ID, Title: "diff status watch test", }) + chat, err = db.UpdateChatLastInjectedContext( + dbauthz.AsChatd(ctx), + database.UpdateChatLastInjectedContextParams{ + ID: chat.ID, + LastInjectedContext: pqtype.NullRawMessage{ + RawMessage: lastInjectedContext, + Valid: true, + }, + }, + ) + require.NoError(t, err) refreshedAt := time.Now().UTC().Truncate(time.Second) staleAt := refreshedAt.Add(time.Hour) - _, err := db.UpsertChatDiffStatusReference( + _, err = db.UpsertChatDiffStatusReference( dbauthz.AsSystemRestricted(ctx), database.UpsertChatDiffStatusReferenceParams{ ChatID: chat.ID, @@ -2017,36 +2035,16 @@ func TestWatchChats(t *testing.T) { ) require.NoError(t, err) + storedChat, err := client.GetChat(ctx, chat.ID) + require.NoError(t, err) + require.NotEmpty(t, storedChat.LastInjectedContext) + // Open the watch WebSocket. conn, err := client.Dial(ctx, "/api/experimental/chats/watch", nil) require.NoError(t, err) defer conn.Close(websocket.StatusNormalClosure, "done") - // Publish a diff_status_change event via pubsub, - // mimicking what PublishDiffStatusChange does after - // it reads the diff status from the DB. - dbStatus, err := db.GetChatDiffStatusByChatID(dbauthz.AsSystemRestricted(ctx), chat.ID) - require.NoError(t, err) - sdkDiffStatus := db2sdk.ChatDiffStatus(chat.ID, &dbStatus) - event := codersdk.ChatWatchEvent{ - Kind: codersdk.ChatWatchEventKindDiffStatusChange, - Chat: codersdk.Chat{ - ID: chat.ID, - OwnerID: chat.OwnerID, - Title: chat.Title, - Status: codersdk.ChatStatus(chat.Status), - CreatedAt: chat.CreatedAt, - UpdatedAt: chat.UpdatedAt, - DiffStatus: &sdkDiffStatus, - }, - } - payload, err := json.Marshal(event) - require.NoError(t, err) - - // A single publish is sufficient because the subscription - // is active before websocket.Accept (and thus before Dial - // returns). This serves as a regression test for the fix. - err = api.Pubsub.Publish(coderdpubsub.ChatWatchEventChannel(user.UserID), payload) + err = chatDaemon.PublishDiffStatusChange(dbauthz.AsChatd(ctx), chat.ID) require.NoError(t, err) var received codersdk.ChatWatchEvent @@ -2071,6 +2069,7 @@ func TestWatchChats(t *testing.T) { require.EqualValues(t, 42, ds.Additions) require.EqualValues(t, 7, ds.Deletions) require.EqualValues(t, 5, ds.ChangedFiles) + require.Empty(t, received.Chat.LastInjectedContext) }) t.Run("ArchiveAndUnarchiveEmitEventsForDescendants", func(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index af448f23c3..adfcc55610 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -3493,23 +3493,29 @@ func (p *Server) publishChatPubsubEvents(chats []database.Chat, kind codersdk.Ch } } +// chatWatchEventSDKChat builds the chat embedded in ChatWatchEvent +// notifications. These payloads travel through PostgreSQL NOTIFY, so +// omit fields that can grow large and that watch consumers already read +// from the REST chat endpoint. +func chatWatchEventSDKChat(chat database.Chat, diffStatus *codersdk.ChatDiffStatus) codersdk.Chat { + sdkChat := db2sdk.Chat(chat, nil, nil) + sdkChat.Files = nil + sdkChat.LastInjectedContext = nil + if diffStatus != nil { + sdkChat.DiffStatus = diffStatus + } + return sdkChat +} + // publishChatPubsubEvent broadcasts a chat lifecycle event via PostgreSQL // pubsub so that all replicas can push updates to watching clients. func (p *Server) publishChatPubsubEvent(chat database.Chat, kind codersdk.ChatWatchEventKind, diffStatus *codersdk.ChatDiffStatus) { if p.pubsub == nil { return } - // diffStatus is applied below. File metadata is intentionally - // omitted from pubsub events to avoid an extra DB query per - // publish. Clients must merge pubsub updates, not replace - // cached file metadata. - sdkChat := db2sdk.Chat(chat, nil, nil) - if diffStatus != nil { - sdkChat.DiffStatus = diffStatus - } event := codersdk.ChatWatchEvent{ Kind: kind, - Chat: sdkChat, + Chat: chatWatchEventSDKChat(chat, diffStatus), } payload, err := json.Marshal(event) if err != nil { diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index e627556d3d..8b2be67c7a 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -12,7 +12,6 @@ import ( "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/database" - "github.com/coder/coder/v2/coderd/database/db2sdk" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" @@ -514,7 +513,7 @@ func (s *taskStarter) publishWatchWithRetry( func publishChatWatchEvent(pubsub chatWorkerPubsub, chat database.Chat, kind codersdk.ChatWatchEventKind) error { event := codersdk.ChatWatchEvent{ Kind: kind, - Chat: db2sdk.Chat(chat, nil, nil), + Chat: chatWatchEventSDKChat(chat, nil), } payload, err := json.Marshal(event) if err != nil {