From 548390ff4b1951c905ca18cb05231a6028457419 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Wed, 8 Jul 2026 07:53:30 +0000 Subject: [PATCH 1/4] fix(coderd): stop manual title generation from writing to chat_messages Manual title generation (RegenerateChatTitle/ProposeChatTitle) recorded token cost by inserting a hidden assistant message into chat_messages and immediately soft-deleting it. Both statements fire AFTER-STATEMENT triggers that sync chats.history_version to snapshot_version. This was the only chat_messages writer outside the chatstate state machine, so it bumped history_version without advancing snapshot_version or publishing a state update. An in-flight generation task captures history_version at spawn and verifies it via a fence at commit. The out-of-band bump made that fence fail, exiting the task via the non-retryable path that performs no cleanup because it assumes a replacement task exists. None is spawned, so the chat stayed running forever and the UI showed "Thinking" indefinitely. Remove the chatd-side accounting path entirely. AI Gateway (aibridge) already records title-call usage independently in aibridge_interceptions and aibridge_token_usages, so the hidden-message accounting was redundant. persistManualTitle now only performs the optimistic title write. This intentionally removes title-generation cost and tokens from chatd's own chat-level cost surfaces (GetChatCostSummary and the spend-limit query paths); that usage now lives only in AI Gateway data. --- coderd/database/querier.go | 6 + coderd/database/queries.sql.go | 6 + coderd/database/queries/chats.sql | 6 + coderd/exp_chats_test.go | 80 +++++ coderd/x/chatd/ARCHITECTURE.md | 2 + coderd/x/chatd/chatd.go | 274 +++--------------- coderd/x/chatd/chatd_internal_test.go | 13 +- coderd/x/chatd/chatd_test.go | 4 +- coderd/x/chatd/quickgen.go | 10 +- coderd/x/chatd/quickgen_internal_test.go | 11 +- .../x/chatd/title_override_internal_test.go | 6 +- 11 files changed, 149 insertions(+), 269 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 218688f7b40..96b603f7b5a 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1041,6 +1041,12 @@ type sqlcQuerier interface { // with concurrent FinalizeStale under READ COMMITTED isolation. InsertChatDebugStep(ctx context.Context, arg InsertChatDebugStepParams) (ChatDebugStep, error) InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error) + // WARNING: All chat_messages writes must go through chatstate + // transitions. AFTER-STATEMENT triggers sync chats.history_version to + // snapshot_version on any chat_messages insert/update, so an + // out-of-band write (even of a hidden or soft-deleted row) breaks the + // history_version fence of an in-flight generation task, killing it + // without a replacement and leaving the chat stuck in running. InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error) // Legacy queue insertion path. When no caller-supplied creator exists, diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index cb2663638c7..9a764d7a475 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -9889,6 +9889,12 @@ type InsertChatMessagesParams struct { RuntimeMs []int64 `db:"runtime_ms" json:"runtime_ms"` } +// WARNING: All chat_messages writes must go through chatstate +// transitions. AFTER-STATEMENT triggers sync chats.history_version to +// snapshot_version on any chat_messages insert/update, so an +// out-of-band write (even of a hidden or soft-deleted row) breaks the +// history_version fence of an in-flight generation task, killing it +// without a replacement and leaving the chat stuck in running. func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) { rows, err := q.db.QueryContext(ctx, insertChatMessages, arg.ChatID, diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 97d2eeb9f7d..3c1648bb89e 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -793,6 +793,12 @@ SELECT * FROM chats_expanded; -- name: InsertChatMessages :many +-- WARNING: All chat_messages writes must go through chatstate +-- transitions. AFTER-STATEMENT triggers sync chats.history_version to +-- snapshot_version on any chat_messages insert/update, so an +-- out-of-band write (even of a hidden or soft-deleted row) breaks the +-- history_version fence of an in-flight generation task, killing it +-- without a replacement and leaving the chat stuck in running. WITH batch AS ( SELECT ( diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 4b93cfa1b8a..9c2d4368997 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -9085,6 +9085,48 @@ func TestRegenerateChatTitle(t *testing.T) { require.Equal(t, "Test Chat", updated.Title) }) + t.Run("DoesNotBumpHistoryVersion", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + user := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createTitleGenerationModelConfig(t, client) + + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + LastModelConfigID: modelConfig.ID, + Title: "history fence chat", + }) + seedManualTitleSourceMessage(t, db, chat, modelConfig.ID) + + // Reproduce the state that broke in production: a state-machine + // transition bumped snapshot_version after the last history + // write, so history_version lags behind. An in-flight generation + // task holds the lagging history_version as its commit fence. + // Any write to chat_messages here would fire the AFTER-STATEMENT + // triggers, sync history_version to snapshot_version, and kill + // that task without a replacement, leaving the chat stuck in + // running. + _, err := db.LockChatAndBumpSnapshotVersion(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + + before, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.NotEqual(t, before.SnapshotVersion, before.HistoryVersion, + "setup must leave history_version lagging snapshot_version") + + updated, err := client.RegenerateChatTitle(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, "Test Chat", updated.Title) + + after, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Equal(t, before.HistoryVersion, after.HistoryVersion, + "manual title regeneration must not touch chat_messages") + }) + t.Run("NoDefaultModelConfig", func(t *testing.T) { t.Parallel() @@ -9256,6 +9298,44 @@ func TestProposeChatTitle(t *testing.T) { require.True(t, persisted.UpdatedAt.Equal(before.UpdatedAt)) }) + t.Run("DoesNotBumpHistoryVersion", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + user := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createTitleGenerationModelConfig(t, client) + + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: user.OrganizationID, + OwnerID: user.UserID, + LastModelConfigID: modelConfig.ID, + Title: "history fence chat", + }) + seedManualTitleSourceMessage(t, db, chat, modelConfig.ID) + + // See the matching TestRegenerateChatTitle subtest: with + // history_version lagging snapshot_version, any chat_messages + // write here would sync history_version and break an in-flight + // generation task's commit fence. + _, err := db.LockChatAndBumpSnapshotVersion(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + + before, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.NotEqual(t, before.SnapshotVersion, before.HistoryVersion, + "setup must leave history_version lagging snapshot_version") + + resp, err := client.ProposeChatTitle(ctx, chat.ID) + require.NoError(t, err) + require.Equal(t, "Test Chat", resp.Title) + + after, err := db.GetChatByID(dbauthz.AsSystemRestricted(ctx), chat.ID) + require.NoError(t, err) + require.Equal(t, before.HistoryVersion, after.HistoryVersion, + "title proposal must not touch chat_messages") + }) + t.Run("NoDefaultModelConfig", func(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index d1519231ff6..f8965949337 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -261,6 +261,8 @@ Each row in `chat_messages` has a `revision` column. It stores the `chats.snapsh Message revision triggers depend on the transition invariant that `snapshot_version` is allocated immediately after the chat row is locked and before any message mutation happens. Runtime code must not assign `chat_messages.revision` directly. +All `chat_messages` writes must go through state machine transitions. Because the triggers fire on every insert or meaningful update, an out-of-band write (even of a hidden or soft-deleted row) advances `history_version` without a corresponding state update, which breaks the `history_version` fence of an in-flight generation task and kills it without a replacement. Manual title endpoints (`RegenerateChatTitle`/`ProposeChatTitle`) used to violate this by inserting and soft-deleting an accounting message for token usage; they no longer write to `chat_messages` at all. Title-generation usage is tracked independently by AI Gateway. + A `BEFORE INSERT` trigger assigns the current chat `snapshot_version` to the inserted message row and records the same value as the chat's latest history version: ```sql diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index ccc313a89ee..9feef216858 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -40,7 +40,6 @@ import ( "github.com/coder/coder/v2/coderd/webpush" "github.com/coder/coder/v2/coderd/workspacestats" "github.com/coder/coder/v2/coderd/x/chatd/chatadvisor" - "github.com/coder/coder/v2/coderd/x/chatd/chatcost" "github.com/coder/coder/v2/coderd/x/chatd/chatdebug" "github.com/coder/coder/v2/coderd/x/chatd/chaterror" "github.com/coder/coder/v2/coderd/x/chatd/chatloop" @@ -2129,21 +2128,6 @@ func (p *Server) ReconcileInvalidStateChat( const manualTitleMessageWindowLimit = 50 -type manualTitleCandidateResult struct { - title string - modelConfig database.ChatModelConfig - usage fantasy.Usage - activeAPIKeyID string - hasMessages bool -} - -type manualTitleGenerationError struct { - cause error - modelConfig database.ChatModelConfig - usage fantasy.Usage - activeAPIKeyID string -} - // generatedChatTitle carries the title produced by the detached // automatic title-generation goroutine. maybeGenerateChatTitle stores // the generated title here so tests can observe it without a database @@ -2177,14 +2161,6 @@ func (t *generatedChatTitle) Load() (string, bool) { return t.title, true } -func (e *manualTitleGenerationError) Error() string { - return e.cause.Error() -} - -func (e *manualTitleGenerationError) Unwrap() error { - return e.cause -} - // RegenerateChatTitle regenerates a chat title from the chat's visible // messages, persists it when it changes, and broadcasts the update. func (p *Server) RegenerateChatTitle( @@ -2195,15 +2171,11 @@ func (p *Server) RegenerateChatTitle( // keeping chat ownership authorization at the HTTP layer. //nolint:gocritic // Non-admin users need chatd-scoped config reads here. chatdCtx := dbauthz.AsChatd(ctx) - updatedChat, err := p.regenerateChatTitleWithStore( + return p.regenerateChatTitleWithStore( chatdCtx, p.db, chat, ) - if err != nil { - return database.Chat{}, p.recordManualTitleGenerationFailure(ctx, chat, err) - } - return updatedChat, nil } // RenameChatTitle persists a user-supplied chat title. @@ -2243,57 +2215,20 @@ func (p *Server) ProposeChatTitle( ) (string, error) { //nolint:gocritic // Non-admin users need chatd-scoped config reads here. chatdCtx := dbauthz.AsChatd(ctx) - title, err := p.proposeChatTitleWithStore(chatdCtx, p.db, chat) - if err != nil { - return "", p.recordManualTitleGenerationFailure(ctx, chat, err) - } - return title, nil + return p.generateManualTitleCandidate(chatdCtx, p.db, chat) } -func (p *Server) recordManualTitleGenerationFailure( - ctx context.Context, - chat database.Chat, - err error, -) error { - var generationErr *manualTitleGenerationError - if !errors.As(err, &generationErr) { - return err - } - - //nolint:gocritic // Failure accounting still needs chatd-scoped config reads. - recordCtx, recordCancel := context.WithTimeout( - dbauthz.AsChatd(context.WithoutCancel(ctx)), - 5*time.Second, - ) - defer recordCancel() - if _, _, recordErr := recordManualTitleUsage( - recordCtx, - p.db, - chat, - generationErr.modelConfig, - generationErr.usage, - generationErr.activeAPIKeyID, - "", - ); recordErr != nil { - return errors.Join( - generationErr, - xerrors.Errorf("record manual title usage: %w", recordErr), - ) - } - return generationErr -} - -// generateManualTitleCandidate performs only model generation and returns the -// candidate plus accounting metadata. Endpoint-specific commit paths are -// responsible for recording usage and deciding whether to persist the title. +// generateManualTitleCandidate generates a title candidate from the chat's +// visible messages. It returns "" when the chat has no messages to summarize. +// Endpoint-specific commit paths decide whether to persist the title. // The context may carry the caller's delegated API key for manual title routes. func (p *Server) generateManualTitleCandidate( ctx context.Context, store database.Store, chat database.Chat, -) (manualTitleCandidateResult, error) { +) (string, error) { if limitErr := p.checkUsageLimit(ctx, store, chat.OwnerID, uuid.NullUUID{UUID: chat.OrganizationID, Valid: true}); limitErr != nil { - return manualTitleCandidateResult{}, limitErr + return "", limitErr } headMessages, err := store.GetChatMessagesByChatIDAscPaginated( @@ -2305,7 +2240,7 @@ func (p *Server) generateManualTitleCandidate( }, ) if err != nil { - return manualTitleCandidateResult{}, xerrors.Errorf("get head chat messages: %w", err) + return "", xerrors.Errorf("get head chat messages: %w", err) } tailMessages, err := store.GetChatMessagesByChatIDDescPaginated( ctx, @@ -2316,15 +2251,15 @@ func (p *Server) generateManualTitleCandidate( }, ) if err != nil { - return manualTitleCandidateResult{}, xerrors.Errorf("get tail chat messages: %w", err) + return "", xerrors.Errorf("get tail chat messages: %w", err) } messages := mergeManualTitleMessages(headMessages, tailMessages) if len(messages) == 0 { - return manualTitleCandidateResult{}, nil + return "", nil } pasteText, err := titlePasteText(ctx, store, messages) if err != nil { - return manualTitleCandidateResult{}, xerrors.Errorf("get pasted-text attachments for manual title: %w", err) + return "", xerrors.Errorf("get pasted-text attachments for manual title: %w", err) } modelOpts := modelBuildOptionsFromMessages(messages) // Manual title routes can run over messages that lack API key attribution. @@ -2336,13 +2271,8 @@ func (p *Server) generateManualTitleCandidate( } model, modelConfig, err := p.resolveManualTitleModel(ctx, store, chat, modelOpts) - result := manualTitleCandidateResult{ - modelConfig: modelConfig, - activeAPIKeyID: modelOpts.ActiveAPIKeyID, - hasMessages: true, - } if err != nil { - return result, err + return "", err } titleCtx := ctx @@ -2360,7 +2290,7 @@ func (p *Server) generateManualTitleCandidate( ) } - title, usage, err := generateManualTitle( + title, err := generateManualTitle( titleCtx, messages, pasteText, @@ -2368,51 +2298,11 @@ func (p *Server) generateManualTitleCandidate( p.titleGenerationProviderOptions(ctx, titleModel, modelConfig), ) finishDebugRun(err) - result.title = title - result.usage = usage - if err != nil { - wrappedErr := xerrors.Errorf("generate manual title: %w", err) - if usage == (fantasy.Usage{}) { - return result, wrappedErr - } - return result, &manualTitleGenerationError{ - cause: wrappedErr, - modelConfig: modelConfig, - usage: usage, - activeAPIKeyID: modelOpts.ActiveAPIKeyID, - } - } - - return result, nil -} - -func (p *Server) proposeChatTitleWithStore( - ctx context.Context, - store database.Store, - chat database.Chat, -) (string, error) { - result, err := p.generateManualTitleCandidate(ctx, store, chat) if err != nil { - return "", err - } - if !result.hasMessages { - return "", nil + return "", xerrors.Errorf("generate manual title: %w", err) } - recordCtx, recordCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - defer recordCancel() - if _, _, recordErr := recordManualTitleUsage( - recordCtx, - store, - chat, - result.modelConfig, - result.usage, - result.activeAPIKeyID, - "", - ); recordErr != nil { - return "", xerrors.Errorf("record manual title usage: %w", recordErr) - } - return result.title, nil + return title, nil } func (p *Server) regenerateChatTitleWithStore( @@ -2420,31 +2310,22 @@ func (p *Server) regenerateChatTitleWithStore( store database.Store, chat database.Chat, ) (database.Chat, error) { - result, err := p.generateManualTitleCandidate(ctx, store, chat) + title, err := p.generateManualTitleCandidate(ctx, store, chat) if err != nil { return database.Chat{}, err } - if !result.hasMessages { + if title == "" { return chat, nil } - recordCtx, recordCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) - defer recordCancel() + // Generation already happened; don't let a client disconnect drop the + // title write. + persistCtx, persistCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer persistCancel() - updatedChat, wroteTitle, recordErr := recordManualTitleUsage( - recordCtx, - store, - chat, - result.modelConfig, - result.usage, - result.activeAPIKeyID, - result.title, - ) - if recordErr != nil { - if result.title != "" { - return database.Chat{}, xerrors.Errorf("record manual title usage and update chat title: %w", recordErr) - } - return database.Chat{}, xerrors.Errorf("record manual title usage: %w", recordErr) + updatedChat, wroteTitle, err := persistManualTitle(persistCtx, store, chat, title) + if err != nil { + return database.Chat{}, xerrors.Errorf("update chat title: %w", err) } // Publish only when this regeneration wrote the title. When a // concurrent rename won the race, the rename path already published @@ -2754,116 +2635,29 @@ func mergeManualTitleMessages( return merged } -func fantasyUsageToChatMessageUsage(usage fantasy.Usage) codersdk.ChatMessageUsage { - var chatUsage codersdk.ChatMessageUsage - if usage.InputTokens != 0 { - chatUsage.InputTokens = ptr.Ref(usage.InputTokens) - } - if usage.OutputTokens != 0 { - chatUsage.OutputTokens = ptr.Ref(usage.OutputTokens) - } - if usage.ReasoningTokens != 0 { - chatUsage.ReasoningTokens = ptr.Ref(usage.ReasoningTokens) - } - if usage.CacheCreationTokens != 0 { - chatUsage.CacheCreationTokens = ptr.Ref(usage.CacheCreationTokens) - } - if usage.CacheReadTokens != 0 { - chatUsage.CacheReadTokens = ptr.Ref(usage.CacheReadTokens) - } - return chatUsage -} - -// recordManualTitleUsage stores token accounting for a manual title -// generation and, when newTitle is set, persists it only if the chat -// title still matches the caller's snapshot. The returned bool reports -// whether the title was actually written; it is false when newTitle is -// empty, when a concurrent writer changed the title first, or when -// newTitle matches the current title. -func recordManualTitleUsage( +// persistManualTitle writes newTitle only if the chat title still +// matches the caller's snapshot. The returned bool reports whether the +// title was actually written; it is false when a concurrent writer +// changed the title first or when newTitle matches the current title. +// Token usage for manual title generation is not recorded here; AI +// Gateway tracks it independently, and writing to chat_messages outside +// the chatstate state machine would break in-flight task fences. +func persistManualTitle( ctx context.Context, store database.Store, chat database.Chat, - modelConfig database.ChatModelConfig, - usage fantasy.Usage, - activeAPIKeyID string, newTitle string, ) (database.Chat, bool, error) { - hasUsage := usage != (fantasy.Usage{}) - if !hasUsage && newTitle == "" { - return chat, false, nil - } - - var totalCostMicros *int64 - if hasUsage { - callConfig := codersdk.ChatModelCallConfig{} - if len(modelConfig.Options) > 0 { - if err := json.Unmarshal(modelConfig.Options, &callConfig); err != nil { - return database.Chat{}, false, xerrors.Errorf("parse model call config: %w", err) - } - } - totalCostMicros = chatcost.CalculateTotalCostMicros( - fantasyUsageToChatMessageUsage(usage), - callConfig.Cost, - ) - } - - // Use a valid empty JSON array for the content column. - // MarshalParts returns a null NullRawMessage for empty - // slices, which becomes an empty string that PostgreSQL - // rejects as invalid JSON. - content := "[]" - updatedChat := chat wroteTitle := false err := store.InTx(func(tx database.Store) error { lockedChat, err := tx.GetChatByIDForUpdate(ctx, chat.ID) if err != nil { - return xerrors.Errorf("lock chat for manual title usage: %w", err) + return xerrors.Errorf("lock chat for manual title persist: %w", err) } updatedChat = lockedChat wroteTitle = false - if hasUsage { - messages, err := tx.InsertChatMessages(ctx, database.InsertChatMessagesParams{ - ChatID: chat.ID, - CreatedBy: []uuid.UUID{chat.OwnerID}, - APIKeyID: []string{activeAPIKeyID}, - ModelConfigID: []uuid.UUID{modelConfig.ID}, - ReasoningEffort: []string{""}, - Role: []database.ChatMessageRole{database.ChatMessageRoleAssistant}, - Content: []string{content}, - ContentVersion: []int16{chatprompt.CurrentContentVersion}, - Visibility: []database.ChatMessageVisibility{database.ChatMessageVisibilityModel}, - InputTokens: []int64{usage.InputTokens}, - OutputTokens: []int64{usage.OutputTokens}, - TotalTokens: []int64{usage.TotalTokens}, - ReasoningTokens: []int64{usage.ReasoningTokens}, - CacheCreationTokens: []int64{usage.CacheCreationTokens}, - CacheReadTokens: []int64{usage.CacheReadTokens}, - ContextLimit: []int64{modelConfig.ContextLimit}, - Compressed: []bool{false}, - TotalCostMicros: []int64{ptr.NilToDefault(totalCostMicros, 0)}, - RuntimeMs: []int64{0}, - }) - if err != nil { - return xerrors.Errorf("insert manual title usage message: %w", err) - } - if len(messages) != 1 { - return xerrors.Errorf("expected 1 manual title usage message, got %d", len(messages)) - } - if err := tx.SoftDeleteChatMessageByID(ctx, messages[0].ID); err != nil { - return xerrors.Errorf("soft delete manual title usage message: %w", err) - } - if lockedChat.LastModelConfigID != modelConfig.ID { - if _, err := tx.UpdateChatLastModelConfigByID(ctx, database.UpdateChatLastModelConfigByIDParams{ - ID: chat.ID, - LastModelConfigID: lockedChat.LastModelConfigID, - }); err != nil { - return xerrors.Errorf("restore chat model config after manual title usage: %w", err) - } - } - } - if newTitle != "" && lockedChat.Title == chat.Title && newTitle != lockedChat.Title { + if lockedChat.Title == chat.Title && newTitle != lockedChat.Title { updatedChat, err = tx.UpdateChatByID(ctx, database.UpdateChatByIDParams{ ID: chat.ID, Title: newTitle, diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index 98adc25ccef..81c190f86ab 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -895,15 +895,6 @@ func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) { ) usageTx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(chat, nil) - usageTx.EXPECT().InsertChatMessages(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatMessagesParams{})).DoAndReturn( - func(_ context.Context, arg database.InsertChatMessagesParams) ([]database.ChatMessage, error) { - require.Equal(t, []uuid.UUID{ownerID}, arg.CreatedBy) - require.Equal(t, []uuid.UUID{modelConfigID}, arg.ModelConfigID) - require.Equal(t, []string{"[]"}, arg.Content) - return []database.ChatMessage{{ID: 91}}, nil - }, - ) - usageTx.EXPECT().SoftDeleteChatMessageByID(gomock.Any(), int64(91)).Return(nil) usageTx.EXPECT().UpdateChatByID(gomock.Any(), database.UpdateChatByIDParams{ ID: chatID, Title: wantTitle, @@ -924,7 +915,7 @@ func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) { } } -// With no request-level locking, recordManualTitleUsage's re-read under +// With no request-level locking, persistManualTitle's re-read under // GetChatByIDForUpdate is the only protection against clobbering a title // that changed while the model call ran. The strict mock has no // UpdateChatByID expectation, so any persist attempt fails the test. @@ -1048,8 +1039,6 @@ func TestRegenerateChatTitle_SkipsPersistWhenTitleChangedConcurrently(t *testing ) usageTx.EXPECT().GetChatByIDForUpdate(gomock.Any(), chatID).Return(landedChat, nil) - usageTx.EXPECT().InsertChatMessages(gomock.Any(), gomock.AssignableToTypeOf(database.InsertChatMessagesParams{})).Return([]database.ChatMessage{{ID: 91}}, nil) - usageTx.EXPECT().SoftDeleteChatMessageByID(gomock.Any(), int64(91)).Return(nil) gotChat, err := server.RegenerateChatTitle(ctx, chat) require.NoError(t, err) diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index 300ed4c647d..2b4bf34441b 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -8446,6 +8446,8 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { require.Equal(t, message.ID, runs[0].HistoryTipMessageID.Int64) } if !tt.wantErr { + // Title generation must not write accounting rows to + // chat_messages; usage is tracked by AI Gateway. var usageMessages int err = rawDB.QueryRowContext( ctx, @@ -8453,7 +8455,7 @@ func TestProposeChatTitle_DebugRun(t *testing.T) { chat.ID, ).Scan(&usageMessages) require.NoError(t, err) - require.Equal(t, 1, usageMessages) + require.Equal(t, 0, usageMessages) } }) } diff --git a/coderd/x/chatd/quickgen.go b/coderd/x/chatd/quickgen.go index ab63d86e495..47b65567465 100644 --- a/coderd/x/chatd/quickgen.go +++ b/coderd/x/chatd/quickgen.go @@ -919,7 +919,7 @@ func generateManualTitle( pasteText map[uuid.UUID]string, fallbackModel fantasy.LanguageModel, providerOptions fantasy.ProviderOptions, -) (string, fantasy.Usage, error) { +) (string, error) { turns := extractManualTitleTurns(messages, pasteText) selected := selectManualTitleTurnIndexes(turns) @@ -927,7 +927,7 @@ func generateManualTitle( return turn.role == string(database.ChatMessageRoleUser) }) if firstUserIndex == -1 { - return "", fantasy.Usage{}, nil + return "", nil } firstUserText := truncateRunes(turns[firstUserIndex].text, maxLatestUserMessageRunes) @@ -946,7 +946,7 @@ func generateManualTitle( userInput = strings.TrimSpace(firstUserText) } - title, usage, err := generateStructuredTitleWithUsage( + title, _, err := generateStructuredTitleWithUsage( titleCtx, fallbackModel, providerOptions, @@ -954,10 +954,10 @@ func generateManualTitle( userInput, ) if err != nil { - return "", usage, err + return "", err } - return title, usage, nil + return title, nil } const turnStatusLabelPrompt = "You write compact chat status labels for a sidebar or push notification. " + diff --git a/coderd/x/chatd/quickgen_internal_test.go b/coderd/x/chatd/quickgen_internal_test.go index cdfa5dbdfaa..d96eceec4cf 100644 --- a/coderd/x/chatd/quickgen_internal_test.go +++ b/coderd/x/chatd/quickgen_internal_test.go @@ -705,7 +705,7 @@ func Test_generateManualTitle_UsesTimeout(t *testing.T) { }, } - title, _, err := generateManualTitle( + title, err := generateManualTitle( context.Background(), messages, nil, @@ -743,7 +743,7 @@ func Test_generateManualTitle_TruncatesFirstUserInput(t *testing.T) { }, } - _, _, err := generateManualTitle( + _, err := generateManualTitle( context.Background(), messages, nil, @@ -753,7 +753,7 @@ func Test_generateManualTitle_TruncatesFirstUserInput(t *testing.T) { require.NoError(t, err) } -func Test_generateManualTitle_ReturnsUsageForEmptyNormalizedTitle(t *testing.T) { +func Test_generateManualTitle_ErrorsOnEmptyNormalizedTitle(t *testing.T) { t.Parallel() messages := []database.ChatMessage{ @@ -778,7 +778,7 @@ func Test_generateManualTitle_ReturnsUsageForEmptyNormalizedTitle(t *testing.T) }, } - _, usage, err := generateManualTitle( + _, err := generateManualTitle( context.Background(), messages, nil, @@ -786,9 +786,6 @@ func Test_generateManualTitle_ReturnsUsageForEmptyNormalizedTitle(t *testing.T) nil, ) require.ErrorContains(t, err, "generated title was empty") - require.Equal(t, int64(11), usage.InputTokens) - require.Equal(t, int64(7), usage.OutputTokens) - require.Equal(t, int64(18), usage.TotalTokens) } func Test_selectPreferredConfiguredShortTextModelConfig(t *testing.T) { diff --git a/coderd/x/chatd/title_override_internal_test.go b/coderd/x/chatd/title_override_internal_test.go index 17a32c5db43..cdd301455f9 100644 --- a/coderd/x/chatd/title_override_internal_test.go +++ b/coderd/x/chatd/title_override_internal_test.go @@ -642,15 +642,13 @@ func TestGenerateManualTitleCandidate_ActiveAPIKeyIDFallback(t *testing.T) { server := titleOverrideTestServer(db, logger) server.aibridgeTransportFactory = aibridgeTestFactoryPointer(factory) - result, err := server.generateManualTitleCandidate(ctx, db, chat) + title, err := server.generateManualTitleCandidate(ctx, db, chat) if tt.wantErrContains != "" { require.ErrorContains(t, err, tt.wantErrContains) return } require.NoError(t, err) - require.Equal(t, wantTitle, result.title) - require.True(t, result.hasMessages) - require.Equal(t, tt.wantAPIKeyID, result.activeAPIKeyID) + require.Equal(t, wantTitle, title) require.Equal(t, tt.wantAPIKeyID, testutil.RequireReceive(ctx, t, seenAPIKeyID)) }) } From 5bb4d53a5e27e4c22501f7ee253f0a5504915239 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Wed, 8 Jul 2026 09:32:54 +0000 Subject: [PATCH 2/4] chore(coderd): consolidate chat_messages write invariant into ARCHITECTURE.md Drop the redundant WARNING comment on InsertChatMessages (and its generated propagation into querier.go/queries.sql.go) and the standalone ARCHITECTURE.md paragraph. The rule that every chat_messages write must go through a state machine transition now lives as a single sentence alongside the existing runtime-code guardrail in the message revision section. --- coderd/database/querier.go | 6 ------ coderd/database/queries.sql.go | 6 ------ coderd/database/queries/chats.sql | 6 ------ coderd/x/chatd/ARCHITECTURE.md | 4 +--- 4 files changed, 1 insertion(+), 21 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 96b603f7b5a..218688f7b40 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1041,12 +1041,6 @@ type sqlcQuerier interface { // with concurrent FinalizeStale under READ COMMITTED isolation. InsertChatDebugStep(ctx context.Context, arg InsertChatDebugStepParams) (ChatDebugStep, error) InsertChatFile(ctx context.Context, arg InsertChatFileParams) (InsertChatFileRow, error) - // WARNING: All chat_messages writes must go through chatstate - // transitions. AFTER-STATEMENT triggers sync chats.history_version to - // snapshot_version on any chat_messages insert/update, so an - // out-of-band write (even of a hidden or soft-deleted row) breaks the - // history_version fence of an in-flight generation task, killing it - // without a replacement and leaving the chat stuck in running. InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error) // Legacy queue insertion path. When no caller-supplied creator exists, diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 9a764d7a475..cb2663638c7 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -9889,12 +9889,6 @@ type InsertChatMessagesParams struct { RuntimeMs []int64 `db:"runtime_ms" json:"runtime_ms"` } -// WARNING: All chat_messages writes must go through chatstate -// transitions. AFTER-STATEMENT triggers sync chats.history_version to -// snapshot_version on any chat_messages insert/update, so an -// out-of-band write (even of a hidden or soft-deleted row) breaks the -// history_version fence of an in-flight generation task, killing it -// without a replacement and leaving the chat stuck in running. func (q *sqlQuerier) InsertChatMessages(ctx context.Context, arg InsertChatMessagesParams) ([]ChatMessage, error) { rows, err := q.db.QueryContext(ctx, insertChatMessages, arg.ChatID, diff --git a/coderd/database/queries/chats.sql b/coderd/database/queries/chats.sql index 3c1648bb89e..97d2eeb9f7d 100644 --- a/coderd/database/queries/chats.sql +++ b/coderd/database/queries/chats.sql @@ -793,12 +793,6 @@ SELECT * FROM chats_expanded; -- name: InsertChatMessages :many --- WARNING: All chat_messages writes must go through chatstate --- transitions. AFTER-STATEMENT triggers sync chats.history_version to --- snapshot_version on any chat_messages insert/update, so an --- out-of-band write (even of a hidden or soft-deleted row) breaks the --- history_version fence of an in-flight generation task, killing it --- without a replacement and leaving the chat stuck in running. WITH batch AS ( SELECT ( diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index f8965949337..73e6f5eafcf 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -259,9 +259,7 @@ Each row in `chat_messages` has a `revision` column. It stores the `chats.snapsh `chats.history_version` stores the latest `snapshot_version` in which chat message history changed. It starts at `0`, remains unchanged for non-history transitions, and is set to the current `snapshot_version` whenever a message is inserted or meaningfully updated. A newly created chat starts with `snapshot_version = 1`; because `Create` inserts initial history in that snapshot, the created chat's `history_version` becomes `1`. No-op message updates do not advance message `revision`, advance `history_version`, or reset `generation_attempt`. Whenever `history_version` changes, `generation_attempt` is reset to `0`; generation attempts are scoped to the current history version. -Message revision triggers depend on the transition invariant that `snapshot_version` is allocated immediately after the chat row is locked and before any message mutation happens. Runtime code must not assign `chat_messages.revision` directly. - -All `chat_messages` writes must go through state machine transitions. Because the triggers fire on every insert or meaningful update, an out-of-band write (even of a hidden or soft-deleted row) advances `history_version` without a corresponding state update, which breaks the `history_version` fence of an in-flight generation task and kills it without a replacement. Manual title endpoints (`RegenerateChatTitle`/`ProposeChatTitle`) used to violate this by inserting and soft-deleting an accounting message for token usage; they no longer write to `chat_messages` at all. Title-generation usage is tracked independently by AI Gateway. +Message revision triggers depend on the transition invariant that `snapshot_version` is allocated immediately after the chat row is locked and before any message mutation happens. Runtime code must not assign `chat_messages.revision` directly, and every `chat_messages` insert or update must go through a state machine transition: the triggers advance `history_version` on any write, so an out-of-band write (even of a hidden or soft-deleted row) moves `history_version` without a matching `snapshot_version` bump and breaks the fence of an in-flight generation task. A `BEFORE INSERT` trigger assigns the current chat `snapshot_version` to the inserted message row and records the same value as the chat's latest history version: From e684fbed21727f3bddddfce55469753e5dfe9a14 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Wed, 8 Jul 2026 09:57:56 +0000 Subject: [PATCH 3/4] chore(coderd): trim regression test comments --- coderd/exp_chats_test.go | 16 ++++------------ 1 file changed, 4 insertions(+), 12 deletions(-) diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index 9c2d4368997..f8631572833 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -9101,14 +9101,9 @@ func TestRegenerateChatTitle(t *testing.T) { }) seedManualTitleSourceMessage(t, db, chat, modelConfig.ID) - // Reproduce the state that broke in production: a state-machine - // transition bumped snapshot_version after the last history - // write, so history_version lags behind. An in-flight generation - // task holds the lagging history_version as its commit fence. - // Any write to chat_messages here would fire the AFTER-STATEMENT - // triggers, sync history_version to snapshot_version, and kill - // that task without a replacement, leaving the chat stuck in - // running. + // Leave history_version lagging snapshot_version, as when a + // generation task is in flight. A chat_messages write here would + // sync it and break that task's commit fence. _, err := db.LockChatAndBumpSnapshotVersion(dbauthz.AsSystemRestricted(ctx), chat.ID) require.NoError(t, err) @@ -9314,10 +9309,7 @@ func TestProposeChatTitle(t *testing.T) { }) seedManualTitleSourceMessage(t, db, chat, modelConfig.ID) - // See the matching TestRegenerateChatTitle subtest: with - // history_version lagging snapshot_version, any chat_messages - // write here would sync history_version and break an in-flight - // generation task's commit fence. + // See the matching TestRegenerateChatTitle subtest. _, err := db.LockChatAndBumpSnapshotVersion(dbauthz.AsSystemRestricted(ctx), chat.ID) require.NoError(t, err) From 240b118365cb2946e26588e8ce2c4d585478a4e5 Mon Sep 17 00:00:00 2001 From: Ethan Dickson Date: Mon, 13 Jul 2026 05:59:11 +0000 Subject: [PATCH 4/4] chore(coderd/x/chatd): name manual-title persist timeout constant --- coderd/x/chatd/chatd.go | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 9feef216858..3e4d17d9446 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -68,6 +68,7 @@ const ( homeInstructionLookupTimeout = 5 * time.Second workspaceDialValidationDelay = 5 * time.Second turnStatusLabelWriteTimeout = 5 * time.Second + manualTitlePersistTimeout = 5 * time.Second // defaultDialTimeout matches the timeout used by ~8 other // server-side AgentConn callers. defaultDialTimeout = 30 * time.Second @@ -2320,7 +2321,7 @@ func (p *Server) regenerateChatTitleWithStore( // Generation already happened; don't let a client disconnect drop the // title write. - persistCtx, persistCancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + persistCtx, persistCancel := context.WithTimeout(context.WithoutCancel(ctx), manualTitlePersistTimeout) defer persistCancel() updatedChat, wroteTitle, err := persistManualTitle(persistCtx, store, chat, title)