From f32592ec6b24f83a1dd932477a77d55886b0a7e1 Mon Sep 17 00:00:00 2001 From: Ethan <39577870+ethanndickson@users.noreply.github.com> Date: Mon, 24 Aug 2026 10:30:01 +0100 Subject: [PATCH] feat!: scope chat model configs to organizations (#27955) ## Summary Scope chat model configurations and runtime model resolution to organizations. Each model configuration now belongs to exactly one organization, each organization can have its own default model, and Chatd only uses models that belong to the chat's organization. Depends on #27968. ## Problem Chat model configurations are currently deployment-wide even though chats belong to organizations. A default model, explicitly selected model, or configured model override can therefore reference the same model configuration from every organization. This prevents model catalogs and defaults from being isolated by organization and allows a chat to retain or select a model that its organization does not own. Existing model configurations also have no organization ownership or organization-scoped ACL data, and the database enforces a single live default across the entire deployment. ## Fix Add required organization ownership and ACL storage to chat model configurations. The migration assigns every existing configuration to the default organization, seeds organization-wide read access, and replaces the deployment-wide default constraint with a per-organization constraint. Scope Chatd lookup, admission, fallback, and queue-promotion behavior to the chat's organization: - New chats use the usable default model from their organization. - Explicit model selections from another organization are rejected. - Deployment or personal overrides that point to a foreign model are not used for the chat. - Historical cross-organization model references remain unchanged in the database, but runtime generation falls back to a usable local default. - Queue promotion resolves unavailable, disabled, missing, or foreign stored models to the usable local default and fails when no usable local default exists. The existing model configuration management and availability routes remain compatibility routes for the default organization in this layer. A later layer in the stack adds organization-scoped management routes and updates clients to use them. ## Breaking change This changes chat model behavior from deployment-wide to organization-scoped. Existing configurations are migrated only to the default organization; they are not copied into every organization. As a result: - A non-default organization no longer inherits the deployment's existing default model. - Creating a chat in an organization without a usable local default fails with `400: No default chat model config is configured.` - Explicitly selecting a configuration owned by another organization fails with `400: Model config does not belong to the specified organization.` - Existing chats whose stored model belongs to another organization can generate only if a usable local default is available; otherwise generation fails. Deployments using Agents in non-default organizations must configure a usable default chat model separately for each organization after the organization-scoped management layers in this stack are available. _This pull request description was generated by Coder Agents._ --- coderd/database/check_constraint.go | 2 + coderd/database/dbauthz/dbauthz.go | 19 +- coderd/database/dbauthz/dbauthz_test.go | 32 +- coderd/database/dbgen/dbgen.go | 16 + coderd/database/dbmetrics/querymetrics.go | 16 +- coderd/database/dbmock/dbmock.go | 31 +- coderd/database/dump.sql | 14 +- coderd/database/foreign_key_constraint.go | 1 + coderd/database/lock.go | 1 + ...80_chat_model_config_organization.down.sql | 27 + ...0580_chat_model_config_organization.up.sql | 40 ++ coderd/database/migrations/migrate_test.go | 532 ++++++++---------- ...0580_chat_model_config_organization.up.sql | 91 +++ coderd/database/models.go | 3 + coderd/database/querier.go | 5 +- coderd/database/querier_test.go | 118 +++- coderd/database/queries.sql.go | 130 ++++- coderd/database/queries/chatmodelconfigs.sql | 36 +- coderd/database/sqlc.yaml | 6 + coderd/database/unique_constraint.go | 2 +- coderd/exp_chats.go | 261 +++++++-- coderd/exp_chats_chatstate_test.go | 4 +- coderd/exp_chats_test.go | 439 ++++++++++++++- coderd/x/chatd/ARCHITECTURE.md | 6 + coderd/x/chatd/advisor_internal_test.go | 30 + coderd/x/chatd/chatd.go | 215 +++++-- coderd/x/chatd/chatd_helpers_test.go | 7 +- coderd/x/chatd/chatd_internal_test.go | 255 ++++++++- coderd/x/chatd/chatd_retry_test.go | 9 +- coderd/x/chatd/chatd_test.go | 293 +++++++++- coderd/x/chatd/chatstate/errors.go | 4 + coderd/x/chatd/chatstate/machine_test.go | 3 +- coderd/x/chatd/chatstate/transitions.go | 71 ++- coderd/x/chatd/chatstate/transitions_test.go | 177 ++++++ coderd/x/chatd/compaction_override.go | 4 +- .../compaction_override_internal_test.go | 22 + coderd/x/chatd/configcache.go | 64 ++- coderd/x/chatd/configcache_internal_test.go | 195 +++++-- .../generation_preparer_internal_test.go | 33 +- coderd/x/chatd/helpers_test.go | 3 +- coderd/x/chatd/quickgen.go | 2 +- coderd/x/chatd/quickgen_internal_test.go | 4 +- coderd/x/chatd/subagent.go | 64 ++- coderd/x/chatd/subagent_catalog.go | 2 + coderd/x/chatd/subagent_internal_test.go | 283 ++++++++-- coderd/x/chatd/tasks.go | 4 +- coderd/x/chatd/tasks_test.go | 5 +- coderd/x/chatd/title_override.go | 8 +- .../x/chatd/title_override_internal_test.go | 80 ++- coderd/x/chatd/toolinput_internal_test.go | 5 +- coderd/x/chatd/turn_summary_internal_test.go | 6 + enterprise/coderd/exp_chats_test.go | 47 +- 52 files changed, 2991 insertions(+), 736 deletions(-) create mode 100644 coderd/database/migrations/000580_chat_model_config_organization.down.sql create mode 100644 coderd/database/migrations/000580_chat_model_config_organization.up.sql create mode 100644 coderd/database/migrations/testdata/fixtures/000580_chat_model_config_organization.up.sql diff --git a/coderd/database/check_constraint.go b/coderd/database/check_constraint.go index 6bdb82ca4e6..139ff171c97 100644 --- a/coderd/database/check_constraint.go +++ b/coderd/database/check_constraint.go @@ -25,6 +25,8 @@ const ( CheckChatModelConfigsAIProviderRequiredWhenActive CheckConstraint = "chat_model_configs_ai_provider_required_when_active" // chat_model_configs CheckChatModelConfigsCompressionThresholdCheck CheckConstraint = "chat_model_configs_compression_threshold_check" // chat_model_configs CheckChatModelConfigsContextLimitCheck CheckConstraint = "chat_model_configs_context_limit_check" // chat_model_configs + CheckChatModelConfigsGroupAclIsObject CheckConstraint = "chat_model_configs_group_acl_is_object" // chat_model_configs + CheckChatModelConfigsUserAclIsObject CheckConstraint = "chat_model_configs_user_acl_is_object" // chat_model_configs CheckChatUsageLimitConfigDefaultLimitMicrosCheck CheckConstraint = "chat_usage_limit_config_default_limit_micros_check" // chat_usage_limit_config CheckChatUsageLimitConfigPeriodCheck CheckConstraint = "chat_usage_limit_config_period_check" // chat_usage_limit_config CheckChatUsageLimitConfigSingletonCheck CheckConstraint = "chat_usage_limit_config_singleton_check" // chat_usage_limit_config diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 616a2fedc31..68eab17b701 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -823,6 +823,9 @@ var ( // user's personal data; token writes use the per-user // AsChatdTokenOwner subject instead. rbac.ResourceUser.Type: {policy.ActionReadPersonal}, + // TODO(mafredri): remove this organization read after CODAGT-709 M3. + // exp_chats and chatd use it for the pre-cutover default-org fallback. + rbac.ResourceOrganization.Type: {policy.ActionRead}, }), User: []rbac.Permission{}, ByOrgID: map[string]rbac.OrgPermissions{}, @@ -3751,7 +3754,7 @@ func (q *querier) GetDatabaseNow(ctx context.Context) (time.Time, error) { return q.db.GetDatabaseNow(ctx) } -func (q *querier) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { +func (q *querier) GetDefaultChatModelConfig(ctx context.Context, organizationID uuid.UUID) (database.ChatModelConfig, error) { // Reading the default model config is needed for chat creation. // TODO(CODAGT-161): scope this check when org context is available. // This function has no org context to scope the check, and @@ -3761,7 +3764,7 @@ func (q *querier) GetDefaultChatModelConfig(ctx context.Context) (database.ChatM if _, ok := ActorFromContext(ctx); !ok { return database.ChatModelConfig{}, ErrNoActor } - return q.db.GetDefaultChatModelConfig(ctx) + return q.db.GetDefaultChatModelConfig(ctx, organizationID) } func (q *querier) GetDefaultOrganization(ctx context.Context) (database.Organization, error) { @@ -3810,6 +3813,14 @@ func (q *querier) GetEnabledChatModelConfigs(ctx context.Context) ([]database.Ge return q.db.GetEnabledChatModelConfigs(ctx) } +func (q *querier) GetEnabledChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + // TODO(CODAGT-161): authorize against an organization-scoped chat model config resource. + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil { + return nil, err + } + return q.db.GetEnabledChatModelConfigsByOrganization(ctx, organizationID) +} + func (q *querier) GetEnabledMCPServerConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.MCPServerConfig, error) { return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetEnabledMCPServerConfigsByOrganization)(ctx, organizationID) } @@ -7302,11 +7313,11 @@ func (q *querier) UnpinChatByID(ctx context.Context, id uuid.UUID) error { return q.db.UnpinChatByID(ctx, id) } -func (q *querier) UnsetDefaultChatModelConfigs(ctx context.Context) error { +func (q *querier) UnsetDefaultChatModelConfigs(ctx context.Context, organizationID uuid.UUID) error { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil { return err } - return q.db.UnsetDefaultChatModelConfigs(ctx) + return q.db.UnsetDefaultChatModelConfigs(ctx, organizationID) } func (q *querier) UpdateAIBridgeInterceptionEnded(ctx context.Context, params database.UpdateAIBridgeInterceptionEndedParams) (database.AIBridgeInterception, error) { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index a2c8d0733fd..83b7c150b20 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -1080,8 +1080,8 @@ func (s *MethodTestSuite) TestChats() { })) s.Run("GetDefaultChatModelConfig", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { config := testutil.Fake(s.T(), faker, database.ChatModelConfig{}) - dbm.EXPECT().GetDefaultChatModelConfig(gomock.Any()).Return(config, nil).AnyTimes() - check.Asserts().Returns(config) + dbm.EXPECT().GetDefaultChatModelConfig(gomock.Any(), config.OrganizationID).Return(config, nil).AnyTimes() + check.Args(config.OrganizationID).Asserts().Returns(config) })) s.Run("GetChatModelConfigs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { configA := testutil.Fake(s.T(), faker, database.ChatModelConfig{}) @@ -1203,6 +1203,14 @@ func (s *MethodTestSuite) TestChats() { check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns([]database.GetEnabledChatModelConfigsRow{rowA, rowB}) })) + s.Run("GetEnabledChatModelConfigsByOrganization", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + orgID := uuid.New() + rowA := testutil.Fake(s.T(), faker, database.GetEnabledChatModelConfigsByOrganizationRow{}) + rowB := testutil.Fake(s.T(), faker, database.GetEnabledChatModelConfigsByOrganizationRow{}) + dbm.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), orgID).Return([]database.GetEnabledChatModelConfigsByOrganizationRow{rowA, rowB}, nil).AnyTimes() + check.Args(orgID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns([]database.GetEnabledChatModelConfigsByOrganizationRow{rowA, rowB}) + })) + s.Run("GetStaleChats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { threshold := dbtime.Now() chats := []database.Chat{testutil.Fake(s.T(), faker, database.Chat{})} @@ -1527,8 +1535,9 @@ func (s *MethodTestSuite) TestChats() { check.Args(arg).Asserts(chat, policy.ActionUpdate).Returns(updatedChat) })) s.Run("UnsetDefaultChatModelConfigs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { - dbm.EXPECT().UnsetDefaultChatModelConfigs(gomock.Any()).Return(nil).AnyTimes() - check.Args().Asserts(rbac.ResourceSystem, policy.ActionUpdate) + orgID := uuid.New() + dbm.EXPECT().UnsetDefaultChatModelConfigs(gomock.Any(), orgID).Return(nil).AnyTimes() + check.Args(orgID).Asserts(rbac.ResourceSystem, policy.ActionUpdate) })) s.Run("UpsertChatDiffStatus", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { chat := testutil.Fake(s.T(), faker, database.Chat{}) @@ -7620,10 +7629,8 @@ func TestAsChatd(t *testing.T) { err = auth.Authorize(ctx, actor, policy.ActionUpdate, rbac.ResourceDeploymentConfig) require.Error(t, err, "deployment config update should not be allowed") - // Pin the complete ResourceUser action set: read_personal (user - // chat custom prompts) only. Token refresh persistence uses the - // per-user AsChatdTokenOwner subject, so a future site-wide - // personal-write grant fails here. + // Pin the complete ResourceUser action set: read_personal only. + // Token refresh persistence uses the per-user AsChatdTokenOwner subject. for _, action := range rbac.ResourceUser.AvailableActions() { err := auth.Authorize(ctx, actor, action, rbac.ResourceUser) if action == policy.ActionReadPersonal { @@ -7632,6 +7639,10 @@ func TestAsChatd(t *testing.T) { require.Error(t, err, "user %s should be denied", action) } } + + // Organization read is temporarily needed for the pre-cutover default-org fallback. + err = auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceOrganization) + require.NoError(t, err, "organization read should be allowed") }) t.Run("DeniedActions", func(t *testing.T) { @@ -7652,11 +7663,6 @@ func TestAsChatd(t *testing.T) { // Cannot access provisioner daemons. err = auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceProvisionerDaemon) require.Error(t, err, "provisioner daemon read should be denied") - - // Cannot access organizations; MCP server config resolution is - // strictly org-scoped and needs no organization reads. - err = auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceOrganization) - require.Error(t, err, "organization read should be denied") }) } diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 792f5fb9416..e54bc35ef22 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -178,6 +178,19 @@ func ChatModelConfig(t testing.TB, db database.Store, seed database.ChatModelCon } aiProviderID = uuid.NullUUID{UUID: provider.ID, Valid: true} } + organizationID := seed.OrganizationID + if organizationID == uuid.Nil { + defaultOrg, err := db.GetDefaultOrganization(genCtx) + require.NoError(t, err, "get default organization") + organizationID = defaultOrg.ID + } + groupACL := seed.GroupACL + if groupACL == nil { + groupACL = database.ChatACL{ + organizationID.String(): {Permissions: []policy.Action{policy.ActionRead}}, + } + } + userACL := seed.UserACL params := database.InsertChatModelConfigParams{ Model: takeFirst(seed.Model, "gpt-4o-mini"), DisplayName: takeFirst(seed.DisplayName, "Test Model"), @@ -189,6 +202,9 @@ func ChatModelConfig(t testing.TB, db database.Store, seed database.ChatModelCon CompressionThreshold: takeFirst(seed.CompressionThreshold, defaultChatModelCompressionThreshold), Options: takeFirstSlice(seed.Options, json.RawMessage(`{}`)), AIProviderID: aiProviderID, + OrganizationID: organizationID, + GroupACL: groupACL, + UserACL: userACL, } for _, fn := range munge { fn(¶ms) diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 765d747eee2..c697cbf2a66 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1936,9 +1936,9 @@ func (m queryMetricsStore) GetDatabaseNow(ctx context.Context) (time.Time, error return r0, r1 } -func (m queryMetricsStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { +func (m queryMetricsStore) GetDefaultChatModelConfig(ctx context.Context, organizationID uuid.UUID) (database.ChatModelConfig, error) { start := time.Now() - r0, r1 := m.s.GetDefaultChatModelConfig(ctx) + r0, r1 := m.s.GetDefaultChatModelConfig(ctx, organizationID) m.queryLatencies.WithLabelValues("GetDefaultChatModelConfig").Observe(time.Since(start).Seconds()) m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetDefaultChatModelConfig").Inc() return r0, r1 @@ -2016,6 +2016,14 @@ func (m queryMetricsStore) GetEnabledChatModelConfigs(ctx context.Context) ([]da return r0, r1 } +func (m queryMetricsStore) GetEnabledChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + start := time.Now() + r0, r1 := m.s.GetEnabledChatModelConfigsByOrganization(ctx, organizationID) + m.queryLatencies.WithLabelValues("GetEnabledChatModelConfigsByOrganization").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetEnabledChatModelConfigsByOrganization").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetEnabledMCPServerConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.MCPServerConfig, error) { start := time.Now() r0, r1 := m.s.GetEnabledMCPServerConfigsByOrganization(ctx, organizationID) @@ -5168,9 +5176,9 @@ func (m queryMetricsStore) UnpinChatByID(ctx context.Context, id uuid.UUID) erro return r0 } -func (m queryMetricsStore) UnsetDefaultChatModelConfigs(ctx context.Context) error { +func (m queryMetricsStore) UnsetDefaultChatModelConfigs(ctx context.Context, organizationID uuid.UUID) error { start := time.Now() - r0 := m.s.UnsetDefaultChatModelConfigs(ctx) + r0 := m.s.UnsetDefaultChatModelConfigs(ctx, organizationID) m.queryLatencies.WithLabelValues("UnsetDefaultChatModelConfigs").Observe(time.Since(start).Seconds()) m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UnsetDefaultChatModelConfigs").Inc() return r0 diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index da8dafdd559..e1765317489 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -3616,18 +3616,18 @@ func (mr *MockStoreMockRecorder) GetDatabaseNow(ctx any) *gomock.Call { } // GetDefaultChatModelConfig mocks base method. -func (m *MockStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { +func (m *MockStore) GetDefaultChatModelConfig(ctx context.Context, organizationID uuid.UUID) (database.ChatModelConfig, error) { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetDefaultChatModelConfig", ctx) + ret := m.ctrl.Call(m, "GetDefaultChatModelConfig", ctx, organizationID) ret0, _ := ret[0].(database.ChatModelConfig) ret1, _ := ret[1].(error) return ret0, ret1 } // GetDefaultChatModelConfig indicates an expected call of GetDefaultChatModelConfig. -func (mr *MockStoreMockRecorder) GetDefaultChatModelConfig(ctx any) *gomock.Call { +func (mr *MockStoreMockRecorder) GetDefaultChatModelConfig(ctx, organizationID any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDefaultChatModelConfig", reflect.TypeOf((*MockStore)(nil).GetDefaultChatModelConfig), ctx) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDefaultChatModelConfig", reflect.TypeOf((*MockStore)(nil).GetDefaultChatModelConfig), ctx, organizationID) } // GetDefaultOrganization mocks base method. @@ -3765,6 +3765,21 @@ func (mr *MockStoreMockRecorder) GetEnabledChatModelConfigs(ctx any) *gomock.Cal return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnabledChatModelConfigs", reflect.TypeOf((*MockStore)(nil).GetEnabledChatModelConfigs), ctx) } +// GetEnabledChatModelConfigsByOrganization mocks base method. +func (m *MockStore) GetEnabledChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetEnabledChatModelConfigsByOrganization", ctx, organizationID) + ret0, _ := ret[0].([]database.GetEnabledChatModelConfigsByOrganizationRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetEnabledChatModelConfigsByOrganization indicates an expected call of GetEnabledChatModelConfigsByOrganization. +func (mr *MockStoreMockRecorder) GetEnabledChatModelConfigsByOrganization(ctx, organizationID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnabledChatModelConfigsByOrganization", reflect.TypeOf((*MockStore)(nil).GetEnabledChatModelConfigsByOrganization), ctx, organizationID) +} + // GetEnabledMCPServerConfigsByOrganization mocks base method. func (m *MockStore) GetEnabledMCPServerConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.MCPServerConfig, error) { m.ctrl.T.Helper() @@ -9793,17 +9808,17 @@ func (mr *MockStoreMockRecorder) UnpinChatByID(ctx, id any) *gomock.Call { } // UnsetDefaultChatModelConfigs mocks base method. -func (m *MockStore) UnsetDefaultChatModelConfigs(ctx context.Context) error { +func (m *MockStore) UnsetDefaultChatModelConfigs(ctx context.Context, organizationID uuid.UUID) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UnsetDefaultChatModelConfigs", ctx) + ret := m.ctrl.Call(m, "UnsetDefaultChatModelConfigs", ctx, organizationID) ret0, _ := ret[0].(error) return ret0 } // UnsetDefaultChatModelConfigs indicates an expected call of UnsetDefaultChatModelConfigs. -func (mr *MockStoreMockRecorder) UnsetDefaultChatModelConfigs(ctx any) *gomock.Call { +func (mr *MockStoreMockRecorder) UnsetDefaultChatModelConfigs(ctx, organizationID any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UnsetDefaultChatModelConfigs", reflect.TypeOf((*MockStore)(nil).UnsetDefaultChatModelConfigs), ctx) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UnsetDefaultChatModelConfigs", reflect.TypeOf((*MockStore)(nil).UnsetDefaultChatModelConfigs), ctx, organizationID) } // UpdateAIBridgeInterceptionEnded mocks base method. diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index f649f5e06af..caf363b1eff 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -2077,9 +2077,14 @@ CREATE TABLE chat_model_configs ( compression_threshold integer NOT NULL, options jsonb DEFAULT '{}'::jsonb NOT NULL, ai_provider_id uuid, + organization_id uuid NOT NULL, + group_acl jsonb DEFAULT '{}'::jsonb NOT NULL, + user_acl jsonb DEFAULT '{}'::jsonb NOT NULL, CONSTRAINT chat_model_configs_ai_provider_required_when_active CHECK (((deleted = true) OR (ai_provider_id IS NOT NULL))), CONSTRAINT chat_model_configs_compression_threshold_check CHECK (((compression_threshold >= 0) AND (compression_threshold <= 100))), - CONSTRAINT chat_model_configs_context_limit_check CHECK ((context_limit > 0)) + CONSTRAINT chat_model_configs_context_limit_check CHECK ((context_limit > 0)), + CONSTRAINT chat_model_configs_group_acl_is_object CHECK ((jsonb_typeof(group_acl) = 'object'::text)), + CONSTRAINT chat_model_configs_user_acl_is_object CHECK ((jsonb_typeof(user_acl) = 'object'::text)) ); CREATE SEQUENCE chat_queued_messages_position_seq @@ -4874,7 +4879,9 @@ CREATE INDEX idx_chat_model_configs_ai_provider_id ON chat_model_configs USING b CREATE INDEX idx_chat_model_configs_enabled ON chat_model_configs USING btree (enabled); -CREATE UNIQUE INDEX idx_chat_model_configs_single_default ON chat_model_configs USING btree ((1)) WHERE ((is_default = true) AND (deleted = false)); +CREATE INDEX idx_chat_model_configs_organization_id ON chat_model_configs USING btree (organization_id); + +CREATE UNIQUE INDEX idx_chat_model_configs_single_default ON chat_model_configs USING btree (organization_id) WHERE ((is_default = true) AND (deleted = false)); CREATE INDEX idx_chat_queued_messages_chat_id ON chat_queued_messages USING btree (chat_id); @@ -5256,6 +5263,9 @@ ALTER TABLE ONLY chat_model_configs ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_created_by_fkey FOREIGN KEY (created_by) REFERENCES users(id); +ALTER TABLE ONLY chat_model_configs + ADD CONSTRAINT chat_model_configs_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; + ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_updated_by_fkey FOREIGN KEY (updated_by) REFERENCES users(id); diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index f760c4c426f..aa3cab22ad5 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -28,6 +28,7 @@ const ( ForeignKeyChatMessagesModelConfigID ForeignKeyConstraint = "chat_messages_model_config_id_fkey" // ALTER TABLE ONLY chat_messages ADD CONSTRAINT chat_messages_model_config_id_fkey FOREIGN KEY (model_config_id) REFERENCES chat_model_configs(id); ForeignKeyChatModelConfigsAIProviderID ForeignKeyConstraint = "chat_model_configs_ai_provider_id_fkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_ai_provider_id_fkey FOREIGN KEY (ai_provider_id) REFERENCES ai_providers(id); ForeignKeyChatModelConfigsCreatedBy ForeignKeyConstraint = "chat_model_configs_created_by_fkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_created_by_fkey FOREIGN KEY (created_by) REFERENCES users(id); + ForeignKeyChatModelConfigsOrganizationID ForeignKeyConstraint = "chat_model_configs_organization_id_fkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_organization_id_fkey FOREIGN KEY (organization_id) REFERENCES organizations(id) ON DELETE CASCADE; ForeignKeyChatModelConfigsUpdatedBy ForeignKeyConstraint = "chat_model_configs_updated_by_fkey" // ALTER TABLE ONLY chat_model_configs ADD CONSTRAINT chat_model_configs_updated_by_fkey FOREIGN KEY (updated_by) REFERENCES users(id); ForeignKeyChatQueuedMessagesChatID ForeignKeyConstraint = "chat_queued_messages_chat_id_fkey" // ALTER TABLE ONLY chat_queued_messages ADD CONSTRAINT chat_queued_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE; ForeignKeyChatsAgentID ForeignKeyConstraint = "chats_agent_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_agent_id_fkey FOREIGN KEY (agent_id) REFERENCES workspace_agents(id) ON DELETE SET NULL; diff --git a/coderd/database/lock.go b/coderd/database/lock.go index 84675889120..7bea74ff12c 100644 --- a/coderd/database/lock.go +++ b/coderd/database/lock.go @@ -16,6 +16,7 @@ const ( LockIDReconcileSystemRoles LockIDBoundaryUsageStats LockIDAIProvidersEnvSeed + // Deprecated: Reserved to prevent reuse. Do not use at runtime. LockIDChatModelConfigWrites LockIDChatCapacityAdmission ) diff --git a/coderd/database/migrations/000580_chat_model_config_organization.down.sql b/coderd/database/migrations/000580_chat_model_config_organization.down.sql new file mode 100644 index 00000000000..819f40a4c74 --- /dev/null +++ b/coderd/database/migrations/000580_chat_model_config_organization.down.sql @@ -0,0 +1,27 @@ +ALTER TABLE chat_model_configs + DROP CONSTRAINT chat_model_configs_user_acl_is_object, + DROP CONSTRAINT chat_model_configs_group_acl_is_object, + DROP COLUMN user_acl, + DROP COLUMN group_acl; + +-- Rollback permanently demotes defaults outside the default organization. +-- Reapplying the up migration does not restore those default selections. +-- The old schema permits only one deployment-wide default. Preserve the +-- default organization's selection and demote defaults from other orgs. +UPDATE chat_model_configs +SET is_default = false, + updated_at = NOW() +WHERE is_default = true + AND deleted = false + AND organization_id IS DISTINCT FROM ( + SELECT id FROM organizations WHERE is_default = true LIMIT 1 + ); + +DROP INDEX idx_chat_model_configs_single_default; +CREATE UNIQUE INDEX idx_chat_model_configs_single_default + ON chat_model_configs ((1)) + WHERE is_default = true AND deleted = false; + +DROP INDEX idx_chat_model_configs_organization_id; + +ALTER TABLE chat_model_configs DROP COLUMN organization_id; diff --git a/coderd/database/migrations/000580_chat_model_config_organization.up.sql b/coderd/database/migrations/000580_chat_model_config_organization.up.sql new file mode 100644 index 00000000000..81b4f7829d5 --- /dev/null +++ b/coderd/database/migrations/000580_chat_model_config_organization.up.sql @@ -0,0 +1,40 @@ +-- Assign every chat model config to an organization. Existing rows belong to +-- the default organization so current deployments keep their model catalog. + +-- Step 1: Add nullable column with FK. +ALTER TABLE chat_model_configs + ADD COLUMN organization_id UUID REFERENCES organizations(id) ON DELETE CASCADE; + +-- Step 2: Backfill every existing row to the default org. +UPDATE chat_model_configs +SET organization_id = (SELECT id FROM organizations WHERE is_default = true LIMIT 1); + +-- Step 3: Enforce NOT NULL going forward. +ALTER TABLE chat_model_configs ALTER COLUMN organization_id SET NOT NULL; + +-- Step 4: Index for efficient lookups by organization. +CREATE INDEX idx_chat_model_configs_organization_id ON chat_model_configs (organization_id); + +-- Step 5: Replace the deployment-wide single-default index with a +-- per-org one. The predicate is unchanged. +DROP INDEX idx_chat_model_configs_single_default; +CREATE UNIQUE INDEX idx_chat_model_configs_single_default + ON chat_model_configs (organization_id) + WHERE is_default = true AND deleted = false; + +-- Step 6: ACL columns (nested chat/workspace shape). They are inert +-- until the RBAC stage evaluates them. +ALTER TABLE chat_model_configs + ADD COLUMN group_acl jsonb DEFAULT '{}'::jsonb NOT NULL, + ADD COLUMN user_acl jsonb DEFAULT '{}'::jsonb NOT NULL, + ADD CONSTRAINT chat_model_configs_group_acl_is_object CHECK (jsonb_typeof(group_acl) = 'object'::text), + ADD CONSTRAINT chat_model_configs_user_acl_is_object CHECK (jsonb_typeof(user_acl) = 'object'::text); + +-- Step 7: Seed every existing row with the everyone-in-org read entry. +-- The Everyone group of an organization always has the organization's +-- own ID (see 000058), so the group_acl key is the org ID itself. +UPDATE chat_model_configs +SET group_acl = jsonb_build_object( + organization_id::text, + jsonb_build_object('permissions', jsonb_build_array('read')) +); diff --git a/coderd/database/migrations/migrate_test.go b/coderd/database/migrations/migrate_test.go index ad543a6d615..3827e58e078 100644 --- a/coderd/database/migrations/migrate_test.go +++ b/coderd/database/migrations/migrate_test.go @@ -3,6 +3,7 @@ package migrations_test import ( "context" "database/sql" + "encoding/json" "fmt" "os" "path/filepath" @@ -117,21 +118,6 @@ func testSQLDB(t testing.TB) *sql.DB { return db } -func stepMigrationsUpTo(t *testing.T, next func() (version uint, more bool, err error), target uint) { - t.Helper() - - for { - version, more, err := next() - require.NoError(t, err) - if !more { - t.Fatalf("migration %d not found", target) - } - if version == target { - return - } - } -} - // paralleltest linter doesn't correctly handle table-driven tests (https://github.com/kunwardeep/paralleltest/issues/8) // nolint:paralleltest func TestCheckLatestVersion(t *testing.T) { @@ -3008,324 +2994,264 @@ func TestMigration000566OAuth2AuthMethodBackfill(t *testing.T) { "the backfill aligns the declaration to what is enforced, so the enforced value must be unchanged") } -func TestMigration000574MCPServerConfigsOrganizationID(t *testing.T) { +func TestMigration000580ChatModelConfigOrganization(t *testing.T) { t.Parallel() - const priorMigrationVersion = 573 + const migrationVersion = 580 + const previousMigrationVersion = 579 sqlDB := testSQLDB(t) - next, err := migrations.Stepper(sqlDB) - require.NoError(t, err) - stepMigrationsUpTo(t, next, priorMigrationVersion) - ctx := testutil.Context(t, testutil.WaitSuperLong) - now := time.Now().UTC().Truncate(time.Microsecond) - - var defaultOrgID uuid.UUID - err = sqlDB.QueryRowContext(ctx, `SELECT id FROM organizations WHERE is_default = true`).Scan(&defaultOrgID) - require.NoError(t, err) - - otherOrgID := uuid.New() - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO organizations ( - id, name, display_name, description, icon, created_at, updated_at, - is_default, deleted, default_org_member_roles - ) VALUES ($1, 'migration-568-org', 'Migration 568 Org', '', '', $2, $2, false, false, '{}') - `, otherOrgID, now) - require.NoError(t, err) - - userID := uuid.New() - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO users ( - id, username, email, hashed_password, created_at, updated_at, - status, rbac_roles, login_type - ) VALUES ($1, 'migration-568-user', 'migration-568@example.com', ''::bytea, $2, $2, 'active', '{}', 'password') - `, userID, now) - require.NoError(t, err) - - const keyDigest = "migration-568-key" - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO dbcrypt_keys (number, active_key_digest, test) - VALUES (568000, $1, 'migration-568-test') - `, keyDigest) - require.NoError(t, err) - - providerID := uuid.New() - modelConfigID := uuid.New() - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO ai_providers ( - id, type, name, display_name, enabled, base_url, created_at, updated_at - ) VALUES ($1, 'openai', 'migration-568-provider', 'Migration 568 Provider', true, 'https://provider.example.com', $2, $2) - `, providerID, now) - require.NoError(t, err) - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO chat_model_configs ( - id, model, display_name, ai_provider_id, context_limit, - compression_threshold, created_at, updated_at - ) VALUES ($1, 'migration-568-model', 'Migration 568 Model', $2, 128000, 70, $3, $3) - `, modelConfigID, providerID, now) - require.NoError(t, err) - - type configSeed struct { - id uuid.UUID - slug string - authType string - oauth2ClientID string - oauth2ClientSecret string - oauth2ClientSecretKeyID sql.NullString - oauth2AuthURL string - oauth2TokenURL string - oauth2RevocationURL string - oauth2Scopes string - apiKeyHeader string - apiKeyValue string - apiKeyValueKeyID sql.NullString - customHeaders string - customHeadersKeyID sql.NullString - } - keyID := sql.NullString{String: keyDigest, Valid: true} - configs := []configSeed{ - {id: uuid.New(), slug: "migration-568-none", authType: "none", apiKeyHeader: "Authorization", customHeaders: "{}"}, - // Every credential column carries ciphertext to prove the backfill - // leaves rows byte-identical apart from organization_id. - { - id: uuid.New(), - slug: "migration-568-oauth2", - authType: "oauth2", - oauth2ClientID: "oauth-client-id", - oauth2ClientSecret: "oauth-secret-ciphertext", - oauth2ClientSecretKeyID: keyID, - oauth2AuthURL: "https://oauth.example.com/authorize", - oauth2TokenURL: "https://oauth.example.com/token", - oauth2RevocationURL: "https://oauth.example.com/revoke", - oauth2Scopes: "openid profile", - apiKeyHeader: "X-API-Key", - apiKeyValue: "api-key-ciphertext", - apiKeyValueKeyID: keyID, - customHeaders: "custom-headers-ciphertext", - customHeadersKeyID: keyID, - }, - } - - originalJSON := make(map[uuid.UUID]string, len(configs)) - for _, config := range configs { - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO mcp_server_configs ( - id, display_name, slug, description, url, auth_type, - oauth2_client_id, oauth2_client_secret, oauth2_client_secret_key_id, - oauth2_auth_url, oauth2_token_url, oauth2_revocation_url, oauth2_scopes, - api_key_header, api_key_value, api_key_value_key_id, - custom_headers, custom_headers_key_id, availability, enabled, - created_by, updated_by, created_at, updated_at - ) VALUES ( - $1, $2, $3, $4, $5, $6, - $7, $8, $9, $10, $11, $12, $13, - $14, $15, $16, $17, $18, 'default_on', true, - $19, $19, $20, $20 - ) - `, - config.id, "Migration 568 "+config.authType, config.slug, "migration 568 config", "https://mcp.example.com/"+config.slug, config.authType, - config.oauth2ClientID, config.oauth2ClientSecret, config.oauth2ClientSecretKeyID, - config.oauth2AuthURL, config.oauth2TokenURL, config.oauth2RevocationURL, config.oauth2Scopes, - config.apiKeyHeader, config.apiKeyValue, config.apiKeyValueKeyID, - config.customHeaders, config.customHeadersKeyID, userID, now, - ) - require.NoError(t, err) - - var rowJSON string - err = sqlDB.QueryRowContext(ctx, `SELECT to_jsonb(config) FROM mcp_server_configs AS config WHERE id = $1`, config.id).Scan(&rowJSON) - require.NoError(t, err) - originalJSON[config.id] = rowJSON - } - - oauthConfigID := configs[1].id - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO mcp_server_user_tokens ( - id, mcp_server_config_id, user_id, access_token, access_token_key_id, - refresh_token, refresh_token_key_id, created_at, updated_at - ) VALUES ($1, $2, $3, 'access-token-ciphertext', $4, 'refresh-token-ciphertext', $4, $5, $5) - `, uuid.New(), oauthConfigID, userID, keyDigest, now) - require.NoError(t, err) - - chatID := uuid.New() - chatConfigIDs := []uuid.UUID{configs[1].id, configs[0].id} - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO chats ( - id, owner_id, organization_id, last_model_config_id, title, - mcp_server_ids, created_at, updated_at - ) VALUES ($1, $2, $3, $4, 'Migration 568 Chat', $5, $6, $6) - `, chatID, userID, otherOrgID, modelConfigID, pq.Array(chatConfigIDs), now) - require.NoError(t, err) - - version, _, err := next() - require.NoError(t, err) - require.EqualValues(t, 574, version) - - var totalConfigs int - err = sqlDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM mcp_server_configs`).Scan(&totalConfigs) - require.NoError(t, err) - require.Equal(t, len(configs), totalConfigs) - - for _, config := range configs { - var gotJSON string - var organizationID uuid.UUID - err = sqlDB.QueryRowContext(ctx, ` - SELECT to_jsonb(config) - 'organization_id', organization_id - FROM mcp_server_configs AS config - WHERE id = $1 - `, config.id).Scan(&gotJSON, &organizationID) - require.NoError(t, err) - require.JSONEq(t, originalJSON[config.id], gotJSON) - require.Equal(t, defaultOrgID, organizationID) - } - - var tokenCount int - var tokenConfigID uuid.UUID - err = sqlDB.QueryRowContext(ctx, `SELECT COUNT(*), MIN(mcp_server_config_id::text)::uuid FROM mcp_server_user_tokens`).Scan(&tokenCount, &tokenConfigID) - require.NoError(t, err) - require.Equal(t, 1, tokenCount) - require.Equal(t, oauthConfigID, tokenConfigID) - - var slugConstraint string - err = sqlDB.QueryRowContext(ctx, ` - SELECT pg_get_constraintdef(oid) - FROM pg_constraint - WHERE conname = 'mcp_server_configs_organization_id_slug_key' - `).Scan(&slugConstraint) - require.NoError(t, err) - require.Equal(t, "UNIQUE (organization_id, slug)", slugConstraint) - - getChatIDs := func(t *testing.T, chatID uuid.UUID) []uuid.UUID { - t.Helper() - var ids []uuid.UUID - err := sqlDB.QueryRowContext(ctx, `SELECT mcp_server_ids FROM chats WHERE id = $1`, chatID).Scan(pq.Array(&ids)) - require.NoError(t, err) - return ids - } - require.Equal(t, chatConfigIDs, getChatIDs(t, chatID)) - - // An organization-created config referenced by a chat exercises the down - // sweep: the config is deleted and its chat references removed. - orgLocalConfigID := uuid.New() - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO mcp_server_configs ( - id, organization_id, display_name, slug, url, auth_type - ) VALUES ($1, $2, 'Org-local config', 'migration-568-org-local', 'https://mcp.example.com/org-local', 'none') - `, orgLocalConfigID, otherOrgID) - require.NoError(t, err) - _, err = sqlDB.ExecContext(ctx, ` - UPDATE chats SET mcp_server_ids = $2 WHERE id = $1 - `, chatID, pq.Array([]uuid.UUID{orgLocalConfigID})) - require.NoError(t, err) - require.Equal(t, []uuid.UUID{orgLocalConfigID}, getChatIDs(t, chatID)) - - downSQL, err := os.ReadFile("000574_mcp_server_configs_organization_id.down.sql") - require.NoError(t, err) - _, err = sqlDB.ExecContext(ctx, string(downSQL)) - require.NoError(t, err) - - err = sqlDB.QueryRowContext(ctx, `SELECT COUNT(*) FROM mcp_server_configs`).Scan(&totalConfigs) - require.NoError(t, err) - require.Equal(t, len(configs), totalConfigs) - - require.Empty(t, getChatIDs(t, chatID)) - var danglingIDs int - err = sqlDB.QueryRowContext(ctx, ` - SELECT COUNT(*) - FROM chats AS chat - CROSS JOIN LATERAL unnest(chat.mcp_server_ids) AS config_id - WHERE NOT EXISTS (SELECT 1 FROM mcp_server_configs WHERE id = config_id) - `).Scan(&danglingIDs) - require.NoError(t, err) - require.Zero(t, danglingIDs) -} - -func TestMigration000576MCPServerConfigACL(t *testing.T) { - t.Parallel() - - const priorMigrationVersion = 575 - - sqlDB := testSQLDB(t) + // Migrate up to the migration before the org-scoping migration. next, err := migrations.Stepper(sqlDB) require.NoError(t, err) for { version, more, err := next() require.NoError(t, err) if !more { - t.Fatalf("migration %d not found", priorMigrationVersion) + t.Fatalf("migration %d not found", previousMigrationVersion) } - if version == priorMigrationVersion { + if version == previousMigrationVersion { break } } ctx := testutil.Context(t, testutil.WaitSuperLong) + + providerID := uuid.New() + defaultConfigID := uuid.New() + plainConfigID := uuid.New() + deletedConfigID := uuid.New() + historyOrgID := uuid.New() + chatID := uuid.New() + messageID := int64(578001) + queuedMessageID := int64(578002) now := time.Now().UTC().Truncate(time.Microsecond) - var defaultOrgID uuid.UUID - err = sqlDB.QueryRowContext(ctx, `SELECT id FROM organizations WHERE is_default = true`).Scan(&defaultOrgID) - require.NoError(t, err) - - orgID := uuid.New() - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO organizations ( - id, name, display_name, description, icon, created_at, updated_at, - is_default, deleted, default_org_member_roles - ) VALUES ($1, 'migration-570-org', 'Migration 570 Org', '', '', $2, $2, false, false, '{}') - `, orgID, now) - require.NoError(t, err) - - configIDs := []uuid.UUID{uuid.New(), uuid.New()} - orgIDs := []uuid.UUID{defaultOrgID, orgID} - for i, configID := range configIDs { - _, err = sqlDB.ExecContext(ctx, ` - INSERT INTO mcp_server_configs ( - id, organization_id, display_name, slug, description, url, auth_type, - availability, enabled, created_at, updated_at - ) VALUES ($1, $2, $3, $4, 'unchanged', $5, 'none', 'default_on', true, $6, $6) - `, configID, orgIDs[i], fmt.Sprintf("Migration 570 Config %d", i), fmt.Sprintf("migration-570-config-%d", i), fmt.Sprintf("https://mcp.example.com/%d", i), now) - require.NoError(t, err) + fixtures := []struct { + query string + args []any + }{ + { + `INSERT INTO ai_providers (id, type, name, enabled, base_url, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7)`, + []any{providerID, "openai", "openai-578", true, "https://api.openai.com/v1", now, now}, + }, + // The deployment's single live default config. + { + `INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, + []any{defaultConfigID, "gpt-5.2", "Default 578", true, true, 200000, 70, providerID, now, now}, + }, + // A live non-default config. + { + `INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)`, + []any{plainConfigID, "gpt-5.2-mini", "Plain 578", true, false, 128000, 70, providerID, now, now}, + }, + // A soft-deleted config is backfilled like any row. + { + `INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, deleted, deleted_at, context_limit, compression_threshold, ai_provider_id, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)`, + []any{deletedConfigID, "gpt-4-legacy", "Deleted 578", false, false, true, now, 128000, 70, providerID, now, now}, + }, + { + `INSERT INTO organizations (id, name, description, display_name, default_org_member_roles, created_at, updated_at) + VALUES ($1, $2, '', '', $3, $4, $4)`, + []any{historyOrgID, "history-org-578", pq.StringArray{}, now}, + }, + // A non-default organization chat keeps the pre-migration model ID. + { + `INSERT INTO chats (id, owner_id, last_model_config_id, title, organization_id, created_at, updated_at) + SELECT $1, u.id, $2, $3, $4, $5, $5 + FROM users u + ORDER BY u.created_at, u.id + LIMIT 1`, + []any{chatID, plainConfigID, "History 578", historyOrgID, now}, + }, + { + `INSERT INTO chat_messages (id, chat_id, model_config_id, role, content, content_version, created_by, created_at) + SELECT $1, $2, $3, 'user', '{"type":"text","text":"history"}'::jsonb, 1, owner_id, $4 + FROM chats WHERE id = $2`, + []any{messageID, chatID, plainConfigID, now}, + }, + { + `INSERT INTO chat_queued_messages (id, chat_id, model_config_id, content, created_by, created_at) + SELECT $1, $2, $3, '{"type":"text","text":"queued"}'::jsonb, owner_id, $4 + FROM chats WHERE id = $2`, + []any{queuedMessageID, chatID, plainConfigID, now}, + }, } + tx, err := sqlDB.BeginTx(ctx, nil) + require.NoError(t, err) + defer tx.Rollback() + for i, f := range fixtures { + _, err := tx.ExecContext(ctx, f.query, f.args...) + require.NoError(t, err, "fixture %d", i) + } + require.NoError(t, tx.Commit()) + + // Run the migration. version, _, err := next() require.NoError(t, err) - require.EqualValues(t, 576, version) + require.EqualValues(t, migrationVersion, version) - for i, configID := range configIDs { - var displayName, description string - var enabled bool - var groupACL, userACL string - err = sqlDB.QueryRowContext(ctx, ` - SELECT display_name, description, enabled, group_acl::text, user_acl::text - FROM mcp_server_configs WHERE id = $1 - `, configID).Scan(&displayName, &description, &enabled, &groupACL, &userACL) + var defaultOrgID uuid.UUID + err = sqlDB.QueryRowContext(ctx, + "SELECT id FROM organizations WHERE is_default = true", + ).Scan(&defaultOrgID) + require.NoError(t, err) + + // Every row, including the soft-deleted one, is backfilled to the + // default org. + var backfilled int + err = sqlDB.QueryRowContext(ctx, + "SELECT COUNT(*) FROM chat_model_configs WHERE organization_id = $1", defaultOrgID, + ).Scan(&backfilled) + require.NoError(t, err) + require.Equal(t, 3, backfilled, "all existing configs should move only to the default org") + + var totalConfigs int + err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_model_configs").Scan(&totalConfigs) + require.NoError(t, err) + require.Equal(t, 3, totalConfigs, "the migration must not copy configs to other organizations") + + var notNullViolation bool + err = sqlDB.QueryRowContext(ctx, + "SELECT EXISTS(SELECT 1 FROM chat_model_configs WHERE organization_id IS NULL)", + ).Scan(¬NullViolation) + require.NoError(t, err) + require.False(t, notNullViolation, "no row may keep a NULL organization_id") + + var chatModelConfigID, messageModelConfigID, queuedModelConfigID uuid.UUID + err = sqlDB.QueryRowContext(ctx, + "SELECT last_model_config_id FROM chats WHERE id = $1", chatID, + ).Scan(&chatModelConfigID) + require.NoError(t, err) + err = sqlDB.QueryRowContext(ctx, + "SELECT model_config_id FROM chat_messages WHERE id = $1", messageID, + ).Scan(&messageModelConfigID) + require.NoError(t, err) + err = sqlDB.QueryRowContext(ctx, + "SELECT model_config_id FROM chat_queued_messages WHERE id = $1", queuedMessageID, + ).Scan(&queuedModelConfigID) + require.NoError(t, err) + require.Equal(t, plainConfigID, chatModelConfigID) + require.Equal(t, plainConfigID, messageModelConfigID) + require.Equal(t, plainConfigID, queuedModelConfigID) + + // Every row's group_acl is seeded with the everyone-in-org read + // entry keyed by the org ID (the Everyone group shares the org ID). + seededACL := map[string]any{ + defaultOrgID.String(): map[string]any{"permissions": []string{"read"}}, + } + for _, id := range []uuid.UUID{defaultConfigID, plainConfigID, deletedConfigID} { + var groupACL []byte + err = sqlDB.QueryRowContext(ctx, + "SELECT group_acl FROM chat_model_configs WHERE id = $1", id, + ).Scan(&groupACL) require.NoError(t, err) - require.Equal(t, fmt.Sprintf("Migration 570 Config %d", i), displayName) - require.Equal(t, "unchanged", description) - require.True(t, enabled) - require.JSONEq(t, fmt.Sprintf(`{%q:{"permissions":["read"]}}`, orgIDs[i].String()), groupACL) - require.JSONEq(t, `{}`, userACL) + require.JSONEq(t, string(mustJSON(t, seededACL)), string(groupACL), + "group_acl should carry the everyone-in-org read entry") } - for _, column := range []string{"group_acl", "user_acl"} { - var defaultValue string - err = sqlDB.QueryRowContext(ctx, ` - SELECT column_default FROM information_schema.columns - WHERE table_name = 'mcp_server_configs' AND column_name = $1 - `, column).Scan(&defaultValue) - require.NoError(t, err) - require.Equal(t, "'{}'::jsonb", defaultValue) - } + // The single-default index is now keyed per organization: the index + // definition references organization_id and keeps its predicate. + var indexDef string + err = sqlDB.QueryRowContext(ctx, + "SELECT indexdef FROM pg_indexes WHERE indexname = 'idx_chat_model_configs_single_default'", + ).Scan(&indexDef) + require.NoError(t, err) + require.Contains(t, indexDef, "organization_id") + require.Contains(t, indexDef, "is_default = true") + require.Contains(t, indexDef, "deleted = false") - downSQL, err := os.ReadFile("000576_mcp_server_config_acl.down.sql") + // The org lookup index exists. + var orgIndexExists bool + err = sqlDB.QueryRowContext(ctx, + "SELECT EXISTS(SELECT 1 FROM pg_indexes WHERE indexname = 'idx_chat_model_configs_organization_id')", + ).Scan(&orgIndexExists) + require.NoError(t, err) + require.True(t, orgIndexExists) + + // The default org can host only one live default: a second insert + // violates the per-org partial unique index. + _, err = sqlDB.ExecContext(ctx, + `INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, organization_id, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`, + uuid.New(), "gpt-5.2-alt", "Second Default 578", true, true, 200000, 70, providerID, defaultOrgID, now, now, + ) + require.Error(t, err, "a second live default in the same org must be rejected") + require.Contains(t, err.Error(), "idx_chat_model_configs_single_default") + + // A second org can host its own live default. + secondOrgID := uuid.New() + secondOrgDefaultConfigID := uuid.New() + _, err = sqlDB.ExecContext(ctx, + `INSERT INTO organizations (id, name, description, display_name, default_org_member_roles, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7)`, + secondOrgID, "second-org-578", "", "", pq.StringArray{}, now, now, + ) + require.NoError(t, err) + _, err = sqlDB.ExecContext(ctx, + `INSERT INTO chat_model_configs (id, model, display_name, enabled, is_default, context_limit, compression_threshold, ai_provider_id, organization_id, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)`, + secondOrgDefaultConfigID, "gpt-5.2-org2", "Second Org Default 578", true, true, 200000, 70, providerID, secondOrgID, now, now, + ) + require.NoError(t, err, "each org can host its own live default") + + // The ACL object CHECKs reject non-object values. + _, err = sqlDB.ExecContext(ctx, + `INSERT INTO chat_model_configs (id, model, display_name, enabled, context_limit, compression_threshold, ai_provider_id, organization_id, group_acl, created_at, updated_at) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, '[]'::jsonb, $9, $10)`, + uuid.New(), "bad-acl", "Bad ACL 578", true, 128000, 70, providerID, secondOrgID, now, now, + ) + require.Error(t, err, "non-object group_acl must be rejected") + + var rowCountBeforeDown int + err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_model_configs").Scan(&rowCountBeforeDown) + require.NoError(t, err) + + downSQL, err := os.ReadFile("000580_chat_model_config_organization.down.sql") require.NoError(t, err) _, err = sqlDB.ExecContext(ctx, string(downSQL)) require.NoError(t, err) - var aclColumns int - err = sqlDB.QueryRowContext(ctx, ` - SELECT COUNT(*) FROM information_schema.columns - WHERE table_name = 'mcp_server_configs' AND column_name IN ('group_acl', 'user_acl') - `).Scan(&aclColumns) + + var defaultOrgIsDefault bool + err = sqlDB.QueryRowContext(ctx, + "SELECT is_default FROM chat_model_configs WHERE id = $1", defaultConfigID, + ).Scan(&defaultOrgIsDefault) + require.NoError(t, err) + require.True(t, defaultOrgIsDefault) + + var secondOrgIsDefault bool + err = sqlDB.QueryRowContext(ctx, + "SELECT is_default FROM chat_model_configs WHERE id = $1", secondOrgDefaultConfigID, + ).Scan(&secondOrgIsDefault) + require.NoError(t, err) + require.False(t, secondOrgIsDefault) + + var rowCountAfterDown int + err = sqlDB.QueryRowContext(ctx, "SELECT COUNT(*) FROM chat_model_configs").Scan(&rowCountAfterDown) + require.NoError(t, err) + require.Equal(t, rowCountBeforeDown, rowCountAfterDown) + + err = sqlDB.QueryRowContext(ctx, + "SELECT indexdef FROM pg_indexes WHERE indexname = 'idx_chat_model_configs_single_default'", + ).Scan(&indexDef) + require.NoError(t, err) + require.Contains(t, indexDef, "((1))") + + upSQL, err := os.ReadFile("000580_chat_model_config_organization.up.sql") + require.NoError(t, err) + _, err = sqlDB.ExecContext(ctx, string(upSQL)) require.NoError(t, err) - require.Zero(t, aclColumns) +} + +func mustJSON(t *testing.T, v any) []byte { + t.Helper() + raw, err := json.Marshal(v) + require.NoError(t, err) + return raw } diff --git a/coderd/database/migrations/testdata/fixtures/000580_chat_model_config_organization.up.sql b/coderd/database/migrations/testdata/fixtures/000580_chat_model_config_organization.up.sql new file mode 100644 index 00000000000..21936acb8d4 --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000580_chat_model_config_organization.up.sql @@ -0,0 +1,91 @@ +-- Fixture for 000580 (org-scope chat model configs). Fixtures apply at +-- the post-migration schema (migrate_test.go applies fixture N right +-- after migration N runs), so rows carry the final shape: an explicit +-- organization_id and the seeded everyone-in-org group_acl entry (the +-- Everyone group shares the org ID). +-- 1. the default org's single live default config, +-- 2. a live non-default config in the default org, +-- 3. a soft-deleted config (outside the partial default index). +INSERT INTO ai_providers ( + id, + type, + name, + enabled, + base_url, + created_at, + updated_at +) VALUES ( + 'a52c6f0e-7d4b-4e1a-9c3f-2b8d5e6f7a8b', + 'openai', + 'fixture-openai-580', + TRUE, + 'https://api.openai.com/v1', + '2024-01-01 00:00:00+00', + '2024-01-01 00:00:00+00' +); + +INSERT INTO chat_model_configs ( + id, + model, + display_name, + enabled, + is_default, + context_limit, + compression_threshold, + ai_provider_id, + organization_id, + group_acl, + created_at, + updated_at +) +SELECT + v.id, + v.model, + v.display_name, + v.enabled, + v.is_default, + v.context_limit, + v.compression_threshold, + 'a52c6f0e-7d4b-4e1a-9c3f-2b8d5e6f7a8b', + o.id, + jsonb_build_object(o.id::text, jsonb_build_object('permissions', jsonb_build_array('read'))), + '2024-01-01 00:00:00+00', + '2024-01-01 00:00:00+00' +FROM + organizations o, + ( + VALUES + ( + '580c0001-0000-4000-8000-000000000001'::uuid, + 'gpt-5.2', + 'Fixture Default 580', + TRUE, + TRUE, + 200000, + 70 + ), + ( + '580c0002-0000-4000-8000-000000000002'::uuid, + 'gpt-5.2-mini', + 'Fixture Non Default 580', + TRUE, + FALSE, + 128000, + 70 + ), + ( + '580c0003-0000-4000-8000-000000000003'::uuid, + 'gpt-4-legacy', + 'Fixture Deleted 580', + FALSE, + FALSE, + 128000, + 70 + ) + ) AS v (id, model, display_name, enabled, is_default, context_limit, compression_threshold) +WHERE + o.is_default = TRUE; + +UPDATE chat_model_configs +SET deleted = TRUE, deleted_at = '2024-06-01 00:00:00+00' +WHERE id = '580c0003-0000-4000-8000-000000000003'; diff --git a/coderd/database/models.go b/coderd/database/models.go index 6cea2c3f8a7..93d19300a74 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -5238,6 +5238,9 @@ type ChatModelConfig struct { CompressionThreshold int32 `db:"compression_threshold" json:"compression_threshold"` Options json.RawMessage `db:"options" json:"options"` AIProviderID uuid.NullUUID `db:"ai_provider_id" json:"ai_provider_id"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + GroupACL ChatACL `db:"group_acl" json:"group_acl"` + UserACL ChatACL `db:"user_acl" json:"user_acl"` } type ChatQueuedMessage struct { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index aac29a4c55d..1b34833f8d0 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -552,7 +552,7 @@ type sqlcQuerier interface { // record deadlines or heartbeats rely on a clock that is consistent // with the database rather than the caller's local clock. GetDatabaseNow(ctx context.Context) (time.Time, error) - GetDefaultChatModelConfig(ctx context.Context) (ChatModelConfig, error) + GetDefaultChatModelConfig(ctx context.Context, organizationID uuid.UUID) (ChatModelConfig, error) GetDefaultOrganization(ctx context.Context) (Organization, error) GetDefaultProxyConfig(ctx context.Context) (GetDefaultProxyConfigRow, error) GetDeploymentID(ctx context.Context) (string, error) @@ -564,6 +564,7 @@ type sqlcQuerier interface { // Check both to ensure the selected config is actually usable. GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error) GetEnabledChatModelConfigs(ctx context.Context) ([]GetEnabledChatModelConfigsRow, error) + GetEnabledChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]GetEnabledChatModelConfigsByOrganizationRow, error) GetEnabledMCPServerConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]MCPServerConfig, error) GetEnabledMCPServerConfigsByOrganizationAndIDs(ctx context.Context, arg GetEnabledMCPServerConfigsByOrganizationAndIDsParams) ([]MCPServerConfig, error) // GetExternalAgentTokensByTemplateID returns the auth tokens for all @@ -1420,7 +1421,7 @@ type sqlcQuerier interface { // re-authenticate under a new OIDC provider. UnlinkOIDCUsersByIssuerMismatch(ctx context.Context, expectedPrefix string) (int64, error) UnpinChatByID(ctx context.Context, id uuid.UUID) error - UnsetDefaultChatModelConfigs(ctx context.Context) error + UnsetDefaultChatModelConfigs(ctx context.Context, organizationID uuid.UUID) error UpdateAIBridgeInterceptionEnded(ctx context.Context, arg UpdateAIBridgeInterceptionEndedParams) (AIBridgeInterception, error) // Records heartbeat liveness for an active Gateway DRPC session. The database sets the // timestamp so it stays consistent regardless of clock drift between API diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index dd2299dcc4d..81526b9baf3 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -12429,6 +12429,60 @@ func TestGetEnabledChatModelConfigsUsesAIProviders(t *testing.T) { require.ErrorIs(t, err, sql.ErrNoRows) } +func TestGetEnabledChatModelConfigsByOrganization(t *testing.T) { + t.Parallel() + + store, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitMedium) + defaultOrg, err := store.GetDefaultOrganization(ctx) + require.NoError(t, err) + provider := dbgen.AIProvider(t, store, database.AIProvider{ + Type: database.AIProviderTypeOpenrouter, + Name: "effective-openrouter-" + uuid.NewString(), + }) + defaultConfig := dbgen.ChatModelConfig(t, store, database.ChatModelConfig{ + Model: "default-model-" + uuid.NewString(), + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: defaultOrg.ID, + IsDefault: true, + }) + localOrg := dbgen.Organization(t, store, database.Organization{}) + localDefault := dbgen.ChatModelConfig(t, store, database.ChatModelConfig{ + Model: "local-default-" + uuid.NewString(), + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: localOrg.ID, + IsDefault: true, + }) + localConfig := dbgen.ChatModelConfig(t, store, database.ChatModelConfig{ + Model: "local-model-" + uuid.NewString(), + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: localOrg.ID, + }) + thirdOrg := dbgen.Organization(t, store, database.Organization{}) + thirdConfig := dbgen.ChatModelConfig(t, store, database.ChatModelConfig{ + Model: "third-model-" + uuid.NewString(), + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: thirdOrg.ID, + }) + + rows, err := store.GetEnabledChatModelConfigsByOrganization(ctx, localOrg.ID) + require.NoError(t, err) + require.Len(t, rows, 2) + ids := make([]uuid.UUID, 0, len(rows)) + for _, row := range rows { + require.Equal(t, localOrg.ID, row.ChatModelConfig.OrganizationID) + ids = append(ids, row.ChatModelConfig.ID) + } + require.ElementsMatch(t, []uuid.UUID{localDefault.ID, localConfig.ID}, ids) + require.NotContains(t, ids, defaultConfig.ID) + require.NotContains(t, ids, thirdConfig.ID) + + emptyOrg := dbgen.Organization(t, store, database.Organization{}) + rows, err = store.GetEnabledChatModelConfigsByOrganization(ctx, emptyOrg.ID) + require.NoError(t, err) + require.Empty(t, rows) +} + func insertChatModelConfigForTest( ctx context.Context, t testing.TB, @@ -12437,32 +12491,46 @@ func insertChatModelConfigForTest( params database.InsertChatModelConfigParams, ) (database.ChatModelConfig, error) { t.Helper() - if params.AIProviderID.Valid { - return store.InsertChatModelConfig(ctx, params) - } - providerName := providerType - if providerName == "" { - providerName = "openai" - } - providers, err := store.GetAIProviders(ctx, database.GetAIProvidersParams{IncludeDisabled: true}) - if err != nil { - return database.ChatModelConfig{}, err - } - var provider database.AIProvider - for _, candidate := range providers { - if candidate.Type != database.AIProviderType(providerName) { - continue + if !params.AIProviderID.Valid { + providerName := providerType + if providerName == "" { + providerName = "openai" } - if provider.ID == uuid.Nil || candidate.CreatedAt.After(provider.CreatedAt) { - provider = candidate + providers, err := store.GetAIProviders(ctx, database.GetAIProvidersParams{IncludeDisabled: true}) + if err != nil { + return database.ChatModelConfig{}, err + } + var provider database.AIProvider + for _, candidate := range providers { + if candidate.Type != database.AIProviderType(providerName) { + continue + } + if provider.ID == uuid.Nil || candidate.CreatedAt.After(provider.CreatedAt) { + provider = candidate + } + } + if provider.ID == uuid.Nil { + provider = dbgen.AIProvider(t, store, database.AIProvider{ + Type: database.AIProviderType(providerName), + }) + } + params.AIProviderID = uuid.NullUUID{UUID: provider.ID, Valid: true} + } + if params.OrganizationID == uuid.Nil { + defaultOrg, err := store.GetDefaultOrganization(ctx) + if err != nil { + return database.ChatModelConfig{}, err + } + params.OrganizationID = defaultOrg.ID + } + if params.GroupACL == nil { + params.GroupACL = database.ChatACL{ + params.OrganizationID.String(): {Permissions: []policy.Action{policy.ActionRead}}, } } - if provider.ID == uuid.Nil { - provider = dbgen.AIProvider(t, store, database.AIProvider{ - Type: database.AIProviderType(providerName), - }) + if params.UserACL == nil { + params.UserACL = database.ChatACL{} } - params.AIProviderID = uuid.NullUUID{UUID: provider.ID, Valid: true} return store.InsertChatModelConfig(ctx, params) } @@ -17681,6 +17749,9 @@ func TestGetChatsFilter(t *testing.T) { ContextLimit: 128000, CompressionThreshold: 80, Options: json.RawMessage(`{}`), + OrganizationID: org.ID, + GroupACL: database.ChatACL{}, + UserACL: database.ChatACL{}, }) require.NoError(t, err) @@ -17973,6 +18044,9 @@ func TestGetChatsSearch(t *testing.T) { ContextLimit: 128000, CompressionThreshold: 80, Options: json.RawMessage(`{}`), + OrganizationID: org.ID, + GroupACL: database.ChatACL{}, + UserACL: database.ChatACL{}, }) require.NoError(t, err) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 17e7731bf8f..a9929822837 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -6157,7 +6157,7 @@ func (q *sqlQuerier) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID const getChatModelConfigByID = `-- name: GetChatModelConfigByID :one SELECT - id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id + id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id, organization_id, group_acl, user_acl FROM chat_model_configs WHERE @@ -6184,13 +6184,16 @@ func (q *sqlQuerier) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) ( &i.CompressionThreshold, &i.Options, &i.AIProviderID, + &i.OrganizationID, + &i.GroupACL, + &i.UserACL, ) return i, err } const getChatModelConfigs = `-- name: GetChatModelConfigs :many SELECT - cmc.id, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options, cmc.ai_provider_id + cmc.id, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options, cmc.ai_provider_id, cmc.organization_id, cmc.group_acl, cmc.user_acl FROM chat_model_configs cmc LEFT JOIN @@ -6229,6 +6232,9 @@ func (q *sqlQuerier) GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig &i.CompressionThreshold, &i.Options, &i.AIProviderID, + &i.OrganizationID, + &i.GroupACL, + &i.UserACL, ); err != nil { return nil, err } @@ -6245,16 +6251,17 @@ func (q *sqlQuerier) GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig const getDefaultChatModelConfig = `-- name: GetDefaultChatModelConfig :one SELECT - id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id + id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id, organization_id, group_acl, user_acl FROM chat_model_configs WHERE is_default = TRUE AND deleted = FALSE + AND organization_id = $1::uuid ` -func (q *sqlQuerier) GetDefaultChatModelConfig(ctx context.Context) (ChatModelConfig, error) { - row := q.db.QueryRowContext(ctx, getDefaultChatModelConfig) +func (q *sqlQuerier) GetDefaultChatModelConfig(ctx context.Context, organizationID uuid.UUID) (ChatModelConfig, error) { + row := q.db.QueryRowContext(ctx, getDefaultChatModelConfig, organizationID) var i ChatModelConfig err := row.Scan( &i.ID, @@ -6272,13 +6279,16 @@ func (q *sqlQuerier) GetDefaultChatModelConfig(ctx context.Context) (ChatModelCo &i.CompressionThreshold, &i.Options, &i.AIProviderID, + &i.OrganizationID, + &i.GroupACL, + &i.UserACL, ) return i, err } const getEnabledChatModelConfigByID = `-- name: GetEnabledChatModelConfigByID :one SELECT - cmc.id, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options, cmc.ai_provider_id + cmc.id, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options, cmc.ai_provider_id, cmc.organization_id, cmc.group_acl, cmc.user_acl FROM chat_model_configs cmc JOIN @@ -6312,13 +6322,16 @@ func (q *sqlQuerier) GetEnabledChatModelConfigByID(ctx context.Context, id uuid. &i.CompressionThreshold, &i.Options, &i.AIProviderID, + &i.OrganizationID, + &i.GroupACL, + &i.UserACL, ) return i, err } const getEnabledChatModelConfigs = `-- name: GetEnabledChatModelConfigs :many SELECT - cmc.id, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options, cmc.ai_provider_id, + cmc.id, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options, cmc.ai_provider_id, cmc.organization_id, cmc.group_acl, cmc.user_acl, ap.type::text AS provider FROM chat_model_configs cmc @@ -6366,6 +6379,78 @@ func (q *sqlQuerier) GetEnabledChatModelConfigs(ctx context.Context) ([]GetEnabl &i.ChatModelConfig.CompressionThreshold, &i.ChatModelConfig.Options, &i.ChatModelConfig.AIProviderID, + &i.ChatModelConfig.OrganizationID, + &i.ChatModelConfig.GroupACL, + &i.ChatModelConfig.UserACL, + &i.Provider, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const getEnabledChatModelConfigsByOrganization = `-- name: GetEnabledChatModelConfigsByOrganization :many +SELECT + cmc.id, cmc.model, cmc.display_name, cmc.created_by, cmc.updated_by, cmc.enabled, cmc.is_default, cmc.deleted, cmc.deleted_at, cmc.created_at, cmc.updated_at, cmc.context_limit, cmc.compression_threshold, cmc.options, cmc.ai_provider_id, cmc.organization_id, cmc.group_acl, cmc.user_acl, + ap.type::text AS provider +FROM + chat_model_configs cmc +JOIN + ai_providers ap ON ap.id = cmc.ai_provider_id +WHERE + cmc.organization_id = $1::uuid + AND cmc.enabled = TRUE + AND cmc.deleted = FALSE + AND ap.enabled = TRUE + AND ap.deleted = FALSE +ORDER BY + ap.type::text ASC, + cmc.model ASC, + cmc.updated_at DESC, + cmc.id DESC +` + +type GetEnabledChatModelConfigsByOrganizationRow struct { + ChatModelConfig ChatModelConfig `db:"chat_model_config" json:"chat_model_config"` + Provider string `db:"provider" json:"provider"` +} + +func (q *sqlQuerier) GetEnabledChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]GetEnabledChatModelConfigsByOrganizationRow, error) { + rows, err := q.db.QueryContext(ctx, getEnabledChatModelConfigsByOrganization, organizationID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetEnabledChatModelConfigsByOrganizationRow + for rows.Next() { + var i GetEnabledChatModelConfigsByOrganizationRow + if err := rows.Scan( + &i.ChatModelConfig.ID, + &i.ChatModelConfig.Model, + &i.ChatModelConfig.DisplayName, + &i.ChatModelConfig.CreatedBy, + &i.ChatModelConfig.UpdatedBy, + &i.ChatModelConfig.Enabled, + &i.ChatModelConfig.IsDefault, + &i.ChatModelConfig.Deleted, + &i.ChatModelConfig.DeletedAt, + &i.ChatModelConfig.CreatedAt, + &i.ChatModelConfig.UpdatedAt, + &i.ChatModelConfig.ContextLimit, + &i.ChatModelConfig.CompressionThreshold, + &i.ChatModelConfig.Options, + &i.ChatModelConfig.AIProviderID, + &i.ChatModelConfig.OrganizationID, + &i.ChatModelConfig.GroupACL, + &i.ChatModelConfig.UserACL, &i.Provider, ); err != nil { return nil, err @@ -6392,7 +6477,10 @@ INSERT INTO chat_model_configs ( context_limit, compression_threshold, options, - ai_provider_id + ai_provider_id, + organization_id, + group_acl, + user_acl ) VALUES ( $1::text, $2::text, @@ -6403,10 +6491,13 @@ INSERT INTO chat_model_configs ( $7::bigint, $8::integer, $9::jsonb, - $10::uuid + $10::uuid, + $11::uuid, + $12, + $13 ) RETURNING - id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id + id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id, organization_id, group_acl, user_acl ` type InsertChatModelConfigParams struct { @@ -6420,6 +6511,9 @@ type InsertChatModelConfigParams struct { CompressionThreshold int32 `db:"compression_threshold" json:"compression_threshold"` Options json.RawMessage `db:"options" json:"options"` AIProviderID uuid.NullUUID `db:"ai_provider_id" json:"ai_provider_id"` + OrganizationID uuid.UUID `db:"organization_id" json:"organization_id"` + GroupACL ChatACL `db:"group_acl" json:"group_acl"` + UserACL ChatACL `db:"user_acl" json:"user_acl"` } func (q *sqlQuerier) InsertChatModelConfig(ctx context.Context, arg InsertChatModelConfigParams) (ChatModelConfig, error) { @@ -6434,6 +6528,9 @@ func (q *sqlQuerier) InsertChatModelConfig(ctx context.Context, arg InsertChatMo arg.CompressionThreshold, arg.Options, arg.AIProviderID, + arg.OrganizationID, + arg.GroupACL, + arg.UserACL, ) var i ChatModelConfig err := row.Scan( @@ -6452,6 +6549,9 @@ func (q *sqlQuerier) InsertChatModelConfig(ctx context.Context, arg InsertChatMo &i.CompressionThreshold, &i.Options, &i.AIProviderID, + &i.OrganizationID, + &i.GroupACL, + &i.UserACL, ) return i, err } @@ -6465,10 +6565,11 @@ SET WHERE is_default = TRUE AND deleted = FALSE + AND organization_id = $1::uuid ` -func (q *sqlQuerier) UnsetDefaultChatModelConfigs(ctx context.Context) error { - _, err := q.db.ExecContext(ctx, unsetDefaultChatModelConfigs) +func (q *sqlQuerier) UnsetDefaultChatModelConfigs(ctx context.Context, organizationID uuid.UUID) error { + _, err := q.db.ExecContext(ctx, unsetDefaultChatModelConfigs, organizationID) return err } @@ -6490,7 +6591,7 @@ WHERE id = $10::uuid AND deleted = FALSE RETURNING - id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id + id, model, display_name, created_by, updated_by, enabled, is_default, deleted, deleted_at, created_at, updated_at, context_limit, compression_threshold, options, ai_provider_id, organization_id, group_acl, user_acl ` type UpdateChatModelConfigParams struct { @@ -6536,6 +6637,9 @@ func (q *sqlQuerier) UpdateChatModelConfig(ctx context.Context, arg UpdateChatMo &i.CompressionThreshold, &i.Options, &i.AIProviderID, + &i.OrganizationID, + &i.GroupACL, + &i.UserACL, ) return i, err } diff --git a/coderd/database/queries/chatmodelconfigs.sql b/coderd/database/queries/chatmodelconfigs.sql index 1e90da264fc..7c64fbcef83 100644 --- a/coderd/database/queries/chatmodelconfigs.sql +++ b/coderd/database/queries/chatmodelconfigs.sql @@ -14,7 +14,8 @@ FROM chat_model_configs WHERE is_default = TRUE - AND deleted = FALSE; + AND deleted = FALSE + AND organization_id = @organization_id::uuid; -- name: GetChatModelConfigs :many SELECT @@ -50,6 +51,26 @@ ORDER BY cmc.updated_at DESC, cmc.id DESC; +-- name: GetEnabledChatModelConfigsByOrganization :many +SELECT + sqlc.embed(cmc), + ap.type::text AS provider +FROM + chat_model_configs cmc +JOIN + ai_providers ap ON ap.id = cmc.ai_provider_id +WHERE + cmc.organization_id = @organization_id::uuid + AND cmc.enabled = TRUE + AND cmc.deleted = FALSE + AND ap.enabled = TRUE + AND ap.deleted = FALSE +ORDER BY + ap.type::text ASC, + cmc.model ASC, + cmc.updated_at DESC, + cmc.id DESC; + -- name: GetEnabledChatModelConfigByID :one SELECT cmc.* @@ -77,7 +98,10 @@ INSERT INTO chat_model_configs ( context_limit, compression_threshold, options, - ai_provider_id + ai_provider_id, + organization_id, + group_acl, + user_acl ) VALUES ( @model::text, @display_name::text, @@ -88,7 +112,10 @@ INSERT INTO chat_model_configs ( @context_limit::bigint, @compression_threshold::integer, @options::jsonb, - sqlc.narg('ai_provider_id')::uuid + sqlc.narg('ai_provider_id')::uuid, + @organization_id::uuid, + @group_acl, + @user_acl ) RETURNING *; @@ -121,7 +148,8 @@ SET updated_at = NOW() WHERE is_default = TRUE - AND deleted = FALSE; + AND deleted = FALSE + AND organization_id = @organization_id::uuid; -- name: DeleteChatModelConfigByID :one UPDATE diff --git a/coderd/database/sqlc.yaml b/coderd/database/sqlc.yaml index 22976e4812a..2267a838160 100644 --- a/coderd/database/sqlc.yaml +++ b/coderd/database/sqlc.yaml @@ -83,6 +83,12 @@ sql: - column: "chats.group_acl" go_type: type: "ChatACL" + - column: "chat_model_configs.user_acl" + go_type: + type: "ChatACL" + - column: "chat_model_configs.group_acl" + go_type: + type: "ChatACL" - column: "mcp_server_configs.user_acl" go_type: type: "ChatACL" diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 87550acf8c0..faa400addbb 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -153,7 +153,7 @@ const ( UniqueIndexAPIKeyName UniqueConstraint = "idx_api_key_name" // CREATE UNIQUE INDEX idx_api_key_name ON api_keys USING btree (user_id, token_name) WHERE (login_type = 'token'::login_type); UniqueIndexChatDebugRunsIDChat UniqueConstraint = "idx_chat_debug_runs_id_chat" // CREATE UNIQUE INDEX idx_chat_debug_runs_id_chat ON chat_debug_runs USING btree (id, chat_id); UniqueIndexChatDebugStepsRunStep UniqueConstraint = "idx_chat_debug_steps_run_step" // CREATE UNIQUE INDEX idx_chat_debug_steps_run_step ON chat_debug_steps USING btree (run_id, step_number); - UniqueIndexChatModelConfigsSingleDefault UniqueConstraint = "idx_chat_model_configs_single_default" // CREATE UNIQUE INDEX idx_chat_model_configs_single_default ON chat_model_configs USING btree ((1)) WHERE ((is_default = true) AND (deleted = false)); + UniqueIndexChatModelConfigsSingleDefault UniqueConstraint = "idx_chat_model_configs_single_default" // CREATE UNIQUE INDEX idx_chat_model_configs_single_default ON chat_model_configs USING btree (organization_id) WHERE ((is_default = true) AND (deleted = false)); UniqueIndexConnectionLogsConnectionIDWorkspaceIDAgentName UniqueConstraint = "idx_connection_logs_connection_id_workspace_id_agent_name" // CREATE UNIQUE INDEX idx_connection_logs_connection_id_workspace_id_agent_name ON connection_logs USING btree (connection_id, workspace_id, agent_name); UniqueIndexCustomRolesNameLowerOrganizationID UniqueConstraint = "idx_custom_roles_name_lower_organization_id" // CREATE UNIQUE INDEX idx_custom_roles_name_lower_organization_id ON custom_roles USING btree (lower(name), COALESCE(organization_id, '00000000-0000-0000-0000-000000000000'::uuid)); UniqueIndexOrganizationNameLower UniqueConstraint = "idx_organization_name_lower" // CREATE UNIQUE INDEX idx_organization_name_lower ON organizations USING btree (lower(name)) WHERE (deleted = false); diff --git a/coderd/exp_chats.go b/coderd/exp_chats.go index 4e503abf549..45e5fcc3546 100644 --- a/coderd/exp_chats.go +++ b/coderd/exp_chats.go @@ -680,6 +680,25 @@ func lookupEnabledChatModelConfigByID( return db.GetEnabledChatModelConfigByID(dbauthz.AsChatd(ctx), id) } +func lookupEnabledDefaultOrgChatModelConfigByID( + ctx context.Context, + db database.Store, + id uuid.UUID, +) (database.ChatModelConfig, error) { + modelConfig, err := lookupEnabledChatModelConfigByID(ctx, db, id) + if err != nil { + return database.ChatModelConfig{}, err + } + defaultOrg, err := defaultOrganizationForChatModelConfigs(ctx, db) + if err != nil { + return database.ChatModelConfig{}, err + } + if modelConfig.OrganizationID != defaultOrg.ID { + return database.ChatModelConfig{}, sql.ErrNoRows + } + return modelConfig, nil +} + func parseChatModelCallConfig(options json.RawMessage) (*codersdk.ChatModelCallConfig, error) { callConfig := &codersdk.ChatModelCallConfig{} if len(options) == 0 { @@ -746,7 +765,7 @@ func validateChatModelOverride( Message: "Invalid model_config_id.", } } - modelConfig, err := lookupEnabledChatModelConfigByID(ctx, db, *id) + modelConfig, err := lookupEnabledDefaultOrgChatModelConfigByID(ctx, db, *id) if err != nil { if xerrors.Is(err, sql.ErrNoRows) { return http.StatusBadRequest, &codersdk.Response{ @@ -991,7 +1010,7 @@ func (api *API) chatPersonalModelOverrideDeploymentDefaults( type userChatModelAvailability struct { configuredProviders []chatprovider.ConfiguredProvider configuredModels []chatprovider.ConfiguredModel - enabledModels []database.GetEnabledChatModelConfigsRow + enabledModels []database.GetEnabledChatModelConfigsByOrganizationRow providerStatus map[string]chatprovider.ProviderAvailability providerStatusByID map[uuid.UUID]chatprovider.ProviderAvailability enabledProviderNames map[string]struct{} @@ -1008,6 +1027,7 @@ const ( chatModelConfigUnavailableModelNotFoundOrDisabled chatModelConfigUnavailableReason = "model_not_found_or_disabled" chatModelConfigUnavailableProviderDisabled chatModelConfigUnavailableReason = "provider_disabled" chatModelConfigUnavailableCredentialsMissing chatModelConfigUnavailableReason = "credentials_missing" + chatModelConfigUnavailableOutsideOrganization chatModelConfigUnavailableReason = "outside_organization" ) // getUserChatProviderAvailability returns the enabled chat providers and models @@ -1016,6 +1036,7 @@ const ( func (api *API) getUserChatProviderAvailability( ctx context.Context, userID uuid.UUID, + organizationID uuid.UUID, ) (userChatModelAvailability, error) { //nolint:gocritic // Chatd context is required to read enabled chat config. chatdCtx := dbauthz.AsChatd(ctx) @@ -1023,7 +1044,7 @@ func (api *API) getUserChatProviderAvailability( if err != nil { return userChatModelAvailability{}, err } - enabledModels, err := api.Database.GetEnabledChatModelConfigs(chatdCtx) + enabledModels, err := api.Database.GetEnabledChatModelConfigsByOrganization(chatdCtx, organizationID) if err != nil { return userChatModelAvailability{}, err } @@ -1145,6 +1166,7 @@ func (api *API) getUserChatProviderAvailability( func (api *API) userCanUseChatModelConfig( ctx context.Context, userID uuid.UUID, + organizationID uuid.UUID, modelConfigID uuid.UUID, ) (database.ChatModelConfig, chatModelConfigUnavailableReason, error) { if modelConfigID == uuid.Nil { @@ -1161,11 +1183,14 @@ func (api *API) userCanUseChatModelConfig( } return database.ChatModelConfig{}, chatModelConfigAvailable, err } + if model.OrganizationID != organizationID { + return database.ChatModelConfig{}, chatModelConfigUnavailableOutsideOrganization, nil + } if !model.Enabled { return database.ChatModelConfig{}, chatModelConfigUnavailableModelNotFoundOrDisabled, nil } - availability, err := api.getUserChatProviderAvailability(ctx, userID) + availability, err := api.getUserChatProviderAvailability(ctx, userID, organizationID) if err != nil { return database.ChatModelConfig{}, chatModelConfigAvailable, err } @@ -1192,9 +1217,10 @@ func (api *API) userCanUseChatModelConfig( func (api *API) validateUserChatModelConfigAvailable( ctx context.Context, userID uuid.UUID, + organizationID uuid.UUID, modelConfigID uuid.UUID, ) (database.ChatModelConfig, int, *codersdk.Response) { - modelConfig, reason, err := api.userCanUseChatModelConfig(ctx, userID, modelConfigID) + modelConfig, reason, err := api.userCanUseChatModelConfig(ctx, userID, organizationID, modelConfigID) if err != nil { return database.ChatModelConfig{}, http.StatusInternalServerError, &codersdk.Response{ Message: "Internal error validating model config override.", @@ -1204,7 +1230,8 @@ func (api *API) validateUserChatModelConfigAvailable( switch reason { case chatModelConfigAvailable: return modelConfig, 0, nil - case chatModelConfigUnavailableModelNotFoundOrDisabled: + case chatModelConfigUnavailableModelNotFoundOrDisabled, + chatModelConfigUnavailableOutsideOrganization: return database.ChatModelConfig{}, http.StatusBadRequest, &codersdk.Response{ Message: "Invalid model_config_id: model config not found or disabled.", } @@ -1235,12 +1262,13 @@ func (api *API) validateUserChatModelConfigAvailable( func (api *API) validateExplicitChatModelConfigAvailable( ctx context.Context, userID uuid.UUID, + organizationID uuid.UUID, modelConfigID uuid.UUID, ) (int, *codersdk.Response) { if modelConfigID == uuid.Nil { return 0, nil } - _, status, resp := api.validateUserChatModelConfigAvailable(ctx, userID, modelConfigID) + _, status, resp := api.validateUserChatModelConfigAvailable(ctx, userID, organizationID, modelConfigID) return status, resp } @@ -1582,7 +1610,15 @@ func (api *API) postChats(rw http.ResponseWriter, r *http.Request) { func (api *API) listChatModels(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() apiKey := httpmw.APIKey(r) - availability, err := api.getUserChatProviderAvailability(ctx, apiKey.UserID) + defaultOrganization, err := defaultOrganizationForChatModelConfigs(ctx, api.Database) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to load the default organization.", + Detail: err.Error(), + }) + return + } + availability, err := api.getUserChatProviderAvailability(ctx, apiKey.UserID, defaultOrganization.ID) if err != nil { httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ Message: "Failed to load chat model configuration.", @@ -2834,7 +2870,7 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) { if req.ModelConfigID != nil { modelConfigID = *req.ModelConfigID } - if status, resp := api.validateExplicitChatModelConfigAvailable(ctx, apiKey.UserID, modelConfigID); resp != nil { + if status, resp := api.validateExplicitChatModelConfigAvailable(ctx, apiKey.UserID, chat.OrganizationID, modelConfigID); resp != nil { httpapi.Write(ctx, rw, status, *resp) return } @@ -3009,7 +3045,7 @@ func (api *API) patchChatMessage(rw http.ResponseWriter, r *http.Request) { if req.ModelConfigID != nil { editModelConfigID = *req.ModelConfigID } - if status, resp := api.validateExplicitChatModelConfigAvailable(ctx, apiKey.UserID, editModelConfigID); resp != nil { + if status, resp := api.validateExplicitChatModelConfigAvailable(ctx, apiKey.UserID, chat.OrganizationID, editModelConfigID); resp != nil { httpapi.Write(ctx, rw, status, *resp) return } @@ -4403,7 +4439,20 @@ func (api *API) resolveCreateChatModelConfigID( Message: "Invalid model config ID.", } } - if _, status, resp := api.validateUserChatModelConfigAvailable(ctx, userID, *req.ModelConfigID); resp != nil { + _, reason, err := api.userCanUseChatModelConfig(ctx, userID, req.OrganizationID, *req.ModelConfigID) + if err != nil { + return uuid.Nil, nil, http.StatusInternalServerError, &codersdk.Response{ + Message: "Failed to resolve chat model config.", + Detail: err.Error(), + } + } + if reason == chatModelConfigUnavailableOutsideOrganization { + return uuid.Nil, nil, http.StatusBadRequest, &codersdk.Response{ + Message: "Model config does not belong to the specified organization.", + } + } + if reason != chatModelConfigAvailable { + _, status, resp := api.validateUserChatModelConfigAvailable(ctx, userID, req.OrganizationID, *req.ModelConfigID) return uuid.Nil, nil, status, resp } return *req.ModelConfigID, nil, 0, nil @@ -4417,7 +4466,7 @@ func (api *API) resolveCreateChatModelConfigID( } } if !personalOverridesEnabled { - id, status, resp := api.defaultCreateChatModelConfigID(ctx) + id, status, resp := api.defaultCreateChatModelConfigID(ctx, req.OrganizationID) return id, nil, status, resp } @@ -4446,12 +4495,12 @@ func (api *API) resolveCreateChatModelConfigID( } switch parsed.Mode { case codersdk.ChatPersonalModelOverrideModeChatDefault: - // For root context, chat_default and the defensive default - // case both fall through to the deployment default model below. + // Root chat_default uses the chat organization default. case codersdk.ChatPersonalModelOverrideModeModel: _, reason, err := api.userCanUseChatModelConfig( ctx, userID, + req.OrganizationID, parsed.ModelConfigID, ) if err != nil { @@ -4480,14 +4529,29 @@ func (api *API) resolveCreateChatModelConfigID( } } - id, status, resp := api.defaultCreateChatModelConfigID(ctx) + id, status, resp := api.defaultCreateChatModelConfigID(ctx, req.OrganizationID) return id, nil, status, resp } +// defaultOrganizationForChatModelConfigs resolves the organization used by +// deployment-wide model configuration routes and overrides. +func defaultOrganizationForChatModelConfigs( + ctx context.Context, + db database.Store, +) (database.Organization, error) { + //nolint:gocritic // Deployment-wide routes need internal default organization metadata. + defaultOrg, err := db.GetDefaultOrganization(dbauthz.AsSystemRestricted(ctx)) + if err != nil { + return database.Organization{}, xerrors.Errorf("get default organization: %w", err) + } + return defaultOrg, nil +} + func (api *API) defaultCreateChatModelConfigID( ctx context.Context, + organizationID uuid.UUID, ) (uuid.UUID, int, *codersdk.Response) { - defaultModelConfig, err := api.Database.GetDefaultChatModelConfig(ctx) + defaultModelConfig, err := api.Database.GetDefaultChatModelConfig(ctx, organizationID) if err != nil { if xerrors.Is(err, sql.ErrNoRows) { return uuid.Nil, http.StatusBadRequest, &codersdk.Response{ @@ -4502,7 +4566,8 @@ func (api *API) defaultCreateChatModelConfigID( // The resolved default may itself be disabled or under a disabled // provider. - if _, err := lookupEnabledChatModelConfigByID(ctx, api.Database, defaultModelConfig.ID); err != nil { + enabledDefault, err := lookupEnabledChatModelConfigByID(ctx, api.Database, defaultModelConfig.ID) + if err != nil { if xerrors.Is(err, sql.ErrNoRows) { return uuid.Nil, http.StatusBadRequest, &codersdk.Response{ Message: "No default chat model config is configured.", @@ -4515,6 +4580,12 @@ func (api *API) defaultCreateChatModelConfigID( } } + if enabledDefault.OrganizationID != organizationID { + return uuid.Nil, http.StatusBadRequest, &codersdk.Response{ + Message: "No default chat model config is configured.", + } + } + return defaultModelConfig.ID, 0, nil } @@ -5253,7 +5324,20 @@ func (api *API) putUserChatPersonalModelOverride(rw http.ResponseWriter, r *http }) return } - modelConfig, status, resp := api.validateUserChatModelConfigAvailable(ctx, apiKey.UserID, parsedModelConfigID) + defaultOrg, err := defaultOrganizationForChatModelConfigs(ctx, api.Database) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to resolve default organization.", + Detail: err.Error(), + }) + return + } + modelConfig, status, resp := api.validateUserChatModelConfigAvailable( + ctx, + apiKey.UserID, + defaultOrg.ID, + parsedModelConfigID, + ) if resp != nil { httpapi.Write(ctx, rw, status, *resp) return @@ -6845,18 +6929,35 @@ func (*API) deleteUserChatProviderKey(rw http.ResponseWriter, r *http.Request) { func (api *API) listChatModelConfigs(rw http.ResponseWriter, r *http.Request) { ctx := r.Context() - // Admin users can see all model configs (including disabled ones) - // for management purposes. Non-admin users see only enabled - // configs, which is sufficient for using the chat feature. + defaultOrg, err := defaultOrganizationForChatModelConfigs(ctx, api.Database) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to resolve default organization.", + Detail: err.Error(), + }) + return + } + + // Admin users can see disabled configs. Non-admin users see only enabled + // configs. These deployment-wide routes manage the default organization. isAdmin := api.Authorize(r, policy.ActionRead, rbac.ResourceDeploymentConfig) var configs []database.ChatModelConfig - var err error if isAdmin { - configs, err = api.Database.GetChatModelConfigs(ctx) + allConfigs, listErr := api.Database.GetChatModelConfigs(ctx) + err = listErr + configs = make([]database.ChatModelConfig, 0, len(allConfigs)) + for _, config := range allConfigs { + if config.OrganizationID == defaultOrg.ID { + configs = append(configs, config) + } + } } else { - //nolint:gocritic // All authenticated users need to read enabled model configs to use the chat feature. - rows, rowsErr := api.Database.GetEnabledChatModelConfigs(dbauthz.AsChatd(ctx)) + //nolint:gocritic // All authenticated users need the shared enabled models. + rows, rowsErr := api.Database.GetEnabledChatModelConfigsByOrganization( + dbauthz.AsChatd(ctx), + defaultOrg.ID, + ) err = rowsErr configs = make([]database.ChatModelConfig, 0, len(rows)) for _, row := range rows { @@ -6900,17 +7001,25 @@ func validateChatModelConfigProviderModel(aiProvider database.AIProvider, model } // inChatModelConfigWriteTx runs fn in a transaction that holds the advisory -// lock serializing chat model config writes. All writes to the table must go -// 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. +// lock serializing chat model config writes for one organization. All writes +// to the table must go through this helper so concurrent writers cannot act +// on stale reads, e.g. two creates on an empty organization 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 { +func (api *API) inChatModelConfigWriteTx( + ctx context.Context, + organizationID uuid.UUID, + 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 { + // organization_id is immutable. UpdateChatModelConfig cannot change it, + // so a lock key from a pre-lock read stays correct. + lockID := database.GenLockID("chat_model_config_writes:" + organizationID.String()) + if err := tx.AcquireLock(ctx, lockID); err != nil { return xerrors.Errorf("acquire chat model config write lock: %w", err) } return fn(tx) @@ -7005,6 +7114,18 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { return } + defaultOrg, err := defaultOrganizationForChatModelConfigs(ctx, api.Database) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to resolve default organization.", + Detail: err.Error(), + }) + return + } + everyoneReadACL := database.ChatACL{ + defaultOrg.ID.String(): {Permissions: []policy.Action{policy.ActionRead}}, + } + insertParams := database.InsertChatModelConfigParams{ Model: model, DisplayName: strings.TrimSpace(req.DisplayName), @@ -7016,10 +7137,13 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { AIProviderID: aiProviderID, CreatedBy: uuid.NullUUID{UUID: apiKey.UserID, Valid: apiKey.UserID != uuid.Nil}, UpdatedBy: uuid.NullUUID{UUID: apiKey.UserID, Valid: apiKey.UserID != uuid.Nil}, + OrganizationID: defaultOrg.ID, + GroupACL: everyoneReadACL, + UserACL: database.ChatACL{}, } var inserted database.ChatModelConfig - err = api.inChatModelConfigWriteTx(ctx, func(tx database.Store) error { + err = api.inChatModelConfigWriteTx(ctx, insertParams.OrganizationID, func(tx database.Store) error { //nolint:gocritic // The route already authorized chat model config updates. lockedAIProvider, err := tx.GetAIProviderByIDForReferenceLock(dbauthz.AsChatd(ctx), insertParams.AIProviderID.UUID) if err != nil { @@ -7037,7 +7161,7 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { insertAsDefault := isDefault if !insertAsDefault { - _, err := tx.GetDefaultChatModelConfig(ctx) + _, err := tx.GetDefaultChatModelConfig(ctx, insertParams.OrganizationID) switch { case err == nil: // A default already exists. @@ -7049,7 +7173,7 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { } if insertAsDefault { - if err := tx.UnsetDefaultChatModelConfigs(ctx); err != nil { + if err := tx.UnsetDefaultChatModelConfigs(ctx, insertParams.OrganizationID); err != nil { return xerrors.Errorf("unset default model configs: %w", err) } } @@ -7061,7 +7185,7 @@ func (api *API) createChatModelConfig(rw http.ResponseWriter, r *http.Request) { } inserted = config - if err := ensureDefaultChatModelConfig(ctx, tx); err != nil { + if err := ensureDefaultChatModelConfig(ctx, tx, insertParams.OrganizationID); err != nil { return err } @@ -7114,7 +7238,8 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { return } - if _, err := api.Database.GetChatModelConfigByID(ctx, modelConfigID); err != nil { + existing, err := api.Database.GetChatModelConfigByID(ctx, modelConfigID) + if err != nil { if httpapi.Is404Error(err) { httpapi.ResourceNotFound(rw) return @@ -7125,6 +7250,18 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { }) return } + defaultOrg, err := defaultOrganizationForChatModelConfigs(ctx, api.Database) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to resolve default organization.", + Detail: err.Error(), + }) + return + } + if existing.OrganizationID != defaultOrg.ID { + httpapi.ResourceNotFound(rw) + return + } var req codersdk.UpdateChatModelConfigRequest if !httpapi.Read(ctx, rw, r, &req) { @@ -7166,7 +7303,7 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { } var updated database.ChatModelConfig - err := api.inChatModelConfigWriteTx(ctx, func(tx database.Store) error { + err = api.inChatModelConfigWriteTx(ctx, existing.OrganizationID, 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. @@ -7246,7 +7383,7 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { setAsDefault := updateParams.IsDefault && !lockedExisting.IsDefault if setAsDefault { - if err := tx.UnsetDefaultChatModelConfigs(ctx); err != nil { + if err := tx.UnsetDefaultChatModelConfigs(ctx, lockedExisting.OrganizationID); err != nil { return xerrors.Errorf("unset default model configs: %w", err) } } @@ -7267,6 +7404,7 @@ func (api *API) updateChatModelConfig(rw http.ResponseWriter, r *http.Request) { if err := ensureDefaultChatModelConfig( ctx, tx, + lockedExisting.OrganizationID, excludeConfigID, ); err != nil { return err @@ -7323,7 +7461,8 @@ func (api *API) deleteChatModelConfig(rw http.ResponseWriter, r *http.Request) { return } - if _, err := api.Database.GetChatModelConfigByID(ctx, modelConfigID); err != nil { + existing, err := api.Database.GetChatModelConfigByID(ctx, modelConfigID) + if err != nil { if httpapi.Is404Error(err) { httpapi.ResourceNotFound(rw) return @@ -7334,15 +7473,27 @@ func (api *API) deleteChatModelConfig(rw http.ResponseWriter, r *http.Request) { }) return } + defaultOrg, err := defaultOrganizationForChatModelConfigs(ctx, api.Database) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Failed to resolve default organization.", + Detail: err.Error(), + }) + return + } + if existing.OrganizationID != defaultOrg.ID { + httpapi.ResourceNotFound(rw) + return + } - if err := api.inChatModelConfigWriteTx(ctx, func(tx database.Store) error { + if err := api.inChatModelConfigWriteTx(ctx, existing.OrganizationID, func(tx database.Store) error { if _, err := tx.DeleteChatModelConfigByID(ctx, modelConfigID); err != nil { if xerrors.Is(err, sql.ErrNoRows) { return errChatModelConfigNotFound } return err } - return ensureDefaultChatModelConfig(ctx, tx) + return ensureDefaultChatModelConfig(ctx, tx, existing.OrganizationID) }); err != nil { if xerrors.Is(err, errChatModelConfigNotFound) { httpapi.ResourceNotFound(rw) @@ -7360,12 +7511,15 @@ func (api *API) deleteChatModelConfig(rw http.ResponseWriter, r *http.Request) { rw.WriteHeader(http.StatusNoContent) } +// ensureDefaultChatModelConfig preserves one default for every organization +// that owns at least one chat model config. func ensureDefaultChatModelConfig( ctx context.Context, tx database.Store, + organizationID uuid.UUID, excludedConfigIDs ...uuid.UUID, ) error { - _, err := tx.GetDefaultChatModelConfig(ctx) + _, err := tx.GetDefaultChatModelConfig(ctx, organizationID) switch { case err == nil: return nil @@ -7377,7 +7531,13 @@ func ensureDefaultChatModelConfig( if err != nil { return xerrors.Errorf("list chat model configs: %w", err) } - if len(modelConfigs) == 0 { + orgModelConfigs := make([]database.ChatModelConfig, 0, len(modelConfigs)) + for _, config := range modelConfigs { + if config.OrganizationID == organizationID { + orgModelConfigs = append(orgModelConfigs, config) + } + } + if len(orgModelConfigs) == 0 { return nil } @@ -7386,7 +7546,7 @@ func ensureDefaultChatModelConfig( // omitted-model chat creation. Fall back to any non-excluded config // when no usable candidate exists. //nolint:gocritic // Candidate usability depends on deployment-wide provider state, not the caller's permissions. - enabledRows, err := tx.GetEnabledChatModelConfigs(dbauthz.AsChatd(ctx)) + enabledRows, err := tx.GetEnabledChatModelConfigsByOrganization(dbauthz.AsChatd(ctx), organizationID) if err != nil { return xerrors.Errorf("list enabled chat model configs: %w", err) } @@ -7403,10 +7563,9 @@ func ensureDefaultChatModelConfig( excluded[configID] = struct{}{} } - candidateConfig := modelConfigs[0] var selected *database.ChatModelConfig - for i := range modelConfigs { - config := &modelConfigs[i] + for i := range orgModelConfigs { + config := &orgModelConfigs[i] if _, skip := excluded[config.ID]; skip { continue } @@ -7418,11 +7577,13 @@ func ensureDefaultChatModelConfig( break } } - if selected != nil { - candidateConfig = *selected + if selected == nil { + // Re-promoting the excluded sole candidate preserves the required default. + selected = &orgModelConfigs[0] } + candidateConfig := *selected - if err := tx.UnsetDefaultChatModelConfigs(ctx); err != nil { + if err := tx.UnsetDefaultChatModelConfigs(ctx, organizationID); err != nil { return xerrors.Errorf("unset default model configs: %w", err) } diff --git a/coderd/exp_chats_chatstate_test.go b/coderd/exp_chats_chatstate_test.go index e89be9a9d2b..424cc5c8e1e 100644 --- a/coderd/exp_chats_chatstate_test.go +++ b/coderd/exp_chats_chatstate_test.go @@ -308,8 +308,8 @@ func TestPromoteChatQueuedMessageMissingReturns404(t *testing.T) { require.Equal(t, http.StatusNotFound, res.StatusCode) } -// TestPromoteChatQueuedMessageEmptyQueueReturnsConflict verifies the -// state-conflict 409 path when the chat has no queued messages. +// TestPromoteChatQueuedMessageEmptyQueueReturnsConflict verifies +// the state-conflict 409 path when the chat has no queued messages. func TestPromoteChatQueuedMessageEmptyQueueReturnsConflict(t *testing.T) { t.Parallel() diff --git a/coderd/exp_chats_test.go b/coderd/exp_chats_test.go index ec50029fb4f..6614afb9ba1 100644 --- a/coderd/exp_chats_test.go +++ b/coderd/exp_chats_test.go @@ -604,6 +604,124 @@ func insertAssistantMessage( func TestPostChats(t *testing.T) { t.Parallel() + t.Run("NonDefaultOrgDoesNotUseDeploymentDefault", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + _ = createChatModelConfig(t, client) + + org := dbgen.Organization(t, db, database.Organization{IsDefault: false}) + memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, org.ID, rbac.ScopedRoleAgentsAccess(org.ID)) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + + _, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: org.ID, + Content: []codersdk.ChatInputPart{ + { + Type: codersdk.ChatInputPartTypeText, + Text: "hello from a non-default org", + }, + }, + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "No default chat model config is configured.", sdkErr.Message) + }) + + t.Run("NonDefaultOrgUsesLocalDefault", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + deploymentDefault := createChatModelConfig(t, client) + org := dbgen.Organization(t, db, database.Organization{IsDefault: false}) + localDefault := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: deploymentDefault.AIProviderID, Valid: true}, + OrganizationID: org.ID, + IsDefault: true, + }) + memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, org.ID, rbac.ScopedRoleAgentsAccess(org.ID)) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + + chat, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: org.ID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "hello with a local default", + }}, + }) + require.NoError(t, err) + require.Equal(t, localDefault.ID, chat.LastModelConfigID) + require.NotEqual(t, deploymentDefault.ID, chat.LastModelConfigID) + }) + + t.Run("RejectsExplicitForeignModel", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + foreignModel := createChatModelConfig(t, client) + org := dbgen.Organization(t, db, database.Organization{IsDefault: false}) + memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, org.ID, rbac.ScopedRoleAgentsAccess(org.ID)) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + _, err := client.UpdateAIProvider(ctx, foreignModel.AIProviderID.String(), codersdk.UpdateAIProviderRequest{ + Enabled: ptr.Ref(false), + }) + require.NoError(t, err) + + _, err = memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: org.ID, + ModelConfigID: &foreignModel.ID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "reject the foreign model", + }}, + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "Model config does not belong to the specified organization.", sdkErr.Message) + }) + + t.Run("ForeignPersonalOverrideUsesLocalDefault", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + foreignModel := createChatModelConfig(t, client) + org := dbgen.Organization(t, db, database.Organization{IsDefault: false}) + localDefault := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: foreignModel.AIProviderID, Valid: true}, + OrganizationID: org.ID, + IsDefault: true, + }) + memberClientRaw, member := coderdtest.CreateAnotherUser(t, client.Client, org.ID, rbac.ScopedRoleAgentsAccess(org.ID)) + memberClient := codersdk.NewExperimentalClient(memberClientRaw) + err := client.UpdateChatPersonalModelOverridesAdminSettings(ctx, codersdk.UpdateChatPersonalModelOverridesAdminSettingsRequest{ + AllowUsers: true, + }) + require.NoError(t, err) + err = db.UpsertUserChatPersonalModelOverride(dbauthz.AsSystemRestricted(ctx), database.UpsertUserChatPersonalModelOverrideParams{ + UserID: member.ID, + Key: chatd.ChatPersonalModelOverrideKey(codersdk.ChatPersonalModelOverrideContextRoot), + Value: "model:" + foreignModel.ID.String(), + }) + require.NoError(t, err) + + chat, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ + OrganizationID: org.ID, + Content: []codersdk.ChatInputPart{{ + Type: codersdk.ChatInputPartTypeText, + Text: "ignore the foreign personal override", + }}, + }) + require.NoError(t, err) + require.Equal(t, localDefault.ID, chat.LastModelConfigID) + require.NotEqual(t, foreignModel.ID, chat.LastModelConfigID) + }) + t.Run("Success", func(t *testing.T) { t.Parallel() @@ -675,7 +793,7 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) firstUser := coderdtest.CreateFirstUser(t, client.Client) - _ = createChatModelConfig(t, client) + deploymentDefault := createChatModelConfig(t, client) defaultOrgConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ OrganizationID: firstUser.OrganizationID, @@ -683,6 +801,11 @@ func TestPostChats(t *testing.T) { }) secondOrg := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: deploymentDefault.AIProviderID, Valid: true}, + OrganizationID: secondOrg.ID, + IsDefault: true, + }) memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, secondOrg.ID, rbac.ScopedRoleAgentsAccess(secondOrg.ID)) memberClient := codersdk.NewExperimentalClient(memberClientRaw) @@ -707,9 +830,14 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) coderdtest.CreateFirstUser(t, client.Client) - _ = createChatModelConfig(t, client) + deploymentDefault := createChatModelConfig(t, client) secondOrg := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: deploymentDefault.AIProviderID, Valid: true}, + OrganizationID: secondOrg.ID, + IsDefault: true, + }) orgConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ OrganizationID: secondOrg.ID, Enabled: true, @@ -885,7 +1013,7 @@ func TestPostChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) client, db := newChatClientWithDatabase(t) coderdtest.CreateFirstUser(t, client.Client) - _ = createChatModelConfig(t, client) + deploymentDefault := createChatModelConfig(t, client) thirdOrg := dbgen.Organization(t, db, database.Organization{}) thirdOrgConfig := dbgen.MCPServerConfig(t, db, database.MCPServerConfig{ @@ -894,6 +1022,11 @@ func TestPostChats(t *testing.T) { }) secondOrg := dbgen.Organization(t, db, database.Organization{}) + _ = dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: deploymentDefault.AIProviderID, Valid: true}, + OrganizationID: secondOrg.ID, + IsDefault: true, + }) memberClientRaw, _ := coderdtest.CreateAnotherUser(t, client.Client, secondOrg.ID, rbac.ScopedRoleAgentsAccess(secondOrg.ID)) memberClient := codersdk.NewExperimentalClient(memberClientRaw) @@ -2554,6 +2687,40 @@ func TestListChatModels(t *testing.T) { require.True(t, googleProvider.Available) }) + t.Run("IncludesOnlyDefaultOrganizationModels", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + provider := createAIProviderForTest(t, client, "openai-compat", "test-api-key") + contextLimit := int64(4096) + defaultConfig, err := client.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{ + AIProviderID: &provider.ID, + Model: "default-" + uuid.NewString(), + ContextLimit: &contextLimit, + }) + require.NoError(t, err) + otherOrganization := dbgen.Organization(t, db, database.Organization{IsDefault: false}) + otherConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "other-" + uuid.NewString(), + OrganizationID: otherOrganization.ID, + }) + + models, err := client.ListChatModels(ctx) + require.NoError(t, err) + + var catalogModels []string + for _, availableProvider := range models.Providers { + for _, model := range availableProvider.Models { + catalogModels = append(catalogModels, model.Model) + } + } + require.Contains(t, catalogModels, defaultConfig.Model) + require.NotContains(t, catalogModels, otherConfig.Model) + }) + t.Run("DisabledProvidersAndModelsAreFilteredOut", func(t *testing.T) { t.Parallel() @@ -4484,6 +4651,27 @@ func TestListChatModelConfigs(t *testing.T) { require.Equal(t, enabledConfig.ID, memberConfigs[0].ID) }) + t.Run("ManagementListExcludesOtherOrganizations", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + defaultConfig := createChatModelConfig(t, client) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + otherConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: otherOrg.ID, + AIProviderID: uuid.NullUUID{UUID: defaultConfig.AIProviderID, Valid: true}, + }) + + configs, err := client.ListChatModelConfigs(ctx) + require.NoError(t, err) + require.Contains(t, configs, defaultConfig) + for _, config := range configs { + require.NotEqual(t, otherConfig.ID, config.ID) + } + }) + t.Run("SuccessForOrganizationMember", func(t *testing.T) { t.Parallel() @@ -4622,6 +4810,90 @@ func TestCreateChatModelConfig(t *testing.T) { require.Equal(t, []uuid.UUID{claimed.ID}, defaults) }) + t.Run("SuccessCustomDeploymentConfigRole", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + // dbauthz rejects site-permission roles on the authorized + // querier; seed the role through the raw store, as persisted + // site custom roles exist in production deployments. + rawDB, pubsub := dbtestutil.NewDB(t) + client := newChatClient(t, func(opts *coderdtest.Options) { + opts.Database = rawDB + opts.Pubsub = pubsub + }) + _ = coderdtest.CreateFirstUser(t, client.Client) + nonDefaultOrg := dbgen.Organization(t, rawDB, database.Organization{IsDefault: false}) + + // A custom site role holding only deployment-config permissions + // passes the route gate but cannot read the default organization + // object; the default-org resolution is an internal step and must + // not require caller-held organization read. The owning org is a + // label on the role; the site permissions apply deployment-wide. + role, err := rawDB.InsertCustomRole(ctx, database.InsertCustomRoleParams{ + Name: testutil.GetRandomName(t), + DisplayName: "Deployment Config Admin", + OrganizationID: uuid.NullUUID{UUID: nonDefaultOrg.ID, Valid: true}, + SitePermissions: database.CustomRolePermissions{ + { + ResourceType: rbac.ResourceDeploymentConfig.Type, + Action: policy.ActionRead, + }, + { + ResourceType: rbac.ResourceDeploymentConfig.Type, + Action: policy.ActionUpdate, + }, + }, + }) + require.NoError(t, err) + + aiProvider := createAIProviderForTest(t, client, "openai", "test-api-key") + // An existing deployment default keeps the create off the + // self-promotion path, which is a separate system-scoped + // authorization requirement. + _ = createChatModelConfig(t, client) + + // A member of a non-default org cannot read the default + // organization object; the implicit organization-member role + // grants organization read only within its own org. + adminClientRaw, adminUser := coderdtest.CreateAnotherUser( + t, + client.Client, + nonDefaultOrg.ID, + ) + _, err = client.UpdateOrganizationMemberRoles( + ctx, + nonDefaultOrg.ID, + adminUser.ID.String(), + codersdk.UpdateRoles{Roles: []string{role.Name}}, + ) + require.NoError(t, err) + adminClient := codersdk.NewExperimentalClient(adminClientRaw) + + contextLimit := int64(4096) + modelConfig, err := adminClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{ + AIProviderID: &aiProvider.ID, + Model: "gpt-4o-mini", + ContextLimit: &contextLimit, + }) + require.NoError(t, err) + + defaultOrg, err := rawDB.GetDefaultOrganization(ctx) + require.NoError(t, err) + row, err := rawDB.GetChatModelConfigByID(ctx, modelConfig.ID) + require.NoError(t, err) + require.Equal(t, defaultOrg.ID, row.OrganizationID) + require.False(t, row.IsDefault) + require.Equal( + t, + database.ChatACL{ + defaultOrg.ID.String(): {Permissions: []policy.Action{policy.ActionRead}}, + }, + row.GroupACL, + ) + require.Equal(t, database.ChatACL{}, row.UserACL) + }) + t.Run("ReasoningEffortStored", func(t *testing.T) { t.Parallel() @@ -4925,6 +5197,83 @@ func TestUpdateChatModelConfig(t *testing.T) { require.Len(t, configs, 1) }) + t.Run("RejectsOtherOrganizationConfig", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + defaultConfig := createChatModelConfig(t, client) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + otherConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: otherOrg.ID, + AIProviderID: uuid.NullUUID{UUID: defaultConfig.AIProviderID, Valid: true}, + }) + + _, err := client.UpdateChatModelConfig(ctx, otherConfig.ID, codersdk.UpdateChatModelConfigRequest{ + DisplayName: "not allowed", + }) + requireSDKError(t, err, http.StatusNotFound) + + stored, err := db.GetChatModelConfigByID(dbauthz.AsSystemRestricted(ctx), otherConfig.ID) + require.NoError(t, err) + require.NotEqual(t, "not allowed", stored.DisplayName) + }) + + t.Run("SoleConfigRemainsDefaultWhenDemoted", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + modelConfig := createChatModelConfig(t, client) + + updated, err := client.UpdateChatModelConfig(ctx, modelConfig.ID, codersdk.UpdateChatModelConfigRequest{ + IsDefault: ptr.Ref(false), + }) + require.NoError(t, err) + require.True(t, updated.IsDefault) + + defaultConfig, err := db.GetDefaultChatModelConfig(dbauthz.AsSystemRestricted(ctx), firstUser.OrganizationID) + require.NoError(t, err) + require.Equal(t, modelConfig.ID, defaultConfig.ID) + }) + + t.Run("PromotionLeavesOtherOrganizationDefaultUnchanged", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + defaultConfig := createChatModelConfig(t, client) + candidateConfig := createAdditionalChatModelConfig( + t, + client, + coderdtest.TestChatProviderOpenAICompat, + "promotion-candidate", + ) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + otherDefault := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "other-org-default-" + uuid.NewString(), + AIProviderID: uuid.NullUUID{UUID: defaultConfig.AIProviderID, Valid: true}, + OrganizationID: otherOrg.ID, + IsDefault: true, + }) + + updated, err := client.UpdateChatModelConfig(ctx, candidateConfig.ID, codersdk.UpdateChatModelConfigRequest{ + IsDefault: ptr.Ref(true), + }) + require.NoError(t, err) + require.True(t, updated.IsDefault) + + orgDefault, err := db.GetDefaultChatModelConfig(dbauthz.AsSystemRestricted(ctx), firstUser.OrganizationID) + require.NoError(t, err) + require.Equal(t, candidateConfig.ID, orgDefault.ID) + unchangedOtherDefault, err := db.GetDefaultChatModelConfig(dbauthz.AsSystemRestricted(ctx), otherOrg.ID) + require.NoError(t, err) + require.Equal(t, otherDefault.ID, unchangedOtherDefault.ID) + }) + t.Run("UnchangedProviderWithoutAIProviderID", func(t *testing.T) { t.Parallel() @@ -5506,6 +5855,57 @@ func TestDeleteChatModelConfig(t *testing.T) { } }) + t.Run("RejectsOtherOrganizationConfig", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + _ = coderdtest.CreateFirstUser(t, client.Client) + defaultConfig := createChatModelConfig(t, client) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + otherConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: otherOrg.ID, + AIProviderID: uuid.NullUUID{UUID: defaultConfig.AIProviderID, Valid: true}, + }) + + err := client.DeleteChatModelConfig(ctx, otherConfig.ID) + requireSDKError(t, err, http.StatusNotFound) + _, err = db.GetChatModelConfigByID(dbauthz.AsSystemRestricted(ctx), otherConfig.ID) + require.NoError(t, err) + }) + + t.Run("DeletionLeavesOtherOrganizationDefaultUnchanged", func(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitLong) + client, db := newChatClientWithDatabase(t) + firstUser := coderdtest.CreateFirstUser(t, client.Client) + defaultConfig := createChatModelConfig(t, client) + candidateConfig := createAdditionalChatModelConfig( + t, + client, + coderdtest.TestChatProviderOpenAICompat, + "delete-candidate", + ) + otherOrg := dbgen.Organization(t, db, database.Organization{}) + otherDefault := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "other-org-default-" + uuid.NewString(), + AIProviderID: uuid.NullUUID{UUID: defaultConfig.AIProviderID, Valid: true}, + OrganizationID: otherOrg.ID, + IsDefault: true, + }) + + err := client.DeleteChatModelConfig(ctx, defaultConfig.ID) + require.NoError(t, err) + + orgDefault, err := db.GetDefaultChatModelConfig(dbauthz.AsSystemRestricted(ctx), firstUser.OrganizationID) + require.NoError(t, err) + require.Equal(t, candidateConfig.ID, orgDefault.ID) + unchangedOtherDefault, err := db.GetDefaultChatModelConfig(dbauthz.AsSystemRestricted(ctx), otherOrg.ID) + require.NoError(t, err) + require.Equal(t, otherDefault.ID, unchangedOtherDefault.ID) + }) + // Deleting the default must not promote a config whose provider is // disabled while a usable candidate exists. t.Run("PromotesUsableConfigOverDisabledProvider", func(t *testing.T) { @@ -14414,6 +14814,12 @@ func TestChatModelOverrides(t *testing.T) { coderdtest.TestChatProviderOpenAICompat, "gpt-4.1-disabled-"+string(setting.context), ) + localOrg := dbgen.Organization(t, db, database.Organization{}) + localModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: localOrg.ID, + AIProviderID: uuid.NullUUID{UUID: openAIModel.AIProviderID, Valid: true}, + Enabled: true, + }) memberClientRaw, _ := coderdtest.CreateAnotherUser(t, adminClient.Client, firstUser.OrganizationID) memberClient := codersdk.NewExperimentalClient(memberClientRaw) @@ -14582,6 +14988,14 @@ func TestChatModelOverrides(t *testing.T) { require.Equal(t, "Invalid model_config_id.", sdkErr.Message) }) + t.Run("ChatOrgModelReturns400", func(t *testing.T) { + ctx := testutil.Context(t, testutil.WaitLong) + + err := putOverride(ctx, adminClient, setting.context, localModel.ID.String()) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Equal(t, "Invalid model_config_id.", sdkErr.Message) + }) + t.Run("NonAdminGETReturns404", func(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) @@ -14747,6 +15161,12 @@ func TestUserChatPersonalModelOverrides(t *testing.T) { require.NotEqual(t, uuid.Nil, provider.ID) require.NotEqual(t, uuid.Nil, disabledProvider.ID) + localOrg := dbgen.Organization(t, db, database.Organization{}) + localModelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: localOrg.ID, + AIProviderID: uuid.NullUUID{UUID: modelProvider.ID, Valid: true}, + Enabled: true, + }) personalOverride := func( resp codersdk.UserChatPersonalModelOverridesResponse, overrideContext codersdk.ChatPersonalModelOverrideContext, @@ -15120,6 +15540,19 @@ func TestUserChatPersonalModelOverrides(t *testing.T) { } }) + t.Run("PUTModelRejectsChatOrgModel", func(t *testing.T) { + for _, overrideContext := range contexts { + rawBefore := getRaw(overrideContext) + err := memberClient.UpdateUserChatPersonalModelOverride(ctx, overrideContext, codersdk.UpdateUserChatPersonalModelOverrideRequest{ + Mode: codersdk.ChatPersonalModelOverrideModeModel, + ModelConfigID: localModelConfig.ID.String(), + }) + sdkErr := requireSDKError(t, err, http.StatusBadRequest) + require.Contains(t, sdkErr.Message, "model config not found or disabled") + require.Equal(t, rawBefore, getRaw(overrideContext)) + } + }) + t.Run("GETMalformedStoredValueFallsBackToContextDefault", func(t *testing.T) { upsertRaw(codersdk.ChatPersonalModelOverrideContextRoot, "model:not-a-uuid") diff --git a/coderd/x/chatd/ARCHITECTURE.md b/coderd/x/chatd/ARCHITECTURE.md index a630e93e6dd..a30a41cff24 100644 --- a/coderd/x/chatd/ARCHITECTURE.md +++ b/coderd/x/chatd/ARCHITECTURE.md @@ -418,6 +418,10 @@ EXECUTE FUNCTION sync_chat_retry_state(); ## HTTP endpoints + + + + This section maps the public endpoints that mutate chat state to the transitions they use. ### `POST /api/experimental/chats` @@ -597,6 +601,8 @@ There are 2 notification channels: # Chat worker + + A chat worker lives inside every coderd replica. It acquires chats, calls the LLM API, executes tools, handles interrupts and tool-result waits, and commits completed outcomes through the core state machine. The chat worker is responsible for: diff --git a/coderd/x/chatd/advisor_internal_test.go b/coderd/x/chatd/advisor_internal_test.go index ebce27d5d01..8b3c166885b 100644 --- a/coderd/x/chatd/advisor_internal_test.go +++ b/coderd/x/chatd/advisor_internal_test.go @@ -210,6 +210,36 @@ func TestResolveAdvisorModelOverride(t *testing.T) { requireChatModel(t, p, codersdk.AdvisorConfig{ModelConfigID: uuid.New()}) }) + t.Run("ForeignOrgConfigUsesChatModel", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitShort) + configID := uuid.New() + foreignOrgID := uuid.New() + chatOrgID := uuid.New() + store := &advisorOverrideStubStore{ + getEnabledChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { + return database.ChatModelConfig{ + ID: configID, + OrganizationID: foreignOrgID, + Enabled: true, + }, nil + }, + } + p := newAdvisorTestServer(ctx, t, store) + + resolved, ok, err := p.resolveAdvisorModelOverride( + ctx, + database.Chat{OrganizationID: chatOrgID}, + codersdk.AdvisorConfig{ModelConfigID: configID}, + advisorTestMaxOutputTokens, + modelBuildOptions{}, + logger, + ) + require.NoError(t, err) + require.False(t, ok) + require.False(t, resolved.model.Valid()) + }) + t.Run("InvalidOptionsJSONUsesChatModel", func(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) diff --git a/coderd/x/chatd/chatd.go b/coderd/x/chatd/chatd.go index 3ea82a68deb..4d540f06d8a 100644 --- a/coderd/x/chatd/chatd.go +++ b/coderd/x/chatd/chatd.go @@ -266,6 +266,14 @@ func (p *Server) resolveAdvisorModelOverride( ctx, advisorCfg.ModelConfigID, ) + if err == nil && overrideConfig.OrganizationID != chat.OrganizationID { + logger.Warn( + ctx, + "advisor model config belongs to another organization, continuing with chat model", + slog.F("model_config_id", advisorCfg.ModelConfigID), + ) + return resolvedModelCall{}, false, nil + } if err != nil { if xerrors.Is(err, sql.ErrNoRows) { logger.Warn( @@ -1071,7 +1079,7 @@ var ( ErrChatArchived = xerrors.New("chat is archived") // ErrNoDefaultChatModelConfig indicates no default chat model config // is configured, so chatd cannot resolve a model for the request. - ErrNoDefaultChatModelConfig = xerrors.New("no default chat model config is configured") + ErrNoDefaultChatModelConfig = chatstate.ErrNoDefaultChatModelConfig // ErrNothingToCompact indicates a manual compaction request found // no uncompressed conversation after the latest compaction // boundary, so running a compaction would produce nothing. @@ -1262,7 +1270,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C deploymentPrompt := p.resolveDeploymentSystemPrompt(ctx) if opts.ModelConfigID != uuid.Nil { - if err := requireEnabledChatModelConfig(ctx, p.db, opts.ModelConfigID); err != nil { + if err := requireEnabledChatModelConfig(ctx, p.db, opts.OrganizationID, opts.ModelConfigID); err != nil { return database.Chat{}, err } } @@ -1276,7 +1284,7 @@ func (p *Server) CreateChat(ctx context.Context, opts CreateOptions) (database.C contentParts := opts.InitialUserContent if p.hooks.Enabled() { // Validate model admission before dispatch, matching the insert path. - if err := validateCreateModelConfigID(ctx, p.db, opts.ModelConfigID); err != nil { + if err := validateCreateModelConfigID(ctx, p.db, opts.OrganizationID, opts.ModelConfigID); err != nil { return database.Chat{}, err } turnID := uuid.New() @@ -1584,63 +1592,81 @@ func resolveSendMessageModelConfigID( requested uuid.UUID, ) (uuid.UUID, error) { if requested == uuid.Nil { - return resolveFallbackModelConfigID(ctx, store, chat.LastModelConfigID) + return resolveFallbackModelConfigID(ctx, store, chat.OrganizationID, chat.LastModelConfigID) } - if err := requireEnabledChatModelConfig(ctx, store, requested); err != nil { + if err := requireEnabledChatModelConfig(ctx, store, chat.OrganizationID, requested); err != nil { return uuid.Nil, err } return requested, nil } -// requireEnabledChatModelConfig rechecks enabled state inside the daemon: -// the coderd preflight can race an admin disabling the model or provider. +// requireEnabledChatModelConfig rechecks enabled state inside the daemon. +// The coderd preflight can race an admin disabling the model or provider. func requireEnabledChatModelConfig( ctx context.Context, store database.Store, + organizationID uuid.UUID, modelConfigID uuid.UUID, ) error { chatdCtx := chatdModelConfigLookupContext(ctx) - if _, err := store.GetEnabledChatModelConfigByID(chatdCtx, modelConfigID); err != nil { - if errors.Is(err, sql.ErrNoRows) { - return xerrors.Errorf( - "%w: %s", - ErrInvalidModelConfigID, - modelConfigID, - ) + config, err := store.GetEnabledChatModelConfigByID(chatdCtx, modelConfigID) + if err == nil { + if config.OrganizationID == organizationID { + return nil } + err = sql.ErrNoRows + } + if errors.Is(err, sql.ErrNoRows) { return xerrors.Errorf( - "get requested model config %s: %w", + "%w: %s", + ErrInvalidModelConfigID, modelConfigID, - err, ) } - return nil + return xerrors.Errorf( + "get requested model config %s: %w", + modelConfigID, + err, + ) } -func validateCreateModelConfigID(ctx context.Context, store database.Store, modelConfigID uuid.UUID) error { +func validateCreateModelConfigID( + ctx context.Context, + store database.Store, + organizationID uuid.UUID, + modelConfigID uuid.UUID, +) error { if modelConfigID == uuid.Nil { return xerrors.Errorf("%w: %s", ErrInvalidModelConfigID, modelConfigID) } chatdCtx := chatdModelConfigLookupContext(ctx) - if _, err := store.GetChatModelConfigByID(chatdCtx, modelConfigID); err != nil { - if errors.Is(err, sql.ErrNoRows) { - return xerrors.Errorf("%w: %s", ErrInvalidModelConfigID, modelConfigID) + config, err := store.GetChatModelConfigByID(chatdCtx, modelConfigID) + if err == nil { + if config.OrganizationID == organizationID { + return nil } - return xerrors.Errorf("get requested model config %s: %w", modelConfigID, err) + err = sql.ErrNoRows } - return nil + if errors.Is(err, sql.ErrNoRows) { + return xerrors.Errorf("%w: %s", ErrInvalidModelConfigID, modelConfigID) + } + return xerrors.Errorf("get requested model config %s: %w", modelConfigID, err) } func resolveFallbackModelConfigID( ctx context.Context, store database.Store, + organizationID uuid.UUID, modelConfigID uuid.UUID, ) (uuid.UUID, error) { chatdCtx := chatdModelConfigLookupContext(ctx) if modelConfigID != uuid.Nil { - if _, err := store.GetEnabledChatModelConfigByID(chatdCtx, modelConfigID); err == nil { - return modelConfigID, nil + config, err := store.GetEnabledChatModelConfigByID(chatdCtx, modelConfigID) + if err == nil { + if config.OrganizationID == organizationID { + return modelConfigID, nil + } } else if !errors.Is(err, sql.ErrNoRows) { return uuid.Nil, xerrors.Errorf( "get chat model config %s: %w", @@ -1650,40 +1676,26 @@ func resolveFallbackModelConfigID( } } - defaultConfig, err := store.GetDefaultChatModelConfig(chatdCtx) + defaultConfig, err := effectiveDefaultChatModelConfig(chatdCtx, store, organizationID) if err != nil { if errors.Is(err, sql.ErrNoRows) { return uuid.Nil, ErrNoDefaultChatModelConfig } return uuid.Nil, xerrors.Errorf("get default chat model config: %w", err) } - // The default may itself be disabled or under a disabled provider. - if _, err := store.GetEnabledChatModelConfigByID(chatdCtx, defaultConfig.ID); err != nil { - if errors.Is(err, sql.ErrNoRows) { - return uuid.Nil, xerrors.Errorf( - "%w: default model config %s or its provider is disabled", - ErrNoDefaultChatModelConfig, - defaultConfig.ID, - ) - } - return uuid.Nil, xerrors.Errorf( - "get default chat model config %s: %w", - defaultConfig.ID, - err, - ) - } return defaultConfig.ID, nil } func validateModelConfigOverride( ctx context.Context, store database.Store, + organizationID uuid.UUID, requested uuid.UUID, ) (uuid.NullUUID, error) { if requested == uuid.Nil { return uuid.NullUUID{}, nil } - if err := requireEnabledChatModelConfig(ctx, store, requested); err != nil { + if err := requireEnabledChatModelConfig(ctx, store, organizationID, requested); err != nil { return uuid.NullUUID{}, err } return uuid.NullUUID{UUID: requested, Valid: true}, nil @@ -1706,6 +1718,31 @@ func validateEditTarget(ctx context.Context, store database.Store, chatID uuid.U return nil } +func effectiveDefaultChatModelConfig( + ctx context.Context, + store database.Store, + organizationID uuid.UUID, +) (database.ChatModelConfig, error) { + rows, err := store.GetEnabledChatModelConfigsByOrganization(ctx, organizationID) + if err != nil { + return database.ChatModelConfig{}, err + } + for _, row := range rows { + if row.ChatModelConfig.IsDefault { + return row.ChatModelConfig, nil + } + } + return database.ChatModelConfig{}, sql.ErrNoRows +} + +func enabledChatModelConfigsForOrganization( + ctx context.Context, + store database.Store, + organizationID uuid.UUID, +) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + return store.GetEnabledChatModelConfigsByOrganization(ctx, organizationID) +} + // EditMessage replaces an earlier user message and discards the // active-history suffix through chatstate.EditMessage. Model-config // override validation and usage-limit admission run in the same @@ -1739,7 +1776,7 @@ func (p *Server) EditMessage( if err := validateEditTarget(ctx, p.db, opts.ChatID, opts.EditedMessageID); err != nil { return EditMessageResult{}, err } - if _, err := validateModelConfigOverride(ctx, p.db, opts.ModelConfigID); err != nil { + if _, err := validateModelConfigOverride(ctx, p.db, chat.OrganizationID, opts.ModelConfigID); err != nil { return EditMessageResult{}, err } sessionStartHookResult, err = p.hooks.Trigger(ctx, chathooks.ChatFor(chat, &turnID), chathooks.Message{Source: chathooks.SessionStartSourceClear}, agenthooks.EventSessionStart, dispatch.CapacityClassAdmission) @@ -1799,7 +1836,7 @@ func (p *Server) EditMessage( } editedMsg = target - modelOverride, err := validateModelConfigOverride(ctx, store, opts.ModelConfigID) + modelOverride, err := validateModelConfigOverride(ctx, store, lockedChat.OrganizationID, opts.ModelConfigID) if err != nil { return err } @@ -1811,7 +1848,7 @@ func (p *Server) EditMessage( if target.ModelConfigID.Valid { preserved = target.ModelConfigID.UUID } - resolved, err := resolveFallbackModelConfigID(ctx, store, preserved) + resolved, err := resolveFallbackModelConfigID(ctx, store, lockedChat.OrganizationID, preserved) if err != nil { return err } @@ -2724,6 +2761,9 @@ func (p *Server) resolveManualTitleModel( modelOpts, ) if overrideErr != nil { + if errors.Is(overrideErr, errModelConfigOutsideOrganization) { + return p.resolveDefaultManualTitleModel(ctx, chat, modelOpts) + } if overrideSet { return resolvedModelCall{}, xerrors.Errorf( "resolve manual title generation model override: %w", @@ -2738,7 +2778,7 @@ func (p *Server) resolveManualTitleModel( return overrideResolved, nil } - configs, err := store.GetEnabledChatModelConfigs(ctx) + configs, err := enabledChatModelConfigsForOrganization(ctx, store, chat.OrganizationID) if err != nil { p.logger.Debug(ctx, "failed to list manual title model configs", slog.F("chat_id", chat.ID), @@ -2769,6 +2809,36 @@ func (p *Server) resolveManualTitleModel( return resolved, nil } +func (p *Server) resolveDefaultManualTitleModel( + ctx context.Context, + chat database.Chat, + modelOpts modelBuildOptions, +) (resolvedModelCall, error) { + config, err := p.configCache.DefaultModelConfig(ctx, chat.OrganizationID) + if err != nil { + if errors.Is(err, sql.ErrNoRows) { + return resolvedModelCall{}, ErrNoDefaultChatModelConfig + } + return resolvedModelCall{}, xerrors.Errorf( + "get default manual title model config: %w", + err, + ) + } + resolved, err := p.resolveModelCall(ctx, modelCallSpec{ + purpose: "title", + chat: chat, + explicitConfig: &config, + buildOptions: modelOpts, + }) + if err != nil { + return resolvedModelCall{}, xerrors.Errorf( + "create default manual title model: %w", + err, + ) + } + return resolved, nil +} + func (p *Server) resolveFallbackManualTitleModel( ctx context.Context, chat database.Chat, @@ -4191,31 +4261,60 @@ func (p *Server) resolveUserProviderAPIKeys( return keys, nil } -// resolveModelConfig looks up the chat's model config by its -// LastModelConfigID. If the referenced config no longer exists -// (e.g. it was deleted), it falls back to the default model -// config. Returns an error when no usable config is available. +func (p *Server) resolveModelConfigForOrganization( + ctx context.Context, + organizationID uuid.UUID, + modelConfigID uuid.UUID, +) (database.ChatModelConfig, string, error) { + if modelConfigID == uuid.Nil { + return database.ChatModelConfig{}, "", sql.ErrNoRows + } + modelConfig, err := p.configCache.ModelConfigByID(ctx, modelConfigID) + if err != nil { + return database.ChatModelConfig{}, "", err + } + if modelConfig.OrganizationID != organizationID { + return database.ChatModelConfig{}, "", errModelConfigOutsideOrganization + } + return p.resolveNormalizedProviderForModelConfig(ctx, modelConfig) +} + +// resolveModelConfig looks up the chat's enabled model config by its +// LastModelConfigID. If the referenced config is unavailable or belongs to +// another organization, it falls back to the local default model config. +// Returns an error when no usable local config is available. func (p *Server) resolveModelConfig( ctx context.Context, chat database.Chat, ) (database.ChatModelConfig, error) { if chat.LastModelConfigID != uuid.Nil { - modelConfig, err := p.configCache.ModelConfigByID( - ctx, chat.LastModelConfigID, - ) + modelConfig, err := p.configCache.ModelConfigByID(ctx, chat.LastModelConfigID) if err == nil { - return modelConfig, nil - } - if !xerrors.Is(err, sql.ErrNoRows) { + if modelConfig.Enabled && modelConfig.OrganizationID == chat.OrganizationID && modelConfig.AIProviderID.Valid { + provider, providerErr := p.db.GetAIProviderByID( + chatdModelConfigLookupContext(ctx), + modelConfig.AIProviderID.UUID, + ) + switch { + case providerErr == nil && provider.Enabled: + return modelConfig, nil + case providerErr == nil, xerrors.Is(providerErr, sql.ErrNoRows): + default: + return database.ChatModelConfig{}, xerrors.Errorf( + "get AI provider %s: %w", + modelConfig.AIProviderID.UUID, providerErr, + ) + } + } + } else if !xerrors.Is(err, sql.ErrNoRows) { return database.ChatModelConfig{}, xerrors.Errorf( "get chat model config %s: %w", chat.LastModelConfigID, err, ) } - // Model config was deleted, fall through to default. } - defaultConfig, err := p.configCache.DefaultModelConfig(ctx) + defaultConfig, err := p.configCache.DefaultModelConfig(ctx, chat.OrganizationID) if err != nil { if xerrors.Is(err, sql.ErrNoRows) { return database.ChatModelConfig{}, ErrNoDefaultChatModelConfig diff --git a/coderd/x/chatd/chatd_helpers_test.go b/coderd/x/chatd/chatd_helpers_test.go index a271016ab82..95c93be3ebe 100644 --- a/coderd/x/chatd/chatd_helpers_test.go +++ b/coderd/x/chatd/chatd_helpers_test.go @@ -64,9 +64,10 @@ func seedAnthropicChatDependencies(t *testing.T, db database.Store, baseURL stri }) dbgen.AIProviderKey(t, db, database.AIProviderKey{ProviderID: provider.ID}) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "claude-sonnet-4-20250514", - IsDefault: true, - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "claude-sonnet-4-20250514", + IsDefault: true, + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }) return user, org, model } diff --git a/coderd/x/chatd/chatd_internal_test.go b/coderd/x/chatd/chatd_internal_test.go index b6895d94e55..f22988c00e9 100644 --- a/coderd/x/chatd/chatd_internal_test.go +++ b/coderd/x/chatd/chatd_internal_test.go @@ -1025,9 +1025,11 @@ func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) { activeAPIKeyID := "key-" + uuid.NewString() wantTitle := "Review PR 23633" + organizationID := uuid.New() chat := database.Chat{ ID: chatID, OwnerID: ownerID, + OrganizationID: organizationID, LastModelConfigID: modelConfigID, Status: database.ChatStatusRunning, WorkerID: uuid.NullUUID{UUID: workerID, Valid: true}, @@ -1035,10 +1037,12 @@ func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) { } providerID := uuid.New() modelConfig := database.ChatModelConfig{ - ID: modelConfigID, - Model: "gpt-4o-mini", - ContextLimit: 8192, - AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + ID: modelConfigID, + Model: "gpt-4o-mini", + ContextLimit: 8192, + Enabled: true, + AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + OrganizationID: organizationID, } updatedChat := chat updatedChat.Title = wantTitle @@ -1129,7 +1133,7 @@ func TestRegenerateChatTitle_PersistsAndBroadcasts(t *testing.T) { }, ).Return(nil, nil) db.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", nil) - db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return(nil, nil) + db.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), organizationID).Return(nil, nil) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, opts *database.TxOptions) error { @@ -1184,18 +1188,22 @@ func TestRegenerateChatTitle_SkipsPersistWhenTitleChangedConcurrently(t *testing activeAPIKeyID := "key-" + uuid.NewString() generatedTitle := "Review PR 23633" + organizationID := uuid.New() chat := database.Chat{ ID: chatID, OwnerID: ownerID, + OrganizationID: organizationID, LastModelConfigID: modelConfigID, Status: database.ChatStatusWaiting, Title: chatprompt.FallbackTitle(userPrompt), } modelConfig := database.ChatModelConfig{ - ID: modelConfigID, - Model: "gpt-4o-mini", - ContextLimit: 8192, - AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + ID: modelConfigID, + Model: "gpt-4o-mini", + ContextLimit: 8192, + Enabled: true, + AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + OrganizationID: organizationID, } // Another writer (rename or a second regenerate) landed while the // model call was in flight. @@ -1274,7 +1282,7 @@ func TestRegenerateChatTitle_SkipsPersistWhenTitleChangedConcurrently(t *testing }, ).Return(nil, nil) db.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", nil) - db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return(nil, nil) + db.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), organizationID).Return(nil, nil) db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( func(fn func(database.Store) error, _ *database.TxOptions) error { @@ -3657,6 +3665,92 @@ func TestServer_inflightContext(t *testing.T) { } } +func TestResolveModelConfigProviderLookupError(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + organizationID := uuid.New() + modelConfigID := uuid.New() + providerID := uuid.New() + modelConfig := database.ChatModelConfig{ + ID: modelConfigID, + Enabled: true, + OrganizationID: organizationID, + AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + } + + db.EXPECT().GetChatModelConfigByID(gomock.Any(), modelConfigID).Return(modelConfig, nil) + db.EXPECT().GetAIProviderByID(gomock.Any(), providerID).Return(database.AIProvider{}, sql.ErrConnDone) + + server := &Server{ + db: db, + configCache: newChatConfigCache(context.Background(), db, quartz.NewReal()), + } + _, err := server.resolveModelConfig(ctx, database.Chat{ + OrganizationID: organizationID, + LastModelConfigID: modelConfigID, + }) + require.ErrorIs(t, err, sql.ErrConnDone) + require.ErrorContains(t, err, "get AI provider") +} + +func TestResolveModelConfigOrganizationScope(t *testing.T) { + t.Parallel() + + t.Run("ForeignConfigFallsBackToLocalDefault", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + chatOrg := dbgen.Organization(t, db, database.Organization{}) + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: defaultOrg.ID, + }) + localDefault := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: chatOrg.ID, + IsDefault: true, + }) + server := &Server{ + db: db, + configCache: newChatConfigCache(context.Background(), db, quartz.NewReal()), + } + + got, err := server.resolveModelConfig(ctx, database.Chat{ + OrganizationID: chatOrg.ID, + LastModelConfigID: foreignModel.ID, + }) + require.NoError(t, err) + require.Equal(t, localDefault.ID, got.ID) + }) + + t.Run("ForeignConfigWithoutLocalDefaultFails", func(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + chatOrg := dbgen.Organization(t, db, database.Organization{}) + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + OrganizationID: defaultOrg.ID, + }) + server := &Server{ + db: db, + configCache: newChatConfigCache(context.Background(), db, quartz.NewReal()), + } + + _, err = server.resolveModelConfig(ctx, database.Chat{ + OrganizationID: chatOrg.ID, + LastModelConfigID: foreignModel.ID, + }) + require.ErrorIs(t, err, ErrNoDefaultChatModelConfig) + }) +} + // TestResolveFallbackModelConfigID verifies that admission does not reuse // a disabled last model and rejects a disabled default. func TestResolveFallbackModelConfigID(t *testing.T) { @@ -3667,10 +3761,17 @@ func TestResolveFallbackModelConfigID(t *testing.T) { p.Enabled = enabled }) } - newModelConfig := func(t *testing.T, db database.Store, providerID uuid.UUID, isDefault bool) database.ChatModelConfig { + // The configs under test live in an ad-hoc org so the per-org + // default lookup stays scoped to them, independent of any seeded + // default-org config. + newModelConfigOrg := func(t *testing.T, db database.Store) uuid.UUID { + return dbgen.Organization(t, db, database.Organization{}).ID + } + newModelConfig := func(t *testing.T, db database.Store, orgID, providerID uuid.UUID, isDefault bool) database.ChatModelConfig { return dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, - IsDefault: isDefault, + AIProviderID: uuid.NullUUID{UUID: providerID, Valid: true}, + IsDefault: isDefault, + OrganizationID: orgID, }) } @@ -3679,10 +3780,11 @@ func TestResolveFallbackModelConfigID(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) + orgID := newModelConfigOrg(t, db) provider := newProvider(t, db, true) - lastModel := newModelConfig(t, db, provider.ID, false) + lastModel := newModelConfig(t, db, orgID, provider.ID, false) - resolved, err := resolveFallbackModelConfigID(ctx, db, lastModel.ID) + resolved, err := resolveFallbackModelConfigID(ctx, db, orgID, lastModel.ID) require.NoError(t, err) require.Equal(t, lastModel.ID, resolved) }) @@ -3692,12 +3794,13 @@ func TestResolveFallbackModelConfigID(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) + orgID := newModelConfigOrg(t, db) disabledProvider := newProvider(t, db, false) - lastModel := newModelConfig(t, db, disabledProvider.ID, false) + lastModel := newModelConfig(t, db, orgID, disabledProvider.ID, false) enabledProvider := newProvider(t, db, true) - defaultModel := newModelConfig(t, db, enabledProvider.ID, true) + defaultModel := newModelConfig(t, db, orgID, enabledProvider.ID, true) - resolved, err := resolveFallbackModelConfigID(ctx, db, lastModel.ID) + resolved, err := resolveFallbackModelConfigID(ctx, db, orgID, lastModel.ID) require.NoError(t, err) require.Equal(t, defaultModel.ID, resolved) }) @@ -3707,24 +3810,79 @@ func TestResolveFallbackModelConfigID(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) + orgID := newModelConfigOrg(t, db) provider := newProvider(t, db, true) - defaultModel := newModelConfig(t, db, provider.ID, true) + defaultModel := newModelConfig(t, db, orgID, provider.ID, true) - resolved, err := resolveFallbackModelConfigID(ctx, db, uuid.Nil) + resolved, err := resolveFallbackModelConfigID(ctx, db, orgID, uuid.Nil) require.NoError(t, err) require.Equal(t, defaultModel.ID, resolved) }) + t.Run("NonDefaultOrgWithoutLocalDefaultRejected", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + otherOrgID := newModelConfigOrg(t, db) + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + provider := newProvider(t, db, true) + newModelConfig(t, db, defaultOrg.ID, provider.ID, true) + + _, err = resolveFallbackModelConfigID(ctx, db, otherOrgID, uuid.Nil) + require.ErrorIs(t, err, ErrNoDefaultChatModelConfig) + }) + + t.Run("DisabledLocalDefaultRejected", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + orgID := newModelConfigOrg(t, db) + disabledProvider := newProvider(t, db, true) + disabledLocal := newModelConfig(t, db, orgID, disabledProvider.ID, true) + _, err := db.UpdateChatModelConfig(ctx, database.UpdateChatModelConfigParams{ + ID: disabledLocal.ID, + Model: disabledLocal.Model, + DisplayName: disabledLocal.DisplayName, + Enabled: false, + IsDefault: true, + ContextLimit: disabledLocal.ContextLimit, + CompressionThreshold: disabledLocal.CompressionThreshold, + Options: disabledLocal.Options, + AIProviderID: disabledLocal.AIProviderID, + }) + require.NoError(t, err) + + _, err = resolveFallbackModelConfigID(ctx, db, orgID, uuid.Nil) + require.ErrorIs(t, err, ErrNoDefaultChatModelConfig) + }) + + t.Run("DisabledLocalDefaultProviderRejected", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + orgID := newModelConfigOrg(t, db) + disabledProvider := newProvider(t, db, false) + newModelConfig(t, db, orgID, disabledProvider.ID, true) + + _, err := resolveFallbackModelConfigID(ctx, db, orgID, uuid.Nil) + require.ErrorIs(t, err, ErrNoDefaultChatModelConfig) + }) + t.Run("ProviderDisabledDefaultRejected", func(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) + orgID := newModelConfigOrg(t, db) disabledProvider := newProvider(t, db, false) - lastModel := newModelConfig(t, db, disabledProvider.ID, false) - newModelConfig(t, db, disabledProvider.ID, true) + lastModel := newModelConfig(t, db, orgID, disabledProvider.ID, false) + newModelConfig(t, db, orgID, disabledProvider.ID, true) - _, err := resolveFallbackModelConfigID(ctx, db, lastModel.ID) + _, err := resolveFallbackModelConfigID(ctx, db, orgID, lastModel.ID) require.ErrorIs(t, err, ErrNoDefaultChatModelConfig) }) @@ -3733,14 +3891,54 @@ func TestResolveFallbackModelConfigID(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) + orgID := newModelConfigOrg(t, db) provider := newProvider(t, db, true) - model := newModelConfig(t, db, provider.ID, false) + model := newModelConfig(t, db, orgID, provider.ID, false) - resolved, err := resolveSendMessageModelConfigID(ctx, db, database.Chat{}, model.ID) + resolved, err := resolveSendMessageModelConfigID(ctx, db, database.Chat{OrganizationID: orgID}, model.ID) require.NoError(t, err) require.Equal(t, model.ID, resolved) }) + t.Run("ExplicitDefaultOrgModelRejected", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + chatOrgID := newModelConfigOrg(t, db) + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + provider := newProvider(t, db, true) + model := newModelConfig(t, db, defaultOrg.ID, provider.ID, false) + + _, err = resolveSendMessageModelConfigID( + ctx, + db, + database.Chat{OrganizationID: chatOrgID}, + model.ID, + ) + require.ErrorIs(t, err, ErrInvalidModelConfigID) + }) + + t.Run("ExplicitUnrelatedOrgModelRejected", func(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + chatOrgID := newModelConfigOrg(t, db) + unrelatedOrgID := newModelConfigOrg(t, db) + provider := newProvider(t, db, true) + model := newModelConfig(t, db, unrelatedOrgID, provider.ID, false) + + _, err := resolveSendMessageModelConfigID( + ctx, + db, + database.Chat{OrganizationID: chatOrgID}, + model.ID, + ) + require.ErrorIs(t, err, ErrInvalidModelConfigID) + }) + // An explicit model whose provider was disabled after the coderd // preflight must still be rejected inside the daemon. t.Run("ExplicitProviderDisabledRejected", func(t *testing.T) { @@ -3748,10 +3946,11 @@ func TestResolveFallbackModelConfigID(t *testing.T) { db, _ := dbtestutil.NewDB(t) ctx := testutil.Context(t, testutil.WaitShort) + orgID := newModelConfigOrg(t, db) disabledProvider := newProvider(t, db, false) - model := newModelConfig(t, db, disabledProvider.ID, false) + model := newModelConfig(t, db, orgID, disabledProvider.ID, false) - _, err := resolveSendMessageModelConfigID(ctx, db, database.Chat{}, model.ID) + _, err := resolveSendMessageModelConfigID(ctx, db, database.Chat{OrganizationID: orgID}, model.ID) require.ErrorIs(t, err, ErrInvalidModelConfigID) }) @@ -3764,7 +3963,7 @@ func TestResolveFallbackModelConfigID(t *testing.T) { owner := dbgen.User(t, db, database.User{}) disabledProvider := newProvider(t, db, false) - model := newModelConfig(t, db, disabledProvider.ID, false) + model := newModelConfig(t, db, newModelConfigOrg(t, db), disabledProvider.ID, false) server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) _, err := server.CreateChat(ctx, CreateOptions{ diff --git a/coderd/x/chatd/chatd_retry_test.go b/coderd/x/chatd/chatd_retry_test.go index 51eb8930970..99b88ca2020 100644 --- a/coderd/x/chatd/chatd_retry_test.go +++ b/coderd/x/chatd/chatd_retry_test.go @@ -104,10 +104,11 @@ func TestActiveServer_RetryStreamSilenceTimeoutAndClassification(t *testing.T) { }) user, org, _ := seedChatDependenciesWithProvider(t, db, "openai", openAIURL) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o", - Enabled: true, - CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + Model: "gpt-4o", + Enabled: true, + CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, + OrganizationID: org.ID, }) factory := chattest.NewMockAIBridgeTransport(t, openAIURL) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { diff --git a/coderd/x/chatd/chatd_test.go b/coderd/x/chatd/chatd_test.go index f1a0ac78152..47babace29b 100644 --- a/coderd/x/chatd/chatd_test.go +++ b/coderd/x/chatd/chatd_test.go @@ -5764,9 +5764,10 @@ func TestActiveServer_CompactionModelOverride(t *testing.T) { seedOverrideModel := func(ctx context.Context, t *testing.T, db database.Store, chatModel database.ChatModelConfig, modelName, effort string, contextLimit int64) database.ChatModelConfig { t.Helper() overrideModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: modelName, - AIProviderID: chatModel.AIProviderID, - ContextLimit: contextLimit, + Model: modelName, + AIProviderID: chatModel.AIProviderID, + ContextLimit: contextLimit, + OrganizationID: chatModel.OrganizationID, }) overrideModel = updateChatModelCallConfig(t, db, overrideModel, codersdk.ChatModelCallConfig{ ReasoningEffort: &codersdk.ChatModelReasoningEffortConfig{ @@ -8786,8 +8787,9 @@ func seedChatDependenciesWithProvider( BaseUrl: baseURL, }) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - AIProviderID: uuid.NullUUID{UUID: providerConfig.ID, Valid: true}, - IsDefault: true, + AIProviderID: uuid.NullUUID{UUID: providerConfig.ID, Valid: true}, + IsDefault: true, + OrganizationID: org.ID, }) return user, org, model } @@ -8824,8 +8826,9 @@ func seedChatDependenciesWithProviderPolicy( }) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - AIProviderID: uuid.NullUUID{UUID: providerConfig.ID, Valid: true}, - IsDefault: true, + AIProviderID: uuid.NullUUID{UUID: providerConfig.ID, Valid: true}, + IsDefault: true, + OrganizationID: org.ID, }) return user, org, providerConfig, model @@ -8902,13 +8905,23 @@ func insertChatModelConfigWithCallConfig( Type: database.AIProviderType(provider), }) } + organizations, err := db.GetOrganizationsByUserID(context.Background(), database.GetOrganizationsByUserIDParams{ + UserID: userID, + Deleted: sql.NullBool{ + Bool: false, + Valid: true, + }, + }) + require.NoError(t, err) + require.Len(t, organizations, 1) return dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - AIProviderID: uuid.NullUUID{UUID: aiProvider.ID, Valid: true}, - Model: model, - DisplayName: model, - CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, - Options: options, + AIProviderID: uuid.NullUUID{UUID: aiProvider.ID, Valid: true}, + Model: model, + DisplayName: model, + CreatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + UpdatedBy: uuid.NullUUID{UUID: userID, Valid: true}, + Options: options, + OrganizationID: organizations[0].ID, }) } @@ -10107,9 +10120,10 @@ func seedAIGatewayOpenAITestDependencies( BaseUrl: openAIURL, }) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - IsDefault: true, - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "gpt-4o-mini", + IsDefault: true, + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }) _, err := db.UpsertUserAIProviderKey(context.Background(), database.UpsertUserAIProviderKeyParams{ ID: uuid.New(), @@ -12416,6 +12430,235 @@ func TestEditMessageRejectsUnknownModelConfig(t *testing.T) { require.Equal(t, modelA.ID, storedChat.LastModelConfigID) } +func TestQueuedCompletionResolvesOrganizationModel(t *testing.T) { + t.Parallel() + + t.Run("TurnCompletion", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitSuperLong) + firstRunStarted := make(chan struct{}) + allowFirstRunFinish := make(chan struct{}) + var requestCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if requestCount.Add(1) == 1 { + chunks := make(chan chattest.OpenAIChunk, 1) + go func() { + defer close(chunks) + select { + case <-firstRunStarted: + default: + close(firstRunStarted) + } + <-allowFirstRunFinish + chunks <- chattest.OpenAITextChunks("first done")[0] + }() + return chattest.OpenAIResponse{StreamingChunks: chunks} + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("second done")...) + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + }) + user, org, localDefault := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, localDefault.ID, "first") + testutil.TryReceive(ctx, t, firstRunStarted) + insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + close(allowFirstRunFinish) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + messages := chatMessages(ctx, t, db, chat.ID) + var queuedMessage database.ChatMessage + for _, message := range messages { + if message.Role != database.ChatMessageRoleUser { + continue + } + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parts) == 1 && parts[0].Text == "queued" { + queuedMessage = message + } + } + require.NotZero(t, queuedMessage.ID) + require.Equal(t, localDefault.ID, queuedMessage.ModelConfigID.UUID) + }) + + t.Run("InterruptionCompletion", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitSuperLong) + streamStarted := make(chan struct{}) + var requestCount atomic.Int32 + openAIURL := chattest.NewOpenAI(t, func(req *chattest.OpenAIRequest) chattest.OpenAIResponse { + if !req.Stream { + return chattest.OpenAINonStreamingResponse("title") + } + if requestCount.Add(1) == 1 { + chunks := make(chan chattest.OpenAIChunk, 1) + go func() { + defer close(chunks) + select { + case <-streamStarted: + default: + close(streamStarted) + } + <-req.Context().Done() + }() + return chattest.OpenAIResponse{StreamingChunks: chunks} + } + return chattest.OpenAIStreamingResponse(chattest.OpenAITextChunks("second done")...) + }) + server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { + cfg.AIBridgeTransportFactory = chatAIGatewayTransportFactoryPointer(chattest.NewMockAIBridgeTransport(t, openAIURL)) + }) + user, org, localDefault := seedChatDependenciesWithProvider(t, db, "openai-compat", openAIURL) + foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := createChatThroughServer(ctx, t, db, server, org.ID, user.ID, localDefault.ID, "first") + testutil.TryReceive(ctx, t, streamStarted) + insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + current, err := db.GetChatByID(ctx, chat.ID) + require.NoError(t, err) + _, err = server.InterruptChat(ctx, current) + require.NoError(t, err) + waitForChatStatus(ctx, t, db, chat.ID, database.ChatStatusWaiting) + + messages := chatMessages(ctx, t, db, chat.ID) + var queuedMessage database.ChatMessage + for _, message := range messages { + if message.Role != database.ChatMessageRoleUser { + continue + } + parts, err := chatprompt.ParseContent(message) + require.NoError(t, err) + if len(parts) == 1 && parts[0].Text == "queued" { + queuedMessage = message + } + } + require.NotZero(t, queuedMessage.ID) + require.Equal(t, localDefault.ID, queuedMessage.ModelConfigID.UUID) + }) +} + +func TestQueuedPromotionResolvesOrganizationModel(t *testing.T) { + t.Parallel() + + t.Run("ExplicitPromotion", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + server := newTestServer(t, db, ps, uuid.New()) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, localDefault := seedChatDependencies(t, db) + foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: localDefault.ID, + Title: "resolve explicit queued promotion", + Status: database.ChatStatusError, + }) + queued := insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + + result, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + QueuedMessageID: queued.ID, + }) + require.NoError(t, err) + require.Equal(t, localDefault.ID, result.PromotedMessage.ModelConfigID.UUID) + }) + + t.Run("ErrorStateSendPromotion", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + server := newTestServer(t, db, ps, uuid.New()) + ctx := testutil.Context(t, testutil.WaitLong) + user, org, localDefault := seedChatDependencies(t, db) + foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: localDefault.ID, + Title: "resolve error send promotion", + Status: database.ChatStatusError, + }) + insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + + result, err := server.SendMessage(ctx, chatd.SendMessageOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + Content: []codersdk.ChatMessagePart{codersdk.ChatMessageText("new")}, + BusyBehavior: chatd.SendMessageBusyBehaviorQueue, + }) + require.NoError(t, err) + require.NotEmpty(t, result.InsertedMessages) + promoted := result.InsertedMessages[len(result.InsertedMessages)-1] + require.Equal(t, localDefault.ID, promoted.ModelConfigID.UUID) + }) + + t.Run("NoLocalDefaultKeepsQueuedRow", func(t *testing.T) { + t.Parallel() + + db, ps := dbtestutil.NewDB(t) + server := newTestServer(t, db, ps, uuid.New()) + ctx := testutil.Context(t, testutil.WaitLong) + user := dbgen.User(t, db, database.User{}) + org := dbgen.Organization(t, db, database.Organization{}) + foreignModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{}) + chat := dbgen.Chat(t, db, database.Chat{ + OrganizationID: org.ID, + OwnerID: user.ID, + LastModelConfigID: foreignModel.ID, + Title: "keep unresolved queued row", + Status: database.ChatStatusError, + }) + queued := insertQueuedMessage(ctx, t, db, chat.ID, user.ID, foreignModel.ID, "queued") + + _, err := server.PromoteQueued(ctx, chatd.PromoteQueuedOptions{ + ChatID: chat.ID, + CreatedBy: user.ID, + QueuedMessageID: queued.ID, + }) + require.ErrorIs(t, err, chatd.ErrNoDefaultChatModelConfig) + rows, err := db.GetChatQueuedMessages(ctx, chat.ID) + require.NoError(t, err) + require.Len(t, rows, 1) + require.Equal(t, queued.ID, rows[0].ID) + messages, err := db.GetChatMessagesByChatID(ctx, database.GetChatMessagesByChatIDParams{ChatID: chat.ID}) + require.NoError(t, err) + require.Empty(t, messages) + }) +} + +func insertQueuedMessage( + ctx context.Context, + t *testing.T, + db database.Store, + chatID uuid.UUID, + createdBy uuid.UUID, + modelConfigID uuid.UUID, + text string, +) database.ChatQueuedMessage { + t.Helper() + content, err := chatprompt.MarshalParts([]codersdk.ChatMessagePart{codersdk.ChatMessageText(text)}) + require.NoError(t, err) + queued, err := db.InsertChatQueuedMessageWithCreator(ctx, database.InsertChatQueuedMessageWithCreatorParams{ + ChatID: chatID, + Content: content.RawMessage, + ModelConfigID: uuid.NullUUID{UUID: modelConfigID, Valid: true}, + CreatedBy: createdBy, + }) + require.NoError(t, err) + return queued +} + func TestPromoteQueuedPreservesReasoningEffort(t *testing.T) { t.Parallel() @@ -13235,16 +13478,18 @@ func TestProviderSwitchSanitizesAndRestoresPEToolHistory(t *testing.T) { }) mA := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - DisplayName: "Model A", - Enabled: true, - AIProviderID: uuid.NullUUID{UUID: cpA.ID, Valid: true}, + Model: "gpt-4o-mini", + DisplayName: "Model A", + Enabled: true, + AIProviderID: uuid.NullUUID{UUID: cpA.ID, Valid: true}, + OrganizationID: org.ID, }) mB := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - DisplayName: "Model B", - Enabled: true, - AIProviderID: uuid.NullUUID{UUID: cpB.ID, Valid: true}, + Model: "gpt-4o-mini", + DisplayName: "Model B", + Enabled: true, + AIProviderID: uuid.NullUUID{UUID: cpB.ID, Valid: true}, + OrganizationID: org.ID, }) server := newActiveTestServer(t, db, ps, func(cfg *chatd.Config) { diff --git a/coderd/x/chatd/chatstate/errors.go b/coderd/x/chatd/chatstate/errors.go index 48da8d2708a..af37b12a0ba 100644 --- a/coderd/x/chatd/chatstate/errors.go +++ b/coderd/x/chatd/chatstate/errors.go @@ -49,6 +49,10 @@ var ( // wraps this sentinel. ErrMessageQueueFull = xerrors.New("chat message queue is full") + // ErrNoDefaultChatModelConfig is returned when queued-message + // promotion cannot find an enabled default model in the chat organization. + ErrNoDefaultChatModelConfig = xerrors.New("no default chat model config is configured") + // ErrChatFileCapExceeded reports a [LinkFiles] cap rejection. ErrChatFileCapExceeded = xerrors.New("chat attachment cap exceeded") diff --git a/coderd/x/chatd/chatstate/machine_test.go b/coderd/x/chatd/chatstate/machine_test.go index e89effc9cef..7e3b92dadf7 100644 --- a/coderd/x/chatd/chatstate/machine_test.go +++ b/coderd/x/chatd/chatstate/machine_test.go @@ -50,7 +50,8 @@ func newTestFixture(t *testing.T) *testFixture { BaseUrl: "http://example.invalid", }) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - IsDefault: true, + OrganizationID: org.ID, + IsDefault: true, }) pub := newRecordingPubsub() return &testFixture{ diff --git a/coderd/x/chatd/chatstate/transitions.go b/coderd/x/chatd/chatstate/transitions.go index c2b49ad62ab..d9e7d1e0342 100644 --- a/coderd/x/chatd/chatstate/transitions.go +++ b/coderd/x/chatd/chatstate/transitions.go @@ -12,6 +12,7 @@ import ( "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbauthz" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/codersdk" @@ -281,18 +282,52 @@ func (tx *Tx) insertQueuedMessage(ownerFallback uuid.UUID, m Message) (database. }) } -// messageFromQueuedRow synthesizes a Message from a stored queued row, -// suitable for promoting into active history. -func messageFromQueuedRow(q database.ChatQueuedMessage) Message { +func (tx *Tx) messageFromQueuedRow(chat database.Chat, queued database.ChatQueuedMessage) (Message, error) { + modelConfigID, err := tx.resolveQueuedMessageModelConfigID(chat, queued) + if err != nil { + return Message{}, err + } return Message{ Role: database.ChatMessageRoleUser, - Content: pqtype.NullRawMessage{RawMessage: q.Content, Valid: q.Content != nil}, + Content: pqtype.NullRawMessage{RawMessage: queued.Content, Valid: queued.Content != nil}, Visibility: database.ChatMessageVisibilityBoth, - ModelConfigID: q.ModelConfigID, - ReasoningEffort: q.ReasoningEffort, - CreatedBy: uuid.NullUUID{UUID: q.CreatedBy, Valid: true}, + ModelConfigID: modelConfigID, + ReasoningEffort: queued.ReasoningEffort, + CreatedBy: uuid.NullUUID{UUID: queued.CreatedBy, Valid: true}, ContentVersion: chatprompt.CurrentContentVersion, + }, nil +} + +func (tx *Tx) resolveQueuedMessageModelConfigID( + chat database.Chat, + queued database.ChatQueuedMessage, +) (uuid.NullUUID, error) { + //nolint:gocritic // Queue promotion needs daemon access to deployment model configuration. + ctx := dbauthz.AsChatd(tx.ctx) + if queued.ModelConfigID.Valid { + config, err := tx.store.GetEnabledChatModelConfigByID(ctx, queued.ModelConfigID.UUID) + if err == nil && config.OrganizationID == chat.OrganizationID { + return queued.ModelConfigID, nil + } + if err != nil && !errors.Is(err, sql.ErrNoRows) { + return uuid.NullUUID{}, xerrors.Errorf( + "get chat model config %s: %w", + queued.ModelConfigID.UUID, + err, + ) + } } + + configs, err := tx.store.GetEnabledChatModelConfigsByOrganization(ctx, chat.OrganizationID) + if err != nil { + return uuid.NullUUID{}, xerrors.Errorf("get default chat model config: %w", err) + } + for _, config := range configs { + if config.ChatModelConfig.IsDefault { + return uuid.NullUUID{UUID: config.ChatModelConfig.ID, Valid: true}, nil + } + } + return uuid.NullUUID{}, ErrNoDefaultChatModelConfig } // SetArchivedInput configures [Tx.SetArchived]. @@ -461,7 +496,10 @@ func (tx *Tx) sendMessageE1(chat database.Chat, input SendMessageInput) (SendMes if err != nil { return SendMessageResult{}, err } - promoted := messageFromQueuedRow(head) + promoted, err := tx.messageFromQueuedRow(chat, head) + if err != nil { + return SendMessageResult{}, xerrors.Errorf("resolve promoted queued head: %w", err) + } inserted, err := tx.insertMessages(append(cancels, promoted)) if err != nil { return SendMessageResult{}, xerrors.Errorf("insert promoted queued head: %w", err) @@ -827,7 +865,10 @@ func (tx *Tx) PromoteQueuedMessage(input PromoteQueuedMessageInput) (PromoteQueu if err != nil { return PromoteQueuedMessageResult{}, err } - promotedMsg := messageFromQueuedRow(target) + promotedMsg, err := tx.messageFromQueuedRow(chat, target) + if err != nil { + return PromoteQueuedMessageResult{}, xerrors.Errorf("resolve promoted queued message: %w", err) + } inserted, err := tx.insertMessages(append(cancels, promotedMsg)) if err != nil { return PromoteQueuedMessageResult{}, xerrors.Errorf("insert promoted queued message: %w", err) @@ -1318,7 +1359,10 @@ func (tx *Tx) FinishInterruption(input FinishInterruptionInput) (FinishInterrupt if err != nil { return FinishInterruptionResult{}, xerrors.Errorf("get queue head: %w", err) } - promotedMsg := messageFromQueuedRow(head) + promotedMsg, err := tx.messageFromQueuedRow(chat, head) + if err != nil { + return FinishInterruptionResult{}, xerrors.Errorf("resolve promoted queue head: %w", err) + } insertedHead, err := tx.insertMessages([]Message{promotedMsg}) if err != nil { return FinishInterruptionResult{}, xerrors.Errorf("insert promoted queue head: %w", err) @@ -1350,7 +1394,7 @@ func (tx *Tx) FinishInterruption(input FinishInterruptionInput) (FinishInterrupt }, nil } -// FinishTurnInput is intentionally empty. +// FinishTurnInput configures [Tx.FinishTurn]. type FinishTurnInput struct{} // FinishTurnResult is returned by [Tx.FinishTurn]. @@ -1388,7 +1432,10 @@ func (tx *Tx) FinishTurn(_ FinishTurnInput) (FinishTurnResult, error) { if err != nil { return FinishTurnResult{}, err } - promotedMsg := messageFromQueuedRow(head) + promotedMsg, err := tx.messageFromQueuedRow(chat, head) + if err != nil { + return FinishTurnResult{}, xerrors.Errorf("resolve promoted queue head: %w", err) + } inserted, err := tx.insertMessages(append(cancels, promotedMsg)) if err != nil { return FinishTurnResult{}, xerrors.Errorf("insert promoted queue head: %w", err) diff --git a/coderd/x/chatd/chatstate/transitions_test.go b/coderd/x/chatd/chatstate/transitions_test.go index 6b5bfb1c3af..31dbe23459a 100644 --- a/coderd/x/chatd/chatstate/transitions_test.go +++ b/coderd/x/chatd/chatstate/transitions_test.go @@ -1,14 +1,17 @@ package chatstate_test import ( + "context" "encoding/json" "testing" "github.com/google/uuid" "github.com/sqlc-dev/pqtype" "github.com/stretchr/testify/require" + "golang.org/x/xerrors" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/x/chatd/chatprompt" "github.com/coder/coder/v2/coderd/x/chatd/chatstate" "github.com/coder/coder/v2/codersdk" @@ -101,6 +104,180 @@ func TestCreateChat_AllowsNonFinalUserMessage(t *testing.T) { require.Len(t, res.InitialMessages, 2) } +func TestPromoteQueuedMessageResolvesOrganizationModel(t *testing.T) { + t.Parallel() + + t.Run("local model remains selected", func(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + localModel := dbgen.ChatModelConfig(t, f.DB, database.ChatModelConfig{ + OrganizationID: f.Org.ID, + }) + promoted := promoteQueuedMessageWithModel(t, f, localModel.ID) + require.Equal(t, localModel.ID, promoted.ModelConfigID.UUID) + }) + + t.Run("foreign model falls back to local default", func(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + foreignOrg := dbgen.Organization(t, f.DB, database.Organization{}) + foreignModel := dbgen.ChatModelConfig(t, f.DB, database.ChatModelConfig{ + OrganizationID: foreignOrg.ID, + IsDefault: true, + }) + promoted := promoteQueuedMessageWithModel(t, f, foreignModel.ID) + require.Equal(t, f.Model.ID, promoted.ModelConfigID.UUID) + }) + + t.Run("missing model falls back to local default", func(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + promoted := promoteQueuedMessageWithModel(t, f, uuid.New()) + require.Equal(t, f.Model.ID, promoted.ModelConfigID.UUID) + }) + + t.Run("disabled model falls back to local default", func(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + disabled := dbgen.ChatModelConfig(t, f.DB, database.ChatModelConfig{ + OrganizationID: f.Org.ID, + }) + setChatModelConfigEnabled(t, f.DB, disabled, false) + promoted := promoteQueuedMessageWithModel(t, f, disabled.ID) + require.Equal(t, f.Model.ID, promoted.ModelConfigID.UUID) + }) + + t.Run("missing local default rejects promotion", func(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + setChatModelConfigEnabled(t, f.DB, f.Model, false) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + machine := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + queued := sendQueuedMessage(t, f, machine, "queued without default") + require.NotNil(t, queued.QueuedMessage) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{RawMessage: json.RawMessage(`{"message":"boom"}`), Valid: true}, + }) + return err + })) + + err := machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: queued.QueuedMessage.ID, + }) + return err + }) + require.ErrorIs(t, err, chatstate.ErrNoDefaultChatModelConfig) + require.Equal(t, chatstate.StateE1, f.classify(ctx, t, created.Chat.ID)) + }) + + t.Run("database error rejects promotion", func(t *testing.T) { + t.Parallel() + f := newTestFixture(t) + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + machine := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + queued := sendQueuedMessage(t, f, machine, "queued database error") + require.NotNil(t, queued.QueuedMessage) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{RawMessage: json.RawMessage(`{"message":"boom"}`), Valid: true}, + }) + return err + })) + + queryErr := xerrors.New("model query failed") + failingStore := &modelConfigErrorStore{Store: f.DB, err: queryErr} + machine = chatstate.NewChatMachine(failingStore, f.Pub, created.Chat.ID) + err := machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: queued.QueuedMessage.ID, + }) + return err + }) + require.ErrorIs(t, err, queryErr) + require.Equal(t, chatstate.StateE1, f.classify(ctx, t, created.Chat.ID)) + }) +} + +type modelConfigErrorStore struct { + database.Store + err error +} + +func (s *modelConfigErrorStore) InTx(fn func(database.Store) error, opts *database.TxOptions) error { + return s.Store.InTx(func(store database.Store) error { + return fn(&modelConfigErrorStore{Store: store, err: s.err}) + }, opts) +} + +func (s *modelConfigErrorStore) GetEnabledChatModelConfigByID( + context.Context, + uuid.UUID, +) (database.ChatModelConfig, error) { + return database.ChatModelConfig{}, s.err +} + +func promoteQueuedMessageWithModel( + t *testing.T, + f *testFixture, + modelConfigID uuid.UUID, +) database.ChatMessage { + t.Helper() + ctx := testutil.Context(t, testutil.WaitShort) + created := createTestChat(t, f) + message := userTextMessage("queued model resolution", f.User.ID, modelConfigID) + queued, err := f.DB.InsertChatQueuedMessage(ctx, database.InsertChatQueuedMessageParams{ + ChatID: created.Chat.ID, + Content: message.Content.RawMessage, + ModelConfigID: message.ModelConfigID, + }) + require.NoError(t, err) + machine := chatstate.NewChatMachine(f.DB, f.Pub, created.Chat.ID) + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + _, err := tx.FinishError(chatstate.FinishErrorInput{ + LastError: pqtype.NullRawMessage{RawMessage: json.RawMessage(`{"message":"boom"}`), Valid: true}, + }) + return err + })) + + var result chatstate.PromoteQueuedMessageResult + require.NoError(t, machine.Update(ctx, func(tx *chatstate.Tx, _ database.Store) error { + var err error + result, err = tx.PromoteQueuedMessage(chatstate.PromoteQueuedMessageInput{ + QueuedMessageID: queued.ID, + }) + return err + })) + require.NotNil(t, result.InsertedMessage) + promoted := requireChatMessageByID(ctx, t, f, result.InsertedMessage.ID) + require.True(t, promoted.ModelConfigID.Valid) + return promoted +} + +func setChatModelConfigEnabled( + t *testing.T, + store database.Store, + config database.ChatModelConfig, + enabled bool, +) { + t.Helper() + _, err := store.UpdateChatModelConfig(context.Background(), database.UpdateChatModelConfigParams{ + ID: config.ID, + Model: config.Model, + DisplayName: config.DisplayName, + Enabled: enabled, + IsDefault: config.IsDefault, + ContextLimit: config.ContextLimit, + CompressionThreshold: config.CompressionThreshold, + Options: config.Options, + AIProviderID: config.AIProviderID, + }) + require.NoError(t, err) +} + // Input-specific rejection cases. // // These tests cover the same matrix rows as TestTransitionMatrix_AllCombinations diff --git a/coderd/x/chatd/compaction_override.go b/coderd/x/chatd/compaction_override.go index 00290ac2a4f..047aa46086b 100644 --- a/coderd/x/chatd/compaction_override.go +++ b/coderd/x/chatd/compaction_override.go @@ -64,7 +64,9 @@ func (p *Server) resolveCompactionOverrideConfig( compactionOverrideContext, raw, chat.OwnerID, - p.resolveModelConfigAndNormalizedProvider, + func(ctx context.Context, modelConfigID uuid.UUID) (database.ChatModelConfig, string, error) { + return p.resolveModelConfigForOrganization(ctx, chat.OrganizationID, modelConfigID) + }, func(ctx context.Context, ownerID uuid.UUID, aiProviderID uuid.UUID) (chatprovider.ProviderAPIKeys, error) { return p.resolveUserProviderAPIKeys(ctx, ownerID, aiProviderID) }, diff --git a/coderd/x/chatd/compaction_override_internal_test.go b/coderd/x/chatd/compaction_override_internal_test.go index 46e6cdd5fb6..30ad039f92e 100644 --- a/coderd/x/chatd/compaction_override_internal_test.go +++ b/coderd/x/chatd/compaction_override_internal_test.go @@ -17,6 +17,28 @@ import ( "github.com/coder/coder/v2/testutil" ) +func TestResolveCompactionOverrideConfig_ForeignConfigFallsBack(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + chat, _ := titleOverrideTestChatAndMessages(t) + chat.OrganizationID = uuid.New() + overrideConfig := titleOverrideModelConfig("gpt-4.1", true) + overrideConfig.OrganizationID = uuid.New() + overrideConfig.AIProviderID = uuid.NullUUID{UUID: uuid.New(), Valid: true} + + db.EXPECT().GetChatCompactionModelOverride(gomock.Any()).Return(overrideConfig.ID.String(), nil) + db.EXPECT().GetChatModelConfigByID(gomock.Any(), overrideConfig.ID).Return(overrideConfig, nil) + + server := titleOverrideTestServer(db, logger) + override, err := server.resolveCompactionOverrideConfig(ctx, chat) + require.NoError(t, err) + require.Nil(t, override) +} + func TestResolveCompactionOverrideConfig_Unset(t *testing.T) { t.Parallel() diff --git a/coderd/x/chatd/configcache.go b/coderd/x/chatd/configcache.go index 69470a9473a..ea5a0138f84 100644 --- a/coderd/x/chatd/configcache.go +++ b/coderd/x/chatd/configcache.go @@ -49,10 +49,24 @@ type modelConfigSnapshot struct { generation uint64 } -// cloneModelConfig returns a shallow copy of cfg with Options -// deep-cloned so the cache owns its own backing array. +func cloneChatACL(acl database.ChatACL) database.ChatACL { + if acl == nil { + return nil + } + clone := make(database.ChatACL, len(acl)) + for id, entry := range acl { + entry.Permissions = slices.Clone(entry.Permissions) + clone[id] = entry + } + return clone +} + +// cloneModelConfig returns a copy of cfg with mutable fields cloned so the +// cache owns their backing storage. func cloneModelConfig(cfg database.ChatModelConfig) database.ChatModelConfig { cfg.Options = slices.Clone(cfg.Options) + cfg.GroupACL = cloneChatACL(cfg.GroupACL) + cfg.UserACL = cloneChatACL(cfg.UserACL) return cfg } @@ -81,8 +95,8 @@ type chatConfigCache struct { modelConfigs map[uuid.UUID]cachedModelConfig modelConfigFetches singleflight.Group[string, database.ChatModelConfig] - // Default model config (singleton). - defaultModelConfig *cachedModelConfig + // Default model configs (keyed by organization ID). + defaultModelConfigs map[uuid.UUID]cachedModelConfig defaultModelConfigGeneration uint64 defaultModelConfigFetches singleflight.Group[string, database.ChatModelConfig] @@ -99,10 +113,11 @@ type chatConfigCache struct { func newChatConfigCache(ctx context.Context, db database.Store, clock quartz.Clock) *chatConfigCache { return &chatConfigCache{ - db: db, - clock: clock, - ctx: ctx, - modelConfigs: make(map[uuid.UUID]cachedModelConfig), + db: db, + clock: clock, + ctx: ctx, + modelConfigs: make(map[uuid.UUID]cachedModelConfig), + defaultModelConfigs: make(map[uuid.UUID]cachedModelConfig), userPrompts: tlru.New[uuid.UUID]( tlru.ConstantCost[string], chatConfigUserPromptEntryLimit, @@ -210,7 +225,7 @@ func (c *chatConfigCache) InvalidateProviders() { // provider existence, so flush all model-config state. clear(c.modelConfigs) c.modelTopologyEpoch++ - c.defaultModelConfig = nil + clear(c.defaultModelConfigs) c.defaultModelConfigGeneration++ c.mu.Unlock() } @@ -281,22 +296,23 @@ func (c *chatConfigCache) storeModelConfig(snap modelConfigSnapshot, config data } } -func (c *chatConfigCache) DefaultModelConfig(ctx context.Context) (database.ChatModelConfig, error) { - if config, ok := c.cachedDefaultModelConfig(); ok { +// DefaultModelConfig returns the usable local default model config. +func (c *chatConfigCache) DefaultModelConfig(ctx context.Context, orgID uuid.UUID) (database.ChatModelConfig, error) { + if config, ok := c.cachedDefaultModelConfig(orgID); ok { return config, nil } snap := c.defaultModelConfigSnapshot() - config, err := singleflightDoChan(ctx, &c.defaultModelConfigFetches, fmt.Sprintf("%d:default", snap.epoch), func() (database.ChatModelConfig, error) { - if cached, ok := c.cachedDefaultModelConfig(); ok { + config, err := singleflightDoChan(ctx, &c.defaultModelConfigFetches, fmt.Sprintf("%d:default:%s", snap.epoch, orgID), func() (database.ChatModelConfig, error) { + if cached, ok := c.cachedDefaultModelConfig(orgID); ok { return cached, nil } - fetched, err := c.db.GetDefaultChatModelConfig(c.ctx) + fetched, err := effectiveDefaultChatModelConfig(c.ctx, c.db, orgID) if err != nil { return database.ChatModelConfig{}, err } - c.storeDefaultModelConfig(snap, fetched) + c.storeDefaultModelConfig(snap, orgID, fetched) return cloneModelConfig(fetched), nil }) if err != nil { @@ -306,11 +322,11 @@ func (c *chatConfigCache) DefaultModelConfig(ctx context.Context) (database.Chat return config, nil } -func (c *chatConfigCache) cachedDefaultModelConfig() (database.ChatModelConfig, bool) { +func (c *chatConfigCache) cachedDefaultModelConfig(orgID uuid.UUID) (database.ChatModelConfig, bool) { c.mu.RLock() - entry := c.defaultModelConfig + entry, ok := c.defaultModelConfigs[orgID] c.mu.RUnlock() - if entry == nil { + if !ok { return database.ChatModelConfig{}, false } if c.clock.Now().Before(entry.expiresAt) { @@ -318,8 +334,8 @@ func (c *chatConfigCache) cachedDefaultModelConfig() (database.ChatModelConfig, } c.mu.Lock() - if current := c.defaultModelConfig; current != nil && !c.clock.Now().Before(current.expiresAt) { - c.defaultModelConfig = nil + if current, ok := c.defaultModelConfigs[orgID]; ok && !c.clock.Now().Before(current.expiresAt) { + delete(c.defaultModelConfigs, orgID) } c.mu.Unlock() @@ -336,7 +352,7 @@ func (c *chatConfigCache) defaultModelConfigSnapshot() modelConfigSnapshot { return snap } -func (c *chatConfigCache) storeDefaultModelConfig(snap modelConfigSnapshot, config database.ChatModelConfig) { +func (c *chatConfigCache) storeDefaultModelConfig(snap modelConfigSnapshot, orgID uuid.UUID, config database.ChatModelConfig) { c.mu.Lock() defer c.mu.Unlock() @@ -347,7 +363,7 @@ func (c *chatConfigCache) storeDefaultModelConfig(snap modelConfigSnapshot, conf return } - c.defaultModelConfig = &cachedModelConfig{ + c.defaultModelConfigs[orgID] = cachedModelConfig{ config: cloneModelConfig(config), expiresAt: c.clock.Now().Add(chatConfigModelConfigTTL), } @@ -412,7 +428,9 @@ func (c *chatConfigCache) InvalidateModelConfig(id uuid.UUID) { c.mu.Lock() delete(c.modelConfigs, id) c.modelTopologyEpoch++ - c.defaultModelConfig = nil + // Coarse invalidation: the event does not identify the changed config's + // organization or default status, so every per-org default is dropped. + clear(c.defaultModelConfigs) c.defaultModelConfigGeneration++ c.mu.Unlock() } diff --git a/coderd/x/chatd/configcache_internal_test.go b/coderd/x/chatd/configcache_internal_test.go index 305e2133115..67ef60c26b5 100644 --- a/coderd/x/chatd/configcache_internal_test.go +++ b/coderd/x/chatd/configcache_internal_test.go @@ -16,6 +16,8 @@ import ( "github.com/coder/coder/v2/coderd/database" "github.com/coder/coder/v2/coderd/database/dbtestutil" coderdpubsub "github.com/coder/coder/v2/coderd/pubsub" + "github.com/coder/coder/v2/coderd/rbac/policy" + "github.com/coder/coder/v2/coderd/util/syncmap" "github.com/coder/coder/v2/coderd/x/chatd/chatprovider" "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/testutil" @@ -25,17 +27,17 @@ import ( type stubChatConfigStore struct { database.Store - getAIProviders func(context.Context) ([]database.AIProvider, error) - getChatModelConfigByID func(context.Context, uuid.UUID) (database.ChatModelConfig, error) - getDefaultChatModelConfig func(context.Context) (database.ChatModelConfig, error) - getUserChatCustomPrompt func(context.Context, uuid.UUID) (string, error) - getChatAdvisorConfig func(context.Context) (string, error) + getAIProviders func(context.Context) ([]database.AIProvider, error) + getChatModelConfigByID func(context.Context, uuid.UUID) (database.ChatModelConfig, error) + getEnabledChatModelConfigsByOrganization func(context.Context, uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) + getUserChatCustomPrompt func(context.Context, uuid.UUID) (string, error) + getChatAdvisorConfig func(context.Context) (string, error) - enabledProvidersCalls atomic.Int32 - modelConfigByIDCalls atomic.Int32 - defaultModelConfigCall atomic.Int32 - userPromptCalls atomic.Int32 - advisorConfigCalls atomic.Int32 + enabledProvidersCalls atomic.Int32 + modelConfigByIDCalls atomic.Int32 + effectiveConfigCalls syncmap.Map[uuid.UUID, *atomic.Int32] + userPromptCalls atomic.Int32 + advisorConfigCalls atomic.Int32 } func (s *stubChatConfigStore) GetAIProviders(ctx context.Context, _ database.GetAIProvidersParams) ([]database.AIProvider, error) { @@ -54,12 +56,33 @@ func (s *stubChatConfigStore) GetChatModelConfigByID(ctx context.Context, id uui return s.getChatModelConfigByID(ctx, id) } -func (s *stubChatConfigStore) GetDefaultChatModelConfig(ctx context.Context) (database.ChatModelConfig, error) { - s.defaultModelConfigCall.Add(1) - if s.getDefaultChatModelConfig == nil { - panic("unexpected GetDefaultChatModelConfig call") +func (s *stubChatConfigStore) GetEnabledChatModelConfigsByOrganization( + ctx context.Context, + organizationID uuid.UUID, +) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + counter, _ := s.effectiveConfigCalls.LoadOrStore(organizationID, &atomic.Int32{}) + counter.Add(1) + if s.getEnabledChatModelConfigsByOrganization == nil { + panic("unexpected GetEnabledChatModelConfigsByOrganization call") } - return s.getDefaultChatModelConfig(ctx) + return s.getEnabledChatModelConfigsByOrganization(ctx, organizationID) +} + +func (s *stubChatConfigStore) effectiveConfigCallCount(organizationID uuid.UUID) int32 { + counter, ok := s.effectiveConfigCalls.Load(organizationID) + if !ok { + return 0 + } + return counter.Load() +} + +func (s *stubChatConfigStore) totalEffectiveConfigCalls() int32 { + var total int32 + s.effectiveConfigCalls.Range(func(_ uuid.UUID, counter *atomic.Int32) bool { + total += counter.Load() + return true + }) + return total } func (s *stubChatConfigStore) GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error) { @@ -169,7 +192,7 @@ func TestConfigCache_ModelConfigByID_CacheHit(t *testing.T) { require.Equal(t, int32(1), store.modelConfigByIDCalls.Load()) } -func TestConfigCache_ModelConfigByID_ClonesOptionsForCache(t *testing.T) { +func TestConfigCache_ModelConfigByID_ClonesMutableFieldsForCache(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitShort) @@ -178,6 +201,12 @@ func TestConfigCache_ModelConfigByID_ClonesOptionsForCache(t *testing.T) { const options = `{"temperature":0.1}` config := testChatModelConfig(configID, "model-a") config.Options = []byte(options) + config.GroupACL = database.ChatACL{ + "group": {Permissions: []policy.Action{policy.ActionRead, policy.ActionUpdate}}, + } + config.UserACL = database.ChatACL{ + "user": {Permissions: []policy.Action{policy.ActionRead, policy.ActionDelete}}, + } store := &stubChatConfigStore{ getChatModelConfigByID: func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { return config, nil @@ -189,17 +218,38 @@ func TestConfigCache_ModelConfigByID_ClonesOptionsForCache(t *testing.T) { first, err := cache.ModelConfigByID(ctx, configID) require.NoError(t, err) first.Options[0] = 'x' // mutate singleflight return + groupEntry := first.GroupACL["group"] + groupEntry.Permissions[0] = policy.ActionDelete + first.GroupACL["group"] = groupEntry + delete(first.GroupACL, "group") + userEntry := first.UserACL["user"] + userEntry.Permissions[0] = policy.ActionUpdate + first.UserACL["user"] = userEntry + first.UserACL["first-mutation"] = database.ChatACLEntry{} // Second call is a cache hit. second, err := cache.ModelConfigByID(ctx, configID) require.NoError(t, err) require.Equal(t, options, string(second.Options)) + require.Equal(t, config.GroupACL, second.GroupACL) + require.Equal(t, config.UserACL, second.UserACL) + require.Equal(t, []policy.Action{policy.ActionRead, policy.ActionUpdate}, second.GroupACL["group"].Permissions) + require.Equal(t, []policy.Action{policy.ActionRead, policy.ActionDelete}, second.UserACL["user"].Permissions) second.Options[0] = 'y' // mutate cache-hit return + groupEntry = second.GroupACL["group"] + groupEntry.Permissions[1] = policy.ActionDelete + second.GroupACL["group"] = groupEntry + second.GroupACL["second-mutation"] = database.ChatACLEntry{} + delete(second.UserACL, "user") - // Third call is another cache hit — must be unaffected. + // Third call is another cache hit. It must be unaffected. third, err := cache.ModelConfigByID(ctx, configID) require.NoError(t, err) require.Equal(t, options, string(third.Options)) + require.Equal(t, config.GroupACL, third.GroupACL) + require.Equal(t, config.UserACL, third.UserACL) + require.Equal(t, []policy.Action{policy.ActionRead, policy.ActionUpdate}, third.GroupACL["group"].Permissions) + require.Equal(t, []policy.Action{policy.ActionRead, policy.ActionDelete}, third.UserACL["user"].Permissions) } func TestConfigCache_ModelConfigByID_NotFound(t *testing.T) { @@ -231,30 +281,89 @@ func TestConfigCache_InvalidateModelConfig_CascadesToDefault(t *testing.T) { ctx := testutil.Context(t, testutil.WaitShort) clock := quartz.NewMock(t) configID := uuid.New() + orgID := uuid.New() config := testChatModelConfig(configID, "model-a") store := &stubChatConfigStore{} store.getChatModelConfigByID = func(context.Context, uuid.UUID) (database.ChatModelConfig, error) { return config, nil } - store.getDefaultChatModelConfig = func(context.Context) (database.ChatModelConfig, error) { - call := store.defaultModelConfigCall.Load() - return testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)), nil + store.getEnabledChatModelConfigsByOrganization = func(_ context.Context, organizationID uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + call := store.effectiveConfigCallCount(organizationID) + config := testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)) + config.IsDefault = true + return []database.GetEnabledChatModelConfigsByOrganizationRow{{ + ChatModelConfig: config, + }}, nil } cache := newChatConfigCache(ctx, store, clock) _, err := cache.ModelConfigByID(ctx, configID) require.NoError(t, err) - firstDefault, err := cache.DefaultModelConfig(ctx) + firstDefault, err := cache.DefaultModelConfig(ctx, orgID) require.NoError(t, err) cache.InvalidateModelConfig(configID) - require.Nil(t, cache.defaultModelConfig) - - secondDefault, err := cache.DefaultModelConfig(ctx) + secondDefault, err := cache.DefaultModelConfig(ctx, orgID) require.NoError(t, err) require.NotEqual(t, firstDefault, secondDefault) - require.Equal(t, int32(2), store.defaultModelConfigCall.Load()) + require.Equal(t, int32(2), store.effectiveConfigCallCount(orgID)) +} + +func TestConfigCache_DefaultModelConfig_PerOrgKeying(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + orgA := uuid.New() + orgB := uuid.New() + store := &stubChatConfigStore{} + store.getEnabledChatModelConfigsByOrganization = func(_ context.Context, organizationID uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + config := testChatModelConfig(uuid.New(), "default-"+organizationID.String()) + config.IsDefault = true + return []database.GetEnabledChatModelConfigsByOrganizationRow{{ + ChatModelConfig: config, + }}, nil + } + cache := newChatConfigCache(ctx, store, clock) + + defaultA, err := cache.DefaultModelConfig(ctx, orgA) + require.NoError(t, err) + defaultB, err := cache.DefaultModelConfig(ctx, orgB) + require.NoError(t, err) + require.NotEqual(t, defaultA.ID, defaultB.ID) + + defaultAAgain, err := cache.DefaultModelConfig(ctx, orgA) + require.NoError(t, err) + require.Equal(t, defaultA, defaultAAgain) + require.Equal(t, int32(1), store.effectiveConfigCallCount(orgA)) + require.Equal(t, int32(1), store.effectiveConfigCallCount(orgB)) + + cache.InvalidateModelConfig(uuid.New()) + _, err = cache.DefaultModelConfig(ctx, orgA) + require.NoError(t, err) + _, err = cache.DefaultModelConfig(ctx, orgB) + require.NoError(t, err) + require.Equal(t, int32(2), store.effectiveConfigCallCount(orgA)) + require.Equal(t, int32(2), store.effectiveConfigCallCount(orgB)) +} + +func TestConfigCache_DefaultModelConfig_NoEffectiveDefault(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + clock := quartz.NewMock(t) + store := &stubChatConfigStore{ + getEnabledChatModelConfigsByOrganization: func(context.Context, uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + return []database.GetEnabledChatModelConfigsByOrganizationRow{{ + ChatModelConfig: testChatModelConfig(uuid.New(), "non-default"), + }}, nil + }, + } + cache := newChatConfigCache(ctx, store, clock) + + _, err := cache.DefaultModelConfig(ctx, uuid.New()) + require.ErrorIs(t, err, sql.ErrNoRows) } func TestConfigCache_UserPrompt_NegativeCaching(t *testing.T) { @@ -585,21 +694,26 @@ func TestConfigCache_InvalidateProviders_CascadesToDefaultModelConfig(t *testing ctx := testutil.Context(t, testutil.WaitShort) clock := quartz.NewMock(t) + orgID := uuid.New() store := &stubChatConfigStore{} - store.getDefaultChatModelConfig = func(context.Context) (database.ChatModelConfig, error) { - call := store.defaultModelConfigCall.Load() - return testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)), nil + store.getEnabledChatModelConfigsByOrganization = func(_ context.Context, organizationID uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { + call := store.effectiveConfigCallCount(organizationID) + config := testChatModelConfig(uuid.New(), fmt.Sprintf("default-model-%d", call)) + config.IsDefault = true + return []database.GetEnabledChatModelConfigsByOrganizationRow{{ + ChatModelConfig: config, + }}, nil } cache := newChatConfigCache(ctx, store, clock) - first, err := cache.DefaultModelConfig(ctx) + first, err := cache.DefaultModelConfig(ctx, orgID) require.NoError(t, err) cache.InvalidateProviders() - second, err := cache.DefaultModelConfig(ctx) + second, err := cache.DefaultModelConfig(ctx, orgID) require.NoError(t, err) require.NotEqual(t, first, second) - require.Equal(t, int32(2), store.defaultModelConfigCall.Load()) + require.Equal(t, int32(2), store.effectiveConfigCallCount(orgID)) } func TestConfigCache_InvalidateProviders_BlocksStaleInFlightModelConfig(t *testing.T) { @@ -732,6 +846,7 @@ func TestConfigCache_CallerCancellation(t *testing.T) { configID := uuid.New() userID := uuid.New() + defaultOrgID := uuid.New() methods := []cacheMethod{ { @@ -798,30 +913,34 @@ func TestConfigCache_CallerCancellation(t *testing.T) { name: "DefaultModelConfig", setupBlocked: func(store *stubChatConfigStore, started, release chan struct{}) { var once sync.Once - store.getDefaultChatModelConfig = func(ctx context.Context) (database.ChatModelConfig, error) { + store.getEnabledChatModelConfigsByOrganization = func(ctx context.Context, _ uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { once.Do(func() { close(started) }) select { case <-ctx.Done(): - return database.ChatModelConfig{}, ctx.Err() + return nil, ctx.Err() case <-release: - return testChatModelConfig(uuid.New(), "default"), nil + config := testChatModelConfig(uuid.New(), "default") + config.IsDefault = true + return []database.GetEnabledChatModelConfigsByOrganizationRow{{ + ChatModelConfig: config, + }}, nil } } }, setupCtxSensitive: func(store *stubChatConfigStore, started chan struct{}) { var once sync.Once - store.getDefaultChatModelConfig = func(ctx context.Context) (database.ChatModelConfig, error) { + store.getEnabledChatModelConfigsByOrganization = func(ctx context.Context, _ uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) { once.Do(func() { close(started) }) <-ctx.Done() - return database.ChatModelConfig{}, ctx.Err() + return nil, ctx.Err() } }, call: func(ctx context.Context, cache *chatConfigCache) error { - _, err := cache.DefaultModelConfig(ctx) + _, err := cache.DefaultModelConfig(ctx, defaultOrgID) return err }, storeCalls: func(store *stubChatConfigStore) int32 { - return store.defaultModelConfigCall.Load() + return store.totalEffectiveConfigCalls() }, }, { diff --git a/coderd/x/chatd/generation_preparer_internal_test.go b/coderd/x/chatd/generation_preparer_internal_test.go index 485147e6d98..17a6f94de50 100644 --- a/coderd/x/chatd/generation_preparer_internal_test.go +++ b/coderd/x/chatd/generation_preparer_internal_test.go @@ -117,9 +117,10 @@ func TestPrepareGenerationClampsRequestedReasoningEffortToMax(t *testing.T) { }) require.NoError(t, err) modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - Options: modelConfigRaw, - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "gpt-4o-mini", + Options: modelConfigRaw, + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }, func(p *database.InsertChatModelConfigParams) { p.Enabled = true }) @@ -209,9 +210,10 @@ func TestPrepareGenerationComputerUseIgnoresChatTransportOverride(t *testing.T) }) require.NoError(t, err) modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - Options: modelConfigRaw, - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "gpt-4o-mini", + Options: modelConfigRaw, + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }, func(p *database.InsertChatModelConfigParams) { p.Enabled = true }) @@ -304,8 +306,9 @@ func TestPrepareGenerationSubagentUsesOwnerSyntheticAPIKey(t *testing.T) { Type: database.AIProviderTypeOpenai, }, "test-key") modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }, func(p *database.InsertChatModelConfigParams) { p.Enabled = true }) @@ -385,9 +388,10 @@ func TestDeriveFinalTurnRunResult(t *testing.T) { CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - DisplayName: "gpt-4o-mini", - Options: json.RawMessage(`{"openai_config":{"use_responses_api":false}}`), + Model: "gpt-4o-mini", + DisplayName: "gpt-4o-mini", + Options: json.RawMessage(`{"openai_config":{"use_responses_api":false}}`), + OrganizationID: org.ID, }, func(p *database.InsertChatModelConfigParams) { p.Enabled = true p.IsDefault = true @@ -504,9 +508,10 @@ func TestDeriveFinalTurnRunResult(t *testing.T) { }) provider := insertInternalAIProvider(t, db, database.AIProviderTypeOpenai, "provider-api-key", false) modelCfg := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - DisplayName: "gpt-4o-mini", - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "gpt-4o-mini", + DisplayName: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }) created, err := chatstate.CreateChat(ctx, db, ps, chatstate.CreateChatInput{ diff --git a/coderd/x/chatd/helpers_test.go b/coderd/x/chatd/helpers_test.go index 6747ff66284..3c3e6cbf3f3 100644 --- a/coderd/x/chatd/helpers_test.go +++ b/coderd/x/chatd/helpers_test.go @@ -125,7 +125,8 @@ func newWorkerTestFixture(t *testing.T) *workerTestFixture { BaseUrl: "http://example.invalid", }) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - IsDefault: true, + IsDefault: true, + OrganizationID: org.ID, }) apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) return &workerTestFixture{db: db, pubsub: ps, sqlDB: sqlDB, user: user, org: org, model: model, apiKey: apiKey} diff --git a/coderd/x/chatd/quickgen.go b/coderd/x/chatd/quickgen.go index 9a0ab481e99..5aa47a9015f 100644 --- a/coderd/x/chatd/quickgen.go +++ b/coderd/x/chatd/quickgen.go @@ -145,7 +145,7 @@ type shortTextCandidate struct { } func selectPreferredConfiguredShortTextModelConfig( - configs []database.GetEnabledChatModelConfigsRow, + configs []database.GetEnabledChatModelConfigsByOrganizationRow, ) (database.ChatModelConfig, bool) { for _, preferred := range preferredTitleModels { for _, config := range configs { diff --git a/coderd/x/chatd/quickgen_internal_test.go b/coderd/x/chatd/quickgen_internal_test.go index bd56f3d95ff..bf34f5f78c6 100644 --- a/coderd/x/chatd/quickgen_internal_test.go +++ b/coderd/x/chatd/quickgen_internal_test.go @@ -800,7 +800,7 @@ func Test_selectPreferredConfiguredShortTextModelConfig(t *testing.T) { t.Run("chooses the highest-priority configured lightweight model", func(t *testing.T) { t.Parallel() - configs := []database.GetEnabledChatModelConfigsRow{ + configs := []database.GetEnabledChatModelConfigsByOrganizationRow{ {ChatModelConfig: database.ChatModelConfig{Model: preferredTitleModels[2].model}, Provider: preferredTitleModels[2].provider}, {ChatModelConfig: database.ChatModelConfig{Model: preferredTitleModels[1].model}, Provider: preferredTitleModels[1].provider}, {ChatModelConfig: database.ChatModelConfig{Model: "gpt-4.1"}, Provider: "openai"}, @@ -814,7 +814,7 @@ func Test_selectPreferredConfiguredShortTextModelConfig(t *testing.T) { t.Run("returns false when no preferred lightweight model is configured", func(t *testing.T) { t.Parallel() - got, ok := selectPreferredConfiguredShortTextModelConfig([]database.GetEnabledChatModelConfigsRow{{ + got, ok := selectPreferredConfiguredShortTextModelConfig([]database.GetEnabledChatModelConfigsByOrganizationRow{{ ChatModelConfig: database.ChatModelConfig{Model: "gpt-4.1"}, Provider: "openai", }}) diff --git a/coderd/x/chatd/subagent.go b/coderd/x/chatd/subagent.go index 5c6c8e9e95c..8429fabe881 100644 --- a/coderd/x/chatd/subagent.go +++ b/coderd/x/chatd/subagent.go @@ -56,7 +56,10 @@ type subagentStatusError struct { func (e *subagentStatusError) Error() string { return e.reason } -var errInvalidModelOverrideMetadata = xerrors.New("invalid model override metadata") +var ( + errInvalidModelOverrideMetadata = xerrors.New("invalid model override metadata") + errModelConfigOutsideOrganization = xerrors.Errorf("%w: model config belongs to another organization", sql.ErrNoRows) +) type modelOverrideConfigResolver func( context.Context, @@ -218,11 +221,12 @@ func (p *Server) resolveConfiguredModelOverride( if failureMode == modelOverrideFailureModeHard { label := modelOverrideErrorLabel(overrideContext) switch { - case errors.Is(err, sql.ErrNoRows): + case errors.Is(err, sql.ErrNoRows), errors.Is(err, errModelConfigOutsideOrganization): return database.ChatModelConfig{}, "", parsed.reasoningEffort, true, xerrors.Errorf( - "%s model override is unavailable: %s", + "%s model override is unavailable: %s: %w", label, parsed.modelConfigID, + err, ) case errors.Is(err, errInvalidModelOverrideMetadata): return database.ChatModelConfig{}, "", parsed.reasoningEffort, true, xerrors.Errorf( @@ -242,6 +246,12 @@ func (p *Server) resolveConfiguredModelOverride( } switch { + case errors.Is(err, errModelConfigOutsideOrganization): + p.logger.Info(ctx, + "model override belongs to another organization, ignoring", + slog.F("override_context", overrideContext), + slog.F("model_config_id", parsed.modelConfigID), + ) case errors.Is(err, sql.ErrNoRows): p.logger.Info(ctx, "model override is unavailable, ignoring", @@ -296,6 +306,7 @@ func (p *Server) resolveConfiguredModelOverride( func (p *Server) resolvePersonalSubagentModelConfigID( ctx context.Context, ownerID uuid.UUID, + organizationID uuid.UUID, overrideContext codersdk.ChatModelOverrideContext, ) (uuid.UUID, *string, bool, error) { personalContext, err := personalModelOverrideContextForSubagent(overrideContext) @@ -326,7 +337,7 @@ func (p *Server) resolvePersonalSubagentModelConfigID( ) if parsed.Malformed { p.logger.Debug(ctx, - "personal model override is malformed, using deployment default", + "personal model override is malformed, using chat organization default", slog.F("override_context", overrideContext), slog.F("owner_id", ownerID), slog.F("raw_model_config_id", strings.TrimSpace(raw)), @@ -341,6 +352,7 @@ func (p *Server) resolvePersonalSubagentModelConfigID( ctx, overrideContext, ownerID, + organizationID, parsed.ModelConfigID, ) if err != nil { @@ -351,7 +363,7 @@ func (p *Server) resolvePersonalSubagentModelConfigID( } default: p.logger.Warn(ctx, - "unsupported personal model override mode, using deployment default", + "unsupported personal model override mode, using chat organization default", slog.F("override_context", overrideContext), slog.F("owner_id", ownerID), slog.F("mode", parsed.Mode), @@ -365,24 +377,33 @@ func (p *Server) resolvePersonalModelOverride( ctx context.Context, overrideContext codersdk.ChatModelOverrideContext, ownerID uuid.UUID, + organizationID uuid.UUID, modelConfigID uuid.UUID, ) (database.ChatModelConfig, bool, error) { - modelConfig, providerName, err := p.resolveModelConfigAndNormalizedProvider( + modelConfig, providerName, err := p.resolveModelConfigForOrganization( ctx, + organizationID, modelConfigID, ) if err != nil { switch { + case errors.Is(err, errModelConfigOutsideOrganization): + p.logger.Debug(ctx, + "personal model override belongs to another organization, using chat organization default", + slog.F("override_context", overrideContext), + slog.F("owner_id", ownerID), + slog.F("model_config_id", modelConfigID), + ) case xerrors.Is(err, sql.ErrNoRows): p.logger.Debug(ctx, - "personal model override is unavailable, using deployment default", + "personal model override is unavailable, using chat organization default", slog.F("override_context", overrideContext), slog.F("owner_id", ownerID), slog.F("model_config_id", modelConfigID), ) case errors.Is(err, errInvalidModelOverrideMetadata): p.logger.Debug(ctx, - "personal model override metadata is invalid, using deployment default", + "personal model override metadata is invalid, using chat organization default", slog.F("override_context", overrideContext), slog.F("owner_id", ownerID), slog.F("model_config_id", modelConfigID), @@ -390,7 +411,7 @@ func (p *Server) resolvePersonalModelOverride( ) default: p.logger.Warn(ctx, - "failed to resolve personal model override, using deployment default", + "failed to resolve personal model override, using chat organization default", slog.F("override_context", overrideContext), slog.F("owner_id", ownerID), slog.F("model_config_id", modelConfigID), @@ -408,7 +429,7 @@ func (p *Server) resolvePersonalModelOverride( } if !userCanUseProviderKeys(providerKeys, providerName) { p.logger.Debug(ctx, - "personal model override credentials are unavailable, using deployment default", + "personal model override credentials are unavailable, using chat organization default", slog.F("override_context", overrideContext), slog.F("owner_id", ownerID), slog.F("model_config_id", modelConfigID), @@ -422,6 +443,7 @@ func (p *Server) resolvePersonalModelOverride( func (p *Server) resolveSubagentModelConfigID( ctx context.Context, ownerID uuid.UUID, + organizationID uuid.UUID, overrideContext codersdk.ChatModelOverrideContext, ) (uuid.UUID, *string, error) { //nolint:gocritic // Chatd needs its scoped config and user-data access here. @@ -437,6 +459,7 @@ func (p *Server) resolveSubagentModelConfigID( modelConfigID, reasoningEffort, resolved, err := p.resolvePersonalSubagentModelConfigID( chatdCtx, ownerID, + organizationID, overrideContext, ) if err != nil { @@ -460,7 +483,9 @@ func (p *Server) resolveSubagentModelConfigID( string(overrideContext), raw, ownerID, - p.resolveModelConfigAndNormalizedProvider, + func(ctx context.Context, modelConfigID uuid.UUID) (database.ChatModelConfig, string, error) { + return p.resolveModelConfigForOrganization(ctx, organizationID, modelConfigID) + }, p.resolveUserProviderAPIKeys, modelOverrideFailureModeSoft, ) @@ -491,6 +516,13 @@ func (p *Server) resolveModelConfigAndNormalizedProvider( if err != nil { return database.ChatModelConfig{}, "", err } + return p.resolveNormalizedProviderForModelConfig(ctx, modelConfig) +} + +func (p *Server) resolveNormalizedProviderForModelConfig( + ctx context.Context, + modelConfig database.ChatModelConfig, +) (database.ChatModelConfig, string, error) { if !modelConfig.Enabled { return database.ChatModelConfig{}, "", sql.ErrNoRows } @@ -518,6 +550,7 @@ func (p *Server) resolveModelConfigAndNormalizedProvider( func (p *Server) resolveExplicitSpawnOverrides( ctx context.Context, ownerID uuid.UUID, + organizationID uuid.UUID, args spawnAgentArgs, ) (*uuid.UUID, *string, error) { var explicitModelConfigID *uuid.UUID @@ -531,8 +564,9 @@ func (p *Server) resolveExplicitSpawnOverrides( } //nolint:gocritic // Chatd needs its scoped config and user-data access here. chatdCtx := dbauthz.AsChatd(ctx) - modelConfig, providerName, err := p.resolveModelConfigAndNormalizedProvider( + modelConfig, providerName, err := p.resolveModelConfigForOrganization( chatdCtx, + organizationID, modelConfigID, ) if err != nil { @@ -594,10 +628,11 @@ func (p *Server) resolveExplicitSpawnOverrides( func (p *Server) listSpawnableModelConfigs( ctx context.Context, ownerID uuid.UUID, + organizationID uuid.UUID, ) ([]map[string]any, error) { //nolint:gocritic // Chatd needs its scoped config and user-data access here. chatdCtx := dbauthz.AsChatd(ctx) - rows, err := p.db.GetEnabledChatModelConfigs(chatdCtx) + rows, err := enabledChatModelConfigsForOrganization(chatdCtx, p.db, organizationID) if err != nil { return nil, xerrors.Errorf("get enabled chat model configs: %w", err) } @@ -703,6 +738,7 @@ func (p *Server) subagentTools( explicitModelConfigID, explicitReasoningEffort, err := p.resolveExplicitSpawnOverrides( ctx, parent.OwnerID, + parent.OrganizationID, args, ) if err != nil { @@ -775,7 +811,7 @@ func (p *Server) subagentTools( return fantasy.NewTextErrorResponse(err.Error()), nil } - models, err := p.listSpawnableModelConfigs(ctx, parent.OwnerID) + models, err := p.listSpawnableModelConfigs(ctx, parent.OwnerID, parent.OrganizationID) if err != nil { p.logger.Warn(ctx, "failed to list spawnable model configs", slog.F("chat_id", parent.ID), diff --git a/coderd/x/chatd/subagent_catalog.go b/coderd/x/chatd/subagent_catalog.go index a79bf81fa5d..843f07a3fa4 100644 --- a/coderd/x/chatd/subagent_catalog.go +++ b/coderd/x/chatd/subagent_catalog.go @@ -65,6 +65,7 @@ func allSubagentDefinitions() []subagentDefinition { modelConfigID, reasoningEffort, err := p.resolveSubagentModelConfigID( ctx, parent.OwnerID, + parent.OrganizationID, codersdk.ChatModelOverrideContextGeneral, ) if err != nil { @@ -90,6 +91,7 @@ func allSubagentDefinitions() []subagentDefinition { resolvedModelConfigID, resolvedReasoningEffort, err := p.resolveSubagentModelConfigID( ctx, turnParent.OwnerID, + turnParent.OrganizationID, codersdk.ChatModelOverrideContextExplore, ) if err != nil { diff --git a/coderd/x/chatd/subagent_internal_test.go b/coderd/x/chatd/subagent_internal_test.go index 7879bd74d24..4bd7a34a038 100644 --- a/coderd/x/chatd/subagent_internal_test.go +++ b/coderd/x/chatd/subagent_internal_test.go @@ -230,8 +230,9 @@ func seedInternalChatDeps( }) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, - IsDefault: true, + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + IsDefault: true, + OrganizationID: org.ID, }) return user, org, model @@ -519,15 +520,16 @@ func TestResolveUserProviderAPIKeys_AIProvider(t *testing.T) { }) } -func TestResolveChatModel_AIProviderDisabled(t *testing.T) { +func TestResolveModelConfig_DisabledProviderFallsBackToLocalDefault(t *testing.T) { t.Parallel() ctx := chatdTestContext(t) db, ps := dbtestutil.NewDB(t) - user, org, _ := seedInternalChatDeps(t, db) + user, org, localDefault := seedInternalChatDeps(t, db) provider := insertInternalAIProvider(t, db, database.AIProviderTypeOpenai, "provider-api-key", false) modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", + Model: "gpt-4o-mini", + OrganizationID: org.ID, AIProviderID: uuid.NullUUID{ UUID: provider.ID, Valid: true, @@ -540,12 +542,9 @@ func TestResolveChatModel_AIProviderDisabled(t *testing.T) { LastModelConfigID: modelConfig.ID, }) - resolved, err := server.resolveModelCall(ctx, modelCallSpec{ - purpose: "standard_turn", - chat: chat, - }) - require.ErrorContains(t, err, "is disabled") - require.Equal(t, resolvedModelCall{}, resolved) + config, err := server.resolveModelConfig(ctx, chat) + require.NoError(t, err) + require.Equal(t, localDefault.ID, config.ID) } func TestResolveUserProviderAPIKeys_PreservesAnthropicKeyFromDBProvider(t *testing.T) { @@ -591,12 +590,14 @@ func TestResolveUserProviderAPIKeys_PreservesAnthropicKeyFromDBProvider(t *testi func insertInternalChatModelConfig( t *testing.T, db database.Store, + orgID uuid.UUID, model string, enabled bool, ) database.ChatModelConfig { return insertInternalChatModelConfigForProvider( t, db, + orgID, "openai", model, enabled, @@ -633,6 +634,7 @@ func insertInternalChatProvider( func insertInternalChatModelConfigForProvider( t *testing.T, db database.Store, + orgID uuid.UUID, provider string, model string, enabled bool, @@ -641,6 +643,7 @@ func insertInternalChatModelConfigForProvider( return insertInternalChatModelConfigWithOptions( t, db, + orgID, provider, model, enabled, @@ -651,6 +654,7 @@ func insertInternalChatModelConfigForProvider( func insertInternalChatModelConfigWithOptions( t *testing.T, db database.Store, + orgID uuid.UUID, provider string, model string, enabled bool, @@ -678,10 +682,11 @@ func insertInternalChatModelConfigWithOptions( }) } modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - AIProviderID: uuid.NullUUID{UUID: aiProvider.ID, Valid: true}, - Model: model, - DisplayName: model, - Options: options, + AIProviderID: uuid.NullUUID{UUID: aiProvider.ID, Valid: true}, + Model: model, + DisplayName: model, + Options: options, + OrganizationID: orgID, }, func(p *database.InsertChatModelConfigParams) { p.Enabled = enabled }) @@ -1045,7 +1050,7 @@ func TestSpawnAgent_GeneralUsesConfiguredModelOverride(t *testing.T) { ctx := chatdTestContext(t) user, org, model := seedInternalChatDeps(t, db) overrideModel := insertInternalChatModelConfig( - t, db, "general-override-"+uuid.NewString(), true, + t, db, org.ID, "general-override-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatGeneralModelOverride(ctx, overrideModel.ID.String())) parentChat := createInternalParentChat( @@ -1071,7 +1076,7 @@ func TestSpawnAgent_GeneralHonorsPersonalModelOverrides(t *testing.T) { name string enablePersonalOverride bool personalRaw func(database.ChatModelConfig) string - personalModel func(context.Context, *testing.T, database.Store, uuid.UUID) database.ChatModelConfig + personalModel func(context.Context, *testing.T, database.Store, uuid.UUID, uuid.UUID) database.ChatModelConfig wantModelID func( database.ChatModelConfig, database.ChatModelConfig, @@ -1132,11 +1137,13 @@ func TestSpawnAgent_GeneralHonorsPersonalModelOverrides(t *testing.T) { ctx context.Context, t *testing.T, db database.Store, + orgID uuid.UUID, userID uuid.UUID, ) database.ChatModelConfig { return insertInternalChatModelConfig( t, db, + orgID, "general-personal-disabled-"+uuid.NewString(), false, ) @@ -1156,6 +1163,7 @@ func TestSpawnAgent_GeneralHonorsPersonalModelOverrides(t *testing.T) { ctx context.Context, t *testing.T, db database.Store, + orgID uuid.UUID, userID uuid.UUID, ) database.ChatModelConfig { insertInternalChatProvider( @@ -1171,6 +1179,7 @@ func TestSpawnAgent_GeneralHonorsPersonalModelOverrides(t *testing.T) { return insertInternalChatModelConfigForProvider( t, db, + orgID, "openai-compat", "gpt-4o-mini", true, @@ -1208,6 +1217,7 @@ func TestSpawnAgent_GeneralHonorsPersonalModelOverrides(t *testing.T) { deploymentModel := insertInternalChatModelConfig( t, db, + org.ID, "general-deployment-"+uuid.NewString(), true, ) @@ -1215,11 +1225,12 @@ func TestSpawnAgent_GeneralHonorsPersonalModelOverrides(t *testing.T) { personalModel := insertInternalChatModelConfig( t, db, + org.ID, "general-personal-"+uuid.NewString(), true, ) if tt.personalModel != nil { - personalModel = tt.personalModel(ctx, t, db, user.ID) + personalModel = tt.personalModel(ctx, t, db, org.ID, user.ID) } if tt.enablePersonalOverride { enableInternalChatPersonalModelOverrides(t, db) @@ -1286,6 +1297,7 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenCredentialsUnavailable(t overrideModel := insertInternalChatModelConfigForProvider( t, db, + org.ID, "openai-compat", "gpt-4o-mini", true, @@ -1355,6 +1367,7 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenProviderDisabled(t *testi overrideModel := insertInternalChatModelConfigForProvider( t, db, + org.ID, "openai-compat", "gpt-4o-mini", true, @@ -1389,6 +1402,64 @@ func TestSpawnAgent_GeneralOverrideLogsAndFallsBackWhenProviderDisabled(t *testi ), 1) } +func TestResolveConfiguredModelOverride_OrganizationMismatch(t *testing.T) { + t.Parallel() + + modelConfigID := uuid.New() + ownerID := uuid.New() + resolveModelConfig := func(context.Context, uuid.UUID) (database.ChatModelConfig, string, error) { + return database.ChatModelConfig{}, "", errModelConfigOutsideOrganization + } + resolveProviderKeys := func(context.Context, uuid.UUID, uuid.UUID) (chatprovider.ProviderAPIKeys, error) { + return chatprovider.ProviderAPIKeys{}, nil + } + + t.Run("Soft", func(t *testing.T) { + t.Parallel() + + logSink := &subagentTestLogSink{} + server := &Server{logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}).AppendSinks(logSink)} + modelConfig, providerName, reasoningEffort, ok, err := server.resolveConfiguredModelOverride( + chatdTestContext(t), + "general", + modelConfigID.String(), + ownerID, + resolveModelConfig, + resolveProviderKeys, + modelOverrideFailureModeSoft, + ) + require.NoError(t, err) + require.Empty(t, modelConfig) + require.Empty(t, providerName) + require.Nil(t, reasoningEffort) + require.False(t, ok) + require.Len(t, logSink.entriesAtLevelWithMessage( + slog.LevelInfo, + "model override belongs to another organization, ignoring", + ), 1) + }) + + t.Run("Hard", func(t *testing.T) { + t.Parallel() + + server := &Server{logger: slogtest.Make(t, &slogtest.Options{IgnoreErrors: true})} + modelConfig, providerName, reasoningEffort, ok, err := server.resolveConfiguredModelOverride( + chatdTestContext(t), + "title", + modelConfigID.String(), + ownerID, + resolveModelConfig, + resolveProviderKeys, + modelOverrideFailureModeHard, + ) + require.ErrorContains(t, err, "title model override is unavailable") + require.Empty(t, modelConfig) + require.Empty(t, providerName) + require.Nil(t, reasoningEffort) + require.True(t, ok) + }) +} + func TestResolveConfiguredModelOverride_AcceptsAmbientCredentialsProvider( t *testing.T, ) { @@ -1474,7 +1545,7 @@ func TestCreateChildSubagentChat_OverrideWorksWhenParentHasNoModel(t *testing.T) ctx := chatdTestContext(t) user, org, model := seedInternalChatDeps(t, db) overrideModel := insertInternalChatModelConfig( - t, db, "override-no-parent-model-"+uuid.NewString(), true, + t, db, org.ID, "override-no-parent-model-"+uuid.NewString(), true, ) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, model.ID, "parent-no-model", @@ -1510,7 +1581,7 @@ func TestSpawnAgent_ExplicitModelConfigID(t *testing.T) { ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) selectedModel := insertInternalChatModelConfig( - t, db, "spawn-explicit-"+uuid.NewString(), true, + t, db, org.ID, "spawn-explicit-"+uuid.NewString(), true, ) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explicit-model", @@ -1564,7 +1635,7 @@ func TestSpawnAgent_ExplicitModelSkipsConfiguredOverrides(t *testing.T) { user, org, parentModel := seedInternalChatDeps(t, db) enableInternalChatPersonalModelOverrides(t, db) selectedModel := insertInternalChatModelConfig( - t, db, "selected-skip-overrides-"+uuid.NewString(), true, + t, db, org.ID, "selected-skip-overrides-"+uuid.NewString(), true, ) parentChat := createInternalParentChat( ctx, t, server, store, org.ID, user.ID, parentModel.ID, "parent-skip-overrides", @@ -1591,13 +1662,13 @@ func TestSpawnAgent_ExplicitModelBeatsConfiguredOverrides(t *testing.T) { ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) deploymentModel := insertInternalChatModelConfig( - t, db, "general-deployment-"+uuid.NewString(), true, + t, db, org.ID, "general-deployment-"+uuid.NewString(), true, ) personalModel := insertInternalChatModelConfig( - t, db, "general-personal-"+uuid.NewString(), true, + t, db, org.ID, "general-personal-"+uuid.NewString(), true, ) selectedModel := insertInternalChatModelConfig( - t, db, "general-selected-"+uuid.NewString(), true, + t, db, org.ID, "general-selected-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatGeneralModelOverride( ctx, deploymentModel.ID.String()+":medium", @@ -1677,7 +1748,7 @@ func TestSpawnAgent_ExplicitReasoningEffort(t *testing.T) { ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) selectedModel := insertInternalChatModelConfig( - t, db, "combined-"+uuid.NewString(), true, + t, db, org.ID, "combined-"+uuid.NewString(), true, ) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-combined", @@ -1708,7 +1779,7 @@ func TestSpawnAgent_ExplicitReasoningEffort(t *testing.T) { ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) deploymentModel := insertInternalChatModelConfig( - t, db, "general-deployment-"+uuid.NewString(), true, + t, db, org.ID, "general-deployment-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatGeneralModelOverride( ctx, deploymentModel.ID.String()+":low", @@ -1738,28 +1809,28 @@ func TestSpawnAgent_ModelArgsRejections(t *testing.T) { tests := []struct { name string - args func(ctx context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs + args func(ctx context.Context, t *testing.T, db database.Store, orgID uuid.UUID, user database.User) spawnAgentArgs wantError string }{ { name: "MalformedModelConfigID", - args: func(context.Context, *testing.T, database.Store, database.User) spawnAgentArgs { + args: func(context.Context, *testing.T, database.Store, uuid.UUID, database.User) spawnAgentArgs { return spawnAgentArgs{ModelConfigID: "not-a-uuid"} }, wantError: "invalid model_config_id: must be a valid UUID", }, { name: "UnknownModelConfigID", - args: func(context.Context, *testing.T, database.Store, database.User) spawnAgentArgs { + args: func(context.Context, *testing.T, database.Store, uuid.UUID, database.User) spawnAgentArgs { return spawnAgentArgs{ModelConfigID: uuid.NewString()} }, wantError: "model_config_id not found or is disabled", }, { name: "DisabledModelConfig", - args: func(_ context.Context, t *testing.T, db database.Store, _ database.User) spawnAgentArgs { + args: func(_ context.Context, t *testing.T, db database.Store, orgID uuid.UUID, _ database.User) spawnAgentArgs { disabled := insertInternalChatModelConfig( - t, db, "reject-disabled-"+uuid.NewString(), false, + t, db, orgID, "reject-disabled-"+uuid.NewString(), false, ) return spawnAgentArgs{ModelConfigID: disabled.ID.String()} }, @@ -1767,7 +1838,7 @@ func TestSpawnAgent_ModelArgsRejections(t *testing.T) { }, { name: "DisabledProvider", - args: func(_ context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs { + args: func(_ context.Context, t *testing.T, db database.Store, orgID uuid.UUID, user database.User) spawnAgentArgs { dbgen.ChatProvider(t, db, database.ChatProvider{ Provider: "openai-compat", CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, @@ -1775,7 +1846,7 @@ func TestSpawnAgent_ModelArgsRejections(t *testing.T) { p.Enabled = false }) config := insertInternalChatModelConfigForProvider( - t, db, "openai-compat", "gpt-4o-mini", true, + t, db, orgID, "openai-compat", "gpt-4o-mini", true, ) return spawnAgentArgs{ModelConfigID: config.ID.String()} }, @@ -1783,13 +1854,13 @@ func TestSpawnAgent_ModelArgsRejections(t *testing.T) { }, { name: "DeletedProvider", - args: func(_ context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs { + args: func(_ context.Context, t *testing.T, db database.Store, orgID uuid.UUID, user database.User) spawnAgentArgs { provider := dbgen.ChatProvider(t, db, database.ChatProvider{ Provider: "openai-compat", CreatedBy: uuid.NullUUID{UUID: user.ID, Valid: true}, }) config := insertInternalChatModelConfigForProvider( - t, db, "openai-compat", "gpt-4o-mini", true, + t, db, orgID, "openai-compat", "gpt-4o-mini", true, ) require.NoError(t, db.DeleteAIProviderByID( systemRestrictedTestContext(t), provider.ID, @@ -1800,12 +1871,12 @@ func TestSpawnAgent_ModelArgsRejections(t *testing.T) { }, { name: "MissingCredentials", - args: func(_ context.Context, t *testing.T, db database.Store, user database.User) spawnAgentArgs { + args: func(_ context.Context, t *testing.T, db database.Store, orgID uuid.UUID, user database.User) spawnAgentArgs { insertInternalChatProvider( t, db, user.ID, "openai-compat", "", false, true, false, ) config := insertInternalChatModelConfigForProvider( - t, db, "openai-compat", "gpt-4o-mini", true, + t, db, orgID, "openai-compat", "gpt-4o-mini", true, ) return spawnAgentArgs{ModelConfigID: config.ID.String()} }, @@ -1813,7 +1884,7 @@ func TestSpawnAgent_ModelArgsRejections(t *testing.T) { }, { name: "InvalidReasoningEffort", - args: func(context.Context, *testing.T, database.Store, database.User) spawnAgentArgs { + args: func(context.Context, *testing.T, database.Store, uuid.UUID, database.User) spawnAgentArgs { return spawnAgentArgs{ReasoningEffort: "ultra"} }, wantError: "invalid reasoning_effort: must be one of", @@ -1833,7 +1904,7 @@ func TestSpawnAgent_ModelArgsRejections(t *testing.T) { ctx, t, server, db, org.ID, user.ID, model.ID, "parent-rejection", ) - args := tt.args(ctx, t, db, user) + args := tt.args(ctx, t, db, org.ID, user) args.Type = subagentTypeGeneral args.Prompt = "delegate work" @@ -1917,17 +1988,18 @@ func TestListSubagentModels(t *testing.T) { effortModel := insertInternalChatModelConfigWithOptions( t, db, + org.ID, "openai", "gpt-effort-"+uuid.NewString(), true, json.RawMessage(`{"reasoning_effort":{"default":"medium","max":"high"}}`), ) disabledModel := insertInternalChatModelConfig( - t, db, "list-disabled-"+uuid.NewString(), false, + t, db, org.ID, "list-disabled-"+uuid.NewString(), false, ) insertInternalChatProvider(t, db, user.ID, "anthropic", "", false, true, false) noCredsModel := insertInternalChatModelConfigForProvider( - t, db, "anthropic", "claude-sonnet-4-5", true, + t, db, org.ID, "anthropic", "claude-sonnet-4-5", true, ) dbgen.ChatProvider(t, db, database.ChatProvider{ Provider: "openai-compat", @@ -1936,7 +2008,7 @@ func TestListSubagentModels(t *testing.T) { p.Enabled = false }) disabledProviderModel := insertInternalChatModelConfigForProvider( - t, db, "openai-compat", "gpt-4o-mini", true, + t, db, org.ID, "openai-compat", "gpt-4o-mini", true, ) parentChat := createInternalParentChat( @@ -2030,7 +2102,7 @@ func TestSpawnAgent_ExploreUsesConfiguredModelOverride(t *testing.T) { ctx := chatdTestContext(t) user, org, model := seedInternalChatDeps(t, db) overrideModel := insertInternalChatModelConfig( - t, db, "explore-override-"+uuid.NewString(), true, + t, db, org.ID, "explore-override-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatExploreModelOverride(ctx, overrideModel.ID.String())) parentChat := createInternalParentChat( @@ -2068,7 +2140,7 @@ func TestSpawnAgent_ExploreFallsBackToCurrentTurnModel(t *testing.T) { ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - t, db, "explore-current-turn-"+uuid.NewString(), true, + t, db, org.ID, "explore-current-turn-"+uuid.NewString(), true, ) parentChat := createInternalParentChat( ctx, t, server, db, org.ID, user.ID, parentModel.ID, "parent-explore-fallback", @@ -2098,7 +2170,7 @@ func TestSpawnAgent_ExploreHonorsPersonalModelOverrides(t *testing.T) { name string enablePersonalOverride bool personalRaw func(database.ChatModelConfig) string - personalModel func(context.Context, *testing.T, database.Store, uuid.UUID) database.ChatModelConfig + personalModel func(context.Context, *testing.T, database.Store, uuid.UUID, uuid.UUID) database.ChatModelConfig wantModelID func( database.ChatModelConfig, database.ChatModelConfig, @@ -2160,11 +2232,13 @@ func TestSpawnAgent_ExploreHonorsPersonalModelOverrides(t *testing.T) { ctx context.Context, t *testing.T, db database.Store, + orgID uuid.UUID, userID uuid.UUID, ) database.ChatModelConfig { return insertInternalChatModelConfig( t, db, + orgID, "explore-personal-disabled-"+uuid.NewString(), false, ) @@ -2184,6 +2258,7 @@ func TestSpawnAgent_ExploreHonorsPersonalModelOverrides(t *testing.T) { ctx context.Context, t *testing.T, db database.Store, + orgID uuid.UUID, userID uuid.UUID, ) database.ChatModelConfig { insertInternalChatProvider( @@ -2199,6 +2274,7 @@ func TestSpawnAgent_ExploreHonorsPersonalModelOverrides(t *testing.T) { return insertInternalChatModelConfigForProvider( t, db, + orgID, "openai-compat", "gpt-4o-mini", true, @@ -2236,12 +2312,14 @@ func TestSpawnAgent_ExploreHonorsPersonalModelOverrides(t *testing.T) { currentTurnModel := insertInternalChatModelConfig( t, db, + org.ID, "explore-current-turn-"+uuid.NewString(), true, ) deploymentModel := insertInternalChatModelConfig( t, db, + org.ID, "explore-deployment-"+uuid.NewString(), true, ) @@ -2249,11 +2327,12 @@ func TestSpawnAgent_ExploreHonorsPersonalModelOverrides(t *testing.T) { personalModel := insertInternalChatModelConfig( t, db, + org.ID, "explore-personal-"+uuid.NewString(), true, ) if tt.personalModel != nil { - personalModel = tt.personalModel(ctx, t, db, user.ID) + personalModel = tt.personalModel(ctx, t, db, org.ID, user.ID) } if tt.enablePersonalOverride { enableInternalChatPersonalModelOverrides(t, db) @@ -2545,7 +2624,7 @@ func TestSpawnAgent_ExploreFallsBackOnInvalidUUID(t *testing.T) { ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - t, db, "explore-invalid-override-"+uuid.NewString(), true, + t, db, org.ID, "explore-invalid-override-"+uuid.NewString(), true, ) require.NoError(t, db.UpsertChatExploreModelOverride(ctx, "not-a-uuid")) parentChat := createInternalParentChat( @@ -2577,10 +2656,10 @@ func TestSpawnAgent_ExploreFallsBackWhenOverrideIsUnavailable(t *testing.T) { ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - t, db, "explore-fallback-current-"+uuid.NewString(), true, + t, db, org.ID, "explore-fallback-current-"+uuid.NewString(), true, ) disabledModel := insertInternalChatModelConfig( - t, db, "explore-disabled-"+uuid.NewString(), false, + t, db, org.ID, "explore-disabled-"+uuid.NewString(), false, ) require.NoError(t, db.UpsertChatExploreModelOverride(ctx, disabledModel.ID.String())) parentChat := createInternalParentChat( @@ -2612,7 +2691,7 @@ func TestSpawnAgent_ExploreFallsBackWhenOverrideCredentialsAreUnavailable(t *tes ctx := chatdTestContext(t) user, org, parentModel := seedInternalChatDeps(t, db) currentTurnModel := insertInternalChatModelConfig( - t, db, "explore-missing-user-key-current-"+uuid.NewString(), true, + t, db, org.ID, "explore-missing-user-key-current-"+uuid.NewString(), true, ) overrideProvider := dbgen.ChatProvider(t, db, database.ChatProvider{ Provider: "openai-compat", @@ -2896,6 +2975,7 @@ func TestSpawnAgent_ComputerUseRejectsMissingConfiguredProvider(t *testing.T) { model := insertInternalChatModelConfigForProvider( t, db, + org.ID, string(codersdk.ChatComputerUseProviderOpenAI), "gpt-4o-mini", true, @@ -4186,8 +4266,9 @@ func TestAwaitSubagentCompletion(t *testing.T) { BaseUrl: providerServer.URL, }) model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }) parent, child := createParentChildChats(ctx, t, server, user, org, model) @@ -4797,3 +4878,101 @@ func TestListAgents(t *testing.T) { require.Contains(t, msg, "only available on root chats") }) } + +func TestEnabledChatModelConfigsForOrganization(t *testing.T) { + t.Parallel() + + db, _ := dbtestutil.NewDB(t) + ctx := testutil.Context(t, testutil.WaitShort) + + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + localOrg := dbgen.Organization(t, db, database.Organization{}) + thirdOrg := dbgen.Organization(t, db, database.Organization{}) + provider := dbgen.AIProviderWithOptionalKey(t, db, database.AIProvider{ + Type: database.AIProviderTypeOpenai, + }, "test-key") + defaultOrgConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: defaultOrg.ID, + IsDefault: true, + }) + localDefault := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: localOrg.ID, + IsDefault: true, + }) + localConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: localOrg.ID, + }) + thirdConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: thirdOrg.ID, + }) + + rows, err := enabledChatModelConfigsForOrganization(ctx, db, localOrg.ID) + require.NoError(t, err) + require.Len(t, rows, 2) + ids := make([]uuid.UUID, 0, len(rows)) + for _, row := range rows { + require.Equal(t, localOrg.ID, row.ChatModelConfig.OrganizationID) + ids = append(ids, row.ChatModelConfig.ID) + } + require.ElementsMatch(t, []uuid.UUID{localDefault.ID, localConfig.ID}, ids) + require.NotContains(t, ids, defaultOrgConfig.ID) + require.NotContains(t, ids, thirdConfig.ID) + + emptyOrg := dbgen.Organization(t, db, database.Organization{}) + rows, err = enabledChatModelConfigsForOrganization(ctx, db, emptyOrg.ID) + require.NoError(t, err) + require.Empty(t, rows) +} + +func TestListSubagentModels_NonDefaultOrgExcludesDefaultOrgConfigs(t *testing.T) { + t.Parallel() + db, ps := dbtestutil.NewDB(t) + server := newInternalTestServer(t, db, ps, chatprovider.ProviderAPIKeys{}) + + ctx := chatdTestContext(t) + user, org, localModel := seedInternalChatDeps(t, db) + + defaultOrg, err := db.GetDefaultOrganization(ctx) + require.NoError(t, err) + defaultOrgProvider := dbgen.AIProviderWithOptionalKey(t, db, database.AIProvider{ + Type: database.AIProviderTypeOpenai, + }, "test-key") + defaultOrgModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + AIProviderID: uuid.NullUUID{UUID: defaultOrgProvider.ID, Valid: true}, + OrganizationID: defaultOrg.ID, + }) + + parentChat := createInternalParentChat( + ctx, t, server, db, org.ID, user.ID, localModel.ID, "parent-list-local-configs", + ) + + resp := runSubagentTool( + ctx, + t, + server, + parentChat, + parentChat.LastModelConfigID, + listSubagentModelsToolName, + struct{}{}, + ) + require.False(t, resp.IsError, "expected success but got: %s", resp.Content) + + var result struct { + Models []struct { + ModelConfigID string `json:"model_config_id"` + } `json:"models"` + } + require.NoError(t, json.Unmarshal([]byte(resp.Content), &result)) + + ids := make([]string, 0, len(result.Models)) + for _, model := range result.Models { + ids = append(ids, model.ModelConfigID) + } + require.Contains(t, ids, localModel.ID.String()) + require.NotContains(t, ids, defaultOrgModel.ID.String()) +} diff --git a/coderd/x/chatd/tasks.go b/coderd/x/chatd/tasks.go index a3738b886db..6bfc90ce94a 100644 --- a/coderd/x/chatd/tasks.go +++ b/coderd/x/chatd/tasks.go @@ -309,7 +309,9 @@ func (s *taskStarter) StartInterrupt(ctx context.Context, input chatWorkerTaskSt if len(committedCancels) > 0 { messages = append(append([]chatstate.Message{}, partialMessages...), committedCancels...) } - if _, err := tx.FinishInterruption(chatstate.FinishInterruptionInput{PartialMessages: messages}); err != nil { + if _, err := tx.FinishInterruption(chatstate.FinishInterruptionInput{ + PartialMessages: messages, + }); err != nil { return xerrors.Errorf("finish interruption: %w", err) } committed, err = store.GetChatByID(ctx, input.ChatID) diff --git a/coderd/x/chatd/tasks_test.go b/coderd/x/chatd/tasks_test.go index a1265b270b9..9531cff927c 100644 --- a/coderd/x/chatd/tasks_test.go +++ b/coderd/x/chatd/tasks_test.go @@ -959,7 +959,10 @@ func newTaskTestFixture(t *testing.T) *taskTestFixture { DisplayName: "openai", BaseUrl: "http://example.invalid", }) - model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{IsDefault: true}) + model := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + IsDefault: true, + OrganizationID: org.ID, + }) apiKey, _ := dbgen.APIKey(t, db, database.APIKey{UserID: user.ID}) return &taskTestFixture{db: db, pubsub: newTaskRecordingPubsub(ps), rawPS: ps, sqlDB: sqlDB, user: user, org: org, model: model, apiKey: apiKey} } diff --git a/coderd/x/chatd/title_override.go b/coderd/x/chatd/title_override.go index c8cbc6ddf9c..7b10541cae3 100644 --- a/coderd/x/chatd/title_override.go +++ b/coderd/x/chatd/title_override.go @@ -2,6 +2,7 @@ package chatd import ( "context" + "errors" "strings" "github.com/google/uuid" @@ -74,13 +75,18 @@ func (p *Server) resolveTitleGenerationModelOverride( titleGenerationOverrideContext, raw, chat.OwnerID, - p.resolveModelConfigAndNormalizedProvider, + func(ctx context.Context, modelConfigID uuid.UUID) (database.ChatModelConfig, string, error) { + return p.resolveModelConfigForOrganization(ctx, chat.OrganizationID, modelConfigID) + }, func(ctx context.Context, ownerID uuid.UUID, aiProviderID uuid.UUID) (chatprovider.ProviderAPIKeys, error) { return p.resolveUserProviderAPIKeys(ctx, ownerID, aiProviderID) }, modelOverrideFailureModeHard, ) if err != nil { + if errors.Is(err, errModelConfigOutsideOrganization) { + return resolvedModelCall{}, false, err + } return resolvedModelCall{}, overrideSet, err } if !overrideSet { diff --git a/coderd/x/chatd/title_override_internal_test.go b/coderd/x/chatd/title_override_internal_test.go index 0d0bda18796..779a96f974b 100644 --- a/coderd/x/chatd/title_override_internal_test.go +++ b/coderd/x/chatd/title_override_internal_test.go @@ -367,6 +367,55 @@ func TestMaybeGenerateChatTitle_TitleGenerationOverrideCallFailureSkipsFallback( require.False(t, ok) } +func TestResolveManualTitleModel_TitleGenerationOverrideIgnoresForeignConfig(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + chat, _ := titleOverrideTestChatAndMessages(t) + chat.OrganizationID = uuid.New() + overrideConfig := titleOverrideModelConfig("gpt-4.1", true) + overrideConfig.OrganizationID = uuid.New() + overrideProviderID := uuid.New() + overrideConfig.AIProviderID = uuid.NullUUID{UUID: overrideProviderID, Valid: true} + localProviderID := uuid.New() + localPreferred := database.ChatModelConfig{ + ID: uuid.New(), + Model: preferredTitleModels[0].model, + Enabled: true, + OrganizationID: chat.OrganizationID, + } + localDefault := database.ChatModelConfig{ + ID: uuid.New(), + AIProviderID: uuid.NullUUID{UUID: localProviderID, Valid: true}, + Model: "local-default-model", + Enabled: true, + IsDefault: true, + OrganizationID: chat.OrganizationID, + } + + db.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return(overrideConfig.ID.String(), nil) + db.EXPECT().GetChatModelConfigByID(gomock.Any(), overrideConfig.ID).Return(overrideConfig, nil) + db.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), chat.OrganizationID).Return([]database.GetEnabledChatModelConfigsByOrganizationRow{ + {ChatModelConfig: localPreferred, Provider: preferredTitleModels[0].provider}, + {ChatModelConfig: localDefault, Provider: "openai"}, + }, nil).AnyTimes() + db.EXPECT().GetAIProviderByID(gomock.Any(), localProviderID).Return(aibridgeTestAIProvider(localProviderID, "primary-openai", database.AIProviderTypeOpenai), nil).AnyTimes() + + server := titleOverrideTestServer(db, logger) + resolved, err := server.resolveManualTitleModel( + ctx, + db, + chat, + modelBuildOptions{ActiveAPIKeyID: uuid.NewString()}, + ) + require.NoError(t, err) + require.True(t, resolved.model.Valid()) + require.Equal(t, localDefault, resolved.dbConfig) +} + func TestResolveManualTitleModel_TitleGenerationOverrideUnset(t *testing.T) { t.Parallel() @@ -384,7 +433,7 @@ func TestResolveManualTitleModel_TitleGenerationOverrideUnset(t *testing.T) { } db.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", nil) - db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return([]database.GetEnabledChatModelConfigsRow{ + db.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), chat.OrganizationID).Return([]database.GetEnabledChatModelConfigsByOrganizationRow{ {ChatModelConfig: database.ChatModelConfig{Model: "gpt-4.1", Enabled: true}, Provider: "openai"}, {ChatModelConfig: preferredConfig, Provider: preferredTitleModels[1].provider}, }, nil) @@ -402,6 +451,30 @@ func TestResolveManualTitleModel_TitleGenerationOverrideUnset(t *testing.T) { require.Equal(t, preferredConfig, resolved.dbConfig) } +func TestResolveManualTitleModel_NonDefaultOrgDoesNotUseDefaultOrgConfigs(t *testing.T) { + t.Parallel() + + ctx := testutil.Context(t, testutil.WaitShort) + ctrl := gomock.NewController(t) + db := dbmock.NewMockStore(ctrl) + logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: true}) + chat, _ := titleOverrideTestChatAndMessages(t) + chat.OrganizationID = uuid.New() + + db.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", nil) + db.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), chat.OrganizationID).Return(nil, nil).AnyTimes() + + server := titleOverrideTestServer(db, logger) + resolved, err := server.resolveManualTitleModel( + ctx, + db, + chat, + modelBuildOptions{ActiveAPIKeyID: uuid.NewString()}, + ) + require.ErrorIs(t, err, ErrNoDefaultChatModelConfig) + require.Equal(t, resolvedModelCall{}, resolved) +} + func TestResolveManualTitleModel_TitleGenerationOverrideUnsetAIProvider(t *testing.T) { t.Parallel() @@ -430,7 +503,7 @@ func TestResolveManualTitleModel_TitleGenerationOverrideUnsetAIProvider(t *testi } db.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", nil) - db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return([]database.GetEnabledChatModelConfigsRow{ + db.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), chat.OrganizationID).Return([]database.GetEnabledChatModelConfigsByOrganizationRow{ {ChatModelConfig: preferredConfig, Provider: preferredTitleModels[1].provider}, }, nil) db.EXPECT().GetAIProviderByID(gomock.Any(), providerID).Return(provider, nil) @@ -468,7 +541,7 @@ func TestResolveManualTitleModel_TitleGenerationOverrideReadDBError(t *testing.T } db.EXPECT().GetChatTitleGenerationModelOverride(gomock.Any()).Return("", sql.ErrConnDone) - db.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return([]database.GetEnabledChatModelConfigsRow{ + db.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), chat.OrganizationID).Return([]database.GetEnabledChatModelConfigsByOrganizationRow{ {ChatModelConfig: database.ChatModelConfig{Model: "gpt-4.1", Enabled: true}, Provider: "openai"}, {ChatModelConfig: preferredConfig, Provider: preferredTitleModels[1].provider}, }, nil) @@ -563,6 +636,7 @@ func TestGenerateManualTitleCandidate_UsesSyntheticAPIKey(t *testing.T) { chat, messages := titleOverrideTestChatAndMessages(t) chat.OrganizationID = uuid.New() overrideConfig := titleOverrideModelConfig("gpt-4.1", true) + overrideConfig.OrganizationID = chat.OrganizationID providerID := uuid.New() overrideConfig.AIProviderID = uuid.NullUUID{UUID: providerID, Valid: true} overrideConfig.Options = modelCallSentinelOptions(t, "title-options-sentinel") diff --git a/coderd/x/chatd/toolinput_internal_test.go b/coderd/x/chatd/toolinput_internal_test.go index e19bbe48a9d..70815793c91 100644 --- a/coderd/x/chatd/toolinput_internal_test.go +++ b/coderd/x/chatd/toolinput_internal_test.go @@ -133,8 +133,9 @@ func TestBuiltinToolSchemasDescribeTheirInputs(t *testing.T) { Type: database.AIProviderTypeOpenai, }, "test-key") modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ - Model: "gpt-4o-mini", - AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + Model: "gpt-4o-mini", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: org.ID, }, func(p *database.InsertChatModelConfigParams) { p.Enabled = true }) diff --git a/coderd/x/chatd/turn_summary_internal_test.go b/coderd/x/chatd/turn_summary_internal_test.go index d8494c4ab13..7f133591c20 100644 --- a/coderd/x/chatd/turn_summary_internal_test.go +++ b/coderd/x/chatd/turn_summary_internal_test.go @@ -51,6 +51,9 @@ func TestUpdateLastTurnSummaryRejectsStaleWrites(t *testing.T) { ContextLimit: 128000, CompressionThreshold: 80, Options: json.RawMessage(`{}`), + OrganizationID: org.ID, + GroupACL: database.ChatACL{}, + UserACL: database.ChatACL{}, }) require.NoError(t, err) @@ -146,6 +149,9 @@ func TestSuccessfulChildChatOutcomeStoresReportSummaryWithoutPush(t *testing.T) ContextLimit: 128000, CompressionThreshold: 80, Options: json.RawMessage(`{}`), + OrganizationID: org.ID, + GroupACL: database.ChatACL{}, + UserACL: database.ChatACL{}, }) require.NoError(t, err) diff --git a/enterprise/coderd/exp_chats_test.go b/enterprise/coderd/exp_chats_test.go index 364b82fc682..5f6c3ecc3df 100644 --- a/enterprise/coderd/exp_chats_test.go +++ b/enterprise/coderd/exp_chats_test.go @@ -14,6 +14,7 @@ import ( "github.com/coder/coder/v2/coderd/aibridgedtest" "github.com/coder/coder/v2/coderd/coderdtest" "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/dbgen" "github.com/coder/coder/v2/coderd/database/dbtestutil" "github.com/coder/coder/v2/coderd/rbac" "github.com/coder/coder/v2/coderd/util/ptr" @@ -1074,12 +1075,12 @@ func (p cookieOnlySessionTokenProvider) SetDialOption(opts *websocket.DialOption opts.HTTPHeader.Set("Cookie", cookieName+"="+p.token) } -func TestCreateChatNonDefaultOrg(t *testing.T) { +func TestCreateChatUsesOrganizationLocalModel(t *testing.T) { t.Parallel() ctx := testutil.Context(t, testutil.WaitLong) - client, firstUser := coderdenttest.New(t, &coderdenttest.Options{ + client, db, firstUser := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ Options: &coderdtest.Options{ DeploymentValues: func() *codersdk.DeploymentValues { v := coderdtest.DeploymentValues(t) @@ -1095,29 +1096,36 @@ func TestCreateChatNonDefaultOrg(t *testing.T) { expClient := codersdk.NewExperimentalClient(client) provider := createOpenAIProviderForTest(ctx, t, expClient, "test-key", "https://example.com") - _, err := expClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{ + defaultModel, err := expClient.CreateChatModelConfig(ctx, codersdk.CreateChatModelConfigRequest{ AIProviderID: &provider.ID, Model: "gpt-4o-mini", - DisplayName: "Test Model", + DisplayName: "Default Organization Model", IsDefault: ptr.Ref(true), ContextLimit: ptr.Ref(int64(1000)), CompressionThreshold: ptr.Ref(int32(70)), }) require.NoError(t, err) - // Create a second (non-default) org via the API. secondOrg := coderdenttest.CreateOrganization(t, client, coderdenttest.CreateOrganizationOptions{}) + localModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "gpt-4o-mini-local", + DisplayName: "Second Organization Model", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: secondOrg.ID, + IsDefault: true, + }) + require.NotEqual(t, defaultModel.ID, localModel.ID) - // Create a member with agents-access in both orgs. memberClientRaw, member := coderdtest.CreateAnotherUser( t, client, firstUser.OrganizationID, rbac.ScopedRoleAgentsAccess(firstUser.OrganizationID), rbac.ScopedRoleAgentsAccess(secondOrg.ID), ) memberClient := codersdk.NewExperimentalClient(memberClientRaw) - // Create a chat in the non-default org. + chat, err := memberClient.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: secondOrg.ID, + ModelConfigID: &localModel.ID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -1128,19 +1136,7 @@ func TestCreateChatNonDefaultOrg(t *testing.T) { require.NoError(t, err) require.Equal(t, secondOrg.ID, chat.OrganizationID) require.Equal(t, member.ID, chat.OwnerID) - - // Verify the chat is visible when listing. - chats, err := memberClient.ListChats(ctx, nil) - require.NoError(t, err) - var found bool - for _, c := range chats { - if c.ID == chat.ID { - found = true - require.Equal(t, secondOrg.ID, c.OrganizationID) - break - } - } - require.True(t, found, "chat should be visible in list") + require.Equal(t, localModel.ID, chat.LastModelConfigID) } func TestListChats_OrgAdminOnlySeesOwnChats(t *testing.T) { @@ -1148,7 +1144,7 @@ func TestListChats_OrgAdminOnlySeesOwnChats(t *testing.T) { ctx := testutil.Context(t, testutil.WaitLong) - client, firstUser := coderdenttest.New(t, &coderdenttest.Options{ + client, db, firstUser := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ Options: &coderdtest.Options{ DeploymentValues: func() *codersdk.DeploymentValues { v := coderdtest.DeploymentValues(t) @@ -1176,6 +1172,13 @@ func TestListChats_OrgAdminOnlySeesOwnChats(t *testing.T) { // Create a second (non-default) org. secondOrg := coderdenttest.CreateOrganization(t, client, coderdenttest.CreateOrganizationOptions{}) + localModel := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{ + Model: "gpt-4o-mini-local", + DisplayName: "Second Organization Model", + AIProviderID: uuid.NullUUID{UUID: provider.ID, Valid: true}, + OrganizationID: secondOrg.ID, + IsDefault: true, + }) // Create a member with agents-access in both orgs. memberClientRaw, _ := coderdtest.CreateAnotherUser( @@ -1187,6 +1190,7 @@ func TestListChats_OrgAdminOnlySeesOwnChats(t *testing.T) { // Member creates a chat in the second org. memberChat, err := memberExp.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: secondOrg.ID, + ModelConfigID: &localModel.ID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText, @@ -1207,6 +1211,7 @@ func TestListChats_OrgAdminOnlySeesOwnChats(t *testing.T) { // Admin creates a chat in the second org. adminChat, err := adminExp.CreateChat(ctx, codersdk.CreateChatRequest{ OrganizationID: secondOrg.ID, + ModelConfigID: &localModel.ID, Content: []codersdk.ChatInputPart{ { Type: codersdk.ChatInputPartTypeText,