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,