mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
feat: remove native chat usage limits in favor of AI Gateway budgets (#27329)
## Stack Context This stack makes AI Gateway data and budgets the source of truth for AI spend controls. 1. Re-back the per-chat cost endpoint with AI Gateway data (#27328, merged). 2. **This PR:** remove native chat usage limits. 3. Remove native chat cost tracking and its dedicated admin UI (#27330). ## Summary Removes the native usage-limit API, SDK types, SQL, and chat enforcement for deployment, user, and group chat limits. Compact AI Gateway budget indicators remain in the Agents sidebar, user menu, and group settings. Gateway budget rejections and provider quota failures continue to classify as usage-limit errors, including a 409 response for synchronous title generation. Budget-period labels now use the API's UTC boundaries, so users see the same dates in every browser timezone. The documentation explains the AI Gateway replacement, its licensing requirements, and the differences from native limits. No schema is dropped in this release. The usage-limit table, index, user and group columns, constraints, audit mappings, and generated scan fields remain for mixed-version rolling upgrades. #27600 tracks their removal after the compatibility window. ## Breaking change Native day, week, and month chat spend limits are removed and are not migrated. AI Gateway budgets are month-based, group-scoped with per-user overrides, and require the AI Gateway entitlement. Deployments without that entitlement no longer have chat spend enforcement. > Mux prepared this PR on Mike's behalf.
This commit is contained in:
Michael Suchacz
@@ -2213,20 +2213,6 @@ func (q *querier) DeleteChatQueuedMessageReturningCount(ctx context.Context, arg
|
||||
return q.db.DeleteChatQueuedMessageReturningCount(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return err
|
||||
}
|
||||
return q.db.DeleteChatUsageLimitGroupOverride(ctx, groupID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return err
|
||||
}
|
||||
return q.db.DeleteChatUsageLimitUserOverride(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) DeleteCryptoKey(ctx context.Context, arg database.DeleteCryptoKeyParams) (database.CryptoKey, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionDelete, rbac.ResourceCryptoKey); err != nil {
|
||||
return database.CryptoKey{}, err
|
||||
@@ -3636,27 +3622,6 @@ func (q *querier) GetChatTitleGenerationModelOverride(ctx context.Context) (stri
|
||||
return q.db.GetChatTitleGenerationModelOverride(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatUsageLimitConfig(ctx context.Context) (database.ChatUsageLimitConfig, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return database.ChatUsageLimitConfig{}, err
|
||||
}
|
||||
return q.db.GetChatUsageLimitConfig(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (database.GetChatUsageLimitGroupOverrideRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return database.GetChatUsageLimitGroupOverrideRow{}, err
|
||||
}
|
||||
return q.db.GetChatUsageLimitGroupOverride(ctx, groupID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (database.GetChatUsageLimitUserOverrideRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return database.GetChatUsageLimitUserOverrideRow{}, err
|
||||
}
|
||||
return q.db.GetChatUsageLimitUserOverride(ctx, userID)
|
||||
}
|
||||
|
||||
func (q *querier) GetChatUserPromptsByChatID(ctx context.Context, arg database.GetChatUserPromptsByChatIDParams) ([]database.GetChatUserPromptsByChatIDRow, error) {
|
||||
// Authorize read on the parent chat.
|
||||
_, err := q.GetChatByID(ctx, arg.ChatID)
|
||||
@@ -5179,13 +5144,6 @@ func (q *querier) GetUserChatPersonalModelOverride(ctx context.Context, arg data
|
||||
return q.db.GetUserChatPersonalModelOverride(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserChatSpendInPeriod(ctx context.Context, arg database.GetUserChatSpendInPeriodParams) (int64, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return q.db.GetUserChatSpendInPeriod(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) {
|
||||
user, err := q.db.GetUserByID(ctx, userID)
|
||||
if err != nil {
|
||||
@@ -5216,13 +5174,6 @@ func (q *querier) GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid
|
||||
return fetchWithAction(q.log, q.auth, policy.ActionReadPersonal, q.db.GetUserForChatSyntheticAPIKeyByID)(ctx, id)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserGroupSpendLimit(ctx context.Context, arg database.GetUserGroupSpendLimitParams) (int64, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return q.db.GetUserGroupSpendLimit(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
|
||||
// Used by insights endpoints. Need to check both for auditors and for regular users with template acl perms.
|
||||
if err := q.authorizeContext(ctx, policy.ActionViewInsights, rbac.ResourceTemplate); err != nil {
|
||||
@@ -6908,20 +6859,6 @@ func (q *querier) ListChatContextResourcesByChatID(ctx context.Context, chatID u
|
||||
return q.db.ListChatContextResourcesByChatID(ctx, chatID)
|
||||
}
|
||||
|
||||
func (q *querier) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.ListChatUsageLimitGroupOverrides(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) ListChatUsageLimitOverrides(ctx context.Context) ([]database.ListChatUsageLimitOverridesRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return q.db.ListChatUsageLimitOverrides(ctx)
|
||||
}
|
||||
|
||||
func (q *querier) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) {
|
||||
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ListProvisionerKeysByOrganization)(ctx, organizationID)
|
||||
}
|
||||
@@ -7157,13 +7094,6 @@ func (q *querier) ReorderChatQueuedMessageToHead(ctx context.Context, arg databa
|
||||
return q.db.ReorderChatQueuedMessageToHead(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(arg.UserID.String())); err != nil {
|
||||
return database.ResolveUserChatSpendLimitRow{}, err
|
||||
}
|
||||
return q.db.ResolveUserChatSpendLimit(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
|
||||
return err
|
||||
@@ -9047,27 +8977,6 @@ func (q *querier) UpsertChatTitleGenerationModelOverride(ctx context.Context, va
|
||||
return q.db.UpsertChatTitleGenerationModelOverride(ctx, value)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertChatUsageLimitConfig(ctx context.Context, arg database.UpsertChatUsageLimitConfigParams) (database.ChatUsageLimitConfig, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return database.ChatUsageLimitConfig{}, err
|
||||
}
|
||||
return q.db.UpsertChatUsageLimitConfig(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg database.UpsertChatUsageLimitGroupOverrideParams) (database.UpsertChatUsageLimitGroupOverrideRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return database.UpsertChatUsageLimitGroupOverrideRow{}, err
|
||||
}
|
||||
return q.db.UpsertChatUsageLimitGroupOverride(ctx, arg)
|
||||
}
|
||||
|
||||
func (q *querier) UpsertChatUsageLimitUserOverride(ctx context.Context, arg database.UpsertChatUsageLimitUserOverrideParams) (database.UpsertChatUsageLimitUserOverrideRow, error) {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
return database.UpsertChatUsageLimitUserOverrideRow{}, err
|
||||
}
|
||||
return q.db.UpsertChatUsageLimitUserOverride(ctx, arg)
|
||||
}
|
||||
|
||||
//nolint:revive // Parameter name matches the generated querier interface.
|
||||
func (q *querier) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error {
|
||||
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
|
||||
|
||||
@@ -1688,152 +1688,6 @@ func (s *MethodTestSuite) TestChats() {
|
||||
dbm.EXPECT().UpsertChatWorkspaceTTL(gomock.Any(), "1h").Return(nil).AnyTimes()
|
||||
check.Args("1h").Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("GetUserChatSpendInPeriod", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.GetUserChatSpendInPeriodParams{
|
||||
UserID: uuid.New(),
|
||||
OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
|
||||
|
||||
StartTime: time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC),
|
||||
EndTime: time.Date(2025, 2, 1, 0, 0, 0, 0, time.UTC),
|
||||
}
|
||||
spend := int64(123)
|
||||
dbm.EXPECT().GetUserChatSpendInPeriod(gomock.Any(), arg).Return(spend, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(spend)
|
||||
}))
|
||||
s.Run("GetUserGroupSpendLimit", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.GetUserGroupSpendLimitParams{
|
||||
UserID: uuid.New(),
|
||||
OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
|
||||
}
|
||||
limit := int64(456)
|
||||
dbm.EXPECT().GetUserGroupSpendLimit(gomock.Any(), arg).Return(limit, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(limit)
|
||||
}))
|
||||
|
||||
s.Run("ResolveUserChatSpendLimit", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.ResolveUserChatSpendLimitParams{
|
||||
UserID: uuid.New(),
|
||||
OrganizationID: uuid.NullUUID{UUID: uuid.New(), Valid: true},
|
||||
}
|
||||
row := database.ResolveUserChatSpendLimitRow{EffectiveLimitMicros: 789, LimitSource: "group"}
|
||||
dbm.EXPECT().ResolveUserChatSpendLimit(gomock.Any(), arg).Return(row, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.UserID.String()), policy.ActionRead).Returns(row)
|
||||
}))
|
||||
|
||||
s.Run("GetChatUsageLimitConfig", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
now := dbtime.Now()
|
||||
config := database.ChatUsageLimitConfig{
|
||||
ID: 1,
|
||||
Singleton: true,
|
||||
Enabled: true,
|
||||
DefaultLimitMicros: 1_000_000,
|
||||
Period: "monthly",
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
dbm.EXPECT().GetChatUsageLimitConfig(gomock.Any()).Return(config, nil).AnyTimes()
|
||||
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(config)
|
||||
}))
|
||||
s.Run("GetChatUsageLimitGroupOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
groupID := uuid.New()
|
||||
override := database.GetChatUsageLimitGroupOverrideRow{
|
||||
GroupID: groupID,
|
||||
SpendLimitMicros: sql.NullInt64{Int64: 2_000_000, Valid: true},
|
||||
}
|
||||
dbm.EXPECT().GetChatUsageLimitGroupOverride(gomock.Any(), groupID).Return(override, nil).AnyTimes()
|
||||
check.Args(groupID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(override)
|
||||
}))
|
||||
s.Run("GetChatUsageLimitUserOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
userID := uuid.New()
|
||||
override := database.GetChatUsageLimitUserOverrideRow{
|
||||
UserID: userID,
|
||||
SpendLimitMicros: sql.NullInt64{Int64: 3_000_000, Valid: true},
|
||||
}
|
||||
dbm.EXPECT().GetChatUsageLimitUserOverride(gomock.Any(), userID).Return(override, nil).AnyTimes()
|
||||
check.Args(userID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(override)
|
||||
}))
|
||||
s.Run("ListChatUsageLimitGroupOverrides", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
overrides := []database.ListChatUsageLimitGroupOverridesRow{{
|
||||
GroupID: uuid.New(),
|
||||
GroupName: "group-name",
|
||||
GroupDisplayName: "Group Name",
|
||||
GroupAvatarUrl: "https://example.com/group.png",
|
||||
SpendLimitMicros: sql.NullInt64{Int64: 4_000_000, Valid: true},
|
||||
MemberCount: 5,
|
||||
}}
|
||||
dbm.EXPECT().ListChatUsageLimitGroupOverrides(gomock.Any()).Return(overrides, nil).AnyTimes()
|
||||
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(overrides)
|
||||
}))
|
||||
s.Run("ListChatUsageLimitOverrides", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
overrides := []database.ListChatUsageLimitOverridesRow{{
|
||||
UserID: uuid.New(),
|
||||
Username: "usage-limit-user",
|
||||
Name: "Usage Limit User",
|
||||
AvatarURL: "https://example.com/avatar.png",
|
||||
SpendLimitMicros: sql.NullInt64{Int64: 5_000_000, Valid: true},
|
||||
}}
|
||||
dbm.EXPECT().ListChatUsageLimitOverrides(gomock.Any()).Return(overrides, nil).AnyTimes()
|
||||
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(overrides)
|
||||
}))
|
||||
s.Run("UpsertChatUsageLimitConfig", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
now := dbtime.Now()
|
||||
arg := database.UpsertChatUsageLimitConfigParams{
|
||||
Enabled: true,
|
||||
DefaultLimitMicros: 6_000_000,
|
||||
Period: "monthly",
|
||||
}
|
||||
config := database.ChatUsageLimitConfig{
|
||||
ID: 1,
|
||||
Singleton: true,
|
||||
Enabled: arg.Enabled,
|
||||
DefaultLimitMicros: arg.DefaultLimitMicros,
|
||||
Period: arg.Period,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
}
|
||||
dbm.EXPECT().UpsertChatUsageLimitConfig(gomock.Any(), arg).Return(config, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(config)
|
||||
}))
|
||||
s.Run("UpsertChatUsageLimitGroupOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.UpsertChatUsageLimitGroupOverrideParams{
|
||||
SpendLimitMicros: 7_000_000,
|
||||
GroupID: uuid.New(),
|
||||
}
|
||||
override := database.UpsertChatUsageLimitGroupOverrideRow{
|
||||
GroupID: arg.GroupID,
|
||||
Name: "group",
|
||||
DisplayName: "Group",
|
||||
AvatarURL: "",
|
||||
SpendLimitMicros: sql.NullInt64{Int64: arg.SpendLimitMicros, Valid: true},
|
||||
}
|
||||
dbm.EXPECT().UpsertChatUsageLimitGroupOverride(gomock.Any(), arg).Return(override, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(override)
|
||||
}))
|
||||
s.Run("UpsertChatUsageLimitUserOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
arg := database.UpsertChatUsageLimitUserOverrideParams{
|
||||
SpendLimitMicros: 8_000_000,
|
||||
UserID: uuid.New(),
|
||||
}
|
||||
override := database.UpsertChatUsageLimitUserOverrideRow{
|
||||
UserID: arg.UserID,
|
||||
Username: "user",
|
||||
Name: "User",
|
||||
AvatarURL: "",
|
||||
SpendLimitMicros: sql.NullInt64{Int64: arg.SpendLimitMicros, Valid: true},
|
||||
}
|
||||
dbm.EXPECT().UpsertChatUsageLimitUserOverride(gomock.Any(), arg).Return(override, nil).AnyTimes()
|
||||
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(override)
|
||||
}))
|
||||
s.Run("DeleteChatUsageLimitGroupOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
groupID := uuid.New()
|
||||
dbm.EXPECT().DeleteChatUsageLimitGroupOverride(gomock.Any(), groupID).Return(nil).AnyTimes()
|
||||
check.Args(groupID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("DeleteChatUsageLimitUserOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
userID := uuid.New()
|
||||
dbm.EXPECT().DeleteChatUsageLimitUserOverride(gomock.Any(), userID).Return(nil).AnyTimes()
|
||||
check.Args(userID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
|
||||
}))
|
||||
s.Run("CleanupDeletedMCPServerIDsFromChats", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
|
||||
dbm.EXPECT().CleanupDeletedMCPServerIDsFromChats(gomock.Any()).Return(nil).AnyTimes()
|
||||
check.Args().Asserts(rbac.ResourceChat, policy.ActionUpdate)
|
||||
|
||||
-104
@@ -545,22 +545,6 @@ func (m queryMetricsStore) DeleteChatQueuedMessageReturningCount(ctx context.Con
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.DeleteChatUsageLimitGroupOverride(ctx, groupID)
|
||||
m.queryLatencies.WithLabelValues("DeleteChatUsageLimitGroupOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatUsageLimitGroupOverride").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.DeleteChatUsageLimitUserOverride(ctx, userID)
|
||||
m.queryLatencies.WithLabelValues("DeleteChatUsageLimitUserOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatUsageLimitUserOverride").Inc()
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) DeleteCryptoKey(ctx context.Context, arg database.DeleteCryptoKeyParams) (database.CryptoKey, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.DeleteCryptoKey(ctx, arg)
|
||||
@@ -1849,30 +1833,6 @@ func (m queryMetricsStore) GetChatTitleGenerationModelOverride(ctx context.Conte
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatUsageLimitConfig(ctx context.Context) (database.ChatUsageLimitConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatUsageLimitConfig(ctx)
|
||||
m.queryLatencies.WithLabelValues("GetChatUsageLimitConfig").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitConfig").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (database.GetChatUsageLimitGroupOverrideRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatUsageLimitGroupOverride(ctx, groupID)
|
||||
m.queryLatencies.WithLabelValues("GetChatUsageLimitGroupOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitGroupOverride").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (database.GetChatUsageLimitUserOverrideRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatUsageLimitUserOverride(ctx, userID)
|
||||
m.queryLatencies.WithLabelValues("GetChatUsageLimitUserOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitUserOverride").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetChatUserPromptsByChatID(ctx context.Context, arg database.GetChatUserPromptsByChatIDParams) ([]database.GetChatUserPromptsByChatIDRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetChatUserPromptsByChatID(ctx, arg)
|
||||
@@ -3337,14 +3297,6 @@ func (m queryMetricsStore) GetUserChatPersonalModelOverride(ctx context.Context,
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetUserChatSpendInPeriod(ctx context.Context, arg database.GetUserChatSpendInPeriodParams) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetUserChatSpendInPeriod(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("GetUserChatSpendInPeriod").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserChatSpendInPeriod").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetUserCodeDiffDisplayMode(ctx, userID)
|
||||
@@ -3377,14 +3329,6 @@ func (m queryMetricsStore) GetUserForChatSyntheticAPIKeyByID(ctx context.Context
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetUserGroupSpendLimit(ctx context.Context, userID database.GetUserGroupSpendLimitParams) (int64, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetUserGroupSpendLimit(ctx, userID)
|
||||
m.queryLatencies.WithLabelValues("GetUserGroupSpendLimit").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetUserGroupSpendLimit").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.GetUserLatencyInsights(ctx, arg)
|
||||
@@ -4849,22 +4793,6 @@ func (m queryMetricsStore) ListChatContextResourcesByChatID(ctx context.Context,
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ListChatUsageLimitGroupOverrides(ctx)
|
||||
m.queryLatencies.WithLabelValues("ListChatUsageLimitGroupOverrides").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListChatUsageLimitGroupOverrides").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ListChatUsageLimitOverrides(ctx context.Context) ([]database.ListChatUsageLimitOverridesRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ListChatUsageLimitOverrides(ctx)
|
||||
m.queryLatencies.WithLabelValues("ListChatUsageLimitOverrides").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListChatUsageLimitOverrides").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ListProvisionerKeysByOrganization(ctx, organizationID)
|
||||
@@ -5073,14 +5001,6 @@ func (m queryMetricsStore) ReorderChatQueuedMessageToHead(ctx context.Context, a
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.ResolveUserChatSpendLimit(ctx, userID)
|
||||
m.queryLatencies.WithLabelValues("ResolveUserChatSpendLimit").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ResolveUserChatSpendLimit").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.RevokeDBCryptKey(ctx, activeKeyDigest)
|
||||
@@ -6449,30 +6369,6 @@ func (m queryMetricsStore) UpsertChatTitleGenerationModelOverride(ctx context.Co
|
||||
return r0
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertChatUsageLimitConfig(ctx context.Context, arg database.UpsertChatUsageLimitConfigParams) (database.ChatUsageLimitConfig, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpsertChatUsageLimitConfig(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpsertChatUsageLimitConfig").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatUsageLimitConfig").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg database.UpsertChatUsageLimitGroupOverrideParams) (database.UpsertChatUsageLimitGroupOverrideRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpsertChatUsageLimitGroupOverride(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpsertChatUsageLimitGroupOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatUsageLimitGroupOverride").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertChatUsageLimitUserOverride(ctx context.Context, arg database.UpsertChatUsageLimitUserOverrideParams) (database.UpsertChatUsageLimitUserOverrideRow, error) {
|
||||
start := time.Now()
|
||||
r0, r1 := m.s.UpsertChatUsageLimitUserOverride(ctx, arg)
|
||||
m.queryLatencies.WithLabelValues("UpsertChatUsageLimitUserOverride").Observe(time.Since(start).Seconds())
|
||||
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatUsageLimitUserOverride").Inc()
|
||||
return r0, r1
|
||||
}
|
||||
|
||||
func (m queryMetricsStore) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error {
|
||||
start := time.Now()
|
||||
r0 := m.s.UpsertChatWorkspaceTTL(ctx, workspaceTtl)
|
||||
|
||||
Generated
-193
@@ -891,34 +891,6 @@ func (mr *MockStoreMockRecorder) DeleteChatQueuedMessageReturningCount(ctx, arg
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatQueuedMessageReturningCount", reflect.TypeOf((*MockStore)(nil).DeleteChatQueuedMessageReturningCount), ctx, arg)
|
||||
}
|
||||
|
||||
// DeleteChatUsageLimitGroupOverride mocks base method.
|
||||
func (m *MockStore) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteChatUsageLimitGroupOverride", ctx, groupID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteChatUsageLimitGroupOverride indicates an expected call of DeleteChatUsageLimitGroupOverride.
|
||||
func (mr *MockStoreMockRecorder) DeleteChatUsageLimitGroupOverride(ctx, groupID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatUsageLimitGroupOverride", reflect.TypeOf((*MockStore)(nil).DeleteChatUsageLimitGroupOverride), ctx, groupID)
|
||||
}
|
||||
|
||||
// DeleteChatUsageLimitUserOverride mocks base method.
|
||||
func (m *MockStore) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "DeleteChatUsageLimitUserOverride", ctx, userID)
|
||||
ret0, _ := ret[0].(error)
|
||||
return ret0
|
||||
}
|
||||
|
||||
// DeleteChatUsageLimitUserOverride indicates an expected call of DeleteChatUsageLimitUserOverride.
|
||||
func (mr *MockStoreMockRecorder) DeleteChatUsageLimitUserOverride(ctx, userID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatUsageLimitUserOverride", reflect.TypeOf((*MockStore)(nil).DeleteChatUsageLimitUserOverride), ctx, userID)
|
||||
}
|
||||
|
||||
// DeleteCryptoKey mocks base method.
|
||||
func (m *MockStore) DeleteCryptoKey(ctx context.Context, arg database.DeleteCryptoKeyParams) (database.CryptoKey, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -3418,51 +3390,6 @@ func (mr *MockStoreMockRecorder) GetChatTitleGenerationModelOverride(ctx any) *g
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatTitleGenerationModelOverride", reflect.TypeOf((*MockStore)(nil).GetChatTitleGenerationModelOverride), ctx)
|
||||
}
|
||||
|
||||
// GetChatUsageLimitConfig mocks base method.
|
||||
func (m *MockStore) GetChatUsageLimitConfig(ctx context.Context) (database.ChatUsageLimitConfig, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatUsageLimitConfig", ctx)
|
||||
ret0, _ := ret[0].(database.ChatUsageLimitConfig)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatUsageLimitConfig indicates an expected call of GetChatUsageLimitConfig.
|
||||
func (mr *MockStoreMockRecorder) GetChatUsageLimitConfig(ctx any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitConfig", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitConfig), ctx)
|
||||
}
|
||||
|
||||
// GetChatUsageLimitGroupOverride mocks base method.
|
||||
func (m *MockStore) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (database.GetChatUsageLimitGroupOverrideRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatUsageLimitGroupOverride", ctx, groupID)
|
||||
ret0, _ := ret[0].(database.GetChatUsageLimitGroupOverrideRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatUsageLimitGroupOverride indicates an expected call of GetChatUsageLimitGroupOverride.
|
||||
func (mr *MockStoreMockRecorder) GetChatUsageLimitGroupOverride(ctx, groupID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitGroupOverride", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitGroupOverride), ctx, groupID)
|
||||
}
|
||||
|
||||
// GetChatUsageLimitUserOverride mocks base method.
|
||||
func (m *MockStore) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (database.GetChatUsageLimitUserOverrideRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetChatUsageLimitUserOverride", ctx, userID)
|
||||
ret0, _ := ret[0].(database.GetChatUsageLimitUserOverrideRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetChatUsageLimitUserOverride indicates an expected call of GetChatUsageLimitUserOverride.
|
||||
func (mr *MockStoreMockRecorder) GetChatUsageLimitUserOverride(ctx, userID any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitUserOverride", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitUserOverride), ctx, userID)
|
||||
}
|
||||
|
||||
// GetChatUserPromptsByChatID mocks base method.
|
||||
func (m *MockStore) GetChatUserPromptsByChatID(ctx context.Context, arg database.GetChatUserPromptsByChatIDParams) ([]database.GetChatUserPromptsByChatIDRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -6238,21 +6165,6 @@ func (mr *MockStoreMockRecorder) GetUserChatPersonalModelOverride(ctx, arg any)
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserChatPersonalModelOverride", reflect.TypeOf((*MockStore)(nil).GetUserChatPersonalModelOverride), ctx, arg)
|
||||
}
|
||||
|
||||
// GetUserChatSpendInPeriod mocks base method.
|
||||
func (m *MockStore) GetUserChatSpendInPeriod(ctx context.Context, arg database.GetUserChatSpendInPeriodParams) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetUserChatSpendInPeriod", ctx, arg)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetUserChatSpendInPeriod indicates an expected call of GetUserChatSpendInPeriod.
|
||||
func (mr *MockStoreMockRecorder) GetUserChatSpendInPeriod(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserChatSpendInPeriod", reflect.TypeOf((*MockStore)(nil).GetUserChatSpendInPeriod), ctx, arg)
|
||||
}
|
||||
|
||||
// GetUserCodeDiffDisplayMode mocks base method.
|
||||
func (m *MockStore) GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -6313,21 +6225,6 @@ func (mr *MockStoreMockRecorder) GetUserForChatSyntheticAPIKeyByID(ctx, id any)
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserForChatSyntheticAPIKeyByID", reflect.TypeOf((*MockStore)(nil).GetUserForChatSyntheticAPIKeyByID), ctx, id)
|
||||
}
|
||||
|
||||
// GetUserGroupSpendLimit mocks base method.
|
||||
func (m *MockStore) GetUserGroupSpendLimit(ctx context.Context, arg database.GetUserGroupSpendLimitParams) (int64, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "GetUserGroupSpendLimit", ctx, arg)
|
||||
ret0, _ := ret[0].(int64)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// GetUserGroupSpendLimit indicates an expected call of GetUserGroupSpendLimit.
|
||||
func (mr *MockStoreMockRecorder) GetUserGroupSpendLimit(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserGroupSpendLimit", reflect.TypeOf((*MockStore)(nil).GetUserGroupSpendLimit), ctx, arg)
|
||||
}
|
||||
|
||||
// GetUserLatencyInsights mocks base method.
|
||||
func (m *MockStore) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -9131,36 +9028,6 @@ func (mr *MockStoreMockRecorder) ListChatContextResourcesByChatID(ctx, chatID an
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatContextResourcesByChatID", reflect.TypeOf((*MockStore)(nil).ListChatContextResourcesByChatID), ctx, chatID)
|
||||
}
|
||||
|
||||
// ListChatUsageLimitGroupOverrides mocks base method.
|
||||
func (m *MockStore) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]database.ListChatUsageLimitGroupOverridesRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ListChatUsageLimitGroupOverrides", ctx)
|
||||
ret0, _ := ret[0].([]database.ListChatUsageLimitGroupOverridesRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ListChatUsageLimitGroupOverrides indicates an expected call of ListChatUsageLimitGroupOverrides.
|
||||
func (mr *MockStoreMockRecorder) ListChatUsageLimitGroupOverrides(ctx any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatUsageLimitGroupOverrides", reflect.TypeOf((*MockStore)(nil).ListChatUsageLimitGroupOverrides), ctx)
|
||||
}
|
||||
|
||||
// ListChatUsageLimitOverrides mocks base method.
|
||||
func (m *MockStore) ListChatUsageLimitOverrides(ctx context.Context) ([]database.ListChatUsageLimitOverridesRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ListChatUsageLimitOverrides", ctx)
|
||||
ret0, _ := ret[0].([]database.ListChatUsageLimitOverridesRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ListChatUsageLimitOverrides indicates an expected call of ListChatUsageLimitOverrides.
|
||||
func (mr *MockStoreMockRecorder) ListChatUsageLimitOverrides(ctx any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListChatUsageLimitOverrides", reflect.TypeOf((*MockStore)(nil).ListChatUsageLimitOverrides), ctx)
|
||||
}
|
||||
|
||||
// ListProvisionerKeysByOrganization mocks base method.
|
||||
func (m *MockStore) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -9578,21 +9445,6 @@ func (mr *MockStoreMockRecorder) ReorderChatQueuedMessageToHead(ctx, arg any) *g
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ReorderChatQueuedMessageToHead", reflect.TypeOf((*MockStore)(nil).ReorderChatQueuedMessageToHead), ctx, arg)
|
||||
}
|
||||
|
||||
// ResolveUserChatSpendLimit mocks base method.
|
||||
func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, arg database.ResolveUserChatSpendLimitParams) (database.ResolveUserChatSpendLimitRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "ResolveUserChatSpendLimit", ctx, arg)
|
||||
ret0, _ := ret[0].(database.ResolveUserChatSpendLimitRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// ResolveUserChatSpendLimit indicates an expected call of ResolveUserChatSpendLimit.
|
||||
func (mr *MockStoreMockRecorder) ResolveUserChatSpendLimit(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResolveUserChatSpendLimit", reflect.TypeOf((*MockStore)(nil).ResolveUserChatSpendLimit), ctx, arg)
|
||||
}
|
||||
|
||||
// RevokeDBCryptKey mocks base method.
|
||||
func (m *MockStore) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
|
||||
m.ctrl.T.Helper()
|
||||
@@ -12070,51 +11922,6 @@ func (mr *MockStoreMockRecorder) UpsertChatTitleGenerationModelOverride(ctx, val
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatTitleGenerationModelOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatTitleGenerationModelOverride), ctx, value)
|
||||
}
|
||||
|
||||
// UpsertChatUsageLimitConfig mocks base method.
|
||||
func (m *MockStore) UpsertChatUsageLimitConfig(ctx context.Context, arg database.UpsertChatUsageLimitConfigParams) (database.ChatUsageLimitConfig, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpsertChatUsageLimitConfig", ctx, arg)
|
||||
ret0, _ := ret[0].(database.ChatUsageLimitConfig)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// UpsertChatUsageLimitConfig indicates an expected call of UpsertChatUsageLimitConfig.
|
||||
func (mr *MockStoreMockRecorder) UpsertChatUsageLimitConfig(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatUsageLimitConfig", reflect.TypeOf((*MockStore)(nil).UpsertChatUsageLimitConfig), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertChatUsageLimitGroupOverride mocks base method.
|
||||
func (m *MockStore) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg database.UpsertChatUsageLimitGroupOverrideParams) (database.UpsertChatUsageLimitGroupOverrideRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpsertChatUsageLimitGroupOverride", ctx, arg)
|
||||
ret0, _ := ret[0].(database.UpsertChatUsageLimitGroupOverrideRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// UpsertChatUsageLimitGroupOverride indicates an expected call of UpsertChatUsageLimitGroupOverride.
|
||||
func (mr *MockStoreMockRecorder) UpsertChatUsageLimitGroupOverride(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatUsageLimitGroupOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatUsageLimitGroupOverride), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertChatUsageLimitUserOverride mocks base method.
|
||||
func (m *MockStore) UpsertChatUsageLimitUserOverride(ctx context.Context, arg database.UpsertChatUsageLimitUserOverrideParams) (database.UpsertChatUsageLimitUserOverrideRow, error) {
|
||||
m.ctrl.T.Helper()
|
||||
ret := m.ctrl.Call(m, "UpsertChatUsageLimitUserOverride", ctx, arg)
|
||||
ret0, _ := ret[0].(database.UpsertChatUsageLimitUserOverrideRow)
|
||||
ret1, _ := ret[1].(error)
|
||||
return ret0, ret1
|
||||
}
|
||||
|
||||
// UpsertChatUsageLimitUserOverride indicates an expected call of UpsertChatUsageLimitUserOverride.
|
||||
func (mr *MockStoreMockRecorder) UpsertChatUsageLimitUserOverride(ctx, arg any) *gomock.Call {
|
||||
mr.mock.ctrl.T.Helper()
|
||||
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatUsageLimitUserOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatUsageLimitUserOverride), ctx, arg)
|
||||
}
|
||||
|
||||
// UpsertChatWorkspaceTTL mocks base method.
|
||||
func (m *MockStore) UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error {
|
||||
m.ctrl.T.Helper()
|
||||
|
||||
Generated
-33
@@ -154,8 +154,6 @@ type sqlcQuerier interface {
|
||||
// number of affected rows so callers can detect missing rows without
|
||||
// a follow-up read.
|
||||
DeleteChatQueuedMessageReturningCount(ctx context.Context, arg DeleteChatQueuedMessageReturningCountParams) (int64, error)
|
||||
DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error
|
||||
DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error
|
||||
DeleteCryptoKey(ctx context.Context, arg DeleteCryptoKeyParams) (CryptoKey, error)
|
||||
DeleteCustomRole(ctx context.Context, arg DeleteCustomRoleParams) error
|
||||
DeleteExpiredAPIKeys(ctx context.Context, arg DeleteExpiredAPIKeysParams) (int64, error)
|
||||
@@ -531,9 +529,6 @@ type sqlcQuerier interface {
|
||||
// Returns an empty string when no allowlist has been configured (all templates allowed).
|
||||
GetChatTemplateAllowlist(ctx context.Context) (string, error)
|
||||
GetChatTitleGenerationModelOverride(ctx context.Context) (string, error)
|
||||
GetChatUsageLimitConfig(ctx context.Context) (ChatUsageLimitConfig, error)
|
||||
GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (GetChatUsageLimitGroupOverrideRow, error)
|
||||
GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (GetChatUsageLimitUserOverrideRow, error)
|
||||
// Returns the concatenated text of each user-visible user prompt in a
|
||||
// chat, newest first. Used by the composer to populate the up/down
|
||||
// arrow prompt-history cycle. Non-text parts (tool calls, files,
|
||||
@@ -931,11 +926,6 @@ type sqlcQuerier interface {
|
||||
GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error)
|
||||
GetUserChatDebugLoggingEnabled(ctx context.Context, userID uuid.UUID) (bool, error)
|
||||
GetUserChatPersonalModelOverride(ctx context.Context, arg GetUserChatPersonalModelOverrideParams) (string, error)
|
||||
// Returns the total spend for a user in the given period.
|
||||
// When organization_id is NULL, spend across all organizations is
|
||||
// returned (global behavior). Otherwise only spend within the
|
||||
// specified organization is included.
|
||||
GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error)
|
||||
GetUserCodeDiffDisplayMode(ctx context.Context, userID uuid.UUID) (string, error)
|
||||
GetUserCount(ctx context.Context, includeSystem bool) (int64, error)
|
||||
// Returns the "Everyone" group (id == organization_id) to attribute a user's
|
||||
@@ -944,12 +934,6 @@ type sqlcQuerier interface {
|
||||
// no organization membership.
|
||||
GetUserEveryoneFallbackGroup(ctx context.Context, userID uuid.UUID) (uuid.UUID, error)
|
||||
GetUserForChatSyntheticAPIKeyByID(ctx context.Context, id uuid.UUID) (User, error)
|
||||
// Returns the minimum (most restrictive) group limit for a user.
|
||||
// Returns -1 if no group limits match the specified scope.
|
||||
// When organization_id is NULL, groups across all organizations are
|
||||
// considered (global behavior). Otherwise only groups within the
|
||||
// specified organization are considered.
|
||||
GetUserGroupSpendLimit(ctx context.Context, arg GetUserGroupSpendLimitParams) (int64, error)
|
||||
// GetUserLatencyInsights returns the median and 95th percentile connection
|
||||
// latency that users have experienced. The result can be filtered on
|
||||
// template_ids, meaning only user data from workspaces based on those templates
|
||||
@@ -1303,8 +1287,6 @@ type sqlcQuerier interface {
|
||||
// Lists a chat's pinned context resources, ordered deterministically by
|
||||
// source.
|
||||
ListChatContextResourcesByChatID(ctx context.Context, chatID uuid.UUID) ([]ChatContextResource, error)
|
||||
ListChatUsageLimitGroupOverrides(ctx context.Context) ([]ListChatUsageLimitGroupOverridesRow, error)
|
||||
ListChatUsageLimitOverrides(ctx context.Context) ([]ListChatUsageLimitOverridesRow, error)
|
||||
ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error)
|
||||
ListProvisionerKeysByOrganizationExcludeReserved(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error)
|
||||
ListTasks(ctx context.Context, arg ListTasksParams) ([]Task, error)
|
||||
@@ -1366,18 +1348,6 @@ type sqlcQuerier interface {
|
||||
// Sets the target queued message's position to one less than the
|
||||
// current minimum position for that chat, moving it to the head.
|
||||
ReorderChatQueuedMessageToHead(ctx context.Context, arg ReorderChatQueuedMessageToHeadParams) (int64, error)
|
||||
// Resolves the effective spend limit for a user using the hierarchy:
|
||||
// 1. Individual user override (highest priority, applies globally across
|
||||
// all organizations since it lives on the users table)
|
||||
// 2. Minimum group limit across the user's groups
|
||||
// 3. Global default from config
|
||||
// Returns -1 if limits are not enabled.
|
||||
// When organization_id is NULL, groups across all organizations are
|
||||
// considered (global behavior). Otherwise only groups within the
|
||||
// specified organization are considered.
|
||||
// limit_source indicates which tier won: 'user', 'group', 'default',
|
||||
// or 'disabled'.
|
||||
ResolveUserChatSpendLimit(ctx context.Context, arg ResolveUserChatSpendLimitParams) (ResolveUserChatSpendLimitRow, error)
|
||||
RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error
|
||||
// Note that this selects from the CTE, not the original table. The CTE is named
|
||||
// the same as the original table to trick sqlc into reusing the existing struct
|
||||
@@ -1697,9 +1667,6 @@ type sqlcQuerier interface {
|
||||
UpsertChatSystemPrompt(ctx context.Context, value string) error
|
||||
UpsertChatTemplateAllowlist(ctx context.Context, templateAllowlist string) error
|
||||
UpsertChatTitleGenerationModelOverride(ctx context.Context, value string) error
|
||||
UpsertChatUsageLimitConfig(ctx context.Context, arg UpsertChatUsageLimitConfigParams) (ChatUsageLimitConfig, error)
|
||||
UpsertChatUsageLimitGroupOverride(ctx context.Context, arg UpsertChatUsageLimitGroupOverrideParams) (UpsertChatUsageLimitGroupOverrideRow, error)
|
||||
UpsertChatUsageLimitUserOverride(ctx context.Context, arg UpsertChatUsageLimitUserOverrideParams) (UpsertChatUsageLimitUserOverrideRow, error)
|
||||
UpsertChatWorkspaceTTL(ctx context.Context, workspaceTtl string) error
|
||||
// The default proxy is implied and not actually stored in the database.
|
||||
// So we need to store it's configuration here for display purposes.
|
||||
|
||||
Generated
-390
@@ -7097,24 +7097,6 @@ func (q *sqlQuerier) DeleteChatQueuedMessageReturningCount(ctx context.Context,
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
const deleteChatUsageLimitGroupOverride = `-- name: DeleteChatUsageLimitGroupOverride :exec
|
||||
UPDATE groups SET chat_spend_limit_micros = NULL WHERE id = $1::uuid
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
|
||||
_, err := q.db.ExecContext(ctx, deleteChatUsageLimitGroupOverride, groupID)
|
||||
return err
|
||||
}
|
||||
|
||||
const deleteChatUsageLimitUserOverride = `-- name: DeleteChatUsageLimitUserOverride :exec
|
||||
UPDATE users SET chat_spend_limit_micros = NULL WHERE id = $1::uuid
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) DeleteChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) error {
|
||||
_, err := q.db.ExecContext(ctx, deleteChatUsageLimitUserOverride, userID)
|
||||
return err
|
||||
}
|
||||
|
||||
const deleteOldChats = `-- name: DeleteOldChats :execrows
|
||||
WITH deletable AS (
|
||||
SELECT id
|
||||
@@ -9119,61 +9101,6 @@ func (q *sqlQuerier) GetChatStreamSyncRows(ctx context.Context, ids []uuid.UUID)
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const getChatUsageLimitConfig = `-- name: GetChatUsageLimitConfig :one
|
||||
SELECT id, singleton, enabled, default_limit_micros, period, created_at, updated_at FROM chat_usage_limit_config WHERE singleton = TRUE LIMIT 1
|
||||
`
|
||||
|
||||
func (q *sqlQuerier) GetChatUsageLimitConfig(ctx context.Context) (ChatUsageLimitConfig, error) {
|
||||
row := q.db.QueryRowContext(ctx, getChatUsageLimitConfig)
|
||||
var i ChatUsageLimitConfig
|
||||
err := row.Scan(
|
||||
&i.ID,
|
||||
&i.Singleton,
|
||||
&i.Enabled,
|
||||
&i.DefaultLimitMicros,
|
||||
&i.Period,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getChatUsageLimitGroupOverride = `-- name: GetChatUsageLimitGroupOverride :one
|
||||
SELECT id AS group_id, chat_spend_limit_micros AS spend_limit_micros
|
||||
FROM groups
|
||||
WHERE id = $1::uuid AND chat_spend_limit_micros IS NOT NULL
|
||||
`
|
||||
|
||||
type GetChatUsageLimitGroupOverrideRow struct {
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) GetChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) (GetChatUsageLimitGroupOverrideRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, getChatUsageLimitGroupOverride, groupID)
|
||||
var i GetChatUsageLimitGroupOverrideRow
|
||||
err := row.Scan(&i.GroupID, &i.SpendLimitMicros)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getChatUsageLimitUserOverride = `-- name: GetChatUsageLimitUserOverride :one
|
||||
SELECT id AS user_id, chat_spend_limit_micros AS spend_limit_micros
|
||||
FROM users
|
||||
WHERE id = $1::uuid AND chat_spend_limit_micros IS NOT NULL
|
||||
`
|
||||
|
||||
type GetChatUsageLimitUserOverrideRow struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) GetChatUsageLimitUserOverride(ctx context.Context, userID uuid.UUID) (GetChatUsageLimitUserOverrideRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, getChatUsageLimitUserOverride, userID)
|
||||
var i GetChatUsageLimitUserOverrideRow
|
||||
err := row.Scan(&i.UserID, &i.SpendLimitMicros)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const getChatUserPromptsByChatID = `-- name: GetChatUserPromptsByChatID :many
|
||||
SELECT
|
||||
cm.id,
|
||||
@@ -10346,68 +10273,6 @@ func (q *sqlQuerier) GetTotalChatMessageRuntimeMsInRange(ctx context.Context, ar
|
||||
return total_runtime_ms, err
|
||||
}
|
||||
|
||||
const getUserChatSpendInPeriod = `-- name: GetUserChatSpendInPeriod :one
|
||||
SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros
|
||||
FROM chat_messages cm
|
||||
JOIN chats c ON c.id = cm.chat_id
|
||||
WHERE c.owner_id = $1::uuid
|
||||
AND ($2::uuid IS NULL
|
||||
OR c.organization_id = $2::uuid)
|
||||
AND cm.created_at >= $3::timestamptz
|
||||
AND cm.created_at < $4::timestamptz
|
||||
AND cm.total_cost_micros IS NOT NULL
|
||||
`
|
||||
|
||||
type GetUserChatSpendInPeriodParams struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"`
|
||||
StartTime time.Time `db:"start_time" json:"start_time"`
|
||||
EndTime time.Time `db:"end_time" json:"end_time"`
|
||||
}
|
||||
|
||||
// Returns the total spend for a user in the given period.
|
||||
// When organization_id is NULL, spend across all organizations is
|
||||
// returned (global behavior). Otherwise only spend within the
|
||||
// specified organization is included.
|
||||
func (q *sqlQuerier) GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error) {
|
||||
row := q.db.QueryRowContext(ctx, getUserChatSpendInPeriod,
|
||||
arg.UserID,
|
||||
arg.OrganizationID,
|
||||
arg.StartTime,
|
||||
arg.EndTime,
|
||||
)
|
||||
var total_spend_micros int64
|
||||
err := row.Scan(&total_spend_micros)
|
||||
return total_spend_micros, err
|
||||
}
|
||||
|
||||
const getUserGroupSpendLimit = `-- name: GetUserGroupSpendLimit :one
|
||||
SELECT COALESCE(MIN(g.chat_spend_limit_micros), -1)::bigint AS limit_micros
|
||||
FROM groups g
|
||||
JOIN group_members_expanded gme ON gme.group_id = g.id
|
||||
WHERE gme.user_id = $1::uuid
|
||||
AND ($2::uuid IS NULL
|
||||
OR g.organization_id = $2::uuid)
|
||||
AND g.chat_spend_limit_micros IS NOT NULL
|
||||
`
|
||||
|
||||
type GetUserGroupSpendLimitParams struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"`
|
||||
}
|
||||
|
||||
// Returns the minimum (most restrictive) group limit for a user.
|
||||
// Returns -1 if no group limits match the specified scope.
|
||||
// When organization_id is NULL, groups across all organizations are
|
||||
// considered (global behavior). Otherwise only groups within the
|
||||
// specified organization are considered.
|
||||
func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, arg GetUserGroupSpendLimitParams) (int64, error) {
|
||||
row := q.db.QueryRowContext(ctx, getUserGroupSpendLimit, arg.UserID, arg.OrganizationID)
|
||||
var limit_micros int64
|
||||
err := row.Scan(&limit_micros)
|
||||
return limit_micros, err
|
||||
}
|
||||
|
||||
const hydrateAgentChatsContext = `-- name: HydrateAgentChatsContext :many
|
||||
WITH hydrated AS (
|
||||
UPDATE chats
|
||||
@@ -11138,106 +11003,6 @@ func (q *sqlQuerier) ListChatContextResourcesByChatID(ctx context.Context, chatI
|
||||
return items, nil
|
||||
}
|
||||
|
||||
const listChatUsageLimitGroupOverrides = `-- name: ListChatUsageLimitGroupOverrides :many
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
g.name AS group_name,
|
||||
g.display_name AS group_display_name,
|
||||
g.avatar_url AS group_avatar_url,
|
||||
g.chat_spend_limit_micros AS spend_limit_micros,
|
||||
(SELECT COUNT(*)
|
||||
FROM group_members_expanded gme
|
||||
WHERE gme.group_id = g.id
|
||||
AND gme.user_is_system = FALSE) AS member_count
|
||||
FROM groups g
|
||||
WHERE g.chat_spend_limit_micros IS NOT NULL
|
||||
ORDER BY g.name ASC
|
||||
`
|
||||
|
||||
type ListChatUsageLimitGroupOverridesRow struct {
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
GroupName string `db:"group_name" json:"group_name"`
|
||||
GroupDisplayName string `db:"group_display_name" json:"group_display_name"`
|
||||
GroupAvatarUrl string `db:"group_avatar_url" json:"group_avatar_url"`
|
||||
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
MemberCount int64 `db:"member_count" json:"member_count"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) ListChatUsageLimitGroupOverrides(ctx context.Context) ([]ListChatUsageLimitGroupOverridesRow, error) {
|
||||
rows, err := q.db.QueryContext(ctx, listChatUsageLimitGroupOverrides)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []ListChatUsageLimitGroupOverridesRow
|
||||
for rows.Next() {
|
||||
var i ListChatUsageLimitGroupOverridesRow
|
||||
if err := rows.Scan(
|
||||
&i.GroupID,
|
||||
&i.GroupName,
|
||||
&i.GroupDisplayName,
|
||||
&i.GroupAvatarUrl,
|
||||
&i.SpendLimitMicros,
|
||||
&i.MemberCount,
|
||||
); 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 listChatUsageLimitOverrides = `-- name: ListChatUsageLimitOverrides :many
|
||||
SELECT u.id AS user_id, u.username, u.name, u.avatar_url,
|
||||
u.chat_spend_limit_micros AS spend_limit_micros
|
||||
FROM users u
|
||||
WHERE u.chat_spend_limit_micros IS NOT NULL
|
||||
ORDER BY u.username ASC
|
||||
`
|
||||
|
||||
type ListChatUsageLimitOverridesRow struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
Username string `db:"username" json:"username"`
|
||||
Name string `db:"name" json:"name"`
|
||||
AvatarURL string `db:"avatar_url" json:"avatar_url"`
|
||||
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) ListChatUsageLimitOverrides(ctx context.Context) ([]ListChatUsageLimitOverridesRow, error) {
|
||||
rows, err := q.db.QueryContext(ctx, listChatUsageLimitOverrides)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
var items []ListChatUsageLimitOverridesRow
|
||||
for rows.Next() {
|
||||
var i ListChatUsageLimitOverridesRow
|
||||
if err := rows.Scan(
|
||||
&i.UserID,
|
||||
&i.Username,
|
||||
&i.Name,
|
||||
&i.AvatarURL,
|
||||
&i.SpendLimitMicros,
|
||||
); 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 lockChatAndBumpSnapshotVersion = `-- name: LockChatAndBumpSnapshotVersion :one
|
||||
WITH bumped_chat AS (
|
||||
UPDATE chats
|
||||
@@ -11558,63 +11323,6 @@ func (q *sqlQuerier) ReorderChatQueuedMessageToHead(ctx context.Context, arg Reo
|
||||
return result.RowsAffected()
|
||||
}
|
||||
|
||||
const resolveUserChatSpendLimit = `-- name: ResolveUserChatSpendLimit :one
|
||||
SELECT CASE
|
||||
WHEN NOT cfg.enabled THEN -1
|
||||
WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros
|
||||
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
|
||||
ELSE cfg.default_limit_micros
|
||||
END::bigint AS effective_limit_micros,
|
||||
CASE
|
||||
WHEN NOT cfg.enabled THEN 'disabled'
|
||||
WHEN u.chat_spend_limit_micros IS NOT NULL THEN 'user'
|
||||
WHEN gl.limit_micros IS NOT NULL THEN 'group'
|
||||
ELSE 'default'
|
||||
END AS limit_source
|
||||
FROM chat_usage_limit_config cfg
|
||||
CROSS JOIN users u
|
||||
LEFT JOIN LATERAL (
|
||||
SELECT MIN(g.chat_spend_limit_micros) AS limit_micros
|
||||
FROM groups g
|
||||
JOIN group_members_expanded gme ON gme.group_id = g.id
|
||||
WHERE gme.user_id = $1::uuid
|
||||
AND ($2::uuid IS NULL
|
||||
OR g.organization_id = $2::uuid)
|
||||
AND g.chat_spend_limit_micros IS NOT NULL
|
||||
) gl ON TRUE
|
||||
WHERE u.id = $1::uuid
|
||||
LIMIT 1
|
||||
`
|
||||
|
||||
type ResolveUserChatSpendLimitParams struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
OrganizationID uuid.NullUUID `db:"organization_id" json:"organization_id"`
|
||||
}
|
||||
|
||||
type ResolveUserChatSpendLimitRow struct {
|
||||
EffectiveLimitMicros int64 `db:"effective_limit_micros" json:"effective_limit_micros"`
|
||||
LimitSource string `db:"limit_source" json:"limit_source"`
|
||||
}
|
||||
|
||||
// Resolves the effective spend limit for a user using the hierarchy:
|
||||
// 1. Individual user override (highest priority, applies globally across
|
||||
// all organizations since it lives on the users table)
|
||||
// 2. Minimum group limit across the user's groups
|
||||
// 3. Global default from config
|
||||
//
|
||||
// Returns -1 if limits are not enabled.
|
||||
// When organization_id is NULL, groups across all organizations are
|
||||
// considered (global behavior). Otherwise only groups within the
|
||||
// specified organization are considered.
|
||||
// limit_source indicates which tier won: 'user', 'group', 'default',
|
||||
// or 'disabled'.
|
||||
func (q *sqlQuerier) ResolveUserChatSpendLimit(ctx context.Context, arg ResolveUserChatSpendLimitParams) (ResolveUserChatSpendLimitRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, resolveUserChatSpendLimit, arg.UserID, arg.OrganizationID)
|
||||
var i ResolveUserChatSpendLimitRow
|
||||
err := row.Scan(&i.EffectiveLimitMicros, &i.LimitSource)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const setChatContextSnapshot = `-- name: SetChatContextSnapshot :exec
|
||||
UPDATE chats
|
||||
SET
|
||||
@@ -13823,104 +13531,6 @@ func (q *sqlQuerier) UpsertChatHeartbeat(ctx context.Context, arg UpsertChatHear
|
||||
return err
|
||||
}
|
||||
|
||||
const upsertChatUsageLimitConfig = `-- name: UpsertChatUsageLimitConfig :one
|
||||
INSERT INTO chat_usage_limit_config (singleton, enabled, default_limit_micros, period, updated_at)
|
||||
VALUES (TRUE, $1::boolean, $2::bigint, $3::text, NOW())
|
||||
ON CONFLICT (singleton) DO UPDATE SET
|
||||
enabled = EXCLUDED.enabled,
|
||||
default_limit_micros = EXCLUDED.default_limit_micros,
|
||||
period = EXCLUDED.period,
|
||||
updated_at = NOW()
|
||||
RETURNING id, singleton, enabled, default_limit_micros, period, created_at, updated_at
|
||||
`
|
||||
|
||||
type UpsertChatUsageLimitConfigParams struct {
|
||||
Enabled bool `db:"enabled" json:"enabled"`
|
||||
DefaultLimitMicros int64 `db:"default_limit_micros" json:"default_limit_micros"`
|
||||
Period string `db:"period" json:"period"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpsertChatUsageLimitConfig(ctx context.Context, arg UpsertChatUsageLimitConfigParams) (ChatUsageLimitConfig, error) {
|
||||
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitConfig, arg.Enabled, arg.DefaultLimitMicros, arg.Period)
|
||||
var i ChatUsageLimitConfig
|
||||
err := row.Scan(
|
||||
&i.ID,
|
||||
&i.Singleton,
|
||||
&i.Enabled,
|
||||
&i.DefaultLimitMicros,
|
||||
&i.Period,
|
||||
&i.CreatedAt,
|
||||
&i.UpdatedAt,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const upsertChatUsageLimitGroupOverride = `-- name: UpsertChatUsageLimitGroupOverride :one
|
||||
UPDATE groups
|
||||
SET chat_spend_limit_micros = $1::bigint
|
||||
WHERE id = $2::uuid
|
||||
RETURNING id AS group_id, name, display_name, avatar_url, chat_spend_limit_micros AS spend_limit_micros
|
||||
`
|
||||
|
||||
type UpsertChatUsageLimitGroupOverrideParams struct {
|
||||
SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
}
|
||||
|
||||
type UpsertChatUsageLimitGroupOverrideRow struct {
|
||||
GroupID uuid.UUID `db:"group_id" json:"group_id"`
|
||||
Name string `db:"name" json:"name"`
|
||||
DisplayName string `db:"display_name" json:"display_name"`
|
||||
AvatarURL string `db:"avatar_url" json:"avatar_url"`
|
||||
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg UpsertChatUsageLimitGroupOverrideParams) (UpsertChatUsageLimitGroupOverrideRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitGroupOverride, arg.SpendLimitMicros, arg.GroupID)
|
||||
var i UpsertChatUsageLimitGroupOverrideRow
|
||||
err := row.Scan(
|
||||
&i.GroupID,
|
||||
&i.Name,
|
||||
&i.DisplayName,
|
||||
&i.AvatarURL,
|
||||
&i.SpendLimitMicros,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const upsertChatUsageLimitUserOverride = `-- name: UpsertChatUsageLimitUserOverride :one
|
||||
UPDATE users
|
||||
SET chat_spend_limit_micros = $1::bigint
|
||||
WHERE id = $2::uuid
|
||||
RETURNING id AS user_id, username, name, avatar_url, chat_spend_limit_micros AS spend_limit_micros
|
||||
`
|
||||
|
||||
type UpsertChatUsageLimitUserOverrideParams struct {
|
||||
SpendLimitMicros int64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
}
|
||||
|
||||
type UpsertChatUsageLimitUserOverrideRow struct {
|
||||
UserID uuid.UUID `db:"user_id" json:"user_id"`
|
||||
Username string `db:"username" json:"username"`
|
||||
Name string `db:"name" json:"name"`
|
||||
AvatarURL string `db:"avatar_url" json:"avatar_url"`
|
||||
SpendLimitMicros sql.NullInt64 `db:"spend_limit_micros" json:"spend_limit_micros"`
|
||||
}
|
||||
|
||||
func (q *sqlQuerier) UpsertChatUsageLimitUserOverride(ctx context.Context, arg UpsertChatUsageLimitUserOverrideParams) (UpsertChatUsageLimitUserOverrideRow, error) {
|
||||
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitUserOverride, arg.SpendLimitMicros, arg.UserID)
|
||||
var i UpsertChatUsageLimitUserOverrideRow
|
||||
err := row.Scan(
|
||||
&i.UserID,
|
||||
&i.Username,
|
||||
&i.Name,
|
||||
&i.AvatarURL,
|
||||
&i.SpendLimitMicros,
|
||||
)
|
||||
return i, err
|
||||
}
|
||||
|
||||
const batchUpsertConnectionLogs = `-- name: BatchUpsertConnectionLogs :exec
|
||||
INSERT INTO connection_logs (
|
||||
id, connect_time, organization_id, workspace_owner_id, workspace_id,
|
||||
|
||||
@@ -2434,55 +2434,6 @@ LIMIT
|
||||
OFFSET
|
||||
sqlc.arg('page_offset')::int;
|
||||
|
||||
-- name: GetChatUsageLimitConfig :one
|
||||
SELECT * FROM chat_usage_limit_config WHERE singleton = TRUE LIMIT 1;
|
||||
|
||||
-- name: UpsertChatUsageLimitConfig :one
|
||||
INSERT INTO chat_usage_limit_config (singleton, enabled, default_limit_micros, period, updated_at)
|
||||
VALUES (TRUE, @enabled::boolean, @default_limit_micros::bigint, @period::text, NOW())
|
||||
ON CONFLICT (singleton) DO UPDATE SET
|
||||
enabled = EXCLUDED.enabled,
|
||||
default_limit_micros = EXCLUDED.default_limit_micros,
|
||||
period = EXCLUDED.period,
|
||||
updated_at = NOW()
|
||||
RETURNING *;
|
||||
|
||||
-- name: ListChatUsageLimitOverrides :many
|
||||
SELECT u.id AS user_id, u.username, u.name, u.avatar_url,
|
||||
u.chat_spend_limit_micros AS spend_limit_micros
|
||||
FROM users u
|
||||
WHERE u.chat_spend_limit_micros IS NOT NULL
|
||||
ORDER BY u.username ASC;
|
||||
|
||||
-- name: UpsertChatUsageLimitUserOverride :one
|
||||
UPDATE users
|
||||
SET chat_spend_limit_micros = @spend_limit_micros::bigint
|
||||
WHERE id = @user_id::uuid
|
||||
RETURNING id AS user_id, username, name, avatar_url, chat_spend_limit_micros AS spend_limit_micros;
|
||||
|
||||
-- name: DeleteChatUsageLimitUserOverride :exec
|
||||
UPDATE users SET chat_spend_limit_micros = NULL WHERE id = @user_id::uuid;
|
||||
|
||||
-- name: GetChatUsageLimitUserOverride :one
|
||||
SELECT id AS user_id, chat_spend_limit_micros AS spend_limit_micros
|
||||
FROM users
|
||||
WHERE id = @user_id::uuid AND chat_spend_limit_micros IS NOT NULL;
|
||||
|
||||
-- name: GetUserChatSpendInPeriod :one
|
||||
-- Returns the total spend for a user in the given period.
|
||||
-- When organization_id is NULL, spend across all organizations is
|
||||
-- returned (global behavior). Otherwise only spend within the
|
||||
-- specified organization is included.
|
||||
SELECT COALESCE(SUM(cm.total_cost_micros), 0)::bigint AS total_spend_micros
|
||||
FROM chat_messages cm
|
||||
JOIN chats c ON c.id = cm.chat_id
|
||||
WHERE c.owner_id = @user_id::uuid
|
||||
AND (sqlc.narg('organization_id')::uuid IS NULL
|
||||
OR c.organization_id = sqlc.narg('organization_id')::uuid)
|
||||
AND cm.created_at >= @start_time::timestamptz
|
||||
AND cm.created_at < @end_time::timestamptz
|
||||
AND cm.total_cost_micros IS NOT NULL;
|
||||
|
||||
-- name: GetTotalChatMessageRuntimeMsInRange :one
|
||||
-- Computes hb_agent_runtime_v1 usage event payloads. Deliberately includes
|
||||
-- soft-deleted messages and messages from all chats.
|
||||
@@ -2508,49 +2459,6 @@ WHERE enabled = TRUE
|
||||
)
|
||||
);
|
||||
|
||||
-- name: ListChatUsageLimitGroupOverrides :many
|
||||
SELECT
|
||||
g.id AS group_id,
|
||||
g.name AS group_name,
|
||||
g.display_name AS group_display_name,
|
||||
g.avatar_url AS group_avatar_url,
|
||||
g.chat_spend_limit_micros AS spend_limit_micros,
|
||||
(SELECT COUNT(*)
|
||||
FROM group_members_expanded gme
|
||||
WHERE gme.group_id = g.id
|
||||
AND gme.user_is_system = FALSE) AS member_count
|
||||
FROM groups g
|
||||
WHERE g.chat_spend_limit_micros IS NOT NULL
|
||||
ORDER BY g.name ASC;
|
||||
|
||||
-- name: UpsertChatUsageLimitGroupOverride :one
|
||||
UPDATE groups
|
||||
SET chat_spend_limit_micros = @spend_limit_micros::bigint
|
||||
WHERE id = @group_id::uuid
|
||||
RETURNING id AS group_id, name, display_name, avatar_url, chat_spend_limit_micros AS spend_limit_micros;
|
||||
|
||||
-- name: DeleteChatUsageLimitGroupOverride :exec
|
||||
UPDATE groups SET chat_spend_limit_micros = NULL WHERE id = @group_id::uuid;
|
||||
|
||||
-- name: GetChatUsageLimitGroupOverride :one
|
||||
SELECT id AS group_id, chat_spend_limit_micros AS spend_limit_micros
|
||||
FROM groups
|
||||
WHERE id = @group_id::uuid AND chat_spend_limit_micros IS NOT NULL;
|
||||
|
||||
-- name: GetUserGroupSpendLimit :one
|
||||
-- Returns the minimum (most restrictive) group limit for a user.
|
||||
-- Returns -1 if no group limits match the specified scope.
|
||||
-- When organization_id is NULL, groups across all organizations are
|
||||
-- considered (global behavior). Otherwise only groups within the
|
||||
-- specified organization are considered.
|
||||
SELECT COALESCE(MIN(g.chat_spend_limit_micros), -1)::bigint AS limit_micros
|
||||
FROM groups g
|
||||
JOIN group_members_expanded gme ON gme.group_id = g.id
|
||||
WHERE gme.user_id = @user_id::uuid
|
||||
AND (sqlc.narg('organization_id')::uuid IS NULL
|
||||
OR g.organization_id = sqlc.narg('organization_id')::uuid)
|
||||
AND g.chat_spend_limit_micros IS NOT NULL;
|
||||
|
||||
-- name: GetChatsByWorkspaceIDs :many
|
||||
SELECT *
|
||||
FROM chats_expanded
|
||||
@@ -2558,44 +2466,6 @@ WHERE archived = false
|
||||
AND workspace_id = ANY(@ids::uuid[])
|
||||
ORDER BY workspace_id, updated_at DESC;
|
||||
|
||||
-- name: ResolveUserChatSpendLimit :one
|
||||
-- Resolves the effective spend limit for a user using the hierarchy:
|
||||
-- 1. Individual user override (highest priority, applies globally across
|
||||
-- all organizations since it lives on the users table)
|
||||
-- 2. Minimum group limit across the user's groups
|
||||
-- 3. Global default from config
|
||||
-- Returns -1 if limits are not enabled.
|
||||
-- When organization_id is NULL, groups across all organizations are
|
||||
-- considered (global behavior). Otherwise only groups within the
|
||||
-- specified organization are considered.
|
||||
-- limit_source indicates which tier won: 'user', 'group', 'default',
|
||||
-- or 'disabled'.
|
||||
SELECT CASE
|
||||
WHEN NOT cfg.enabled THEN -1
|
||||
WHEN u.chat_spend_limit_micros IS NOT NULL THEN u.chat_spend_limit_micros
|
||||
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
|
||||
ELSE cfg.default_limit_micros
|
||||
END::bigint AS effective_limit_micros,
|
||||
CASE
|
||||
WHEN NOT cfg.enabled THEN 'disabled'
|
||||
WHEN u.chat_spend_limit_micros IS NOT NULL THEN 'user'
|
||||
WHEN gl.limit_micros IS NOT NULL THEN 'group'
|
||||
ELSE 'default'
|
||||
END AS limit_source
|
||||
FROM chat_usage_limit_config cfg
|
||||
CROSS JOIN users u
|
||||
LEFT JOIN LATERAL (
|
||||
SELECT MIN(g.chat_spend_limit_micros) AS limit_micros
|
||||
FROM groups g
|
||||
JOIN group_members_expanded gme ON gme.group_id = g.id
|
||||
WHERE gme.user_id = @user_id::uuid
|
||||
AND (sqlc.narg('organization_id')::uuid IS NULL
|
||||
OR g.organization_id = sqlc.narg('organization_id')::uuid)
|
||||
AND g.chat_spend_limit_micros IS NOT NULL
|
||||
) gl ON TRUE
|
||||
WHERE u.id = @user_id::uuid
|
||||
LIMIT 1;
|
||||
|
||||
-- name: UpdateChatLastReadMessageID :exec
|
||||
-- Updates the last read message ID for a chat. This is used to track
|
||||
-- which messages the owner has seen, enabling unread indicators.
|
||||
|
||||
Reference in New Issue
Block a user