From c4db03f11ac955a55dc8dd14351d443de48da18a Mon Sep 17 00:00:00 2001 From: Ethan <39577870+ethanndickson@users.noreply.github.com> Date: Tue, 17 Mar 2026 00:44:10 +1100 Subject: [PATCH] perf(coderd/database): skip redundant chat row update in InsertChatMessage (#23111) ## Summary - add an `IS DISTINCT FROM` guard to `InsertChatMessage`'s `updated_chat` CTE so `chats.last_model_config_id` is only rewritten when the incoming `model_config_id` actually changes - regenerate the query layer - add focused regression coverage for the two meaningful behaviors: same-model inserts and real model switches - trim redundant message-field assertions so the new test stays focused on the guard behavior ## Proof this is an improvement This PR reduces work in the hottest chat write query without changing the insert behavior. ### Why the old query did unnecessary work Before this change, `InsertChatMessage` always ran this update whenever `model_config_id` was non-null: ```sql UPDATE chats SET last_model_config_id = sqlc.narg('model_config_id')::uuid WHERE id = @chat_id::uuid AND sqlc.narg('model_config_id')::uuid IS NOT NULL ``` That means the query rewrote the `chats` row even when `chats.last_model_config_id` was already equal to the incoming value. ### What changes in this PR This PR adds: ```sql AND chats.last_model_config_id IS DISTINCT FROM sqlc.narg('model_config_id')::uuid ``` So same-model inserts still insert the message, but they no longer perform a redundant `UPDATE chats`. ### Why this matters on the hot path From the chat scaletest investigation that motivated this change: - `InsertChatMessage` (+ `updated_chat` CTE) was the hottest write query - about **104k calls** - about **0.69 ms average latency** - about **71.8 s total DB execution time** We also verified common callsites where the update is provably redundant: - `CreateChat` inserts the chat with `LastModelConfigID = opts.ModelConfigID`, then immediately inserts initial system/user messages with that same model config - follow-up user messages commonly pass `lockedChat.LastModelConfigID` straight into `InsertChatMessage` - assistant/tool/summary persistence keeps the current model in the common case; only real switches or fallback cases need the chat row update That means a meaningful fraction of executions of the hottest DB write query move from: - **before:** insert message **+** rewrite chat row - **after:** insert message only This should reduce row churn and write contention on `chats`, especially against other chat-row writers like `UpdateChatStatus` and `GetChatByIDForUpdate`. --- coderd/database/querier_test.go | 125 ++++++++++++++++++++++++++++++ coderd/database/queries.sql.go | 1 + coderd/database/queries/chats.sql | 1 + 3 files changed, 127 insertions(+) diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index c186e2e11e..df1e278b83 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -8997,6 +8997,131 @@ func TestInsertWorkspaceAgentDevcontainers(t *testing.T) { } } +func TestInsertChatMessage(t *testing.T) { + t.Parallel() + + insertModelConfig := func( + t *testing.T, + store database.Store, + ctx context.Context, + userID uuid.UUID, + provider string, + model string, + displayName string, + isDefault bool, + ) database.ChatModelConfig { + t.Helper() + + modelConfig, err := store.InsertChatModelConfig(ctx, database.InsertChatModelConfigParams{ + Provider: provider, + Model: model, + DisplayName: displayName, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + Enabled: true, + IsDefault: isDefault, + ContextLimit: 128000, + CompressionThreshold: 80, + Options: json.RawMessage(`{}`), + }) + require.NoError(t, err) + + return modelConfig + } + + setupChat := func(t *testing.T) (database.Store, context.Context, database.User, database.Chat, string, database.ChatModelConfig) { + t.Helper() + + store, _ := dbtestutil.NewDB(t) + ctx := context.Background() + + dbgen.Organization(t, store, database.Organization{}) + user := dbgen.User(t, store, database.User{}) + provider := "openai" + + _, err := store.InsertChatProvider(ctx, database.InsertChatProviderParams{ + Provider: provider, + DisplayName: "OpenAI", + APIKey: "test-key", + Enabled: true, + }) + require.NoError(t, err) + + modelConfigA := insertModelConfig( + t, + store, + ctx, + user.ID, + provider, + "test-model-a-"+uuid.NewString(), + "Test Model A", + true, + ) + + chat, err := store.InsertChat(ctx, database.InsertChatParams{ + OwnerID: user.ID, + LastModelConfigID: modelConfigA.ID, + Title: "test-chat-" + uuid.NewString(), + }) + require.NoError(t, err) + + return store, ctx, user, chat, provider, modelConfigA + } + + insertMessage := func(t *testing.T, store database.Store, ctx context.Context, chatID, userID, modelConfigID uuid.UUID, content string) { + t.Helper() + + _, err := store.InsertChatMessage(ctx, database.InsertChatMessageParams{ + ChatID: chatID, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + Role: database.ChatMessageRoleUser, + ContentVersion: chatprompt.CurrentContentVersion, + Visibility: database.ChatMessageVisibilityBoth, + Content: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(fmt.Sprintf("%q", content)), + Valid: true, + }, + }) + require.NoError(t, err) + } + + t.Run("ModelSwitchUpdatesLastModelConfigID", func(t *testing.T) { + t.Parallel() + + store, ctx, user, chat, provider, modelConfigA := setupChat(t) + modelConfigB := insertModelConfig( + t, + store, + ctx, + user.ID, + provider, + "test-model-b-"+uuid.NewString(), + "Test Model B", + false, + ) + + insertMessage(t, store, ctx, chat.ID, user.ID, modelConfigB.ID, "switch models") + + gotChat, err := store.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, modelConfigA.ID, chat.LastModelConfigID) + require.Equal(t, modelConfigB.ID, gotChat.LastModelConfigID) + }) + + t.Run("SameModelDoesNotBreakAnything", func(t *testing.T) { + t.Parallel() + + store, ctx, user, chat, _, modelConfigA := setupChat(t) + + insertMessage(t, store, ctx, chat.ID, user.ID, modelConfigA.ID, "same model") + + gotChat, err := store.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, modelConfigA.ID, gotChat.LastModelConfigID) + }) +} + func TestGetChatMessagesForPromptByChatID(t *testing.T) { t.Parallel() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 5ca326923c..ecc8e46989 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -4211,6 +4211,7 @@ WITH updated_chat AS ( WHERE id = $1::uuid AND $3::uuid IS NOT NULL + AND chats.last_model_config_id IS DISTINCT FROM $3::uuid ) INSERT INTO chat_messages ( chat_id, diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 6c20b46275..c45a8f056e 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -165,6 +165,7 @@ WITH updated_chat AS ( WHERE id = @chat_id::uuid AND sqlc.narg('model_config_id')::uuid IS NOT NULL + AND chats.last_model_config_id IS DISTINCT FROM sqlc.narg('model_config_id')::uuid ) INSERT INTO chat_messages ( chat_id,