diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 2a776824d77..616a2fedc31 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2217,20 +2217,13 @@ func (q *querier) DeleteChatDebugDataByChatID(ctx context.Context, arg database. return q.db.DeleteChatDebugDataByChatID(ctx, arg) } -func (q *querier) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) error { +func (q *querier) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil { - return err + return uuid.Nil, err } return q.db.DeleteChatModelConfigByID(ctx, id) } -func (q *querier) DeleteChatModelConfigsByAIProviderID(ctx context.Context, aiProviderID uuid.UUID) error { - if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceAIProvider); err != nil { - return err - } - return q.db.DeleteChatModelConfigsByAIProviderID(ctx, aiProviderID) -} - func (q *querier) DeleteChatQueuedMessage(ctx context.Context, arg database.DeleteChatQueuedMessageParams) error { chat, err := q.db.GetChatByID(ctx, arg.ChatID) if err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 73b2c9654d2..a2c8d0733fd 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -704,14 +704,9 @@ func (s *MethodTestSuite) TestChats() { check.Args(msg.ID).Asserts(chat, policy.ActionUpdate).Returns() })) s.Run("DeleteChatModelConfigByID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - id := uuid.New() - dbm.EXPECT().DeleteChatModelConfigByID(gomock.Any(), id).Return(nil).AnyTimes() - check.Args(id).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate) - })) - s.Run("DeleteChatModelConfigsByAIProviderID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - providerID := uuid.New() - dbm.EXPECT().DeleteChatModelConfigsByAIProviderID(gomock.Any(), providerID).Return(nil).AnyTimes() - check.Args(providerID).Asserts(rbac.ResourceAIProvider, policy.ActionDelete) + configID := uuid.New() + dbm.EXPECT().DeleteChatModelConfigByID(gomock.Any(), configID).Return(configID, nil).AnyTimes() + check.Args(configID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(configID) })) s.Run("DeleteChatQueuedMessage", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index ca75449cca4..765d747eee2 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -520,20 +520,12 @@ func (m queryMetricsStore) DeleteChatDebugDataByChatID(ctx context.Context, chat return r0, r1 } -func (m queryMetricsStore) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) error { +func (m queryMetricsStore) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { start := time.Now() - r0 := m.s.DeleteChatModelConfigByID(ctx, id) + r0, r1 := m.s.DeleteChatModelConfigByID(ctx, id) m.queryLatencies.WithLabelValues("DeleteChatModelConfigByID").Observe(time.Since(start).Seconds()) m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatModelConfigByID").Inc() - return r0 -} - -func (m queryMetricsStore) DeleteChatModelConfigsByAIProviderID(ctx context.Context, aiProviderID uuid.UUID) error { - start := time.Now() - r0 := m.s.DeleteChatModelConfigsByAIProviderID(ctx, aiProviderID) - m.queryLatencies.WithLabelValues("DeleteChatModelConfigsByAIProviderID").Observe(time.Since(start).Seconds()) - m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatModelConfigsByAIProviderID").Inc() - return r0 + return r0, r1 } func (m queryMetricsStore) DeleteChatQueuedMessage(ctx context.Context, arg database.DeleteChatQueuedMessageParams) error { diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 328c4d0e8cd..da8dafdd559 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -849,11 +849,12 @@ func (mr *MockStoreMockRecorder) DeleteChatDebugDataByChatID(ctx, arg any) *gomo } // DeleteChatModelConfigByID mocks base method. -func (m *MockStore) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) error { +func (m *MockStore) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "DeleteChatModelConfigByID", ctx, id) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(uuid.UUID) + ret1, _ := ret[1].(error) + return ret0, ret1 } // DeleteChatModelConfigByID indicates an expected call of DeleteChatModelConfigByID. @@ -862,20 +863,6 @@ func (mr *MockStoreMockRecorder) DeleteChatModelConfigByID(ctx, id any) *gomock. return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatModelConfigByID", reflect.TypeOf((*MockStore)(nil).DeleteChatModelConfigByID), ctx, id) } -// DeleteChatModelConfigsByAIProviderID mocks base method. -func (m *MockStore) DeleteChatModelConfigsByAIProviderID(ctx context.Context, aiProviderID uuid.UUID) error { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "DeleteChatModelConfigsByAIProviderID", ctx, aiProviderID) - ret0, _ := ret[0].(error) - return ret0 -} - -// DeleteChatModelConfigsByAIProviderID indicates an expected call of DeleteChatModelConfigsByAIProviderID. -func (mr *MockStoreMockRecorder) DeleteChatModelConfigsByAIProviderID(ctx, aiProviderID any) *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatModelConfigsByAIProviderID", reflect.TypeOf((*MockStore)(nil).DeleteChatModelConfigsByAIProviderID), ctx, aiProviderID) -} - // DeleteChatQueuedMessage mocks base method. func (m *MockStore) DeleteChatQueuedMessage(ctx context.Context, arg database.DeleteChatQueuedMessageParams) error { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index a17b3bc26c3..aac29a4c55d 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -146,8 +146,7 @@ type sqlcQuerier interface { // window (for example, after an unarchive races with a pending // archive-cleanup retry). DeleteChatDebugDataByChatID(ctx context.Context, arg DeleteChatDebugDataByChatIDParams) (int64, error) - DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) error - DeleteChatModelConfigsByAIProviderID(ctx context.Context, aiProviderID uuid.UUID) error + DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) DeleteChatQueuedMessage(ctx context.Context, arg DeleteChatQueuedMessageParams) error // Deletes a queued message, scoped to the parent chat. Returns the // number of affected rows so callers can detect missing rows without diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index c09f38e7313..dd2299dcc4d 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12351,6 +12351,21 @@ func TestInsertWorkspaceAgentDevcontainers(t *testing.T) { } } +func TestDeleteChatModelConfigByID(t *testing.T) { + t.Parallel() + + store, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitMedium) + config := dbgen.ChatModelConfig(t, store, database.ChatModelConfig{}) + + deletedID, err := store.DeleteChatModelConfigByID(ctx, config.ID) + require.NoError(t, err) + require.Equal(t, config.ID, deletedID) + + _, err = store.DeleteChatModelConfigByID(ctx, config.ID) + require.ErrorIs(t, err, sql.ErrNoRows) +} + func TestGetEnabledChatModelConfigsUsesAIProviders(t *testing.T) { t.Parallel() diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 722b9dfe0ae..17e7731bf8f 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6135,7 +6135,7 @@ func (q *sqlQuerier) InsertChatFile(ctx context.Context, arg InsertChatFileParam return i, err } -const deleteChatModelConfigByID = `-- name: DeleteChatModelConfigByID :exec +const deleteChatModelConfigByID = `-- name: DeleteChatModelConfigByID :one UPDATE chat_model_configs SET @@ -6144,28 +6144,15 @@ SET updated_at = NOW() WHERE id = $1::uuid -` - -func (q *sqlQuerier) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) error { - _, err := q.db.ExecContext(ctx, deleteChatModelConfigByID, id) - return err -} - -const deleteChatModelConfigsByAIProviderID = `-- name: DeleteChatModelConfigsByAIProviderID :exec -UPDATE - chat_model_configs -SET - deleted = TRUE, - deleted_at = NOW(), - updated_at = NOW() -WHERE - ai_provider_id = $1::uuid AND deleted = FALSE +RETURNING id ` -func (q *sqlQuerier) DeleteChatModelConfigsByAIProviderID(ctx context.Context, aiProviderID uuid.UUID) error { - _, err := q.db.ExecContext(ctx, deleteChatModelConfigsByAIProviderID, aiProviderID) - return err +func (q *sqlQuerier) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) { + row := q.db.QueryRowContext(ctx, deleteChatModelConfigByID, id) + var id_2 uuid.UUID + err := row.Scan(&id_2) + return id_2, err } const getChatModelConfigByID = `-- name: GetChatModelConfigByID :one diff --git a/coderd/database/queries/chatmodelconfigs.sql b/coderd/database/queries/chatmodelconfigs.sql index ae950839913..1e90da264fc 100644 --- a/coderd/database/queries/chatmodelconfigs.sql +++ b/coderd/database/queries/chatmodelconfigs.sql @@ -123,7 +123,7 @@ WHERE is_default = TRUE AND deleted = FALSE; --- name: DeleteChatModelConfigByID :exec +-- name: DeleteChatModelConfigByID :one UPDATE chat_model_configs SET @@ -131,15 +131,6 @@ SET deleted_at = NOW(), updated_at = NOW() WHERE - id = @id::uuid; - --- name: DeleteChatModelConfigsByAIProviderID :exec -UPDATE - chat_model_configs -SET - deleted = TRUE, - deleted_at = NOW(), - updated_at = NOW() -WHERE - ai_provider_id = @ai_provider_id::uuid - AND deleted = FALSE; + id = @id::uuid + AND deleted = FALSE +RETURNING id; diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 8c7a3fa4f16..4e503abf549 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -4518,6 +4518,20 @@ func (api *API) defaultCreateChatModelConfigID( return defaultModelConfig.ID, 0, nil } +// validateChatCompressionThreshold enforces the chat_model_configs CHECK +// constraint range. +func validateChatCompressionThreshold(threshold int32) error { + if threshold < minChatContextCompressionThreshold || + threshold > maxChatContextCompressionThreshold { + return xerrors.Errorf( + "context_compression_threshold must be between %d and %d", + minChatContextCompressionThreshold, + maxChatContextCompressionThreshold, + ) + } + return nil +} + func normalizeChatCompressionThreshold( requested *int32, fallback int32, @@ -4527,13 +4541,8 @@ func normalizeChatCompressionThreshold( threshold = *requested } - if threshold < minChatContextCompressionThreshold || - threshold > maxChatContextCompressionThreshold { - return 0, xerrors.Errorf( - "context_compression_threshold must be between %d and %d", - minChatContextCompressionThreshold, - maxChatContextCompressionThreshold, - ) + if err := validateChatCompressionThreshold(threshold); err != nil { + return 0, err } return threshold, nil @@ -6689,7 +6698,7 @@ func (api *API) upsertUserAIProviderKey(rw http.ResponseWriter, r *http.Request) return } if !provider.Enabled { - httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{Message: "AI provider is disabled."}) + writeChatProviderPreconditionError(ctx, rw, errChatProviderDisabled) return } var req codersdk.CreateUserAIProviderKeyRequest @@ -6895,13 +6904,17 @@ func validateChatModelConfigProviderModel(aiProvider database.AIProvider, model // through this helper so concurrent writers cannot act on stale reads, e.g. // two creates on an empty deployment both self-promoting to default and // violating the idx_chat_model_configs_single_default unique index. +// +// The transaction must run at ReadCommitted. A snapshot isolation level takes +// the snapshot at the lock statement, so a re-read inside fn would observe +// pre-lock state and reintroduce the lost update this helper prevents. func (api *API) inChatModelConfigWriteTx(ctx context.Context, fn func(tx database.Store) error) error { return api.Database.InTx(func(tx database.Store) error { if err := tx.AcquireLock(ctx, database.LockIDChatModelConfigWrites); err != nil { return xerrors.Errorf("acquire chat model config write lock: %w", err) } return fn(tx) - }, nil) + }, &database.TxOptions{Isolation: sql.LevelReadCommitted}) } func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { @@ -6925,7 +6938,7 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { aiProvider, err := api.Database.GetAIProviderByID(dbauthz.AsChatd(ctx), *req.AIProviderID) if err != nil { if httpapi.Is404Error(err) { - httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{Message: "AI provider is not configured."}) + writeChatProviderPreconditionError(ctx, rw, errChatProviderMissing) return } httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ @@ -6935,7 +6948,7 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { return } if !aiProvider.Enabled { - httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{Message: "AI provider is disabled."}) + writeChatProviderPreconditionError(ctx, rw, errChatProviderDisabled) return } aiProviderID := uuid.NullUUID{UUID: aiProvider.ID, Valid: true} @@ -7011,12 +7024,12 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { lockedAIProvider, err := tx.GetAIProviderByIDForReferenceLock(dbauthz.AsChatd(ctx), insertParams.AIProviderID.UUID) if err != nil { if xerrors.Is(err, sql.ErrNoRows) { - return errChatProviderNotConfigured + return errChatProviderMissing } - return xerrors.Errorf("get AI provider for update: %w", err) + return xerrors.Errorf("get AI provider for create: %w", err) } if !lockedAIProvider.Enabled { - return errChatProviderNotConfigured + return errChatProviderDisabled } if err := validateChatModelConfigProviderModel(lockedAIProvider, insertParams.Model); err != nil { return err @@ -7060,6 +7073,9 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { return nil }) if err != nil { + if writeChatProviderPreconditionError(ctx, rw, err) { + return + } var providerModelErr *chatModelConfigProviderModelError switch { case errors.As(err, &providerModelErr): @@ -7071,12 +7087,6 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { Detail: err.Error(), }) return - case xerrors.Is(err, errChatProviderNotConfigured): - httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{ - Message: "Chat provider is not configured.", - Detail: err.Error(), - }) - return default: httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to create chat model config.", @@ -7104,8 +7114,7 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { return } - existing, err := api.Database.GetChatModelConfigByID(ctx, modelConfigID) - if err != nil { + if _, err := api.Database.GetChatModelConfigByID(ctx, modelConfigID); err != nil { if httpapi.Is404Error(err) { httpapi.ResourceNotFound(rw) return @@ -7122,71 +7131,28 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { return } - aiProviderID := existing.AIProviderID - if req.AIProviderID != nil { - //nolint:gocritic // The route already authorized chat model config updates. - aiProvider, err := api.Database.GetAIProviderByID(dbauthz.AsChatd(ctx), *req.AIProviderID) - if err != nil { - if httpapi.Is404Error(err) { - httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{Message: "AI provider is not configured."}) - return - } - httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ - Message: "Failed to get AI provider.", - Detail: err.Error(), - }) - return - } - if !aiProvider.Enabled { - httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{Message: "AI provider is disabled."}) - return - } - aiProviderID = uuid.NullUUID{UUID: aiProvider.ID, Valid: true} - } - - model := existing.Model - if trimmed := strings.TrimSpace(req.Model); trimmed != "" { - model = trimmed - } - - displayName := existing.DisplayName - if trimmed := strings.TrimSpace(req.DisplayName); trimmed != "" { - displayName = trimmed - } - - enabled := existing.Enabled - if req.Enabled != nil { - enabled = *req.Enabled - } - isDefault := existing.IsDefault - if req.IsDefault != nil { - isDefault = *req.IsDefault + if req.ContextLimit != nil && *req.ContextLimit <= 0 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Context limit must be greater than zero.", + }) + return } - contextLimit := existing.ContextLimit - if req.ContextLimit != nil { - if *req.ContextLimit <= 0 { + // A PATCH that omits the field keeps the stored value, which the + // chat_model_configs CHECK constraint already bounds to 0..100. + var requestedCompressionThreshold *int32 + if req.CompressionThreshold != nil { + if thresholdErr := validateChatCompressionThreshold(*req.CompressionThreshold); thresholdErr != nil { httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Context limit must be greater than zero.", + Message: "Invalid compression threshold.", + Detail: thresholdErr.Error(), }) return } - contextLimit = *req.ContextLimit + requestedCompressionThreshold = req.CompressionThreshold } - compressionThreshold, thresholdErr := normalizeChatCompressionThreshold( - req.CompressionThreshold, - existing.CompressionThreshold, - ) - if thresholdErr != nil { - httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ - Message: "Invalid compression threshold.", - Detail: thresholdErr.Error(), - }) - return - } - - modelConfigRaw := existing.Options + var requestedModelConfig json.RawMessage if req.ModelConfig != nil { encodedModelConfig, modelConfigErr := marshalChatModelCallConfig(req.ModelConfig) if modelConfigErr != nil { @@ -7196,51 +7162,96 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { }) return } - modelConfigRaw = encodedModelConfig + requestedModelConfig = encodedModelConfig } - updateParams := database.UpdateChatModelConfigParams{ - Model: model, - DisplayName: displayName, - Enabled: enabled, - IsDefault: isDefault, - ContextLimit: contextLimit, - CompressionThreshold: compressionThreshold, - Options: modelConfigRaw, - AIProviderID: aiProviderID, - UpdatedBy: uuid.NullUUID{UUID: apiKey.UserID, Valid: apiKey.UserID != uuid.Nil}, - ID: existing.ID, - } - - // Re-derive the provider type under lock when the model or provider changes. - revalidateProviderModel := updateParams.AIProviderID.Valid && (req.AIProviderID != nil || strings.TrimSpace(req.Model) != "") var updated database.ChatModelConfig - err = api.inChatModelConfigWriteTx(ctx, func(tx database.Store) error { + err := api.inChatModelConfigWriteTx(ctx, func(tx database.Store) error { + // The unlocked read above only rejects unknown IDs; a concurrent writer + // can change or delete the row between the two reads, so merge against + // this copy. + lockedExisting, err := tx.GetChatModelConfigByID(ctx, modelConfigID) + if err != nil { + if xerrors.Is(err, sql.ErrNoRows) { + return errChatModelConfigNotFound + } + return xerrors.Errorf("get chat model config for update: %w", err) + } + + model := lockedExisting.Model + if trimmed := strings.TrimSpace(req.Model); trimmed != "" { + model = trimmed + } + displayName := lockedExisting.DisplayName + if trimmed := strings.TrimSpace(req.DisplayName); trimmed != "" { + displayName = trimmed + } + enabled := lockedExisting.Enabled + if req.Enabled != nil { + enabled = *req.Enabled + } + isDefault := lockedExisting.IsDefault + if req.IsDefault != nil { + isDefault = *req.IsDefault + } + contextLimit := lockedExisting.ContextLimit + if req.ContextLimit != nil { + contextLimit = *req.ContextLimit + } + compressionThreshold := lockedExisting.CompressionThreshold + if requestedCompressionThreshold != nil { + compressionThreshold = *requestedCompressionThreshold + } + modelConfigRaw := lockedExisting.Options + if requestedModelConfig != nil { + modelConfigRaw = requestedModelConfig + } + aiProviderID := lockedExisting.AIProviderID + if req.AIProviderID != nil { + aiProviderID = uuid.NullUUID{UUID: *req.AIProviderID, Valid: true} + } + + updateParams := database.UpdateChatModelConfigParams{ + Model: model, + DisplayName: displayName, + Enabled: enabled, + IsDefault: isDefault, + ContextLimit: contextLimit, + CompressionThreshold: compressionThreshold, + Options: modelConfigRaw, + AIProviderID: aiProviderID, + UpdatedBy: uuid.NullUUID{UUID: apiKey.UserID, Valid: apiKey.UserID != uuid.Nil}, + ID: lockedExisting.ID, + } + + // An update that touches neither the provider nor the model cannot + // invalidate the stored provider/model pair. + revalidateProviderModel := updateParams.AIProviderID.Valid && (req.AIProviderID != nil || strings.TrimSpace(req.Model) != "") if revalidateProviderModel { //nolint:gocritic // The route already authorized chat model config updates. aiProvider, err := tx.GetAIProviderByIDForReferenceLock(dbauthz.AsChatd(ctx), updateParams.AIProviderID.UUID) if err != nil { if xerrors.Is(err, sql.ErrNoRows) { - return errChatProviderNotConfigured + return errChatProviderMissing } return xerrors.Errorf("get AI provider for update: %w", err) } if !aiProvider.Enabled { - return errChatProviderNotConfigured + return errChatProviderDisabled } if err := validateChatModelConfigProviderModel(aiProvider, updateParams.Model); err != nil { return err } } - setAsDefault := updateParams.IsDefault && !existing.IsDefault + setAsDefault := updateParams.IsDefault && !lockedExisting.IsDefault if setAsDefault { if err := tx.UnsetDefaultChatModelConfigs(ctx); err != nil { return xerrors.Errorf("unset default model configs: %w", err) } } - _, err := tx.UpdateChatModelConfig(ctx, updateParams) + _, err = tx.UpdateChatModelConfig(ctx, updateParams) if err != nil { if xerrors.Is(err, sql.ErrNoRows) { return errChatModelConfigNotFound @@ -7249,8 +7260,8 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { } excludeConfigID := uuid.Nil - if existing.IsDefault && req.IsDefault != nil && !*req.IsDefault { - excludeConfigID = existing.ID + if lockedExisting.IsDefault && req.IsDefault != nil && !*req.IsDefault { + excludeConfigID = lockedExisting.ID } if err := ensureDefaultChatModelConfig( @@ -7261,18 +7272,17 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { return err } - refreshedConfig, err := tx.GetChatModelConfigByID(ctx, existing.ID) + refreshedConfig, err := tx.GetChatModelConfigByID(ctx, lockedExisting.ID) if err != nil { - if xerrors.Is(err, sql.ErrNoRows) { - // Do not wrap with %w. The outer handler maps target misses to 404. - return xerrors.Errorf("refresh updated chat model config: %v", err) - } return xerrors.Errorf("refresh updated chat model config: %w", err) } updated = refreshedConfig return nil }) if err != nil { + if writeChatProviderPreconditionError(ctx, rw, err) { + return + } var providerModelErr *chatModelConfigProviderModelError switch { case errors.As(err, &providerModelErr): @@ -7284,12 +7294,6 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { Detail: err.Error(), }) return - case xerrors.Is(err, errChatProviderNotConfigured): - httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{ - Message: "Chat provider is not configured.", - Detail: err.Error(), - }) - return case xerrors.Is(err, errChatModelConfigNotFound): httpapi.ResourceNotFound(rw) return @@ -7332,11 +7336,18 @@ func (api *API) deleteChatModelConfig(rw http.ResponseWriter, r *http.Request) { } if err := api.inChatModelConfigWriteTx(ctx, func(tx database.Store) error { - if err := tx.DeleteChatModelConfigByID(ctx, modelConfigID); err != nil { + if _, err := tx.DeleteChatModelConfigByID(ctx, modelConfigID); err != nil { + if xerrors.Is(err, sql.ErrNoRows) { + return errChatModelConfigNotFound + } return err } return ensureDefaultChatModelConfig(ctx, tx) }); err != nil { + if xerrors.Is(err, errChatModelConfigNotFound) { + httpapi.ResourceNotFound(rw) + return + } httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to delete chat model config.", Detail: err.Error(), @@ -7624,9 +7635,24 @@ func validateChatProviderAPIKeySize(apiKey string) error { return nil } +func writeChatProviderPreconditionError(ctx context.Context, rw http.ResponseWriter, err error) bool { + var message string + switch { + case xerrors.Is(err, errChatProviderMissing): + message = "AI provider is not configured." + case xerrors.Is(err, errChatProviderDisabled): + message = "AI provider is disabled." + default: + return false + } + httpapi.Write(ctx, rw, http.StatusPreconditionFailed, codersdk.Response{Message: message}) + return true +} + var ( - errChatModelConfigNotFound = xerrors.New("chat model config not found") - errChatProviderNotConfigured = xerrors.New("chat provider is not configured") + errChatProviderDisabled = xerrors.New("AI provider is disabled") + errChatProviderMissing = xerrors.New("AI provider is not configured") + errChatModelConfigNotFound = xerrors.New("chat model config not found") ) // ChatProviderAPIKeysFromDeploymentValues returns deployment-backed chat diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index cfb8e00b1f9..ec50029fb4f 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -443,44 +443,143 @@ func (s *lockSwitchChatPlanModeInstructionsStore) GetChatPlanModeInstructions(ct return instructions, err } -// failNextUpdateChatModelConfigStore shares its failure state across InTx -// wrappers so tests can force a specific in-transaction model-config update to -// return sql.ErrNoRows. -type failNextUpdateChatModelConfigStore struct { +type chatModelConfigLockedReadMutation struct { + id uuid.UUID + model string +} + +// chatModelConfigHookStore forces a specific in-transaction chat model config +// operation to behave abnormally. Each hook targets one ID and consumes +// itself on first match. The single-consume gate uses atomic.Pointer so a +// concurrent InTx wrapper cannot race with the arming goroutine. +type chatModelConfigHookStore struct { database.Store - failNextUpdateChatModelConfig *atomic.Bool - failNextUpdateChatModelConfigID uuid.UUID + inTx bool + + failNextUpdate *atomic.Pointer[uuid.UUID] + failNextDelete *atomic.Pointer[uuid.UUID] + failProviderReferenceLock *atomic.Pointer[uuid.UUID] + vanishAtLockedRead *atomic.Pointer[uuid.UUID] + mutateAtLockedRead *atomic.Pointer[chatModelConfigLockedReadMutation] } -func newFailNextUpdateChatModelConfigStore(store database.Store) *failNextUpdateChatModelConfigStore { - return &failNextUpdateChatModelConfigStore{ - Store: store, - failNextUpdateChatModelConfig: &atomic.Bool{}, +func newChatModelConfigHookStore(store database.Store) *chatModelConfigHookStore { + return &chatModelConfigHookStore{ + Store: store, + failNextUpdate: &atomic.Pointer[uuid.UUID]{}, + failNextDelete: &atomic.Pointer[uuid.UUID]{}, + failProviderReferenceLock: &atomic.Pointer[uuid.UUID]{}, + vanishAtLockedRead: &atomic.Pointer[uuid.UUID]{}, + mutateAtLockedRead: &atomic.Pointer[chatModelConfigLockedReadMutation]{}, } } -func (s *failNextUpdateChatModelConfigStore) InTx(function func(database.Store) error, txOpts *database.TxOptions) error { +func newChatClientWithModelConfigHookStore(t testing.TB) (*codersdk.ExperimentalClient, *chatModelConfigHookStore) { + t.Helper() + + rawDB, pubsub := dbtestutil.NewDB(t) + store := newChatModelConfigHookStore(rawDB) + client := newChatClient(t, func(opts *coderdtest.Options) { + opts.Database = store + opts.Pubsub = pubsub + }) + return client, store +} + +func (s *chatModelConfigHookStore) armFailNextUpdate(id uuid.UUID) { + s.failNextUpdate.Store(&id) +} + +func (s *chatModelConfigHookStore) armFailNextDelete(id uuid.UUID) { + s.failNextDelete.Store(&id) +} + +func (s *chatModelConfigHookStore) armFailProviderReferenceLock(id uuid.UUID) { + s.failProviderReferenceLock.Store(&id) +} + +func (s *chatModelConfigHookStore) armVanishAtLockedRead(id uuid.UUID) { + s.vanishAtLockedRead.Store(&id) +} + +func (s *chatModelConfigHookStore) armMutateAtLockedRead(id uuid.UUID, model string) { + s.mutateAtLockedRead.Store(&chatModelConfigLockedReadMutation{id: id, model: model}) +} + +func (s *chatModelConfigHookStore) InTx(function func(database.Store) error, txOpts *database.TxOptions) error { return s.Store.InTx(func(tx database.Store) error { - return function(&failNextUpdateChatModelConfigStore{ - Store: tx, - failNextUpdateChatModelConfig: s.failNextUpdateChatModelConfig, - failNextUpdateChatModelConfigID: s.failNextUpdateChatModelConfigID, + return function(&chatModelConfigHookStore{ + Store: tx, + inTx: true, + failNextUpdate: s.failNextUpdate, + failNextDelete: s.failNextDelete, + failProviderReferenceLock: s.failProviderReferenceLock, + vanishAtLockedRead: s.vanishAtLockedRead, + mutateAtLockedRead: s.mutateAtLockedRead, }) }, txOpts) } -func (s *failNextUpdateChatModelConfigStore) UpdateChatModelConfig( +func consumeChatModelConfigHook(hook *atomic.Pointer[uuid.UUID], id uuid.UUID) bool { + target := hook.Load() + if target == nil || *target != id { + return false + } + return hook.CompareAndSwap(target, nil) +} + +func (s *chatModelConfigHookStore) GetAIProviderByIDForReferenceLock( + ctx context.Context, + id uuid.UUID, +) (database.AIProvider, error) { + if consumeChatModelConfigHook(s.failProviderReferenceLock, id) { + return database.AIProvider{}, stderrors.New("forced provider reference lock failure") + } + return s.Store.GetAIProviderByIDForReferenceLock(ctx, id) +} + +func (s *chatModelConfigHookStore) UpdateChatModelConfig( ctx context.Context, arg database.UpdateChatModelConfigParams, ) (database.ChatModelConfig, error) { - if arg.ID == s.failNextUpdateChatModelConfigID && - s.failNextUpdateChatModelConfig.CompareAndSwap(true, false) { + if consumeChatModelConfigHook(s.failNextUpdate, arg.ID) { return database.ChatModelConfig{}, sql.ErrNoRows } return s.Store.UpdateChatModelConfig(ctx, arg) } +func (s *chatModelConfigHookStore) DeleteChatModelConfigByID( + ctx context.Context, + id uuid.UUID, +) (uuid.UUID, error) { + if consumeChatModelConfigHook(s.failNextDelete, id) { + return uuid.Nil, sql.ErrNoRows + } + return s.Store.DeleteChatModelConfigByID(ctx, id) +} + +func (s *chatModelConfigHookStore) GetChatModelConfigByID( + ctx context.Context, + id uuid.UUID, +) (database.ChatModelConfig, error) { + if s.inTx && consumeChatModelConfigHook(s.vanishAtLockedRead, id) { + return database.ChatModelConfig{}, sql.ErrNoRows + } + if s.inTx { + mutation := s.mutateAtLockedRead.Load() + if mutation != nil && mutation.id == id && s.mutateAtLockedRead.CompareAndSwap(mutation, nil) { + row, err := s.Store.GetChatModelConfigByID(ctx, id) + if err != nil { + return row, err + } + row.Model = mutation.model + return row, nil + } + } + return s.Store.GetChatModelConfigByID(ctx, id) +} + func insertAssistantMessage( t *testing.T, db database.Store, @@ -4444,6 +4543,28 @@ func TestCreateChatModelConfig(t *testing.T) { require.Len(t, configs, 1) }) + t.Run("ProviderReferenceLockFailureUsesCreateErrorText", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, store := newChatClientWithModelConfigHookStore(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + aiProvider := createAIProviderForTest(t, client, "openai", "test-api-key") + + store.armFailProviderReferenceLock(aiProvider.ID) + + contextLimit := int64(4096) + _, err := client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{ + AIProviderID: &aiProvider.ID, + Model: "gpt-4o-mini", + ContextLimit: &contextLimit, + }) + sdkErr := requireSDKError(t, err, http.StatusInternalServerError) + require.Equal(t, "Failed to create chat model config.", sdkErr.Message) + require.Contains(t, sdkErr.Detail, "get AI provider for create") + require.NotContains(t, sdkErr.Detail, "get AI provider for update") + }) + t.Run("ConcurrentCreatesElectSingleDefault", func(t *testing.T) { t.Parallel() @@ -5127,20 +5248,11 @@ func TestUpdateChatModelConfig(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - rawDB, pubsub := dbtestutil.NewDB(t) - store := newFailNextUpdateChatModelConfigStore(rawDB) - rawClient, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{ - Database: store, - Pubsub: pubsub, - DeploymentValues: coderdtest.DeploymentValues(t), - }) - aibridgedtest.StartTestAIBridgeDaemon(t.Context(), t, api, nil) - client := codersdk.NewExperimentalClient(rawClient) + client, store := newChatClientWithModelConfigHookStore(t) _ = coderdtest.CreateFirstUser(t, client.Client) modelConfig := createChatModelConfig(t, client) - store.failNextUpdateChatModelConfigID = modelConfig.ID - store.failNextUpdateChatModelConfig.Store(true) + store.armFailNextUpdate(modelConfig.ID) _, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ DisplayName: "missing in tx", @@ -5152,15 +5264,7 @@ func TestUpdateChatModelConfig(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - rawDB, pubsub := dbtestutil.NewDB(t) - store := newFailNextUpdateChatModelConfigStore(rawDB) - rawClient, _, api := coderdtest.NewWithAPI(t, &coderdtest.Options{ - Database: store, - Pubsub: pubsub, - DeploymentValues: coderdtest.DeploymentValues(t), - }) - aibridgedtest.StartTestAIBridgeDaemon(t.Context(), t, api, nil) - client := codersdk.NewExperimentalClient(rawClient) + client, store := newChatClientWithModelConfigHookStore(t) _ = coderdtest.CreateFirstUser(t, client.Client) defaultConfig := createChatModelConfig(t, client) @@ -5176,8 +5280,7 @@ func TestUpdateChatModelConfig(t *testing.T) { }) require.NoError(t, err) - store.failNextUpdateChatModelConfigID = candidateConfig.ID - store.failNextUpdateChatModelConfig.Store(true) + store.armFailNextUpdate(candidateConfig.ID) _, err = client.UpdateChatModelConfig(ctx, defaultConfig.ID, codersdk.UpdateChatModelConfigRequest{ IsDefault: ptr.Ref(false), @@ -5215,6 +5318,135 @@ func TestUpdateChatModelConfig(t *testing.T) { require.Equal(t, "Context limit must be greater than zero.", sdkErr.Message) }) + // A row that disappears between the pre-read and the locked read must + // 404, not silently no-op the update. + t.Run("NotFoundWhenTargetRowDisappearsAtLockedRead", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, store := newChatClientWithModelConfigHookStore(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + store.armVanishAtLockedRead(modelConfig.ID) + + _, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + DisplayName: "vanished before update", + }) + requireSDKError(t, err, http.StatusNotFound) + }) + + t.Run("MergesFromLockedCopy", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, store := newChatClientWithModelConfigHookStore(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + const sentinel = "locked-copy-sentinel-model" + store.armMutateAtLockedRead(modelConfig.ID, sentinel) + + updated, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + DisplayName: "only-display-name", + }) + require.NoError(t, err) + require.Equal(t, sentinel, updated.Model) + require.Equal(t, "only-display-name", updated.DisplayName) + }) + + t.Run("StoredProviderDisabledOnModelOnlyUpdate", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + disabled := false + _, err := client.UpdateAIProvider(ctx, modelConfig.AIProviderID.String(), codersdk.UpdateAIProviderRequest{ + Enabled: &disabled, + }) + require.NoError(t, err) + + _, err = client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + Model: "gpt-4o-different", + }) + sdkErr := requireSDKError(t, err, http.StatusPreconditionFailed) + require.Equal(t, "AI provider is disabled.", sdkErr.Message) + }) + + t.Run("StoredProviderMissingOnModelOnlyUpdate", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + err := client.DeleteAIProvider(ctx, modelConfig.AIProviderID.String()) + require.NoError(t, err) + + _, err = client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + Model: "gpt-4o-different", + }) + sdkErr := requireSDKError(t, err, http.StatusPreconditionFailed) + require.Equal(t, "AI provider is not configured.", sdkErr.Message) + }) + + t.Run("CompressionThresholdOutOfRange", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + _, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + CompressionThreshold: ptr.Ref(int32(150)), + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "Invalid compression threshold.", sdkErr.Message) + }) + + t.Run("CompressionThresholdInRange", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + updated, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + CompressionThreshold: ptr.Ref(int32(55)), + }) + require.NoError(t, err) + require.Equal(t, int32(55), updated.CompressionThreshold) + }) + + t.Run("ModelConfigUpdated", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client := newChatClient(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + updated, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + ModelConfig: &codersdk.ChatModelCallConfig{ + ReasoningEffort: &codersdk.ChatModelReasoningEffortConfig{ + Default: ptr.Ref("high"), + Max: ptr.Ref("high"), + }, + }, + }) + require.NoError(t, err) + require.NotNil(t, updated.ModelConfig) + require.NotNil(t, updated.ModelConfig.ReasoningEffort) + require.NotNil(t, updated.ModelConfig.ReasoningEffort.Default) + require.Equal(t, "high", *updated.ModelConfig.ReasoningEffort.Default) + }) + t.Run("InvalidModelConfigID", func(t *testing.T) { t.Parallel() @@ -5325,6 +5557,20 @@ func TestDeleteChatModelConfig(t *testing.T) { requireSDKError(t, err, http.StatusNotFound) }) + t.Run("NotFoundWhenTargetRowDisappearsAtDelete", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, store := newChatClientWithModelConfigHookStore(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + store.armFailNextDelete(modelConfig.ID) + + err := client.DeleteChatModelConfig(ctx, modelConfig.ID) + requireSDKError(t, err, http.StatusNotFound) + }) + t.Run("InvalidModelConfigID", func(t *testing.T) { t.Parallel() @@ -12604,10 +12850,11 @@ func seedChatWithDeletedModelConfig( Title: "chat without model config", }) seedManualTitleSourceMessage(t, db, chat, modelConfig.ID) - require.NoError(t, db.DeleteChatModelConfigByID( + _, err := db.DeleteChatModelConfigByID( dbauthz.AsSystemRestricted(ctx), modelConfig.ID, - )) + ) + require.NoError(t, err) return chat } diff --git a/coderd/telemetry/telemetry_test.go b/coderd/telemetry/telemetry_test.go index 875882388f4..a03fe05abdb 100644 --- a/coderd/telemetry/telemetry_test.go +++ b/coderd/telemetry/telemetry_test.go @@ -1662,7 +1662,7 @@ func TestChatsTelemetry(t *testing.T) { DisplayName: "Deleted Model", ContextLimit: 100000, }) - err := db.DeleteChatModelConfigByID(ctx, deletedCfg.ID) + _, err := db.DeleteChatModelConfigByID(ctx, deletedCfg.ID) require.NoError(t, err) // Create a root chat with a workspace.