mirror of
https://github.com/coder/coder.git
synced 2026-10-31 11:02:50 +08:00
feat!: enforce organization scope for chat models (#27959)
## Summary Move chat model management and runtime availability to organization-scoped API routes. The API now identifies the organization in every model-management operation, applies chat model RBAC and ACLs within that organization, and returns only models the caller can use there. Depends on #27958. ## Problem Chat model configurations now belong to organizations, but the existing experimental management and discovery APIs are deployment-shaped. Their routes do not identify an organization, SDK methods do not accept an organization ID, and item operations address a model configuration by ID alone. That contract cannot safely expose organization-owned model configurations. It also does not provide the organization-specific provider information, authorization behavior, ACL management, or recovery guidance needed by the settings and chat clients later in the stack. ## Fix Add organization-scoped chat model endpoints under: ```text /api/experimental/organizations/{organization}/chats/models ``` The new API provides: - Organization-scoped model listing and creation. - Model get, update, and delete operations scoped by both organization and model ID. - Model ACL read and update operations. - Organization-scoped runtime availability and provider status. - Redacted provider descriptors for model-management clients. - Chat model RBAC and row ACL enforcement for writes, runtime discovery, and Chatd model selection. - Recovery guidance when a chat's organization has no usable local model. Historical chats can retain a model reference owned by another organization, but runtime selection uses an authorized local default when one is available. Foreign, inaccessible, disabled, and otherwise unusable models are not exposed as usable selections for the organization. The existing collection and availability routes remain as default-organization compatibility endpoints in this layer: ```text GET /api/experimental/chats/model-configs POST /api/experimental/chats/model-configs GET /api/experimental/chats/models ``` These compatibility routes do not provide access to models in non-default organizations. The previous unscoped item update and delete routes are removed because item operations must include the owning organization. ## Breaking change This changes the experimental HTTP API and Go SDK contracts for chat model management. Clients must migrate model-management and availability requests to the organization-scoped routes and provide an organization identifier. In the Go SDK, the following methods now require an organization ID: - `ChatModelAvailability` - `ChatModels` - `CreateChatModel` - `UpdateChatModel` - `DeleteChatModel` `ChatModels` also changes its return value from `[]ChatModel` to `OrganizationChatModelsResponse`, which includes the organization models and the redacted provider descriptors required by management clients. New item and ACL methods likewise require both the organization ID and model ID. The default-organization compatibility routes reduce the immediate impact for collection and availability callers, but they do not preserve the previous item API or SDK method signatures. Generated clients and direct consumers of the experimental API must be updated before adopting this change. _This pull request description was generated by Coder Agents._
This commit is contained in:
@@ -0,0 +1,21 @@
|
||||
package database
|
||||
|
||||
// EffectiveChatModelConfigs contains enabled configs for one organization.
|
||||
type EffectiveChatModelConfigs struct {
|
||||
Configs []GetEnabledChatModelConfigsByOrganizationRow
|
||||
DefaultConfig ChatModelConfig
|
||||
}
|
||||
|
||||
// DeriveEffectiveChatModelConfigs selects the stored default from enabled
|
||||
// configs in one organization.
|
||||
func DeriveEffectiveChatModelConfigs(
|
||||
rows []GetEnabledChatModelConfigsByOrganizationRow,
|
||||
) EffectiveChatModelConfigs {
|
||||
result := EffectiveChatModelConfigs{Configs: rows}
|
||||
for _, row := range rows {
|
||||
if row.ChatModelConfig.IsDefault {
|
||||
result.DefaultConfig = row.ChatModelConfig
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
@@ -34,9 +34,11 @@ const wrapname = "dbauthz.querier"
|
||||
// ErrNoActor is returned if no actor is present in the context.
|
||||
var ErrNoActor = xerrors.Errorf("no authorization actor in context")
|
||||
|
||||
// NotAuthorizedError is a sentinel error that unwraps to sql.ErrNoRows.
|
||||
// This allows the internal error to be read by the caller if needed. Otherwise
|
||||
// it will be handled as a 404.
|
||||
// NotAuthorizedError is a sentinel error wrapping the underlying RBAC denial.
|
||||
// HTTP layers map it to a 404 (via IsUnauthorized) so denials are
|
||||
// indistinguishable from missing resources, but it does NOT unwrap to
|
||||
// sql.ErrNoRows: 'errors.Is(err, sql.ErrNoRows)' is false for it, so callers
|
||||
// distinguishing "not found" from "denied" must use httpapi.Is404Error.
|
||||
type NotAuthorizedError struct {
|
||||
Err error
|
||||
}
|
||||
@@ -57,8 +59,7 @@ func (NotAuthorizedError) IsUnauthorized() bool {
|
||||
return true
|
||||
}
|
||||
|
||||
// Unwrap will always unwrap to a sql.ErrNoRows so the API returns a 404.
|
||||
// So 'errors.Is(err, sql.ErrNoRows)' will always be true.
|
||||
// Unwrap returns the wrapped RBAC denial, never sql.ErrNoRows.
|
||||
func (e NotAuthorizedError) Unwrap() error {
|
||||
return e.Err
|
||||
}
|
||||
@@ -824,8 +825,8 @@ 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.
|
||||
// Organization reads support organization-scoped model
|
||||
// resolution and user-global settings in the default org.
|
||||
rbac.ResourceOrganization.Type: {policy.ActionRead},
|
||||
}),
|
||||
User: []rbac.Permission{},
|
||||
@@ -2222,9 +2223,16 @@ func (q *querier) DeleteChatDebugDataByChatID(ctx context.Context, arg database.
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) (uuid.UUID, error) {
|
||||
// Interim write gate: see InsertChatModelConfig.
|
||||
// TODO(mafredri): swap to ResourceChatModelConfig object delete after CODAGT-709 M3 (org-scoping cutover)
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
config, err := q.db.GetChatModelConfigByID(ctx, id)
|
||||
if err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
if err := q.authorizeContext(ctx, policy.ActionDelete, object); err != nil {
|
||||
return uuid.Nil, err
|
||||
}
|
||||
return q.db.DeleteChatModelConfigByID(ctx, id)
|
||||
@@ -3512,20 +3520,41 @@ func (q *querier) GetChatMessagesForPromptByChatID(ctx context.Context, chatID u
|
||||
}
|
||||
|
||||
func (q *querier) GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
|
||||
// Interim read gate: see InsertChatModelConfig.
|
||||
// TODO(mafredri): swap to ResourceChatModelConfig object read after CODAGT-709 M3 (org-scoping cutover)
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
config, err := q.db.GetChatModelConfigByID(ctx, id)
|
||||
if err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return q.db.GetChatModelConfigByID(ctx, id)
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, object); err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return config, nil
|
||||
}
|
||||
|
||||
func (q *querier) GetChatModelConfigs(ctx context.Context) ([]database.ChatModelConfig, error) {
|
||||
func (q *querier) GetChatModelConfigs(ctx context.Context, organizationID uuid.UUID) ([]database.ChatModelConfig, error) {
|
||||
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceChatModelConfig.Type)
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err)
|
||||
}
|
||||
return q.db.GetAuthorizedChatModelConfigs(ctx, prep)
|
||||
return q.db.GetAuthorizedChatModelConfigs(ctx, organizationID, prep)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ChatModelConfig, error) {
|
||||
// Consumed only by ensureDefaultChatModelConfig's default-election read
|
||||
// inside the write transaction. Every path that reaches the election
|
||||
// already requires update-in-org (insert/update/delete and the election's
|
||||
// own UnsetDefaultChatModelConfigs/UpdateChatModelConfig), so the read
|
||||
// authorizes update-in-org rather than read; a create-only principal never
|
||||
// reaches it because create short-circuits when the inserted config is the
|
||||
// org's first.
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChatModelConfig.InOrg(organizationID)); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetChatModelConfigsByOrganization(ctx, organizationID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatModelConfigsForTelemetry(ctx context.Context) ([]database.GetChatModelConfigsForTelemetryRow, error) {
|
||||
@@ -3761,12 +3790,19 @@ func (q *querier) GetDatabaseNow(ctx context.Context) (time.Time, error) {
|
||||
}
|
||||
|
||||
func (q *querier) GetDefaultChatModelConfig(ctx context.Context, organizationID uuid.UUID) (database.ChatModelConfig, error) {
|
||||
// Chat creation authorizes the organization-scoped chat before this lookup.
|
||||
// TODO(mafredri): remove after CODAGT-709 M3.
|
||||
if _, ok := ActorFromContext(ctx); !ok {
|
||||
return database.ChatModelConfig{}, ErrNoActor
|
||||
config, err := q.db.GetDefaultChatModelConfig(ctx, organizationID)
|
||||
if err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return q.db.GetDefaultChatModelConfig(ctx, organizationID)
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, object); err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return config, nil
|
||||
}
|
||||
|
||||
func (q *querier) GetDefaultOrganization(ctx context.Context) (database.Organization, error) {
|
||||
@@ -3802,28 +3838,49 @@ func (q *querier) GetEligibleProvisionerDaemonsByProvisionerJobIDs(ctx context.C
|
||||
}
|
||||
|
||||
func (q *querier) GetEnabledChatModelConfigByID(ctx context.Context, id uuid.UUID) (database.ChatModelConfig, error) {
|
||||
// TODO(mafredri): authorize the fetched object after CODAGT-709 M3
|
||||
// (org-scoping cutover).
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChatModelConfig); err != nil {
|
||||
config, err := q.db.GetEnabledChatModelConfigByID(ctx, id)
|
||||
if err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return q.db.GetEnabledChatModelConfigByID(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) GetEnabledChatModelConfigs(ctx context.Context) ([]database.GetEnabledChatModelConfigsRow, error) {
|
||||
// TODO(mafredri): replace this deployment-wide query after CODAGT-709 M3
|
||||
// (org-scoping cutover).
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChatModelConfig); err != nil {
|
||||
return nil, err
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, object); err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return q.db.GetEnabledChatModelConfigs(ctx)
|
||||
return config, nil
|
||||
}
|
||||
|
||||
func (q *querier) GetEnabledChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.GetEnabledChatModelConfigsByOrganizationRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChatModelConfig.InOrg(organizationID)); err != nil {
|
||||
rows, err := q.db.GetEnabledChatModelConfigsByOrganization(ctx, organizationID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetEnabledChatModelConfigsByOrganization(ctx, organizationID)
|
||||
act, ok := ActorFromContext(ctx)
|
||||
if !ok {
|
||||
return nil, ErrNoActor
|
||||
}
|
||||
out := make([]database.GetEnabledChatModelConfigsByOrganizationRow, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
config := row.ChatModelConfig
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
err := q.auth.Authorize(ctx, act, policy.ActionRead, object)
|
||||
switch {
|
||||
case err == nil:
|
||||
out = append(out, row)
|
||||
case rbac.IsUnauthorizedError(err):
|
||||
continue
|
||||
default:
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (q *querier) GetEnabledMCPServerConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.MCPServerConfig, error) {
|
||||
@@ -6116,13 +6173,7 @@ func (q *querier) InsertChatMessages(ctx context.Context, arg database.InsertCha
|
||||
}
|
||||
|
||||
func (q *querier) InsertChatModelConfig(ctx context.Context, arg database.InsertChatModelConfigParams) (database.ChatModelConfig, error) {
|
||||
// The routes create configs in the default organization. Organization
|
||||
// administrators must not manage configs that other organizations can read.
|
||||
// TODO(mafredri): swap to ResourceChatModelConfig create-in-org after CODAGT-709 M3 (org-scoping cutover)
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return q.db.InsertChatModelConfig(ctx, arg)
|
||||
return insert(q.log, q.auth, rbac.ResourceChatModelConfig.InOrg(arg.OrganizationID), q.db.InsertChatModelConfig)(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) InsertChatQueuedMessage(ctx context.Context, arg database.InsertChatQueuedMessageParams) (database.ChatQueuedMessage, error) {
|
||||
@@ -7322,9 +7373,7 @@ func (q *querier) UnpinChatByID(ctx context.Context, id uuid.UUID) error {
|
||||
}
|
||||
|
||||
func (q *querier) UnsetDefaultChatModelConfigs(ctx context.Context, organizationID uuid.UUID) error {
|
||||
// Interim write gate: see InsertChatModelConfig.
|
||||
// TODO(mafredri): swap to ResourceChatModelConfig update-in-org after CODAGT-709 M3 (org-scoping cutover)
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceChatModelConfig.InOrg(organizationID)); err != nil {
|
||||
return err
|
||||
}
|
||||
return q.db.UnsetDefaultChatModelConfigs(ctx, organizationID)
|
||||
@@ -7503,9 +7552,16 @@ func (q *querier) UpdateChatMCPServerIDs(ctx context.Context, arg database.Updat
|
||||
}
|
||||
|
||||
func (q *querier) UpdateChatModelConfig(ctx context.Context, arg database.UpdateChatModelConfigParams) (database.ChatModelConfig, error) {
|
||||
// Interim write gate: see InsertChatModelConfig.
|
||||
// TODO(mafredri): swap to ResourceChatModelConfig object update after CODAGT-709 M3 (org-scoping cutover)
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
existing, err := q.db.GetChatModelConfigByID(ctx, arg.ID)
|
||||
if err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(existing.ID).
|
||||
InOrg(existing.OrganizationID).
|
||||
WithACLUserList(existing.UserACL.RBACACL()).
|
||||
WithGroupACL(existing.GroupACL.RBACACL())
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, object); err != nil {
|
||||
return database.ChatModelConfig{}, err
|
||||
}
|
||||
return q.db.UpdateChatModelConfig(ctx, arg)
|
||||
@@ -9478,18 +9534,8 @@ func (q *querier) GetAuthorizedChatsByChatFileID(ctx context.Context, fileID uui
|
||||
return q.db.GetAuthorizedChatsByChatFileID(ctx, fileID, prepared)
|
||||
}
|
||||
|
||||
func (q *querier) GetAuthorizedChatModelConfigs(ctx context.Context, prepared rbac.PreparedAuthorized) ([]database.ChatModelConfig, error) {
|
||||
return q.db.GetAuthorizedChatModelConfigs(ctx, prepared)
|
||||
}
|
||||
|
||||
// GetDefaultChatModelConfigCandidates returns all non-deleted configs.
|
||||
// A filtered result can leave an organization without a default.
|
||||
// TODO(mafredri): remove after CODAGT-709 M3.
|
||||
func (q *querier) GetDefaultChatModelConfigCandidates(ctx context.Context) ([]database.ChatModelConfig, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.GetDefaultChatModelConfigCandidates(ctx)
|
||||
func (q *querier) GetAuthorizedChatModelConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]database.ChatModelConfig, error) {
|
||||
return q.db.GetAuthorizedChatModelConfigs(ctx, organizationID, prepared)
|
||||
}
|
||||
|
||||
func (q *querier) GetAuthorizedMCPServerConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]database.MCPServerConfig, error) {
|
||||
|
||||
@@ -802,10 +802,16 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().SoftDeleteChatMessageByID(gomock.Any(), msg.ID).Return(nil).AnyTimes()
|
||||
check.Args(msg.ID).Asserts(chat, policy.ActionUpdate).Returns()
|
||||
}))
|
||||
s.Run("DeleteChatModelConfigByID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
configID := uuid.New()
|
||||
dbm.EXPECT().DeleteChatModelConfigByID(gomock.Any(), configID).Return(configID, nil).AnyTimes()
|
||||
check.Args(configID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(configID)
|
||||
s.Run("DeleteChatModelConfigByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
config := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
|
||||
dbm.EXPECT().GetChatModelConfigByID(gomock.Any(), config.ID).Return(config, nil).AnyTimes()
|
||||
dbm.EXPECT().DeleteChatModelConfigByID(gomock.Any(), config.ID).Return(config.ID, nil).AnyTimes()
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
check.Args(config.ID).Asserts(object, policy.ActionDelete).Returns(config.ID)
|
||||
}))
|
||||
s.Run("DeleteChatQueuedMessage", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
@@ -1175,30 +1181,40 @@ func (s *MethodTestSuite) TestChats() {
|
||||
s.Run("GetChatModelConfigByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
config := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
|
||||
dbm.EXPECT().GetChatModelConfigByID(gomock.Any(), config.ID).Return(config, nil).AnyTimes()
|
||||
check.Args(config.ID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(config)
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
check.Args(config.ID).Asserts(object, policy.ActionRead).Returns(config)
|
||||
}))
|
||||
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(), config.OrganizationID).Return(config, nil).AnyTimes()
|
||||
check.Args(config.OrganizationID).Asserts().Returns(config)
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
check.Args(config.OrganizationID).Asserts(object, policy.ActionRead).Returns(config)
|
||||
}))
|
||||
s.Run("GetChatModelConfigs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().GetAuthorizedChatModelConfigs(gomock.Any(), gomock.Any()).Return([]database.ChatModelConfig{}, nil).AnyTimes()
|
||||
orgID := uuid.New()
|
||||
dbm.EXPECT().GetAuthorizedChatModelConfigs(gomock.Any(), orgID, gomock.Any()).Return([]database.ChatModelConfig{}, nil).AnyTimes()
|
||||
// No asserts here because SQLFilter.
|
||||
check.Args().Asserts()
|
||||
check.Args(orgID).Asserts()
|
||||
}))
|
||||
s.Run("GetChatModelConfigsByOrganization", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
orgID := uuid.New()
|
||||
dbm.EXPECT().GetChatModelConfigsByOrganization(gomock.Any(), orgID).Return([]database.ChatModelConfig{}, nil).AnyTimes()
|
||||
check.Args(orgID).Asserts(rbac.ResourceChatModelConfig.InOrg(orgID), policy.ActionUpdate).Returns([]database.ChatModelConfig{})
|
||||
}))
|
||||
s.Run("GetAuthorizedChatModelConfigs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().GetAuthorizedChatModelConfigs(gomock.Any(), gomock.Any()).Return([]database.ChatModelConfig{}, nil).AnyTimes()
|
||||
orgID := uuid.New()
|
||||
dbm.EXPECT().GetAuthorizedChatModelConfigs(gomock.Any(), orgID, gomock.Any()).Return([]database.ChatModelConfig{}, nil).AnyTimes()
|
||||
// No asserts here because callers provide the SQL filter.
|
||||
check.Args(emptyPreparedAuthorized{}).Asserts()
|
||||
check.Args(orgID, emptyPreparedAuthorized{}).Asserts()
|
||||
}))
|
||||
s.Run("GetDefaultChatModelConfigCandidates", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
configA := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
|
||||
configB := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
|
||||
dbm.EXPECT().GetDefaultChatModelConfigCandidates(gomock.Any()).Return([]database.ChatModelConfig{configA, configB}, nil).AnyTimes()
|
||||
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns([]database.ChatModelConfig{configA, configB})
|
||||
}))
|
||||
|
||||
s.Run("GetChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
params := database.GetChatsParams{}
|
||||
dbm.EXPECT().GetAuthorizedChats(gomock.Any(), params, gomock.Any()).Return([]database.GetChatsRow{}, nil).AnyTimes()
|
||||
@@ -1303,21 +1319,33 @@ func (s *MethodTestSuite) TestChats() {
|
||||
s.Run("GetEnabledChatModelConfigByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
config := testutil.Fake(s.T(), faker, database.ChatModelConfig{})
|
||||
dbm.EXPECT().GetEnabledChatModelConfigByID(gomock.Any(), config.ID).Return(config, nil).AnyTimes()
|
||||
check.Args(config.ID).Asserts(rbac.ResourceChatModelConfig, policy.ActionRead).Returns(config)
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
check.Args(config.ID).Asserts(object, policy.ActionRead).Returns(config)
|
||||
}))
|
||||
s.Run("GetEnabledChatModelConfigs", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
rowA := testutil.Fake(s.T(), faker, database.GetEnabledChatModelConfigsRow{})
|
||||
rowB := testutil.Fake(s.T(), faker, database.GetEnabledChatModelConfigsRow{})
|
||||
dbm.EXPECT().GetEnabledChatModelConfigs(gomock.Any()).Return([]database.GetEnabledChatModelConfigsRow{rowA, rowB}, nil).AnyTimes()
|
||||
check.Args().Asserts(rbac.ResourceChatModelConfig, 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{})
|
||||
rowA := testutil.Fake(s.T(), faker, database.GetEnabledChatModelConfigsByOrganizationRow{
|
||||
ChatModelConfig: database.ChatModelConfig{OrganizationID: orgID},
|
||||
})
|
||||
rowB := testutil.Fake(s.T(), faker, database.GetEnabledChatModelConfigsByOrganizationRow{
|
||||
ChatModelConfig: database.ChatModelConfig{OrganizationID: orgID},
|
||||
})
|
||||
dbm.EXPECT().GetEnabledChatModelConfigsByOrganization(gomock.Any(), orgID).Return([]database.GetEnabledChatModelConfigsByOrganizationRow{rowA, rowB}, nil).AnyTimes()
|
||||
check.Args(orgID).Asserts(rbac.ResourceChatModelConfig.InOrg(orgID), policy.ActionRead).Returns([]database.GetEnabledChatModelConfigsByOrganizationRow{rowA, rowB})
|
||||
objectA := rbac.ResourceChatModelConfig.
|
||||
WithID(rowA.ChatModelConfig.ID).
|
||||
InOrg(rowA.ChatModelConfig.OrganizationID).
|
||||
WithACLUserList(rowA.ChatModelConfig.UserACL.RBACACL()).
|
||||
WithGroupACL(rowA.ChatModelConfig.GroupACL.RBACACL())
|
||||
objectB := rbac.ResourceChatModelConfig.
|
||||
WithID(rowB.ChatModelConfig.ID).
|
||||
InOrg(rowB.ChatModelConfig.OrganizationID).
|
||||
WithACLUserList(rowB.ChatModelConfig.UserACL.RBACACL()).
|
||||
WithGroupACL(rowB.ChatModelConfig.GroupACL.RBACACL())
|
||||
check.Args(orgID).Asserts(objectA, policy.ActionRead, objectB, policy.ActionRead).Returns([]database.GetEnabledChatModelConfigsByOrganizationRow{rowA, rowB})
|
||||
}))
|
||||
|
||||
s.Run("GetStaleChats", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
@@ -1360,13 +1388,14 @@ func (s *MethodTestSuite) TestChats() {
|
||||
}))
|
||||
s.Run("InsertChatModelConfig", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
arg := database.InsertChatModelConfigParams{
|
||||
Model: "test-model",
|
||||
DisplayName: "Test Model",
|
||||
Enabled: true,
|
||||
Model: "test-model",
|
||||
DisplayName: "Test Model",
|
||||
Enabled: true,
|
||||
OrganizationID: uuid.New(),
|
||||
}
|
||||
config := testutil.Fake(s.T(), faker, database.ChatModelConfig{Model: arg.Model, DisplayName: arg.DisplayName, Enabled: arg.Enabled})
|
||||
dbm.EXPECT().InsertChatModelConfig(gomock.Any(), arg).Return(config, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(config)
|
||||
check.Args(arg).Asserts(rbac.ResourceChatModelConfig.InOrg(arg.OrganizationID), policy.ActionCreate).Returns(config)
|
||||
}))
|
||||
|
||||
s.Run("PopNextQueuedMessage", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
@@ -1594,8 +1623,14 @@ func (s *MethodTestSuite) TestChats() {
|
||||
DisplayName: "Updated Model",
|
||||
Enabled: true,
|
||||
}
|
||||
dbm.EXPECT().GetChatModelConfigByID(gomock.Any(), config.ID).Return(config, nil).AnyTimes()
|
||||
dbm.EXPECT().UpdateChatModelConfig(gomock.Any(), arg).Return(config, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(config)
|
||||
object := rbac.ResourceChatModelConfig.
|
||||
WithID(config.ID).
|
||||
InOrg(config.OrganizationID).
|
||||
WithACLUserList(config.UserACL.RBACACL()).
|
||||
WithGroupACL(config.GroupACL.RBACACL())
|
||||
check.Args(arg).Asserts(object, policy.ActionUpdate).Returns(config)
|
||||
}))
|
||||
|
||||
s.Run("UpdateChatModelConfigACLByID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
@@ -1668,7 +1703,7 @@ func (s *MethodTestSuite) TestChats() {
|
||||
s.Run("UnsetDefaultChatModelConfigs", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
orgID := uuid.New()
|
||||
dbm.EXPECT().UnsetDefaultChatModelConfigs(gomock.Any(), orgID).Return(nil).AnyTimes()
|
||||
check.Args(orgID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
|
||||
check.Args(orgID).Asserts(rbac.ResourceChatModelConfig.InOrg(orgID), policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("UpsertChatDiffStatus", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
|
||||
chat := testutil.Fake(s.T(), faker, database.Chat{})
|
||||
@@ -7760,6 +7795,10 @@ 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")
|
||||
|
||||
// Chatd resolves organization model configs during background work.
|
||||
err = auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceChatModelConfig.InOrg(uuid.New()))
|
||||
require.NoError(t, err, "chat model config read should be allowed")
|
||||
|
||||
// Pin the complete ResourceUser action set: read_personal only.
|
||||
// Token refresh persistence uses the per-user AsChatdTokenOwner subject.
|
||||
for _, action := range rbac.ResourceUser.AvailableActions() {
|
||||
@@ -7771,7 +7810,7 @@ func TestAsChatd(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Organization read is temporarily needed for the pre-cutover default-org fallback.
|
||||
// Organization reads support the pre-cutover default-org fallback.
|
||||
err = auth.Authorize(ctx, actor, policy.ActionRead, rbac.ResourceOrganization)
|
||||
require.NoError(t, err, "organization read should be allowed")
|
||||
})
|
||||
|
||||
+12
-20
@@ -1696,14 +1696,22 @@ func (m queryMetricsStore) GetChatModelConfigByID(ctx context.Context, id uuid.U
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatModelConfigs(ctx context.Context) ([]database.ChatModelConfig, error) {
|
||||
func (m queryMetricsStore) GetChatModelConfigs(ctx context.Context, organizationID uuid.UUID) ([]database.ChatModelConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatModelConfigs(ctx)
|
||||
r0, r1 := m.s.GetChatModelConfigs(ctx, organizationID)
|
||||
m.queryLatencies.WithLabelValues("GetChatModelConfigs").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatModelConfigs").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ChatModelConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatModelConfigsByOrganization(ctx, organizationID)
|
||||
m.queryLatencies.WithLabelValues("GetChatModelConfigsByOrganization").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatModelConfigsByOrganization").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatModelConfigsForTelemetry(ctx context.Context) ([]database.GetChatModelConfigsForTelemetryRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatModelConfigsForTelemetry(ctx)
|
||||
@@ -2008,14 +2016,6 @@ func (m queryMetricsStore) GetEnabledChatModelConfigByID(ctx context.Context, id
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetEnabledChatModelConfigs(ctx context.Context) ([]database.GetEnabledChatModelConfigsRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetEnabledChatModelConfigs(ctx)
|
||||
m.queryLatencies.WithLabelValues("GetEnabledChatModelConfigs").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetEnabledChatModelConfigs").Inc()
|
||||
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)
|
||||
@@ -6840,22 +6840,14 @@ func (m queryMetricsStore) GetAuthorizedChatsByChatFileID(ctx context.Context, f
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetAuthorizedChatModelConfigs(ctx context.Context, prepared rbac.PreparedAuthorized) ([]database.ChatModelConfig, error) {
|
||||
func (m queryMetricsStore) GetAuthorizedChatModelConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]database.ChatModelConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetAuthorizedChatModelConfigs(ctx, prepared)
|
||||
r0, r1 := m.s.GetAuthorizedChatModelConfigs(ctx, organizationID, prepared)
|
||||
m.queryLatencies.WithLabelValues("GetAuthorizedChatModelConfigs").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAuthorizedChatModelConfigs").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetDefaultChatModelConfigCandidates(ctx context.Context) ([]database.ChatModelConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetDefaultChatModelConfigCandidates(ctx)
|
||||
m.queryLatencies.WithLabelValues("GetDefaultChatModelConfigCandidates").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetDefaultChatModelConfigCandidates").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetAuthorizedMCPServerConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]database.MCPServerConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetAuthorizedMCPServerConfigs(ctx, organizationID, prepared)
|
||||
|
||||
Generated
+23
-38
@@ -2476,18 +2476,18 @@ func (mr *MockStoreMockRecorder) GetAuthorizedAuditLogsOffset(ctx, arg, prepared
|
||||
}
|
||||
|
||||
// GetAuthorizedChatModelConfigs mocks base method.
|
||||
func (m *MockStore) GetAuthorizedChatModelConfigs(ctx context.Context, prepared rbac.PreparedAuthorized) ([]database.ChatModelConfig, error) {
|
||||
func (m *MockStore) GetAuthorizedChatModelConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]database.ChatModelConfig, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetAuthorizedChatModelConfigs", ctx, prepared)
|
||||
ret := m.ctrl.Call(m, "GetAuthorizedChatModelConfigs", ctx, organizationID, prepared)
|
||||
ret0, _ := ret[0].([]database.ChatModelConfig)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetAuthorizedChatModelConfigs indicates an expected call of GetAuthorizedChatModelConfigs.
|
||||
func (mr *MockStoreMockRecorder) GetAuthorizedChatModelConfigs(ctx, prepared any) *gomock.Call {
|
||||
func (mr *MockStoreMockRecorder) GetAuthorizedChatModelConfigs(ctx, organizationID, prepared any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuthorizedChatModelConfigs", reflect.TypeOf((*MockStore)(nil).GetAuthorizedChatModelConfigs), ctx, prepared)
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAuthorizedChatModelConfigs", reflect.TypeOf((*MockStore)(nil).GetAuthorizedChatModelConfigs), ctx, organizationID, prepared)
|
||||
}
|
||||
|
||||
// GetAuthorizedChats mocks base method.
|
||||
@@ -3181,18 +3181,33 @@ func (mr *MockStoreMockRecorder) GetChatModelConfigByID(ctx, id any) *gomock.Cal
|
||||
}
|
||||
|
||||
// GetChatModelConfigs mocks base method.
|
||||
func (m *MockStore) GetChatModelConfigs(ctx context.Context) ([]database.ChatModelConfig, error) {
|
||||
func (m *MockStore) GetChatModelConfigs(ctx context.Context, organizationID uuid.UUID) ([]database.ChatModelConfig, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatModelConfigs", ctx)
|
||||
ret := m.ctrl.Call(m, "GetChatModelConfigs", ctx, organizationID)
|
||||
ret0, _ := ret[0].([]database.ChatModelConfig)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatModelConfigs indicates an expected call of GetChatModelConfigs.
|
||||
func (mr *MockStoreMockRecorder) GetChatModelConfigs(ctx any) *gomock.Call {
|
||||
func (mr *MockStoreMockRecorder) GetChatModelConfigs(ctx, organizationID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatModelConfigs", reflect.TypeOf((*MockStore)(nil).GetChatModelConfigs), ctx)
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatModelConfigs", reflect.TypeOf((*MockStore)(nil).GetChatModelConfigs), ctx, organizationID)
|
||||
}
|
||||
|
||||
// GetChatModelConfigsByOrganization mocks base method.
|
||||
func (m *MockStore) GetChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ChatModelConfig, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatModelConfigsByOrganization", ctx, organizationID)
|
||||
ret0, _ := ret[0].([]database.ChatModelConfig)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatModelConfigsByOrganization indicates an expected call of GetChatModelConfigsByOrganization.
|
||||
func (mr *MockStoreMockRecorder) GetChatModelConfigsByOrganization(ctx, organizationID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatModelConfigsByOrganization", reflect.TypeOf((*MockStore)(nil).GetChatModelConfigsByOrganization), ctx, organizationID)
|
||||
}
|
||||
|
||||
// GetChatModelConfigsForTelemetry mocks base method.
|
||||
@@ -3645,21 +3660,6 @@ func (mr *MockStoreMockRecorder) GetDefaultChatModelConfig(ctx, organizationID a
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDefaultChatModelConfig", reflect.TypeOf((*MockStore)(nil).GetDefaultChatModelConfig), ctx, organizationID)
|
||||
}
|
||||
|
||||
// GetDefaultChatModelConfigCandidates mocks base method.
|
||||
func (m *MockStore) GetDefaultChatModelConfigCandidates(ctx context.Context) ([]database.ChatModelConfig, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetDefaultChatModelConfigCandidates", ctx)
|
||||
ret0, _ := ret[0].([]database.ChatModelConfig)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetDefaultChatModelConfigCandidates indicates an expected call of GetDefaultChatModelConfigCandidates.
|
||||
func (mr *MockStoreMockRecorder) GetDefaultChatModelConfigCandidates(ctx any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetDefaultChatModelConfigCandidates", reflect.TypeOf((*MockStore)(nil).GetDefaultChatModelConfigCandidates), ctx)
|
||||
}
|
||||
|
||||
// GetDefaultOrganization mocks base method.
|
||||
func (m *MockStore) GetDefaultOrganization(ctx context.Context) (database.Organization, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -3780,21 +3780,6 @@ func (mr *MockStoreMockRecorder) GetEnabledChatModelConfigByID(ctx, id any) *gom
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetEnabledChatModelConfigByID", reflect.TypeOf((*MockStore)(nil).GetEnabledChatModelConfigByID), ctx, id)
|
||||
}
|
||||
|
||||
// GetEnabledChatModelConfigs mocks base method.
|
||||
func (m *MockStore) GetEnabledChatModelConfigs(ctx context.Context) ([]database.GetEnabledChatModelConfigsRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetEnabledChatModelConfigs", ctx)
|
||||
ret0, _ := ret[0].([]database.GetEnabledChatModelConfigsRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetEnabledChatModelConfigs indicates an expected call of GetEnabledChatModelConfigs.
|
||||
func (mr *MockStoreMockRecorder) GetEnabledChatModelConfigs(ctx any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
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()
|
||||
|
||||
@@ -2234,8 +2234,9 @@ func TestPurgeChatDebugRuns(t *testing.T) {
|
||||
DisplayName: "OpenAI",
|
||||
})
|
||||
modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Model: "test-model",
|
||||
ContextLimit: 8192,
|
||||
Model: "test-model",
|
||||
ContextLimit: 8192,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
return chatDebugDeps{user: user, org: org, modelConfig: modelConfig}
|
||||
}
|
||||
@@ -2460,8 +2461,9 @@ func TestDeleteOldChatFiles(t *testing.T) {
|
||||
DisplayName: "OpenAI",
|
||||
})
|
||||
mc := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Model: "test-model",
|
||||
ContextLimit: 8192,
|
||||
Model: "test-model",
|
||||
ContextLimit: 8192,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
return chatDeps{user: user, org: org, modelConfig: mc}
|
||||
}
|
||||
@@ -3051,8 +3053,9 @@ func TestBackfillChatMessagesSearchTsv(t *testing.T) {
|
||||
DisplayName: "OpenAI",
|
||||
})
|
||||
modelConfig := dbgen.ChatModelConfig(t, db, database.ChatModelConfig{
|
||||
Model: "test-model",
|
||||
ContextLimit: 8192,
|
||||
Model: "test-model",
|
||||
ContextLimit: 8192,
|
||||
OrganizationID: org.ID,
|
||||
})
|
||||
chat := dbgen.Chat(t, db, database.Chat{
|
||||
OrganizationID: org.ID,
|
||||
|
||||
@@ -58,18 +58,10 @@ type customQuerier interface {
|
||||
}
|
||||
|
||||
type chatModelConfigQuerier interface {
|
||||
GetAuthorizedChatModelConfigs(ctx context.Context, prepared rbac.PreparedAuthorized) ([]ChatModelConfig, error)
|
||||
// GetDefaultChatModelConfigCandidates returns every non-deleted config in
|
||||
// every organization without an authorization filter. The caller selects
|
||||
// one organization's promotion candidate from this deployment-wide list.
|
||||
GetDefaultChatModelConfigCandidates(ctx context.Context) ([]ChatModelConfig, error)
|
||||
GetAuthorizedChatModelConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]ChatModelConfig, error)
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) GetDefaultChatModelConfigCandidates(ctx context.Context) ([]ChatModelConfig, error) {
|
||||
return q.GetChatModelConfigs(ctx)
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) GetAuthorizedChatModelConfigs(ctx context.Context, prepared rbac.PreparedAuthorized) ([]ChatModelConfig, error) {
|
||||
func (q *sqlQuerier) GetAuthorizedChatModelConfigs(ctx context.Context, organizationID uuid.UUID, prepared rbac.PreparedAuthorized) ([]ChatModelConfig, error) {
|
||||
authorizedFilter, err := prepared.CompileToSQL(ctx, rbac.ConfigChatModelConfigs())
|
||||
if err != nil {
|
||||
return nil, xerrors.Errorf("compile authorized filter: %w", err)
|
||||
@@ -82,7 +74,7 @@ func (q *sqlQuerier) GetAuthorizedChatModelConfigs(ctx context.Context, prepared
|
||||
|
||||
// The name comment is for metric tracking
|
||||
query := fmt.Sprintf("-- name: GetAuthorizedChatModelConfigs :many\n%s", filtered)
|
||||
rows, err := q.db.QueryContext(ctx, query)
|
||||
rows, err := q.db.QueryContext(ctx, query, organizationID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
Generated
+6
-2
@@ -482,7 +482,12 @@ type sqlcQuerier interface {
|
||||
// results remain after their assistant calls.
|
||||
GetChatMessagesForPromptByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatMessage, error)
|
||||
GetChatModelConfigByID(ctx context.Context, id uuid.UUID) (ChatModelConfig, error)
|
||||
GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error)
|
||||
GetChatModelConfigs(ctx context.Context, organizationID uuid.UUID) ([]ChatModelConfig, error)
|
||||
// All live configs in one organization, unfiltered. Consumed ONLY by
|
||||
// ensureDefaultChatModelConfig's default-election read inside the write
|
||||
// transaction; authorization is the caller's update-in-org check that every
|
||||
// path reaching the election already requires. No @authorize_filter.
|
||||
GetChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ChatModelConfig, error)
|
||||
// deleted = false guarantees ai_provider_id is non-null, so INNER JOIN is safe.
|
||||
GetChatModelConfigsForTelemetry(ctx context.Context) ([]GetChatModelConfigsForTelemetryRow, error)
|
||||
// GetChatPersonalModelOverridesEnabled returns whether users may configure
|
||||
@@ -562,7 +567,6 @@ type sqlcQuerier interface {
|
||||
// Providers can be disabled independently of their model configs.
|
||||
// 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)
|
||||
|
||||
@@ -12483,15 +12483,15 @@ func TestGetEnabledChatModelConfigsUsesAIProviders(t *testing.T) {
|
||||
params.Enabled = false
|
||||
})
|
||||
|
||||
configs, err := store.GetEnabledChatModelConfigs(ctx)
|
||||
configs, err := store.GetEnabledChatModelConfigsByOrganization(ctx, enabledConfig.OrganizationID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, slices.ContainsFunc(configs, func(row database.GetEnabledChatModelConfigsRow) bool {
|
||||
require.True(t, slices.ContainsFunc(configs, func(row database.GetEnabledChatModelConfigsByOrganizationRow) bool {
|
||||
return row.ChatModelConfig.ID == enabledConfig.ID
|
||||
}))
|
||||
require.False(t, slices.ContainsFunc(configs, func(row database.GetEnabledChatModelConfigsRow) bool {
|
||||
require.False(t, slices.ContainsFunc(configs, func(row database.GetEnabledChatModelConfigsByOrganizationRow) bool {
|
||||
return row.ChatModelConfig.ID == disabledProviderConfig.ID
|
||||
}))
|
||||
require.False(t, slices.ContainsFunc(configs, func(row database.GetEnabledChatModelConfigsRow) bool {
|
||||
require.False(t, slices.ContainsFunc(configs, func(row database.GetEnabledChatModelConfigsByOrganizationRow) bool {
|
||||
return row.ChatModelConfig.ID == disabledModelConfig.ID
|
||||
}))
|
||||
|
||||
|
||||
Generated
+63
-70
@@ -6200,6 +6200,7 @@ LEFT JOIN
|
||||
ai_providers ap ON ap.id = cmc.ai_provider_id
|
||||
WHERE
|
||||
cmc.deleted = FALSE
|
||||
AND cmc.organization_id = $1::uuid
|
||||
-- Authorize Filter clause will be injected below in GetAuthorizedChatModelConfigs
|
||||
-- @authorize_filter
|
||||
ORDER BY
|
||||
@@ -6209,8 +6210,68 @@ ORDER BY
|
||||
cmc.id DESC
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) GetChatModelConfigs(ctx context.Context) ([]ChatModelConfig, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getChatModelConfigs)
|
||||
func (q *sqlQuerier) GetChatModelConfigs(ctx context.Context, organizationID uuid.UUID) ([]ChatModelConfig, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getChatModelConfigs, organizationID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []ChatModelConfig
|
||||
for rows.Next() {
|
||||
var i ChatModelConfig
|
||||
if err := rows.Scan(
|
||||
&i.ID,
|
||||
&i.Model,
|
||||
&i.DisplayName,
|
||||
&i.CreatedBy,
|
||||
&i.UpdatedBy,
|
||||
&i.Enabled,
|
||||
&i.IsDefault,
|
||||
&i.Deleted,
|
||||
&i.DeletedAt,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
&i.ContextLimit,
|
||||
&i.CompressionThreshold,
|
||||
&i.Options,
|
||||
&i.AIProviderID,
|
||||
&i.OrganizationID,
|
||||
&i.GroupACL,
|
||||
&i.UserACL,
|
||||
); 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 getChatModelConfigsByOrganization = `-- name: GetChatModelConfigsByOrganization :many
|
||||
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, organization_id, group_acl, user_acl
|
||||
FROM
|
||||
chat_model_configs
|
||||
WHERE
|
||||
organization_id = $1::uuid
|
||||
AND deleted = FALSE
|
||||
ORDER BY
|
||||
model ASC,
|
||||
updated_at DESC,
|
||||
id DESC
|
||||
`
|
||||
|
||||
// All live configs in one organization, unfiltered. Consumed ONLY by
|
||||
// ensureDefaultChatModelConfig's default-election read inside the write
|
||||
// transaction; authorization is the caller's update-in-org check that every
|
||||
// path reaching the election already requires. No @authorize_filter.
|
||||
func (q *sqlQuerier) GetChatModelConfigsByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ChatModelConfig, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getChatModelConfigsByOrganization, organizationID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -6331,74 +6392,6 @@ func (q *sqlQuerier) GetEnabledChatModelConfigByID(ctx context.Context, id uuid.
|
||||
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.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.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 GetEnabledChatModelConfigsRow struct {
|
||||
ChatModelConfig ChatModelConfig `db:"chat_model_config" json:"chat_model_config"`
|
||||
Provider string `db:"provider" json:"provider"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) GetEnabledChatModelConfigs(ctx context.Context) ([]GetEnabledChatModelConfigsRow, error) {
|
||||
rows, err := q.db.QueryContext(ctx, getEnabledChatModelConfigs)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []GetEnabledChatModelConfigsRow
|
||||
for rows.Next() {
|
||||
var i GetEnabledChatModelConfigsRow
|
||||
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
|
||||
}
|
||||
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,
|
||||
|
||||
@@ -26,6 +26,7 @@ LEFT JOIN
|
||||
ai_providers ap ON ap.id = cmc.ai_provider_id
|
||||
WHERE
|
||||
cmc.deleted = FALSE
|
||||
AND cmc.organization_id = @organization_id::uuid
|
||||
-- Authorize Filter clause will be injected below in GetAuthorizedChatModelConfigs
|
||||
-- @authorize_filter
|
||||
ORDER BY
|
||||
@@ -34,24 +35,22 @@ ORDER BY
|
||||
cmc.updated_at DESC,
|
||||
cmc.id DESC;
|
||||
|
||||
-- name: GetEnabledChatModelConfigs :many
|
||||
-- name: GetChatModelConfigsByOrganization :many
|
||||
-- All live configs in one organization, unfiltered. Consumed ONLY by
|
||||
-- ensureDefaultChatModelConfig's default-election read inside the write
|
||||
-- transaction; authorization is the caller's update-in-org check that every
|
||||
-- path reaching the election already requires. No @authorize_filter.
|
||||
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
|
||||
chat_model_configs
|
||||
WHERE
|
||||
cmc.enabled = TRUE
|
||||
AND cmc.deleted = FALSE
|
||||
AND ap.enabled = TRUE
|
||||
AND ap.deleted = FALSE
|
||||
organization_id = @organization_id::uuid
|
||||
AND deleted = FALSE
|
||||
ORDER BY
|
||||
ap.type::text ASC,
|
||||
cmc.model ASC,
|
||||
cmc.updated_at DESC,
|
||||
cmc.id DESC;
|
||||
model ASC,
|
||||
updated_at DESC,
|
||||
id DESC;
|
||||
|
||||
-- name: GetEnabledChatModelConfigsByOrganization :many
|
||||
SELECT
|
||||
|
||||
Reference in New Issue
Block a user