Compare commits

...

24 Commits

Author SHA1 Message Date
Michael Suchacz 713facdad5 test(dbauthz): add accounting entry for CountEnabledModelsWithoutPricing 2026-03-14 09:29:25 +00:00
Michael Suchacz 9a70780820 test(migrations): add fixture for 000438 chat usage limits 2026-03-14 09:29:25 +00:00
Michael Suchacz da3aca6d52 test(dbauthz): add accounting entries for chat usage limit methods 2026-03-14 09:29:25 +00:00
Michael Suchacz 8c15ad5950 fix: renumber chat usage limit migration after rebase 2026-03-14 08:32:44 +00:00
Michael Suchacz b3ad420db2 fix(site): use design system form components in LimitsTab 2026-03-14 08:28:04 +00:00
Michael Suchacz 6369ba1cf8 fix(site): use group autocomplete in limits tab 2026-03-14 08:28:04 +00:00
Michael Suchacz 56e369b0d6 refactor(agents): extract limits tab 2026-03-14 08:28:04 +00:00
Michael Suchacz 6fbbd155e7 fix(site): display usage limit error to chat users 2026-03-14 08:28:04 +00:00
Michael Suchacz 71dca287f9 fix(chatd): address RBAC, errname, and docs review feedback 2026-03-14 08:28:04 +00:00
Michael Suchacz 8156ef36c9 docs(chatd): add multi-period architecture note 2026-03-14 08:28:04 +00:00
Michael Suchacz 796f4ac65d fix(chat): address usage limit review feedback 2026-03-14 08:28:04 +00:00
Michael Suchacz c3ae8228ef refactor(chat): consolidate usage limit resolution 2026-03-14 08:28:04 +00:00
Michael Suchacz 095d18dd10 feat(site): add limits tab story 2026-03-14 08:28:04 +00:00
Michael Suchacz f470475e45 feat(agents): add group spend limit overrides 2026-03-14 08:28:03 +00:00
Michael Suchacz 01010fbafc feat(site): add group override API methods and query mutations 2026-03-14 08:28:03 +00:00
Michael Suchacz fb385c60d0 feat(chat): add group usage limit overrides 2026-03-14 08:28:03 +00:00
Michael Suchacz 2bfc2eb6c7 feat(chats): add group usage limit overrides 2026-03-14 08:28:03 +00:00
Michael Suchacz efcf39e98c fix: remove swagger annotations from experimental chat endpoints
Experimental chat endpoints should not appear in generated API docs.
Follow the existing pattern of using only an EXPERIMENTAL comment.
2026-03-14 08:28:03 +00:00
Michael Suchacz 48de0677a4 refactor: remove owner_id denormalization from chat_messages
The spend-limit query now JOINs through the chats table to resolve
ownership, matching the pattern used by existing cost analytics
queries. This eliminates the owner_id column, the table-wide
backfill, and the covering index from the migration.
2026-03-14 08:28:03 +00:00
Michael Suchacz f67c20b71b fix: address code review findings (logging, query, swagger, week bounds) 2026-03-14 08:28:03 +00:00
Michael Suchacz 9ef9e78bf2 test(chatd): cover ComputePeriodBounds 2026-03-14 08:28:02 +00:00
Michael Suchacz e1c1226aee chore: regenerate artifacts after rebase
Run make gen to produce updated dump.sql, models, querier,
dbmetrics, dbmock, apidoc, API docs, and TypeScript types.
2026-03-14 08:28:02 +00:00
Michael Suchacz 960bafddf2 fix: use JSONB extraction for CountEnabledModelsWithoutPricing query
The pricing data is stored in the options JSONB column at
options->'cost', not as top-level columns on chat_model_configs.
Also adds deleted=FALSE filter to match existing query patterns.
2026-03-14 08:28:02 +00:00
Michael Suchacz 596752c4f2 feat: add agent spend limiting (rebased on main) 2026-03-14 08:28:02 +00:00
30 changed files with 3327 additions and 36 deletions
+48
View File
@@ -172,6 +172,23 @@ var (
errChatTakenByOtherWorker = xerrors.New("chat acquired by another worker")
)
// UsageLimitExceededError indicates the user has exceeded their chat spend
// limit.
type UsageLimitExceededError struct {
LimitMicros int64
ConsumedMicros int64
PeriodEnd time.Time
}
func (e *UsageLimitExceededError) Error() string {
return fmt.Sprintf(
"usage limit exceeded: spent %d of %d micros, resets at %s",
e.ConsumedMicros,
e.LimitMicros,
e.PeriodEnd.Format(time.RFC3339),
)
}
// CreateOptions controls chat creation in the shared chat mutation path.
type CreateOptions struct {
OwnerID uuid.UUID
@@ -389,6 +406,12 @@ func (p *Server) SendMessage(
if err != nil {
return xerrors.Errorf("lock chat: %w", err)
}
// Enforce usage limits before queueing or inserting.
if limitErr := p.checkUsageLimit(ctx, lockedChat.OwnerID); limitErr != nil {
return limitErr
}
modelConfigID := lockedChat.LastModelConfigID
if opts.ModelConfigID != nil {
modelConfigID = *opts.ModelConfigID
@@ -492,6 +515,31 @@ func (p *Server) SendMessage(
return result, nil
}
func (p *Server) checkUsageLimit(ctx context.Context, ownerID uuid.UUID) error {
status, err := ResolveUsageLimitStatus(ctx, p.db, ownerID, time.Now())
if err != nil {
// Fail open: never block chat due to a limit-resolution failure.
p.logger.Warn(ctx, "usage limit check failed, allowing message",
slog.F("owner_id", ownerID),
slog.Error(err),
)
return nil
}
if status == nil {
return nil
}
// Block when current spend reaches or exceeds limit (>= ensures
// the user cannot start new conversations once the limit is hit).
if status.SpendLimitMicros != nil && status.CurrentSpend >= *status.SpendLimitMicros {
return &UsageLimitExceededError{
LimitMicros: *status.SpendLimitMicros,
ConsumedMicros: status.CurrentSpend,
PeriodEnd: status.PeriodEnd,
}
}
return nil
}
// EditMessage updates a user message in-place, truncates all following messages,
// clears queued messages, and moves the chat into pending status.
func (p *Server) EditMessage(
+124
View File
@@ -0,0 +1,124 @@
package chatd
import (
"context"
"database/sql"
"errors"
"time"
"github.com/google/uuid"
"golang.org/x/xerrors"
"github.com/coder/coder/v2/coderd/database"
"github.com/coder/coder/v2/coderd/database/dbauthz"
"github.com/coder/coder/v2/codersdk"
)
// ComputePeriodBounds returns the UTC-aligned start and end bounds for the
// active usage-limit period containing now.
func ComputePeriodBounds(now time.Time, period codersdk.ChatUsageLimitPeriod) (start, end time.Time) {
utcNow := now.UTC()
switch period {
case codersdk.ChatUsageLimitPeriodDay:
start = time.Date(utcNow.Year(), utcNow.Month(), utcNow.Day(), 0, 0, 0, 0, time.UTC)
end = start.AddDate(0, 0, 1)
case codersdk.ChatUsageLimitPeriodWeek:
// Walk backward to Monday of the current ISO week.
// ISO 8601 weeks always start on Monday, so this never
// crosses an ISO-week boundary.
start = time.Date(utcNow.Year(), utcNow.Month(), utcNow.Day(), 0, 0, 0, 0, time.UTC)
for start.Weekday() != time.Monday {
start = start.AddDate(0, 0, -1)
}
end = start.AddDate(0, 0, 7)
case codersdk.ChatUsageLimitPeriodMonth:
start = time.Date(utcNow.Year(), utcNow.Month(), 1, 0, 0, 0, 0, time.UTC)
end = start.AddDate(0, 1, 0)
}
return start, end
}
// Architecture note: today this path enforces one period globally
// (day/week/month) from config.
// To support simultaneous periods, add nullable
// daily/weekly/monthly_limit_micros columns on override tables, where NULL
// means no limit for that period.
// Then scan spend once over the widest active window with conditional SUMs
// for each period and compare each spend/limit pair Go-side, blocking on
// whichever period is tightest.
// ResolveUsageLimitStatus resolves the current usage-limit status for userID.
//
// Note: There is a potential race condition where two concurrent messages
// from the same user can both pass the limit check if processed in
// parallel, allowing brief overage. This is acceptable because:
// - Cost is only known after the LLM API returns.
// - Overage is bounded by message cost × concurrency.
// - Fail-open is the deliberate design choice for this feature.
func ResolveUsageLimitStatus(ctx context.Context, db database.Store, userID uuid.UUID, now time.Time) (*codersdk.ChatUsageLimitStatus, error) {
//nolint:gocritic // AsChatd provides narrowly-scoped daemon access for
// deployment config reads and cross-user chat spend aggregation.
authCtx := dbauthz.AsChatd(ctx)
config, err := db.GetChatUsageLimitConfig(authCtx)
if err != nil {
if errors.Is(err, sql.ErrNoRows) {
return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits.
}
return nil, err
}
if !config.Enabled {
return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits.
}
period, ok := mapDBPeriodToSDK(config.Period)
if !ok {
return nil, xerrors.Errorf("invalid chat usage limit period %q", config.Period)
}
// Resolve effective limit in a single query:
// individual override > group limit > global default.
effectiveLimit, err := db.ResolveUserChatSpendLimit(authCtx, userID)
if err != nil {
return nil, err
}
// -1 means limits are disabled (shouldn't happen since we checked above,
// but handle gracefully).
if effectiveLimit < 0 {
return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits.
}
start, end := ComputePeriodBounds(now, period)
spendTotal, err := db.GetUserChatSpendInPeriod(authCtx, database.GetUserChatSpendInPeriodParams{
UserID: userID,
StartTime: start,
EndTime: end,
})
if err != nil {
return nil, err
}
return &codersdk.ChatUsageLimitStatus{
IsLimited: true,
Period: period,
SpendLimitMicros: &effectiveLimit,
CurrentSpend: spendTotal,
PeriodStart: start,
PeriodEnd: end,
}, nil
}
func mapDBPeriodToSDK(dbPeriod string) (codersdk.ChatUsageLimitPeriod, bool) {
switch dbPeriod {
case string(codersdk.ChatUsageLimitPeriodDay):
return codersdk.ChatUsageLimitPeriodDay, true
case string(codersdk.ChatUsageLimitPeriodWeek):
return codersdk.ChatUsageLimitPeriodWeek, true
case string(codersdk.ChatUsageLimitPeriodMonth):
return codersdk.ChatUsageLimitPeriodMonth, true
default:
return "", false
}
}
+132
View File
@@ -0,0 +1,132 @@
package chatd //nolint:testpackage // Keeps chatd unit tests in the package.
import (
"testing"
"time"
"github.com/coder/coder/v2/codersdk"
)
func TestComputePeriodBounds(t *testing.T) {
t.Parallel()
newYork, err := time.LoadLocation("America/New_York")
if err != nil {
t.Fatalf("load America/New_York: %v", err)
}
tests := []struct {
name string
now time.Time
period codersdk.ChatUsageLimitPeriod
wantStart time.Time
wantEnd time.Time
}{
{
name: "day/mid_day",
now: time.Date(2025, time.June, 15, 14, 30, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "day/midnight_exactly",
now: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "day/end_of_day",
now: time.Date(2025, time.June, 15, 23, 59, 59, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/wednesday",
now: time.Date(2025, time.June, 11, 10, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/monday",
now: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/sunday",
now: time.Date(2025, time.June, 15, 23, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2025, time.June, 9, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
},
{
name: "week/year_boundary",
now: time.Date(2024, time.December, 31, 12, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodWeek,
wantStart: time.Date(2024, time.December, 30, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.January, 6, 0, 0, 0, 0, time.UTC),
},
{
name: "month/mid_month",
now: time.Date(2025, time.June, 15, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.July, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/first_day",
now: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.July, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/last_day",
now: time.Date(2025, time.June, 30, 23, 59, 59, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.June, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.July, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/february",
now: time.Date(2025, time.February, 15, 12, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2025, time.February, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.March, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "month/leap_year_february",
now: time.Date(2024, time.February, 29, 12, 0, 0, 0, time.UTC),
period: codersdk.ChatUsageLimitPeriodMonth,
wantStart: time.Date(2024, time.February, 1, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2024, time.March, 1, 0, 0, 0, 0, time.UTC),
},
{
name: "day/non_utc_timezone",
now: time.Date(2025, time.June, 15, 22, 0, 0, 0, newYork),
period: codersdk.ChatUsageLimitPeriodDay,
wantStart: time.Date(2025, time.June, 16, 0, 0, 0, 0, time.UTC),
wantEnd: time.Date(2025, time.June, 17, 0, 0, 0, 0, time.UTC),
},
}
for _, tc := range tests {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
start, end := ComputePeriodBounds(tc.now, tc.period)
if !start.Equal(tc.wantStart) {
t.Errorf("start: got %v, want %v", start, tc.wantStart)
}
if !end.Equal(tc.wantEnd) {
t.Errorf("end: got %v, want %v", end, tc.wantEnd)
}
})
}
}
+397 -2
View File
@@ -431,7 +431,12 @@ func (api *API) chatCostSummary(rw http.ResponseWriter, r *http.Request) {
chatBreakdowns = append(chatBreakdowns, convertChatCostChatBreakdown(chat))
}
httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatCostSummary{
usageStatus, err := chatd.ResolveUsageLimitStatus(ctx, api.Database, targetUser.ID, time.Now())
if err != nil {
api.Logger.Warn(ctx, "failed to resolve usage limit status", slog.Error(err))
}
response := codersdk.ChatCostSummary{
StartDate: startDate,
EndDate: endDate,
TotalCostMicros: summary.TotalCostMicros,
@@ -443,7 +448,12 @@ func (api *API) chatCostSummary(rw http.ResponseWriter, r *http.Request) {
TotalCacheCreationTokens: summary.TotalCacheCreationTokens,
ByModel: modelBreakdowns,
ByChat: chatBreakdowns,
})
}
if usageStatus != nil {
response.UsageLimit = usageStatus
}
httpapi.Write(ctx, rw, http.StatusOK, response)
}
func (api *API) chatCostUsers(rw http.ResponseWriter, r *http.Request) {
@@ -547,6 +557,371 @@ func (api *API) chatCostUsers(rw http.ResponseWriter, r *http.Request) {
})
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
//
//nolint:revive // HTTP handler writes to ResponseWriter.
func (api *API) getChatUsageLimitConfig(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionRead, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
//nolint:gocritic // Deployment config and override hydration require
// system-restricted access.
systemCtx := dbauthz.AsSystemRestricted(ctx)
config, configErr := api.Database.GetChatUsageLimitConfig(systemCtx)
if configErr != nil && !errors.Is(configErr, sql.ErrNoRows) {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to get chat usage limit config.",
Detail: configErr.Error(),
})
return
}
overrideRows, err := api.Database.ListChatUsageLimitOverrides(systemCtx)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to list chat usage limit overrides.",
Detail: err.Error(),
})
return
}
groupOverrides, err := api.Database.ListChatUsageLimitGroupOverrides(systemCtx)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to list group usage limit overrides.",
Detail: err.Error(),
})
return
}
unpricedModelCount, err := api.Database.CountEnabledModelsWithoutPricing(systemCtx)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to count unpriced chat models.",
Detail: err.Error(),
})
return
}
response := codersdk.ChatUsageLimitConfigResponse{
ChatUsageLimitConfig: codersdk.ChatUsageLimitConfig{},
UnpricedModelCount: unpricedModelCount,
Overrides: make([]codersdk.ChatUsageLimitOverride, 0, len(overrideRows)),
GroupOverrides: make([]codersdk.ChatUsageLimitGroupOverride, 0, len(groupOverrides)),
}
if configErr == nil {
response.Period = codersdk.ChatUsageLimitPeriod(config.Period)
response.UpdatedAt = config.UpdatedAt
if config.Enabled {
response.SpendLimitMicros = ptr.Ref(config.DefaultLimitMicros)
}
}
for _, row := range overrideRows {
response.Overrides = append(response.Overrides, codersdk.ChatUsageLimitOverride{
UserID: row.UserID,
Username: row.Username,
Name: row.Name,
AvatarURL: row.AvatarURL,
SpendLimitMicros: ptr.Ref(row.LimitMicros),
CreatedAt: row.CreatedAt,
UpdatedAt: row.UpdatedAt,
})
}
for _, glo := range groupOverrides {
limitMicros := glo.LimitMicros
response.GroupOverrides = append(response.GroupOverrides, codersdk.ChatUsageLimitGroupOverride{
GroupID: glo.GroupID,
GroupName: glo.GroupName,
GroupDisplayName: glo.GroupDisplayName,
GroupAvatarURL: glo.GroupAvatarUrl,
MemberCount: glo.MemberCount,
SpendLimitMicros: &limitMicros,
CreatedAt: glo.CreatedAt,
UpdatedAt: glo.UpdatedAt,
})
}
httpapi.Write(ctx, rw, http.StatusOK, response)
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) updateChatUsageLimitConfig(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
var req codersdk.ChatUsageLimitConfig
if !httpapi.Read(ctx, rw, r, &req) {
return
}
params := database.UpsertChatUsageLimitConfigParams{
Enabled: false,
DefaultLimitMicros: 0,
Period: "",
}
if req.SpendLimitMicros == nil {
params.Enabled = false
params.DefaultLimitMicros = 0
params.Period = string(req.Period)
if params.Period == "" {
params.Period = string(codersdk.ChatUsageLimitPeriodMonth)
}
} else {
if *req.SpendLimitMicros <= 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit spend limit.",
Detail: "Spend limit must be greater than 0.",
})
return
}
switch req.Period {
case codersdk.ChatUsageLimitPeriodDay, codersdk.ChatUsageLimitPeriodWeek, codersdk.ChatUsageLimitPeriodMonth:
default:
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit period.",
Detail: "Period must be one of: day, week, month.",
})
return
}
params.Enabled = true
params.DefaultLimitMicros = *req.SpendLimitMicros
params.Period = string(req.Period)
}
//nolint:gocritic // Deployment config writes require system-restricted access.
systemCtx := dbauthz.AsSystemRestricted(ctx)
config, err := api.Database.UpsertChatUsageLimitConfig(systemCtx, params)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to update chat usage limit config.",
Detail: err.Error(),
})
return
}
response := codersdk.ChatUsageLimitConfig{
Period: codersdk.ChatUsageLimitPeriod(config.Period),
UpdatedAt: config.UpdatedAt,
}
if config.Enabled {
response.SpendLimitMicros = ptr.Ref(config.DefaultLimitMicros)
}
httpapi.Write(ctx, rw, http.StatusOK, response)
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
//
//nolint:revive // HTTP handler writes to ResponseWriter.
func (api *API) getMyChatUsageLimitStatus(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
status, err := chatd.ResolveUsageLimitStatus(ctx, api.Database, httpmw.APIKey(r).UserID, time.Now())
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to get chat usage limit status.",
Detail: err.Error(),
})
return
}
if status == nil {
status = &codersdk.ChatUsageLimitStatus{IsLimited: false}
}
httpapi.Write(ctx, rw, http.StatusOK, status)
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) upsertChatUsageLimitOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
userID, ok := parseChatUsageLimitUserID(rw, r)
if !ok {
return
}
var req codersdk.UpsertChatUsageLimitOverrideRequest
if !httpapi.Read(ctx, rw, r, &req) {
return
}
if req.SpendLimitMicros == nil || *req.SpendLimitMicros <= 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit override.",
Detail: "Spend limit must be greater than 0.",
})
return
}
//nolint:gocritic // Override writes and user hydration require
// system-restricted access.
systemCtx := dbauthz.AsSystemRestricted(ctx)
override, err := api.Database.UpsertChatUsageLimitOverride(systemCtx, database.UpsertChatUsageLimitOverrideParams{
UserID: userID,
LimitMicros: *req.SpendLimitMicros,
})
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to upsert chat usage limit override.",
Detail: err.Error(),
})
return
}
user, err := api.Database.GetUserByID(systemCtx, userID)
if err != nil {
api.Logger.Warn(ctx, "failed to fetch user for override response",
slog.F("user_id", userID),
slog.Error(err),
)
user = database.User{}
}
httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatUsageLimitOverride{
UserID: override.UserID,
Username: user.Username,
Name: user.Name,
AvatarURL: user.AvatarURL,
SpendLimitMicros: ptr.Ref(override.LimitMicros),
CreatedAt: override.CreatedAt,
UpdatedAt: override.UpdatedAt,
})
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) deleteChatUsageLimitOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
userID, ok := parseChatUsageLimitUserID(rw, r)
if !ok {
return
}
//nolint:gocritic // Override deletion requires system-restricted access.
systemCtx := dbauthz.AsSystemRestricted(ctx)
if err := api.Database.DeleteChatUsageLimitOverride(systemCtx, userID); err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to delete chat usage limit override.",
Detail: err.Error(),
})
return
}
rw.WriteHeader(http.StatusNoContent)
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) upsertChatUsageLimitGroupOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
groupIDStr := chi.URLParam(r, "group")
groupID, err := uuid.Parse(groupIDStr)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid group ID.",
Detail: err.Error(),
})
return
}
var req codersdk.UpsertChatUsageLimitGroupOverrideRequest
if !httpapi.Read(ctx, rw, r, &req) {
return
}
if req.SpendLimitMicros == nil || *req.SpendLimitMicros <= 0 {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Spend limit must be a positive value.",
})
return
}
//nolint:gocritic // Group override writes and group hydration require
// system-restricted access.
systemCtx := dbauthz.AsSystemRestricted(ctx)
override, err := api.Database.UpsertChatUsageLimitGroupOverride(systemCtx, database.UpsertChatUsageLimitGroupOverrideParams{
GroupID: groupID,
LimitMicros: *req.SpendLimitMicros,
})
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to upsert group usage limit override.",
Detail: err.Error(),
})
return
}
group, err := api.Database.GetGroupByID(systemCtx, groupID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to fetch group details.",
Detail: err.Error(),
})
return
}
limitMicros := override.LimitMicros
httpapi.Write(ctx, rw, http.StatusOK, codersdk.ChatUsageLimitGroupOverride{
GroupID: override.GroupID,
GroupName: group.Name,
GroupDisplayName: group.DisplayName,
GroupAvatarURL: group.AvatarURL,
SpendLimitMicros: &limitMicros,
CreatedAt: override.CreatedAt,
UpdatedAt: override.UpdatedAt,
})
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
func (api *API) deleteChatUsageLimitGroupOverride(rw http.ResponseWriter, r *http.Request) {
ctx := r.Context()
if !api.Authorize(r, policy.ActionUpdate, rbac.ResourceDeploymentConfig) {
httpapi.Forbidden(rw)
return
}
groupIDStr := chi.URLParam(r, "group")
groupID, err := uuid.Parse(groupIDStr)
if err != nil {
httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid group ID.",
Detail: err.Error(),
})
return
}
//nolint:gocritic // Group override deletion requires system-restricted access.
systemCtx := dbauthz.AsSystemRestricted(ctx)
err = api.Database.DeleteChatUsageLimitGroupOverride(systemCtx, groupID)
if err != nil {
httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
Message: "Failed to delete group usage limit override.",
Detail: err.Error(),
})
return
}
rw.WriteHeader(http.StatusNoContent)
}
// EXPERIMENTAL: this endpoint is experimental and is subject to change.
//
//nolint:revive // HTTP handler writes to ResponseWriter.
@@ -956,6 +1331,14 @@ func (api *API) postChatMessages(rw http.ResponseWriter, r *http.Request) {
},
)
if sendErr != nil {
var limitErr *chatd.UsageLimitExceededError
if errors.As(sendErr, &limitErr) {
httpapi.Write(ctx, rw, http.StatusConflict, codersdk.Response{
Message: "Chat usage limit exceeded.",
Detail: limitErr.Error(),
})
return
}
if xerrors.Is(sendErr, chatd.ErrMessageQueueFull) {
httpapi.Write(ctx, rw, http.StatusTooManyRequests, codersdk.Response{
Message: "Message queue is full.",
@@ -3405,6 +3788,18 @@ func chatModelConfigToUpdateParams(
}
}
func parseChatUsageLimitUserID(rw http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
userID, err := uuid.Parse(chi.URLParam(r, "user"))
if err != nil {
httpapi.Write(r.Context(), rw, http.StatusBadRequest, codersdk.Response{
Message: "Invalid chat usage limit user ID.",
Detail: err.Error(),
})
return uuid.Nil, false
}
return userID, true
}
func parseChatProviderID(rw http.ResponseWriter, r *http.Request) (uuid.UUID, bool) {
providerID, err := uuid.Parse(chi.URLParam(r, "providerConfig"))
if err != nil {
+13
View File
@@ -1175,6 +1175,19 @@ func New(options *Options) *API {
r.Delete("/", api.deleteChatModelConfig)
})
})
r.Route("/usage-limits", func(r chi.Router) {
r.Get("/", api.getChatUsageLimitConfig)
r.Put("/", api.updateChatUsageLimitConfig)
r.Get("/status", api.getMyChatUsageLimitStatus)
r.Route("/overrides/{user}", func(r chi.Router) {
r.Put("/", api.upsertChatUsageLimitOverride)
r.Delete("/", api.deleteChatUsageLimitOverride)
})
r.Route("/group-overrides/{group}", func(r chi.Router) {
r.Put("/", api.upsertChatUsageLimitGroupOverride)
r.Delete("/", api.deleteChatUsageLimitGroupOverride)
})
})
r.Route("/{chat}", func(r chi.Router) {
r.Use(httpmw.ExtractChatParam(options.Database))
r.Get("/", api.getChat)
+2
View File
@@ -10,6 +10,8 @@ const (
CheckChatModelConfigsCompressionThresholdCheck CheckConstraint = "chat_model_configs_compression_threshold_check" // chat_model_configs
CheckChatModelConfigsContextLimitCheck CheckConstraint = "chat_model_configs_context_limit_check" // chat_model_configs
CheckChatProvidersProviderCheck CheckConstraint = "chat_providers_provider_check" // chat_providers
CheckChatUsageLimitConfigPeriodCheck CheckConstraint = "chat_usage_limit_config_period_check" // chat_usage_limit_config
CheckChatUsageLimitConfigSingletonCheck CheckConstraint = "chat_usage_limit_config_singleton_check" // chat_usage_limit_config
CheckOrganizationIDNotZero CheckConstraint = "organization_id_not_zero" // custom_roles
CheckOneTimePasscodeSet CheckConstraint = "one_time_passcode_set" // users
CheckUsersEmailNotEmpty CheckConstraint = "users_email_not_empty" // users
+98
View File
@@ -1726,6 +1726,13 @@ func (q *querier) CountConnectionLogs(ctx context.Context, arg database.CountCon
return q.db.CountAuthorizedConnectionLogs(ctx, arg, prep)
}
func (q *querier) CountEnabledModelsWithoutPricing(ctx context.Context) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return 0, err
}
return q.db.CountEnabledModelsWithoutPricing(ctx)
}
func (q *querier) CountInProgressPrebuilds(ctx context.Context) ([]database.CountInProgressPrebuildsRow, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceWorkspace.All()); err != nil {
return nil, err
@@ -1854,6 +1861,20 @@ func (q *querier) DeleteChatQueuedMessage(ctx context.Context, arg database.Dele
return q.db.DeleteChatQueuedMessage(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) DeleteChatUsageLimitOverride(ctx context.Context, userID uuid.UUID) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return err
}
return q.db.DeleteChatUsageLimitOverride(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
@@ -2596,6 +2617,27 @@ func (q *querier) GetChatSystemPrompt(ctx context.Context) (string, error) {
return q.db.GetChatSystemPrompt(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) GetChatUsageLimitGroupOverrideByGroupID(ctx context.Context, groupID uuid.UUID) (database.ChatUsageLimitGroupOverride, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return database.ChatUsageLimitGroupOverride{}, err
}
return q.db.GetChatUsageLimitGroupOverrideByGroupID(ctx, groupID)
}
func (q *querier) GetChatUsageLimitOverrideByUserID(ctx context.Context, userID uuid.UUID) (database.ChatUsageLimitOverride, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceDeploymentConfig); err != nil {
return database.ChatUsageLimitOverride{}, err
}
return q.db.GetChatUsageLimitOverrideByUserID(ctx, userID)
}
func (q *querier) GetChatsByOwnerID(ctx context.Context, ownerID database.GetChatsByOwnerIDParams) ([]database.Chat, error) {
return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.GetChatsByOwnerID)(ctx, ownerID)
}
@@ -3750,6 +3792,13 @@ func (q *querier) GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID)
return q.db.GetUserChatCustomPrompt(ctx, userID)
}
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) GetUserCount(ctx context.Context, includeSystem bool) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil {
return 0, err
@@ -3757,6 +3806,13 @@ func (q *querier) GetUserCount(ctx context.Context, includeSystem bool) (int64,
return q.db.GetUserCount(ctx, includeSystem)
}
func (q *querier) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(userID.String())); err != nil {
return 0, err
}
return q.db.GetUserGroupSpendLimit(ctx, userID)
}
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 {
@@ -5115,6 +5171,20 @@ func (q *querier) ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context,
return q.db.ListAIBridgeUserPromptsByInterceptionIDs(ctx, interceptionIDs)
}
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)
}
@@ -5234,6 +5304,13 @@ func (q *querier) RemoveUserFromGroups(ctx context.Context, arg database.RemoveU
return q.db.RemoveUserFromGroups(ctx, arg)
}
func (q *querier) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceChat.WithOwner(userID.String())); err != nil {
return 0, err
}
return q.db.ResolveUserChatSpendLimit(ctx, userID)
}
func (q *querier) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil {
return err
@@ -6469,6 +6546,27 @@ func (q *querier) UpsertChatSystemPrompt(ctx context.Context, value string) erro
return q.db.UpsertChatSystemPrompt(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.ChatUsageLimitGroupOverride, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return database.ChatUsageLimitGroupOverride{}, err
}
return q.db.UpsertChatUsageLimitGroupOverride(ctx, arg)
}
func (q *querier) UpsertChatUsageLimitOverride(ctx context.Context, arg database.UpsertChatUsageLimitOverrideParams) (database.ChatUsageLimitOverride, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceDeploymentConfig); err != nil {
return database.ChatUsageLimitOverride{}, err
}
return q.db.UpsertChatUsageLimitOverride(ctx, arg)
}
func (q *querier) UpsertConnectionLog(ctx context.Context, arg database.UpsertConnectionLogParams) (database.ConnectionLog, error) {
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceConnectionLog); err != nil {
return database.ConnectionLog{}, err
+158
View File
@@ -513,6 +513,10 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().GetChatCostSummary(gomock.Any(), arg).Return(row, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceChat.WithOwner(arg.OwnerID.String()), policy.ActionRead).Returns(row)
}))
s.Run("CountEnabledModelsWithoutPricing", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
dbm.EXPECT().CountEnabledModelsWithoutPricing(gomock.Any()).Return(int64(3), nil).AnyTimes()
check.Args().Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(int64(3))
}))
s.Run("GetChatDiffStatusByChatID", s.Mocked(func(dbm *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
chat := testutil.Fake(s.T(), faker, database.Chat{})
diffStatus := testutil.Fake(s.T(), faker, database.ChatDiffStatus{ChatID: chat.ID})
@@ -833,6 +837,160 @@ func (s *MethodTestSuite) TestChats() {
dbm.EXPECT().UpsertChatSystemPrompt(gomock.Any(), "").Return(nil).AnyTimes()
check.Args("").Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
s.Run("GetUserChatSpendInPeriod", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
arg := database.GetUserChatSpendInPeriodParams{
UserID: uuid.New(),
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) {
userID := uuid.New()
limit := int64(456)
dbm.EXPECT().GetUserGroupSpendLimit(gomock.Any(), userID).Return(limit, nil).AnyTimes()
check.Args(userID).Asserts(rbac.ResourceChat.WithOwner(userID.String()), policy.ActionRead).Returns(limit)
}))
s.Run("ResolveUserChatSpendLimit", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
userID := uuid.New()
limit := int64(789)
dbm.EXPECT().ResolveUserChatSpendLimit(gomock.Any(), userID).Return(limit, nil).AnyTimes()
check.Args(userID).Asserts(rbac.ResourceChat.WithOwner(userID.String()), policy.ActionRead).Returns(limit)
}))
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("GetChatUsageLimitGroupOverrideByGroupID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
now := dbtime.Now()
groupID := uuid.New()
override := database.ChatUsageLimitGroupOverride{
ID: 1,
GroupID: groupID,
LimitMicros: 2_000_000,
CreatedAt: now,
UpdatedAt: now,
}
dbm.EXPECT().GetChatUsageLimitGroupOverrideByGroupID(gomock.Any(), groupID).Return(override, nil).AnyTimes()
check.Args(groupID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionRead).Returns(override)
}))
s.Run("GetChatUsageLimitOverrideByUserID", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
now := dbtime.Now()
userID := uuid.New()
override := database.ChatUsageLimitOverride{
ID: 1,
UserID: userID,
LimitMicros: 3_000_000,
CreatedAt: now,
UpdatedAt: now,
}
dbm.EXPECT().GetChatUsageLimitOverrideByUserID(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) {
now := dbtime.Now()
overrides := []database.ListChatUsageLimitGroupOverridesRow{{
ID: 1,
GroupID: uuid.New(),
LimitMicros: 4_000_000,
CreatedAt: now,
UpdatedAt: now,
GroupName: "group-name",
GroupDisplayName: "Group Name",
GroupAvatarUrl: "https://example.com/group.png",
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) {
now := dbtime.Now()
overrides := []database.ListChatUsageLimitOverridesRow{{
ID: 1,
UserID: uuid.New(),
LimitMicros: 5_000_000,
CreatedAt: now,
UpdatedAt: now,
Username: "usage-limit-user",
Name: "Usage Limit User",
AvatarURL: "https://example.com/avatar.png",
}}
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) {
now := dbtime.Now()
arg := database.UpsertChatUsageLimitGroupOverrideParams{
GroupID: uuid.New(),
LimitMicros: 7_000_000,
}
override := database.ChatUsageLimitGroupOverride{
ID: 1,
GroupID: arg.GroupID,
LimitMicros: arg.LimitMicros,
CreatedAt: now,
UpdatedAt: now,
}
dbm.EXPECT().UpsertChatUsageLimitGroupOverride(gomock.Any(), arg).Return(override, nil).AnyTimes()
check.Args(arg).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate).Returns(override)
}))
s.Run("UpsertChatUsageLimitOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
now := dbtime.Now()
arg := database.UpsertChatUsageLimitOverrideParams{
UserID: uuid.New(),
LimitMicros: 8_000_000,
}
override := database.ChatUsageLimitOverride{
ID: 1,
UserID: arg.UserID,
LimitMicros: arg.LimitMicros,
CreatedAt: now,
UpdatedAt: now,
}
dbm.EXPECT().UpsertChatUsageLimitOverride(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("DeleteChatUsageLimitOverride", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
userID := uuid.New()
dbm.EXPECT().DeleteChatUsageLimitOverride(gomock.Any(), userID).Return(nil).AnyTimes()
check.Args(userID).Asserts(rbac.ResourceDeploymentConfig, policy.ActionUpdate)
}))
}
func (s *MethodTestSuite) TestFile() {
+112
View File
@@ -288,6 +288,14 @@ func (m queryMetricsStore) CountConnectionLogs(ctx context.Context, arg database
return r0, r1
}
func (m queryMetricsStore) CountEnabledModelsWithoutPricing(ctx context.Context) (int64, error) {
start := time.Now()
r0, r1 := m.s.CountEnabledModelsWithoutPricing(ctx)
m.queryLatencies.WithLabelValues("CountEnabledModelsWithoutPricing").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "CountEnabledModelsWithoutPricing").Inc()
return r0, r1
}
func (m queryMetricsStore) CountInProgressPrebuilds(ctx context.Context) ([]database.CountInProgressPrebuildsRow, error) {
start := time.Now()
r0, r1 := m.s.CountInProgressPrebuilds(ctx)
@@ -408,6 +416,22 @@ func (m queryMetricsStore) DeleteChatQueuedMessage(ctx context.Context, arg data
return r0
}
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) DeleteChatUsageLimitOverride(ctx context.Context, userID uuid.UUID) error {
start := time.Now()
r0 := m.s.DeleteChatUsageLimitOverride(ctx, userID)
m.queryLatencies.WithLabelValues("DeleteChatUsageLimitOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "DeleteChatUsageLimitOverride").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)
@@ -1127,6 +1151,30 @@ func (m queryMetricsStore) GetChatSystemPrompt(ctx context.Context) (string, err
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) GetChatUsageLimitGroupOverrideByGroupID(ctx context.Context, groupID uuid.UUID) (database.ChatUsageLimitGroupOverride, error) {
start := time.Now()
r0, r1 := m.s.GetChatUsageLimitGroupOverrideByGroupID(ctx, groupID)
m.queryLatencies.WithLabelValues("GetChatUsageLimitGroupOverrideByGroupID").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitGroupOverrideByGroupID").Inc()
return r0, r1
}
func (m queryMetricsStore) GetChatUsageLimitOverrideByUserID(ctx context.Context, userID uuid.UUID) (database.ChatUsageLimitOverride, error) {
start := time.Now()
r0, r1 := m.s.GetChatUsageLimitOverrideByUserID(ctx, userID)
m.queryLatencies.WithLabelValues("GetChatUsageLimitOverrideByUserID").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetChatUsageLimitOverrideByUserID").Inc()
return r0, r1
}
func (m queryMetricsStore) GetChatsByOwnerID(ctx context.Context, ownerID database.GetChatsByOwnerIDParams) ([]database.Chat, error) {
start := time.Now()
r0, r1 := m.s.GetChatsByOwnerID(ctx, ownerID)
@@ -2255,6 +2303,14 @@ func (m queryMetricsStore) GetUserChatCustomPrompt(ctx context.Context, userID u
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) GetUserCount(ctx context.Context, includeSystem bool) (int64, error) {
start := time.Now()
r0, r1 := m.s.GetUserCount(ctx, includeSystem)
@@ -2263,6 +2319,14 @@ func (m queryMetricsStore) GetUserCount(ctx context.Context, includeSystem bool)
return r0, r1
}
func (m queryMetricsStore) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (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)
@@ -3495,6 +3559,22 @@ func (m queryMetricsStore) ListAIBridgeUserPromptsByInterceptionIDs(ctx 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)
@@ -3607,6 +3687,14 @@ func (m queryMetricsStore) RemoveUserFromGroups(ctx context.Context, arg databas
return r0, r1
}
func (m queryMetricsStore) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, 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)
@@ -4454,6 +4542,30 @@ func (m queryMetricsStore) UpsertChatSystemPrompt(ctx context.Context, value str
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.ChatUsageLimitGroupOverride, 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) UpsertChatUsageLimitOverride(ctx context.Context, arg database.UpsertChatUsageLimitOverrideParams) (database.ChatUsageLimitOverride, error) {
start := time.Now()
r0, r1 := m.s.UpsertChatUsageLimitOverride(ctx, arg)
m.queryLatencies.WithLabelValues("UpsertChatUsageLimitOverride").Observe(time.Since(start).Seconds())
m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpsertChatUsageLimitOverride").Inc()
return r0, r1
}
func (m queryMetricsStore) UpsertConnectionLog(ctx context.Context, arg database.UpsertConnectionLogParams) (database.ConnectionLog, error) {
start := time.Now()
r0, r1 := m.s.UpsertConnectionLog(ctx, arg)
+208
View File
@@ -424,6 +424,21 @@ func (mr *MockStoreMockRecorder) CountConnectionLogs(ctx, arg any) *gomock.Call
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountConnectionLogs", reflect.TypeOf((*MockStore)(nil).CountConnectionLogs), ctx, arg)
}
// CountEnabledModelsWithoutPricing mocks base method.
func (m *MockStore) CountEnabledModelsWithoutPricing(ctx context.Context) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "CountEnabledModelsWithoutPricing", ctx)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// CountEnabledModelsWithoutPricing indicates an expected call of CountEnabledModelsWithoutPricing.
func (mr *MockStoreMockRecorder) CountEnabledModelsWithoutPricing(ctx any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CountEnabledModelsWithoutPricing", reflect.TypeOf((*MockStore)(nil).CountEnabledModelsWithoutPricing), ctx)
}
// CountInProgressPrebuilds mocks base method.
func (m *MockStore) CountInProgressPrebuilds(ctx context.Context) ([]database.CountInProgressPrebuildsRow, error) {
m.ctrl.T.Helper()
@@ -639,6 +654,34 @@ func (mr *MockStoreMockRecorder) DeleteChatQueuedMessage(ctx, arg any) *gomock.C
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatQueuedMessage", reflect.TypeOf((*MockStore)(nil).DeleteChatQueuedMessage), 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)
}
// DeleteChatUsageLimitOverride mocks base method.
func (m *MockStore) DeleteChatUsageLimitOverride(ctx context.Context, userID uuid.UUID) error {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "DeleteChatUsageLimitOverride", ctx, userID)
ret0, _ := ret[0].(error)
return ret0
}
// DeleteChatUsageLimitOverride indicates an expected call of DeleteChatUsageLimitOverride.
func (mr *MockStoreMockRecorder) DeleteChatUsageLimitOverride(ctx, userID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "DeleteChatUsageLimitOverride", reflect.TypeOf((*MockStore)(nil).DeleteChatUsageLimitOverride), ctx, userID)
}
// DeleteCryptoKey mocks base method.
func (m *MockStore) DeleteCryptoKey(ctx context.Context, arg database.DeleteCryptoKeyParams) (database.CryptoKey, error) {
m.ctrl.T.Helper()
@@ -2048,6 +2091,51 @@ func (mr *MockStoreMockRecorder) GetChatSystemPrompt(ctx any) *gomock.Call {
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatSystemPrompt", reflect.TypeOf((*MockStore)(nil).GetChatSystemPrompt), 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)
}
// GetChatUsageLimitGroupOverrideByGroupID mocks base method.
func (m *MockStore) GetChatUsageLimitGroupOverrideByGroupID(ctx context.Context, groupID uuid.UUID) (database.ChatUsageLimitGroupOverride, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetChatUsageLimitGroupOverrideByGroupID", ctx, groupID)
ret0, _ := ret[0].(database.ChatUsageLimitGroupOverride)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetChatUsageLimitGroupOverrideByGroupID indicates an expected call of GetChatUsageLimitGroupOverrideByGroupID.
func (mr *MockStoreMockRecorder) GetChatUsageLimitGroupOverrideByGroupID(ctx, groupID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitGroupOverrideByGroupID", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitGroupOverrideByGroupID), ctx, groupID)
}
// GetChatUsageLimitOverrideByUserID mocks base method.
func (m *MockStore) GetChatUsageLimitOverrideByUserID(ctx context.Context, userID uuid.UUID) (database.ChatUsageLimitOverride, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetChatUsageLimitOverrideByUserID", ctx, userID)
ret0, _ := ret[0].(database.ChatUsageLimitOverride)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetChatUsageLimitOverrideByUserID indicates an expected call of GetChatUsageLimitOverrideByUserID.
func (mr *MockStoreMockRecorder) GetChatUsageLimitOverrideByUserID(ctx, userID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChatUsageLimitOverrideByUserID", reflect.TypeOf((*MockStore)(nil).GetChatUsageLimitOverrideByUserID), ctx, userID)
}
// GetChatsByOwnerID mocks base method.
func (m *MockStore) GetChatsByOwnerID(ctx context.Context, arg database.GetChatsByOwnerIDParams) ([]database.Chat, error) {
m.ctrl.T.Helper()
@@ -4193,6 +4281,21 @@ func (mr *MockStoreMockRecorder) GetUserChatCustomPrompt(ctx, userID any) *gomoc
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserChatCustomPrompt", reflect.TypeOf((*MockStore)(nil).GetUserChatCustomPrompt), ctx, userID)
}
// 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)
}
// GetUserCount mocks base method.
func (m *MockStore) GetUserCount(ctx context.Context, includeSystem bool) (int64, error) {
m.ctrl.T.Helper()
@@ -4208,6 +4311,21 @@ func (mr *MockStoreMockRecorder) GetUserCount(ctx, includeSystem any) *gomock.Ca
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserCount", reflect.TypeOf((*MockStore)(nil).GetUserCount), ctx, includeSystem)
}
// GetUserGroupSpendLimit mocks base method.
func (m *MockStore) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "GetUserGroupSpendLimit", ctx, userID)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// GetUserGroupSpendLimit indicates an expected call of GetUserGroupSpendLimit.
func (mr *MockStoreMockRecorder) GetUserGroupSpendLimit(ctx, userID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetUserGroupSpendLimit", reflect.TypeOf((*MockStore)(nil).GetUserGroupSpendLimit), ctx, userID)
}
// GetUserLatencyInsights mocks base method.
func (m *MockStore) GetUserLatencyInsights(ctx context.Context, arg database.GetUserLatencyInsightsParams) ([]database.GetUserLatencyInsightsRow, error) {
m.ctrl.T.Helper()
@@ -6547,6 +6665,36 @@ func (mr *MockStoreMockRecorder) ListAuthorizedAIBridgeModels(ctx, arg, prepared
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAuthorizedAIBridgeModels", reflect.TypeOf((*MockStore)(nil).ListAuthorizedAIBridgeModels), ctx, arg, prepared)
}
// 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()
@@ -6785,6 +6933,21 @@ func (mr *MockStoreMockRecorder) RemoveUserFromGroups(ctx, arg any) *gomock.Call
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RemoveUserFromGroups", reflect.TypeOf((*MockStore)(nil).RemoveUserFromGroups), ctx, arg)
}
// ResolveUserChatSpendLimit mocks base method.
func (m *MockStore) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "ResolveUserChatSpendLimit", ctx, userID)
ret0, _ := ret[0].(int64)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// ResolveUserChatSpendLimit indicates an expected call of ResolveUserChatSpendLimit.
func (mr *MockStoreMockRecorder) ResolveUserChatSpendLimit(ctx, userID any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ResolveUserChatSpendLimit", reflect.TypeOf((*MockStore)(nil).ResolveUserChatSpendLimit), ctx, userID)
}
// RevokeDBCryptKey mocks base method.
func (m *MockStore) RevokeDBCryptKey(ctx context.Context, activeKeyDigest string) error {
m.ctrl.T.Helper()
@@ -8316,6 +8479,51 @@ func (mr *MockStoreMockRecorder) UpsertChatSystemPrompt(ctx, value any) *gomock.
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatSystemPrompt", reflect.TypeOf((*MockStore)(nil).UpsertChatSystemPrompt), 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.ChatUsageLimitGroupOverride, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpsertChatUsageLimitGroupOverride", ctx, arg)
ret0, _ := ret[0].(database.ChatUsageLimitGroupOverride)
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)
}
// UpsertChatUsageLimitOverride mocks base method.
func (m *MockStore) UpsertChatUsageLimitOverride(ctx context.Context, arg database.UpsertChatUsageLimitOverrideParams) (database.ChatUsageLimitOverride, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "UpsertChatUsageLimitOverride", ctx, arg)
ret0, _ := ret[0].(database.ChatUsageLimitOverride)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// UpsertChatUsageLimitOverride indicates an expected call of UpsertChatUsageLimitOverride.
func (mr *MockStoreMockRecorder) UpsertChatUsageLimitOverride(ctx, arg any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpsertChatUsageLimitOverride", reflect.TypeOf((*MockStore)(nil).UpsertChatUsageLimitOverride), ctx, arg)
}
// UpsertConnectionLog mocks base method.
func (m *MockStore) UpsertConnectionLog(ctx context.Context, arg database.UpsertConnectionLogParams) (database.ConnectionLog, error) {
m.ctrl.T.Helper()
+85
View File
@@ -1302,6 +1302,61 @@ CREATE SEQUENCE chat_queued_messages_id_seq
ALTER SEQUENCE chat_queued_messages_id_seq OWNED BY chat_queued_messages.id;
CREATE TABLE chat_usage_limit_config (
id bigint NOT NULL,
singleton boolean DEFAULT true NOT NULL,
enabled boolean DEFAULT false NOT NULL,
default_limit_micros bigint DEFAULT 0 NOT NULL,
period text DEFAULT 'month'::text NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,
CONSTRAINT chat_usage_limit_config_period_check CHECK ((period = ANY (ARRAY['day'::text, 'week'::text, 'month'::text]))),
CONSTRAINT chat_usage_limit_config_singleton_check CHECK (singleton)
);
CREATE SEQUENCE chat_usage_limit_config_id_seq
START WITH 1
INCREMENT BY 1
NO MINVALUE
NO MAXVALUE
CACHE 1;
ALTER SEQUENCE chat_usage_limit_config_id_seq OWNED BY chat_usage_limit_config.id;
CREATE TABLE chat_usage_limit_group_overrides (
id bigint NOT NULL,
group_id uuid NOT NULL,
limit_micros bigint DEFAULT 0 NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL
);
CREATE SEQUENCE chat_usage_limit_group_overrides_id_seq
START WITH 1
INCREMENT BY 1
NO MINVALUE
NO MAXVALUE
CACHE 1;
ALTER SEQUENCE chat_usage_limit_group_overrides_id_seq OWNED BY chat_usage_limit_group_overrides.id;
CREATE TABLE chat_usage_limit_overrides (
id bigint NOT NULL,
user_id uuid NOT NULL,
limit_micros bigint DEFAULT 0 NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL
);
CREATE SEQUENCE chat_usage_limit_overrides_id_seq
START WITH 1
INCREMENT BY 1
NO MINVALUE
NO MAXVALUE
CACHE 1;
ALTER SEQUENCE chat_usage_limit_overrides_id_seq OWNED BY chat_usage_limit_overrides.id;
CREATE TABLE chats (
id uuid DEFAULT gen_random_uuid() NOT NULL,
owner_id uuid NOT NULL,
@@ -3140,6 +3195,12 @@ ALTER TABLE ONLY chat_messages ALTER COLUMN id SET DEFAULT nextval('chat_message
ALTER TABLE ONLY chat_queued_messages ALTER COLUMN id SET DEFAULT nextval('chat_queued_messages_id_seq'::regclass);
ALTER TABLE ONLY chat_usage_limit_config ALTER COLUMN id SET DEFAULT nextval('chat_usage_limit_config_id_seq'::regclass);
ALTER TABLE ONLY chat_usage_limit_group_overrides ALTER COLUMN id SET DEFAULT nextval('chat_usage_limit_group_overrides_id_seq'::regclass);
ALTER TABLE ONLY chat_usage_limit_overrides ALTER COLUMN id SET DEFAULT nextval('chat_usage_limit_overrides_id_seq'::regclass);
ALTER TABLE ONLY licenses ALTER COLUMN id SET DEFAULT nextval('licenses_id_seq'::regclass);
ALTER TABLE ONLY provisioner_job_logs ALTER COLUMN id SET DEFAULT nextval('provisioner_job_logs_id_seq'::regclass);
@@ -3197,6 +3258,24 @@ ALTER TABLE ONLY chat_providers
ALTER TABLE ONLY chat_queued_messages
ADD CONSTRAINT chat_queued_messages_pkey PRIMARY KEY (id);
ALTER TABLE ONLY chat_usage_limit_config
ADD CONSTRAINT chat_usage_limit_config_pkey PRIMARY KEY (id);
ALTER TABLE ONLY chat_usage_limit_config
ADD CONSTRAINT chat_usage_limit_config_singleton_key UNIQUE (singleton);
ALTER TABLE ONLY chat_usage_limit_group_overrides
ADD CONSTRAINT chat_usage_limit_group_overrides_group_id_key UNIQUE (group_id);
ALTER TABLE ONLY chat_usage_limit_group_overrides
ADD CONSTRAINT chat_usage_limit_group_overrides_pkey PRIMARY KEY (id);
ALTER TABLE ONLY chat_usage_limit_overrides
ADD CONSTRAINT chat_usage_limit_overrides_pkey PRIMARY KEY (id);
ALTER TABLE ONLY chat_usage_limit_overrides
ADD CONSTRAINT chat_usage_limit_overrides_user_id_key UNIQUE (user_id);
ALTER TABLE ONLY chats
ADD CONSTRAINT chats_pkey PRIMARY KEY (id);
@@ -3854,6 +3933,12 @@ ALTER TABLE ONLY chat_providers
ALTER TABLE ONLY chat_queued_messages
ADD CONSTRAINT chat_queued_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
ALTER TABLE ONLY chat_usage_limit_group_overrides
ADD CONSTRAINT chat_usage_limit_group_overrides_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE;
ALTER TABLE ONLY chat_usage_limit_overrides
ADD CONSTRAINT chat_usage_limit_overrides_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE;
ALTER TABLE ONLY chats
ADD CONSTRAINT chats_last_model_config_id_fkey FOREIGN KEY (last_model_config_id) REFERENCES chat_model_configs(id);
@@ -19,6 +19,8 @@ const (
ForeignKeyChatProvidersAPIKeyKeyID ForeignKeyConstraint = "chat_providers_api_key_key_id_fkey" // ALTER TABLE ONLY chat_providers ADD CONSTRAINT chat_providers_api_key_key_id_fkey FOREIGN KEY (api_key_key_id) REFERENCES dbcrypt_keys(active_key_digest);
ForeignKeyChatProvidersCreatedBy ForeignKeyConstraint = "chat_providers_created_by_fkey" // ALTER TABLE ONLY chat_providers ADD CONSTRAINT chat_providers_created_by_fkey FOREIGN KEY (created_by) REFERENCES users(id);
ForeignKeyChatQueuedMessagesChatID ForeignKeyConstraint = "chat_queued_messages_chat_id_fkey" // ALTER TABLE ONLY chat_queued_messages ADD CONSTRAINT chat_queued_messages_chat_id_fkey FOREIGN KEY (chat_id) REFERENCES chats(id) ON DELETE CASCADE;
ForeignKeyChatUsageLimitGroupOverridesGroupID ForeignKeyConstraint = "chat_usage_limit_group_overrides_group_id_fkey" // ALTER TABLE ONLY chat_usage_limit_group_overrides ADD CONSTRAINT chat_usage_limit_group_overrides_group_id_fkey FOREIGN KEY (group_id) REFERENCES groups(id) ON DELETE CASCADE;
ForeignKeyChatUsageLimitOverridesUserID ForeignKeyConstraint = "chat_usage_limit_overrides_user_id_fkey" // ALTER TABLE ONLY chat_usage_limit_overrides ADD CONSTRAINT chat_usage_limit_overrides_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE;
ForeignKeyChatsLastModelConfigID ForeignKeyConstraint = "chats_last_model_config_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_last_model_config_id_fkey FOREIGN KEY (last_model_config_id) REFERENCES chat_model_configs(id);
ForeignKeyChatsOwnerID ForeignKeyConstraint = "chats_owner_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE CASCADE;
ForeignKeyChatsParentChatID ForeignKeyConstraint = "chats_parent_chat_id_fkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_parent_chat_id_fkey FOREIGN KEY (parent_chat_id) REFERENCES chats(id) ON DELETE SET NULL;
@@ -0,0 +1,3 @@
DROP TABLE IF EXISTS chat_usage_limit_group_overrides;
DROP TABLE IF EXISTS chat_usage_limit_overrides;
DROP TABLE IF EXISTS chat_usage_limit_config;
@@ -0,0 +1,38 @@
-- 1. Singleton config table
CREATE TABLE chat_usage_limit_config (
id BIGSERIAL PRIMARY KEY,
-- Only one row allowed (enforced by CHECK).
singleton BOOLEAN NOT NULL DEFAULT TRUE CHECK (singleton),
UNIQUE (singleton),
enabled BOOLEAN NOT NULL DEFAULT FALSE,
-- Limit per user per period, in micro-dollars (1 USD = 1,000,000).
default_limit_micros BIGINT NOT NULL DEFAULT 0,
-- Period length: 'day', 'week', or 'month'.
period TEXT NOT NULL DEFAULT 'month'
CHECK (period IN ('day', 'week', 'month')),
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
-- Seed a single disabled row so reads never return empty.
INSERT INTO chat_usage_limit_config (singleton) VALUES (TRUE);
-- 2. Per-user overrides
CREATE TABLE chat_usage_limit_overrides (
id BIGSERIAL PRIMARY KEY,
user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
limit_micros BIGINT NOT NULL DEFAULT 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
UNIQUE (user_id)
);
-- 3. Per-group overrides
CREATE TABLE chat_usage_limit_group_overrides (
id BIGSERIAL PRIMARY KEY,
group_id UUID NOT NULL REFERENCES groups(id) ON DELETE CASCADE,
limit_micros BIGINT NOT NULL DEFAULT 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
UNIQUE (group_id)
);
@@ -0,0 +1,5 @@
INSERT INTO chat_usage_limit_overrides (user_id, limit_micros, created_at, updated_at) VALUES
('fc1511ef-4fcf-4a3b-98a1-8df64160e35a', 5000000, '2025-01-01 00:00:00+00', '2025-01-01 00:00:00+00');
INSERT INTO chat_usage_limit_group_overrides (group_id, limit_micros, created_at, updated_at) VALUES
('bb640d07-ca8a-4869-b6bc-ae61ebb2fda1', 10000000, '2025-01-01 00:00:00+00', '2025-01-01 00:00:00+00');
+26
View File
@@ -4125,6 +4125,32 @@ type ChatQueuedMessage struct {
CreatedAt time.Time `db:"created_at" json:"created_at"`
}
type ChatUsageLimitConfig struct {
ID int64 `db:"id" json:"id"`
Singleton bool `db:"singleton" json:"singleton"`
Enabled bool `db:"enabled" json:"enabled"`
DefaultLimitMicros int64 `db:"default_limit_micros" json:"default_limit_micros"`
Period string `db:"period" json:"period"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
}
type ChatUsageLimitGroupOverride struct {
ID int64 `db:"id" json:"id"`
GroupID uuid.UUID `db:"group_id" json:"group_id"`
LimitMicros int64 `db:"limit_micros" json:"limit_micros"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
}
type ChatUsageLimitOverride struct {
ID int64 `db:"id" json:"id"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
LimitMicros int64 `db:"limit_micros" json:"limit_micros"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
}
type ConnectionLog struct {
ID uuid.UUID `db:"id" json:"id"`
ConnectTime time.Time `db:"connect_time" json:"connect_time"`
+23
View File
@@ -77,6 +77,9 @@ type sqlcQuerier interface {
CountAIBridgeInterceptions(ctx context.Context, arg CountAIBridgeInterceptionsParams) (int64, error)
CountAuditLogs(ctx context.Context, arg CountAuditLogsParams) (int64, error)
CountConnectionLogs(ctx context.Context, arg CountConnectionLogsParams) (int64, error)
// Counts enabled, non-deleted model configs that lack both input and
// output pricing in their JSONB options.cost configuration.
CountEnabledModelsWithoutPricing(ctx context.Context) (int64, error)
// CountInProgressPrebuilds returns the number of in-progress prebuilds, grouped by preset ID and transition.
// Prebuild considered in-progress if it's in the "pending", "starting", "stopping", or "deleting" state.
CountInProgressPrebuilds(ctx context.Context) ([]CountInProgressPrebuildsRow, error)
@@ -99,6 +102,8 @@ type sqlcQuerier interface {
DeleteChatModelConfigByID(ctx context.Context, id uuid.UUID) error
DeleteChatProviderByID(ctx context.Context, id uuid.UUID) error
DeleteChatQueuedMessage(ctx context.Context, arg DeleteChatQueuedMessageParams) error
DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error
DeleteChatUsageLimitOverride(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)
@@ -242,6 +247,9 @@ type sqlcQuerier interface {
GetChatProviders(ctx context.Context) ([]ChatProvider, error)
GetChatQueuedMessages(ctx context.Context, chatID uuid.UUID) ([]ChatQueuedMessage, error)
GetChatSystemPrompt(ctx context.Context) (string, error)
GetChatUsageLimitConfig(ctx context.Context) (ChatUsageLimitConfig, error)
GetChatUsageLimitGroupOverrideByGroupID(ctx context.Context, groupID uuid.UUID) (ChatUsageLimitGroupOverride, error)
GetChatUsageLimitOverrideByUserID(ctx context.Context, userID uuid.UUID) (ChatUsageLimitOverride, error)
GetChatsByOwnerID(ctx context.Context, arg GetChatsByOwnerIDParams) ([]Chat, error)
GetConnectionLogsOffset(ctx context.Context, arg GetConnectionLogsOffsetParams) ([]GetConnectionLogsOffsetRow, error)
GetCryptoKeyByFeatureAndSequence(ctx context.Context, arg GetCryptoKeyByFeatureAndSequenceParams) (CryptoKey, error)
@@ -494,7 +502,11 @@ type sqlcQuerier interface {
GetUserByEmailOrUsername(ctx context.Context, arg GetUserByEmailOrUsernameParams) (User, error)
GetUserByID(ctx context.Context, id uuid.UUID) (User, error)
GetUserChatCustomPrompt(ctx context.Context, userID uuid.UUID) (string, error)
GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error)
GetUserCount(ctx context.Context, includeSystem bool) (int64, error)
// Returns the minimum (most restrictive) group limit for a user.
// Returns -1 if the user has no group limits applied.
GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (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
@@ -694,6 +706,8 @@ type sqlcQuerier interface {
ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeTokenUsage, error)
ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeToolUsage, error)
ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeUserPrompt, 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)
@@ -714,6 +728,12 @@ type sqlcQuerier interface {
ReduceWorkspaceAgentShareLevelToAuthenticatedByTemplate(ctx context.Context, templateID uuid.UUID) error
RegisterWorkspaceProxy(ctx context.Context, arg RegisterWorkspaceProxyParams) (WorkspaceProxy, error)
RemoveUserFromGroups(ctx context.Context, arg RemoveUserFromGroupsParams) ([]uuid.UUID, error)
// Resolves the effective spend limit for a user using the hierarchy:
// 1. Individual user override (highest priority)
// 2. Minimum group limit across all user's groups
// 3. Global default from config
// Returns -1 if limits are not enabled.
ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, 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
@@ -843,6 +863,9 @@ type sqlcQuerier interface {
UpsertChatDiffStatus(ctx context.Context, arg UpsertChatDiffStatusParams) (ChatDiffStatus, error)
UpsertChatDiffStatusReference(ctx context.Context, arg UpsertChatDiffStatusReferenceParams) (ChatDiffStatus, error)
UpsertChatSystemPrompt(ctx context.Context, value string) error
UpsertChatUsageLimitConfig(ctx context.Context, arg UpsertChatUsageLimitConfigParams) (ChatUsageLimitConfig, error)
UpsertChatUsageLimitGroupOverride(ctx context.Context, arg UpsertChatUsageLimitGroupOverrideParams) (ChatUsageLimitGroupOverride, error)
UpsertChatUsageLimitOverride(ctx context.Context, arg UpsertChatUsageLimitOverrideParams) (ChatUsageLimitOverride, error)
UpsertConnectionLog(ctx context.Context, arg UpsertConnectionLogParams) (ConnectionLog, 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.
+362
View File
@@ -3144,6 +3144,30 @@ func (q *sqlQuerier) BackoffChatDiffStatus(ctx context.Context, arg BackoffChatD
return err
}
const countEnabledModelsWithoutPricing = `-- name: CountEnabledModelsWithoutPricing :one
SELECT COUNT(*)::bigint AS count
FROM chat_model_configs
WHERE enabled = TRUE
AND deleted = FALSE
AND (
options->'cost' IS NULL
OR options->'cost' = 'null'::jsonb
OR (
(options->'cost'->>'input_price_per_million_tokens' IS NULL)
AND (options->'cost'->>'output_price_per_million_tokens' IS NULL)
)
)
`
// Counts enabled, non-deleted model configs that lack both input and
// output pricing in their JSONB options.cost configuration.
func (q *sqlQuerier) CountEnabledModelsWithoutPricing(ctx context.Context) (int64, error) {
row := q.db.QueryRowContext(ctx, countEnabledModelsWithoutPricing)
var count int64
err := row.Scan(&count)
return count, err
}
const deleteAllChatQueuedMessages = `-- name: DeleteAllChatQueuedMessages :exec
DELETE FROM chat_queued_messages WHERE chat_id = $1
`
@@ -3185,6 +3209,24 @@ func (q *sqlQuerier) DeleteChatQueuedMessage(ctx context.Context, arg DeleteChat
return err
}
const deleteChatUsageLimitGroupOverride = `-- name: DeleteChatUsageLimitGroupOverride :exec
DELETE FROM chat_usage_limit_group_overrides WHERE group_id = $1::uuid
`
func (q *sqlQuerier) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
_, err := q.db.ExecContext(ctx, deleteChatUsageLimitGroupOverride, groupID)
return err
}
const deleteChatUsageLimitOverride = `-- name: DeleteChatUsageLimitOverride :exec
DELETE FROM chat_usage_limit_overrides WHERE user_id = $1::uuid
`
func (q *sqlQuerier) DeleteChatUsageLimitOverride(ctx context.Context, userID uuid.UUID) error {
_, err := q.db.ExecContext(ctx, deleteChatUsageLimitOverride, userID)
return err
}
const getChatByID = `-- name: GetChatByID :one
SELECT
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
@@ -3943,6 +3985,59 @@ func (q *sqlQuerier) GetChatQueuedMessages(ctx context.Context, chatID 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 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 getChatUsageLimitGroupOverrideByGroupID = `-- name: GetChatUsageLimitGroupOverrideByGroupID :one
SELECT id, group_id, limit_micros, created_at, updated_at FROM chat_usage_limit_group_overrides WHERE group_id = $1::uuid
`
func (q *sqlQuerier) GetChatUsageLimitGroupOverrideByGroupID(ctx context.Context, groupID uuid.UUID) (ChatUsageLimitGroupOverride, error) {
row := q.db.QueryRowContext(ctx, getChatUsageLimitGroupOverrideByGroupID, groupID)
var i ChatUsageLimitGroupOverride
err := row.Scan(
&i.ID,
&i.GroupID,
&i.LimitMicros,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const getChatUsageLimitOverrideByUserID = `-- name: GetChatUsageLimitOverrideByUserID :one
SELECT id, user_id, limit_micros, created_at, updated_at FROM chat_usage_limit_overrides WHERE user_id = $1::uuid
`
func (q *sqlQuerier) GetChatUsageLimitOverrideByUserID(ctx context.Context, userID uuid.UUID) (ChatUsageLimitOverride, error) {
row := q.db.QueryRowContext(ctx, getChatUsageLimitOverrideByUserID, userID)
var i ChatUsageLimitOverride
err := row.Scan(
&i.ID,
&i.UserID,
&i.LimitMicros,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const getChatsByOwnerID = `-- name: GetChatsByOwnerID :many
SELECT
id, owner_id, workspace_id, title, status, worker_id, started_at, heartbeat_at, created_at, updated_at, parent_chat_id, root_chat_id, last_model_config_id, archived, last_error, mode
@@ -4134,6 +4229,45 @@ func (q *sqlQuerier) GetStaleChats(ctx context.Context, staleThreshold time.Time
return items, nil
}
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 cm.created_at >= $2::timestamptz
AND cm.created_at < $3::timestamptz
AND cm.total_cost_micros IS NOT NULL
`
type GetUserChatSpendInPeriodParams struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
StartTime time.Time `db:"start_time" json:"start_time"`
EndTime time.Time `db:"end_time" json:"end_time"`
}
func (q *sqlQuerier) GetUserChatSpendInPeriod(ctx context.Context, arg GetUserChatSpendInPeriodParams) (int64, error) {
row := q.db.QueryRowContext(ctx, getUserChatSpendInPeriod, arg.UserID, 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(glo.limit_micros), -1)::bigint AS limit_micros
FROM chat_usage_limit_group_overrides glo
JOIN group_members_expanded gme ON gme.group_id = glo.group_id
WHERE gme.user_id = $1::uuid
`
// Returns the minimum (most restrictive) group limit for a user.
// Returns -1 if the user has no group limits applied.
func (q *sqlQuerier) GetUserGroupSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, getUserGroupSpendLimit, userID)
var limit_micros int64
err := row.Scan(&limit_micros)
return limit_micros, err
}
const insertChat = `-- name: InsertChat :one
INSERT INTO chats (
owner_id,
@@ -4332,6 +4466,113 @@ func (q *sqlQuerier) InsertChatQueuedMessage(ctx context.Context, arg InsertChat
return i, err
}
const listChatUsageLimitGroupOverrides = `-- name: ListChatUsageLimitGroupOverrides :many
SELECT
glo.id, glo.group_id, glo.limit_micros, glo.created_at, glo.updated_at,
g.name AS group_name,
g.display_name AS group_display_name,
g.avatar_url AS group_avatar_url,
(SELECT COUNT(*) FROM group_members_expanded gme WHERE gme.group_id = glo.group_id) AS member_count
FROM chat_usage_limit_group_overrides glo
JOIN groups g ON g.id = glo.group_id
ORDER BY g.name ASC
`
type ListChatUsageLimitGroupOverridesRow struct {
ID int64 `db:"id" json:"id"`
GroupID uuid.UUID `db:"group_id" json:"group_id"`
LimitMicros int64 `db:"limit_micros" json:"limit_micros"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
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"`
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.ID,
&i.GroupID,
&i.LimitMicros,
&i.CreatedAt,
&i.UpdatedAt,
&i.GroupName,
&i.GroupDisplayName,
&i.GroupAvatarUrl,
&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 o.id, o.user_id, o.limit_micros, o.created_at, o.updated_at, u.username, u.name, u.avatar_url
FROM chat_usage_limit_overrides o
JOIN users u ON u.id = o.user_id
ORDER BY u.username ASC
`
type ListChatUsageLimitOverridesRow struct {
ID int64 `db:"id" json:"id"`
UserID uuid.UUID `db:"user_id" json:"user_id"`
LimitMicros int64 `db:"limit_micros" json:"limit_micros"`
CreatedAt time.Time `db:"created_at" json:"created_at"`
UpdatedAt time.Time `db:"updated_at" json:"updated_at"`
Username string `db:"username" json:"username"`
Name string `db:"name" json:"name"`
AvatarURL string `db:"avatar_url" json:"avatar_url"`
}
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.ID,
&i.UserID,
&i.LimitMicros,
&i.CreatedAt,
&i.UpdatedAt,
&i.Username,
&i.Name,
&i.AvatarURL,
); 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 popNextQueuedMessage = `-- name: PopNextQueuedMessage :one
DELETE FROM chat_queued_messages
WHERE id = (
@@ -4355,6 +4596,41 @@ func (q *sqlQuerier) PopNextQueuedMessage(ctx context.Context, chatID uuid.UUID)
return i, err
}
const resolveUserChatSpendLimit = `-- name: ResolveUserChatSpendLimit :one
SELECT CASE
-- If limits are disabled, return -1.
WHEN NOT cfg.enabled THEN -1
-- Individual override takes priority.
WHEN uo.limit_micros IS NOT NULL THEN uo.limit_micros
-- Group limit (minimum across all user's groups) is next.
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
-- Fall back to global default.
ELSE cfg.default_limit_micros
END::bigint AS effective_limit_micros
FROM chat_usage_limit_config cfg
LEFT JOIN chat_usage_limit_overrides uo
ON uo.user_id = $1::uuid
LEFT JOIN LATERAL (
SELECT MIN(glo.limit_micros) AS limit_micros
FROM chat_usage_limit_group_overrides glo
JOIN group_members_expanded gme ON gme.group_id = glo.group_id
WHERE gme.user_id = $1::uuid
) gl ON TRUE
LIMIT 1
`
// Resolves the effective spend limit for a user using the hierarchy:
// 1. Individual user override (highest priority)
// 2. Minimum group limit across all user's groups
// 3. Global default from config
// Returns -1 if limits are not enabled.
func (q *sqlQuerier) ResolveUserChatSpendLimit(ctx context.Context, userID uuid.UUID) (int64, error) {
row := q.db.QueryRowContext(ctx, resolveUserChatSpendLimit, userID)
var effective_limit_micros int64
err := row.Scan(&effective_limit_micros)
return effective_limit_micros, err
}
const unarchiveChatByID = `-- name: UnarchiveChatByID :exec
UPDATE chats SET archived = false, updated_at = NOW() WHERE id = $1::uuid
`
@@ -4784,6 +5060,92 @@ func (q *sqlQuerier) UpsertChatDiffStatusReference(ctx context.Context, arg Upse
return i, 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
INSERT INTO chat_usage_limit_group_overrides (group_id, limit_micros, updated_at)
VALUES ($1::uuid, $2::bigint, NOW())
ON CONFLICT (group_id) DO UPDATE SET
limit_micros = EXCLUDED.limit_micros,
updated_at = NOW()
RETURNING id, group_id, limit_micros, created_at, updated_at
`
type UpsertChatUsageLimitGroupOverrideParams struct {
GroupID uuid.UUID `db:"group_id" json:"group_id"`
LimitMicros int64 `db:"limit_micros" json:"limit_micros"`
}
func (q *sqlQuerier) UpsertChatUsageLimitGroupOverride(ctx context.Context, arg UpsertChatUsageLimitGroupOverrideParams) (ChatUsageLimitGroupOverride, error) {
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitGroupOverride, arg.GroupID, arg.LimitMicros)
var i ChatUsageLimitGroupOverride
err := row.Scan(
&i.ID,
&i.GroupID,
&i.LimitMicros,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const upsertChatUsageLimitOverride = `-- name: UpsertChatUsageLimitOverride :one
INSERT INTO chat_usage_limit_overrides (user_id, limit_micros, updated_at)
VALUES ($1::uuid, $2::bigint, NOW())
ON CONFLICT (user_id) DO UPDATE SET
limit_micros = EXCLUDED.limit_micros,
updated_at = NOW()
RETURNING id, user_id, limit_micros, created_at, updated_at
`
type UpsertChatUsageLimitOverrideParams struct {
UserID uuid.UUID `db:"user_id" json:"user_id"`
LimitMicros int64 `db:"limit_micros" json:"limit_micros"`
}
func (q *sqlQuerier) UpsertChatUsageLimitOverride(ctx context.Context, arg UpsertChatUsageLimitOverrideParams) (ChatUsageLimitOverride, error) {
row := q.db.QueryRowContext(ctx, upsertChatUsageLimitOverride, arg.UserID, arg.LimitMicros)
var i ChatUsageLimitOverride
err := row.Scan(
&i.ID,
&i.UserID,
&i.LimitMicros,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const countConnectionLogs = `-- name: CountConnectionLogs :one
SELECT
COUNT(*) AS count
+118
View File
@@ -679,3 +679,121 @@ LIMIT
sqlc.arg('page_limit')::int
OFFSET
sqlc.arg('page_offset')::int;
-- name: GetChatUsageLimitConfig :one
SELECT * FROM chat_usage_limit_config 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 o.*, u.username, u.name, u.avatar_url
FROM chat_usage_limit_overrides o
JOIN users u ON u.id = o.user_id
ORDER BY u.username ASC;
-- name: UpsertChatUsageLimitOverride :one
INSERT INTO chat_usage_limit_overrides (user_id, limit_micros, updated_at)
VALUES (@user_id::uuid, @limit_micros::bigint, NOW())
ON CONFLICT (user_id) DO UPDATE SET
limit_micros = EXCLUDED.limit_micros,
updated_at = NOW()
RETURNING *;
-- name: DeleteChatUsageLimitOverride :exec
DELETE FROM chat_usage_limit_overrides WHERE user_id = @user_id::uuid;
-- name: GetChatUsageLimitOverrideByUserID :one
SELECT * FROM chat_usage_limit_overrides WHERE user_id = @user_id::uuid;
-- 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 = @user_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: CountEnabledModelsWithoutPricing :one
-- Counts enabled, non-deleted model configs that lack both input and
-- output pricing in their JSONB options.cost configuration.
SELECT COUNT(*)::bigint AS count
FROM chat_model_configs
WHERE enabled = TRUE
AND deleted = FALSE
AND (
options->'cost' IS NULL
OR options->'cost' = 'null'::jsonb
OR (
(options->'cost'->>'input_price_per_million_tokens' IS NULL)
AND (options->'cost'->>'output_price_per_million_tokens' IS NULL)
)
);
-- name: ListChatUsageLimitGroupOverrides :many
SELECT
glo.*,
g.name AS group_name,
g.display_name AS group_display_name,
g.avatar_url AS group_avatar_url,
(SELECT COUNT(*) FROM group_members_expanded gme WHERE gme.group_id = glo.group_id) AS member_count
FROM chat_usage_limit_group_overrides glo
JOIN groups g ON g.id = glo.group_id
ORDER BY g.name ASC;
-- name: UpsertChatUsageLimitGroupOverride :one
INSERT INTO chat_usage_limit_group_overrides (group_id, limit_micros, updated_at)
VALUES (@group_id::uuid, @limit_micros::bigint, NOW())
ON CONFLICT (group_id) DO UPDATE SET
limit_micros = EXCLUDED.limit_micros,
updated_at = NOW()
RETURNING *;
-- name: DeleteChatUsageLimitGroupOverride :exec
DELETE FROM chat_usage_limit_group_overrides WHERE group_id = @group_id::uuid;
-- name: GetChatUsageLimitGroupOverrideByGroupID :one
SELECT * FROM chat_usage_limit_group_overrides WHERE group_id = @group_id::uuid;
-- name: GetUserGroupSpendLimit :one
-- Returns the minimum (most restrictive) group limit for a user.
-- Returns -1 if the user has no group limits applied.
SELECT COALESCE(MIN(glo.limit_micros), -1)::bigint AS limit_micros
FROM chat_usage_limit_group_overrides glo
JOIN group_members_expanded gme ON gme.group_id = glo.group_id
WHERE gme.user_id = @user_id::uuid;
-- name: ResolveUserChatSpendLimit :one
-- Resolves the effective spend limit for a user using the hierarchy:
-- 1. Individual user override (highest priority)
-- 2. Minimum group limit across all user's groups
-- 3. Global default from config
-- Returns -1 if limits are not enabled.
SELECT CASE
-- If limits are disabled, return -1.
WHEN NOT cfg.enabled THEN -1
-- Individual override takes priority.
WHEN uo.limit_micros IS NOT NULL THEN uo.limit_micros
-- Group limit (minimum across all user's groups) is next.
WHEN gl.limit_micros IS NOT NULL THEN gl.limit_micros
-- Fall back to global default.
ELSE cfg.default_limit_micros
END::bigint AS effective_limit_micros
FROM chat_usage_limit_config cfg
LEFT JOIN chat_usage_limit_overrides uo
ON uo.user_id = @user_id::uuid
LEFT JOIN LATERAL (
SELECT MIN(glo.limit_micros) AS limit_micros
FROM chat_usage_limit_group_overrides glo
JOIN group_members_expanded gme ON gme.group_id = glo.group_id
WHERE gme.user_id = @user_id::uuid
) gl ON TRUE
LIMIT 1;
+6
View File
@@ -21,6 +21,12 @@ const (
UniqueChatProvidersPkey UniqueConstraint = "chat_providers_pkey" // ALTER TABLE ONLY chat_providers ADD CONSTRAINT chat_providers_pkey PRIMARY KEY (id);
UniqueChatProvidersProviderKey UniqueConstraint = "chat_providers_provider_key" // ALTER TABLE ONLY chat_providers ADD CONSTRAINT chat_providers_provider_key UNIQUE (provider);
UniqueChatQueuedMessagesPkey UniqueConstraint = "chat_queued_messages_pkey" // ALTER TABLE ONLY chat_queued_messages ADD CONSTRAINT chat_queued_messages_pkey PRIMARY KEY (id);
UniqueChatUsageLimitConfigPkey UniqueConstraint = "chat_usage_limit_config_pkey" // ALTER TABLE ONLY chat_usage_limit_config ADD CONSTRAINT chat_usage_limit_config_pkey PRIMARY KEY (id);
UniqueChatUsageLimitConfigSingletonKey UniqueConstraint = "chat_usage_limit_config_singleton_key" // ALTER TABLE ONLY chat_usage_limit_config ADD CONSTRAINT chat_usage_limit_config_singleton_key UNIQUE (singleton);
UniqueChatUsageLimitGroupOverridesGroupIDKey UniqueConstraint = "chat_usage_limit_group_overrides_group_id_key" // ALTER TABLE ONLY chat_usage_limit_group_overrides ADD CONSTRAINT chat_usage_limit_group_overrides_group_id_key UNIQUE (group_id);
UniqueChatUsageLimitGroupOverridesPkey UniqueConstraint = "chat_usage_limit_group_overrides_pkey" // ALTER TABLE ONLY chat_usage_limit_group_overrides ADD CONSTRAINT chat_usage_limit_group_overrides_pkey PRIMARY KEY (id);
UniqueChatUsageLimitOverridesPkey UniqueConstraint = "chat_usage_limit_overrides_pkey" // ALTER TABLE ONLY chat_usage_limit_overrides ADD CONSTRAINT chat_usage_limit_overrides_pkey PRIMARY KEY (id);
UniqueChatUsageLimitOverridesUserIDKey UniqueConstraint = "chat_usage_limit_overrides_user_id_key" // ALTER TABLE ONLY chat_usage_limit_overrides ADD CONSTRAINT chat_usage_limit_overrides_user_id_key UNIQUE (user_id);
UniqueChatsPkey UniqueConstraint = "chats_pkey" // ALTER TABLE ONLY chats ADD CONSTRAINT chats_pkey PRIMARY KEY (id);
UniqueConnectionLogsPkey UniqueConstraint = "connection_logs_pkey" // ALTER TABLE ONLY connection_logs ADD CONSTRAINT connection_logs_pkey PRIMARY KEY (id);
UniqueCryptoKeysPkey UniqueConstraint = "crypto_keys_pkey" // ALTER TABLE ONLY crypto_keys ADD CONSTRAINT crypto_keys_pkey PRIMARY KEY (feature, sequence);
+176
View File
@@ -744,6 +744,7 @@ type ChatCostSummary struct {
TotalCacheCreationTokens int64 `json:"total_cache_creation_tokens"`
ByModel []ChatCostModelBreakdown `json:"by_model"`
ByChat []ChatCostChatBreakdown `json:"by_chat"`
UsageLimit *ChatUsageLimitStatus `json:"usage_limit,omitempty"`
}
// ChatCostModelBreakdown contains per-model cost aggregation.
@@ -795,6 +796,77 @@ type ChatCostUsersResponse struct {
Users []ChatCostUserRollup `json:"users"`
}
// ChatUsageLimitPeriod represents the time window for usage limits.
type ChatUsageLimitPeriod string
const (
ChatUsageLimitPeriodDay ChatUsageLimitPeriod = "day"
ChatUsageLimitPeriodWeek ChatUsageLimitPeriod = "week"
ChatUsageLimitPeriodMonth ChatUsageLimitPeriod = "month"
)
// ChatUsageLimitConfig is the deployment-wide default usage limit config.
type ChatUsageLimitConfig struct {
SpendLimitMicros *int64 `json:"spend_limit_micros"` // nil = unlimited
Period ChatUsageLimitPeriod `json:"period"`
UpdatedAt time.Time `json:"updated_at" format:"date-time"`
}
// ChatUsageLimitOverride is a per-user override of the deployment default.
type ChatUsageLimitOverride struct {
UserID uuid.UUID `json:"user_id" format:"uuid"`
Username string `json:"username"`
Name string `json:"name"`
AvatarURL string `json:"avatar_url"`
SpendLimitMicros *int64 `json:"spend_limit_micros"` // nil = unlimited
CreatedAt time.Time `json:"created_at" format:"date-time"`
UpdatedAt time.Time `json:"updated_at" format:"date-time"`
}
// ChatUsageLimitGroupOverride represents a group-scoped spend limit override.
type ChatUsageLimitGroupOverride struct {
GroupID uuid.UUID `json:"group_id" format:"uuid"`
GroupName string `json:"group_name"`
GroupDisplayName string `json:"group_display_name"`
GroupAvatarURL string `json:"group_avatar_url"`
MemberCount int64 `json:"member_count"`
SpendLimitMicros *int64 `json:"spend_limit_micros"` // nil = unlimited
CreatedAt time.Time `json:"created_at" format:"date-time"`
UpdatedAt time.Time `json:"updated_at" format:"date-time"`
}
// UpsertChatUsageLimitOverrideRequest is the body for creating/updating a
// per-user usage limit override.
type UpsertChatUsageLimitOverrideRequest struct {
SpendLimitMicros *int64 `json:"spend_limit_micros"` // nil = unlimited
}
// UpsertChatUsageLimitGroupOverrideRequest is the request to create or update
// a group-level spend limit override.
type UpsertChatUsageLimitGroupOverrideRequest struct {
SpendLimitMicros *int64 `json:"spend_limit_micros"` // nil = unlimited
}
// ChatUsageLimitStatus represents the current spend status for a user
// within their active limit period.
type ChatUsageLimitStatus struct {
IsLimited bool `json:"is_limited"`
Period ChatUsageLimitPeriod `json:"period"`
SpendLimitMicros *int64 `json:"spend_limit_micros,omitempty"`
CurrentSpend int64 `json:"current_spend"`
PeriodStart time.Time `json:"period_start" format:"date-time"`
PeriodEnd time.Time `json:"period_end" format:"date-time"`
}
// ChatUsageLimitConfigResponse is returned from the admin config endpoint
// and includes the config plus a count of models without pricing.
type ChatUsageLimitConfigResponse struct {
ChatUsageLimitConfig
UnpricedModelCount int64 `json:"unpriced_model_count"`
Overrides []ChatUsageLimitOverride `json:"overrides"`
GroupOverrides []ChatUsageLimitGroupOverride `json:"group_overrides"`
}
// ListChatsOptions are optional parameters for ListChats.
type ListChatsOptions struct {
Query string
@@ -1409,6 +1481,110 @@ func (c *Client) GetChatFile(ctx context.Context, fileID uuid.UUID) ([]byte, str
return data, res.Header.Get("Content-Type"), nil
}
// GetChatUsageLimitConfig returns the deployment-wide chat usage limit config.
func (c *Client) GetChatUsageLimitConfig(ctx context.Context) (ChatUsageLimitConfigResponse, error) {
res, err := c.Request(ctx, http.MethodGet, "/api/experimental/chats/usage-limits", nil)
if err != nil {
return ChatUsageLimitConfigResponse{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return ChatUsageLimitConfigResponse{}, ReadBodyAsError(res)
}
var resp ChatUsageLimitConfigResponse
return resp, json.NewDecoder(res.Body).Decode(&resp)
}
// UpdateChatUsageLimitConfig updates the deployment-wide usage limit config.
func (c *Client) UpdateChatUsageLimitConfig(ctx context.Context, req ChatUsageLimitConfig) (ChatUsageLimitConfig, error) {
res, err := c.Request(ctx, http.MethodPut, "/api/experimental/chats/usage-limits", req)
if err != nil {
return ChatUsageLimitConfig{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return ChatUsageLimitConfig{}, ReadBodyAsError(res)
}
var resp ChatUsageLimitConfig
return resp, json.NewDecoder(res.Body).Decode(&resp)
}
// UpsertChatUsageLimitOverride creates or updates a per-user usage limit override.
func (c *Client) UpsertChatUsageLimitOverride(ctx context.Context, userID uuid.UUID, req UpsertChatUsageLimitOverrideRequest) (ChatUsageLimitOverride, error) {
res, err := c.Request(ctx, http.MethodPut, fmt.Sprintf("/api/experimental/chats/usage-limits/overrides/%s", userID), req)
if err != nil {
return ChatUsageLimitOverride{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return ChatUsageLimitOverride{}, ReadBodyAsError(res)
}
var resp ChatUsageLimitOverride
return resp, json.NewDecoder(res.Body).Decode(&resp)
}
// DeleteChatUsageLimitOverride removes a per-user usage limit override.
func (c *Client) DeleteChatUsageLimitOverride(ctx context.Context, userID uuid.UUID) error {
res, err := c.Request(ctx, http.MethodDelete, fmt.Sprintf("/api/experimental/chats/usage-limits/overrides/%s", userID), nil)
if err != nil {
return err
}
defer res.Body.Close()
if res.StatusCode != http.StatusNoContent {
return ReadBodyAsError(res)
}
return nil
}
// UpsertChatUsageLimitGroupOverride creates or updates a group-level
// spend limit override. EXPERIMENTAL: This API is subject to change.
func (c *Client) UpsertChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID, req UpsertChatUsageLimitGroupOverrideRequest) (ChatUsageLimitGroupOverride, error) {
res, err := c.Request(ctx, http.MethodPut,
fmt.Sprintf("/api/experimental/chats/usage-limits/group-overrides/%s", groupID),
req,
)
if err != nil {
return ChatUsageLimitGroupOverride{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return ChatUsageLimitGroupOverride{}, ReadBodyAsError(res)
}
var override ChatUsageLimitGroupOverride
return override, json.NewDecoder(res.Body).Decode(&override)
}
// DeleteChatUsageLimitGroupOverride removes a group-level spend limit
// override. EXPERIMENTAL: This API is subject to change.
func (c *Client) DeleteChatUsageLimitGroupOverride(ctx context.Context, groupID uuid.UUID) error {
res, err := c.Request(ctx, http.MethodDelete,
fmt.Sprintf("/api/experimental/chats/usage-limits/group-overrides/%s", groupID),
nil,
)
if err != nil {
return err
}
defer res.Body.Close()
if res.StatusCode != http.StatusNoContent {
return ReadBodyAsError(res)
}
return nil
}
// GetMyChatUsageLimitStatus returns the current user's chat usage limit status.
func (c *Client) GetMyChatUsageLimitStatus(ctx context.Context) (ChatUsageLimitStatus, error) {
res, err := c.Request(ctx, http.MethodGet, "/api/experimental/chats/usage-limits/status", nil)
if err != nil {
return ChatUsageLimitStatus{}, err
}
defer res.Body.Close()
if res.StatusCode != http.StatusOK {
return ChatUsageLimitStatus{}, ReadBodyAsError(res)
}
var resp ChatUsageLimitStatus
return resp, json.NewDecoder(res.Body).Decode(&resp)
}
func formatChatStreamResponseError(response Response) string {
message := strings.TrimSpace(response.Message)
detail := strings.TrimSpace(response.Detail)
+60
View File
@@ -3213,6 +3213,66 @@ class ApiMethods {
const response = await this.axios.get<TypesGen.ChatCostUsersResponse>(url);
return response.data;
};
getChatUsageLimitConfig =
async (): Promise<TypesGen.ChatUsageLimitConfigResponse> =>
this.axios
.get<TypesGen.ChatUsageLimitConfigResponse>(
"/api/experimental/chats/usage-limits",
)
.then((res) => res.data);
updateChatUsageLimitConfig = async (
req: TypesGen.ChatUsageLimitConfig,
): Promise<TypesGen.ChatUsageLimitConfigResponse> =>
this.axios
.put<TypesGen.ChatUsageLimitConfigResponse>(
"/api/experimental/chats/usage-limits",
req,
)
.then((res) => res.data);
upsertChatUsageLimitOverride = async (
userID: string,
req: TypesGen.UpsertChatUsageLimitOverrideRequest,
): Promise<TypesGen.ChatUsageLimitOverride> =>
this.axios
.put<TypesGen.ChatUsageLimitOverride>(
`/api/experimental/chats/usage-limits/overrides/${encodeURIComponent(userID)}`,
req,
)
.then((res) => res.data);
deleteChatUsageLimitOverride = async (userID: string): Promise<void> =>
this.axios
.delete(
`/api/experimental/chats/usage-limits/overrides/${encodeURIComponent(userID)}`,
)
.then((res) => res.data);
upsertChatUsageLimitGroupOverride = async (
groupID: string,
req: TypesGen.UpsertChatUsageLimitGroupOverrideRequest,
): Promise<TypesGen.ChatUsageLimitGroupOverride> =>
this.axios
.put<TypesGen.ChatUsageLimitGroupOverride>(
`/api/experimental/chats/usage-limits/group-overrides/${encodeURIComponent(groupID)}`,
req,
)
.then((res) => res.data);
deleteChatUsageLimitGroupOverride = async (groupID: string): Promise<void> =>
this.axios
.delete(
`/api/experimental/chats/usage-limits/group-overrides/${encodeURIComponent(groupID)}`,
)
.then((res) => res.data);
getChatUsageLimitStatus = async (): Promise<TypesGen.ChatUsageLimitStatus> =>
this.axios
.get<TypesGen.ChatUsageLimitStatus>(
"/api/experimental/chats/usage-limits/status",
)
.then((res) => res.data);
}
export type TaskFeedbackRating = "good" | "okay" | "bad";
+89
View File
@@ -425,3 +425,92 @@ export const chatCostUsers = (params?: ChatCostUsersParams) => ({
queryFn: () => API.getChatCostUsers(params),
staleTime: 60_000,
});
/** @public */
export const chatUsageLimitConfigKey = ["chatUsageLimitConfig"] as const;
/** @public */
export const chatUsageLimitConfig = () => ({
queryKey: chatUsageLimitConfigKey,
queryFn: () => API.getChatUsageLimitConfig(),
});
/** @public */
export const chatUsageLimitStatusKey = ["chatUsageLimitStatus"] as const;
/** @public */
export const chatUsageLimitStatus = () => ({
queryKey: chatUsageLimitStatusKey,
queryFn: () => API.getChatUsageLimitStatus(),
});
/** @public */
export const updateChatUsageLimitConfig = (queryClient: QueryClient) => ({
mutationFn: (req: TypesGen.ChatUsageLimitConfig) =>
API.updateChatUsageLimitConfig(req),
onSuccess: async () => {
await queryClient.invalidateQueries({
queryKey: chatUsageLimitConfigKey,
});
},
});
type UpsertChatUsageLimitOverrideMutationArgs = {
userID: string;
req: TypesGen.UpsertChatUsageLimitOverrideRequest;
};
/** @public */
export const upsertChatUsageLimitOverride = (queryClient: QueryClient) => ({
mutationFn: ({ userID, req }: UpsertChatUsageLimitOverrideMutationArgs) =>
API.upsertChatUsageLimitOverride(userID, req),
onSuccess: async () => {
await queryClient.invalidateQueries({
queryKey: chatUsageLimitConfigKey,
});
},
});
/** @public */
export const deleteChatUsageLimitOverride = (queryClient: QueryClient) => ({
mutationFn: (userID: string) => API.deleteChatUsageLimitOverride(userID),
onSuccess: async () => {
await queryClient.invalidateQueries({
queryKey: chatUsageLimitConfigKey,
});
},
});
type UpsertChatUsageLimitGroupOverrideMutationArgs = {
groupID: string;
req: TypesGen.UpsertChatUsageLimitGroupOverrideRequest;
};
/** @public */
export const upsertChatUsageLimitGroupOverride = (
queryClient: QueryClient,
) => ({
mutationFn: ({
groupID,
req,
}: UpsertChatUsageLimitGroupOverrideMutationArgs) =>
API.upsertChatUsageLimitGroupOverride(groupID, req),
onSuccess: async () => {
await queryClient.invalidateQueries({
queryKey: chatUsageLimitConfigKey,
});
},
});
/** @public */
export const deleteChatUsageLimitGroupOverride = (
queryClient: QueryClient,
) => ({
mutationFn: (groupID: string) =>
API.deleteChatUsageLimitGroupOverride(groupID),
onSuccess: async () => {
await queryClient.invalidateQueries({
queryKey: chatUsageLimitConfigKey,
});
},
});
+2 -1
View File
@@ -10,7 +10,8 @@ type GroupSortOrder = "asc" | "desc";
export const groupsQueryKey = ["groups"];
const groups = () => {
/** @public */
export const groups = () => {
return {
queryKey: groupsQueryKey,
queryFn: () => API.getGroups(),
+92
View File
@@ -1119,6 +1119,7 @@ export interface ChatCostSummary {
readonly total_cache_creation_tokens: number;
readonly by_model: readonly ChatCostModelBreakdown[];
readonly by_chat: readonly ChatCostChatBreakdown[];
readonly usage_limit?: ChatUsageLimitStatus;
}
// From codersdk/chats.go
@@ -1769,6 +1770,79 @@ export interface ChatSystemPromptResponse {
readonly system_prompt: string;
}
// From codersdk/chats.go
/**
* ChatUsageLimitConfig is the deployment-wide default usage limit config.
*/
export interface ChatUsageLimitConfig {
readonly spend_limit_micros: number | null; // nil = unlimited
readonly period: ChatUsageLimitPeriod;
readonly updated_at: string;
}
// From codersdk/chats.go
/**
* ChatUsageLimitConfigResponse is returned from the admin config endpoint
* and includes the config plus a count of models without pricing.
*/
export interface ChatUsageLimitConfigResponse extends ChatUsageLimitConfig {
readonly unpriced_model_count: number;
readonly overrides: readonly ChatUsageLimitOverride[];
readonly group_overrides: readonly ChatUsageLimitGroupOverride[];
}
// From codersdk/chats.go
/**
* ChatUsageLimitGroupOverride represents a group-scoped spend limit override.
*/
export interface ChatUsageLimitGroupOverride {
readonly group_id: string;
readonly group_name: string;
readonly group_display_name: string;
readonly group_avatar_url: string;
readonly member_count: number;
readonly spend_limit_micros: number | null; // nil = unlimited
readonly created_at: string;
readonly updated_at: string;
}
// From codersdk/chats.go
/**
* ChatUsageLimitOverride is a per-user override of the deployment default.
*/
export interface ChatUsageLimitOverride {
readonly user_id: string;
readonly username: string;
readonly name: string;
readonly avatar_url: string;
readonly spend_limit_micros: number | null; // nil = unlimited
readonly created_at: string;
readonly updated_at: string;
}
// From codersdk/chats.go
export type ChatUsageLimitPeriod = "day" | "month" | "week";
export const ChatUsageLimitPeriods: ChatUsageLimitPeriod[] = [
"day",
"month",
"week",
];
// From codersdk/chats.go
/**
* ChatUsageLimitStatus represents the current spend status for a user
* within their active limit period.
*/
export interface ChatUsageLimitStatus {
readonly is_limited: boolean;
readonly period: ChatUsageLimitPeriod;
readonly spend_limit_micros?: number;
readonly current_spend: number;
readonly period_start: string;
readonly period_end: string;
}
// From codersdk/client.go
/**
* CoderDesktopTelemetryHeader contains a JSON-encoded representation of Desktop telemetry
@@ -6817,6 +6891,24 @@ export interface UploadResponse {
readonly hash: string;
}
// From codersdk/chats.go
/**
* UpsertChatUsageLimitGroupOverrideRequest is the request to create or update
* a group-level spend limit override.
*/
export interface UpsertChatUsageLimitGroupOverrideRequest {
readonly spend_limit_micros: number | null; // nil = unlimited
}
// From codersdk/chats.go
/**
* UpsertChatUsageLimitOverrideRequest is the body for creating/updating a
* per-user usage limit override.
*/
export interface UpsertChatUsageLimitOverrideRequest {
readonly spend_limit_micros: number | null; // nil = unlimited
}
// From codersdk/workspaceagentportshare.go
export interface UpsertWorkspaceAgentPortShareRequest {
readonly agent_name: string;
+27 -17
View File
@@ -1,4 +1,5 @@
import { API, watchWorkspace } from "api/api";
import { getErrorMessage, isApiError } from "api/errors";
import {
chat,
chatDiffStatus,
@@ -160,8 +161,7 @@ export const AgentDetailTimeline: FC<AgentDetailTimelineProps> = ({
() => buildParsedMessageSections(parsedMessages),
[parsedMessages],
);
const detailErrorMessage =
(chatStatus === "error" ? persistedErrorReason : undefined) || streamError;
const detailErrorMessage = persistedErrorReason || streamError;
const latestMessage = messages[messages.length - 1];
const latestMessageNeedsAssistantResponse =
!latestMessage || latestMessage.role !== "assistant";
@@ -913,22 +913,32 @@ const AgentDetail: FC = () => {
// timeline when the server confirms via the POST response or
// via the SSE stream.
store.clearStreamState();
const response = await sendMutation.mutateAsync(request);
// When the server accepts the message immediately (not
// queued), insert it into the store so it appears in the
// timeline without waiting for the SSE stream.
if (!response.queued && response.message) {
store.upsertDurableMessage(response.message);
}
if (typeof window !== "undefined") {
if (selectedModelConfigID) {
localStorage.setItem(
lastModelConfigIDStorageKey,
selectedModelConfigID,
);
} else {
localStorage.removeItem(lastModelConfigIDStorageKey);
try {
const response = await sendMutation.mutateAsync(request);
// When the server accepts the message immediately (not
// queued), insert it into the store so it appears in the
// timeline without waiting for the SSE stream.
if (!response.queued && response.message) {
store.upsertDurableMessage(response.message);
}
if (typeof window !== "undefined") {
if (selectedModelConfigID) {
localStorage.setItem(
lastModelConfigIDStorageKey,
selectedModelConfigID,
);
} else {
localStorage.removeItem(lastModelConfigIDStorageKey);
}
}
} catch (error) {
if (isApiError(error) && error.response?.status === 409) {
setChatErrorReason(
agentId,
getErrorMessage(error, "Your usage limit has been reached."),
);
}
throw error;
}
};
@@ -10,6 +10,7 @@ import {
TableHeader,
TableRow,
} from "components/Table/Table";
import dayjs from "dayjs";
import { TriangleAlertIcon } from "lucide-react";
import type { FC } from "react";
import { formatCostMicros, formatTokenCount } from "utils/analytics";
@@ -23,6 +24,21 @@ interface ChatCostSummaryViewProps {
emptyMessage: string;
}
const getUsageLimitPeriodLabel = (
period: TypesGen.ChatUsageLimitPeriod,
): string => {
switch (period) {
case "day":
return "Daily";
case "week":
return "Weekly";
case "month":
return "Monthly";
default:
return "";
}
};
export const ChatCostSummaryView: FC<ChatCostSummaryViewProps> = ({
summary,
isLoading,
@@ -60,6 +76,35 @@ export const ChatCostSummaryView: FC<ChatCostSummaryViewProps> = ({
return null;
}
const usageLimit = summary.usage_limit;
const showUsageLimitCard = usageLimit?.is_limited === true;
const usageLimitCurrentSpend = usageLimit?.current_spend ?? 0;
const usageLimitSpendMicros = usageLimit?.spend_limit_micros ?? 0;
const usageLimitPeriodLabel = usageLimit
? getUsageLimitPeriodLabel(usageLimit.period)
: "";
const usageProgressPercentage =
showUsageLimitCard && usageLimitSpendMicros > 0
? Math.min((usageLimitCurrentSpend / usageLimitSpendMicros) * 100, 100)
: 0;
const usageProgressBarClass =
usageProgressPercentage > 90
? "bg-surface-red"
: usageProgressPercentage >= 75
? "bg-surface-orange"
: "bg-surface-green";
const usageLimitExceeded =
showUsageLimitCard && usageLimitCurrentSpend > usageLimitSpendMicros;
const usageLimitStatusText = usageLimitExceeded
? "Limit exceeded"
: `${formatCostMicros(
Math.max(usageLimitSpendMicros - usageLimitCurrentSpend, 0),
)} remaining`;
const usageLimitResetAt =
showUsageLimitCard && usageLimit?.period_end
? dayjs(usageLimit.period_end).format("MMM D, YYYY h:mm A")
: "";
return (
<div className="space-y-6">
<div className="grid grid-cols-2 gap-4 md:grid-cols-3">
@@ -115,6 +160,54 @@ export const ChatCostSummaryView: FC<ChatCostSummaryViewProps> = ({
</div>
</div>
{showUsageLimitCard && usageLimit && (
<div className="rounded-lg border border-border-default bg-surface-secondary p-4">
<div className="flex flex-col gap-4">
<div className="flex flex-col gap-2 md:flex-row md:items-end md:justify-between">
<div>
<p className="text-xs font-medium uppercase tracking-wide text-content-secondary">
{usageLimitPeriodLabel} Spend Limit
</p>
<p className="mt-1 text-2xl font-semibold text-content-primary">
{formatCostMicros(usageLimitCurrentSpend)} /{" "}
{formatCostMicros(usageLimitSpendMicros)}
</p>
</div>
<p className="text-sm text-content-secondary">
{Math.round(usageProgressPercentage)}% used
</p>
</div>
<div
role="progressbar"
aria-label={`${usageLimitPeriodLabel} spend usage`}
aria-valuemin={0}
aria-valuemax={100}
aria-valuenow={Math.round(usageProgressPercentage)}
className="h-2 overflow-hidden rounded-full bg-surface-tertiary"
>
<div
className={`h-full rounded-full ${usageProgressBarClass}`}
style={{ width: `${usageProgressPercentage}%` }}
/>
</div>
<div className="flex flex-col gap-1 text-sm md:flex-row md:items-center md:justify-between">
<p
className={
usageLimitExceeded
? "text-content-destructive"
: "text-content-secondary"
}
>
{usageLimitStatusText}
</p>
<p className="text-content-secondary">
Resets {usageLimitResetAt}
</p>
</div>
</div>
</div>
)}
{summary.unpriced_message_count > 0 && (
<div className="flex items-start gap-3 rounded-lg border border-border-warning bg-surface-warning p-4 text-sm text-content-primary">
<TriangleAlertIcon className="h-5 w-5 shrink-0 text-content-warning" />
@@ -4,7 +4,9 @@ import {
chatModelConfigsKey,
chatModelsKey,
chatProviderConfigsKey,
chatUsageLimitConfigKey,
} from "api/queries/chats";
import { groupsQueryKey } from "api/queries/groups";
import type {
ChatCostSummary,
ChatCostUserRollup,
@@ -12,6 +14,10 @@ import type {
ChatModelConfig,
ChatModelsResponse,
ChatProviderConfig,
ChatUsageLimitConfigResponse,
ChatUsageLimitGroupOverride,
ChatUsageLimitOverride,
Group,
} from "api/typesGenerated";
import {
expect,
@@ -155,6 +161,69 @@ const mockUsageSummary: ChatCostSummary = {
],
};
const mockGroupOverrides: ChatUsageLimitGroupOverride[] = [
{
group_id: "grp-1",
group_name: "engineering",
group_display_name: "Engineering",
group_avatar_url: "",
member_count: 12,
spend_limit_micros: 20_000_000,
created_at: "2026-03-01T00:00:00Z",
updated_at: "2026-03-01T00:00:00Z",
},
];
const mockUserOverrides: ChatUsageLimitOverride[] = [
{
user_id: "user-1",
username: "alice",
name: "Alice Example",
avatar_url: "https://example.com/alice.png",
spend_limit_micros: 50_000_000,
created_at: "2026-03-01T00:00:00Z",
updated_at: "2026-03-01T00:00:00Z",
},
];
const mockLimitConfig: ChatUsageLimitConfigResponse = {
spend_limit_micros: 10_000_000,
period: "month",
updated_at: "2026-03-01T00:00:00Z",
unpriced_model_count: 1,
group_overrides: mockGroupOverrides,
overrides: mockUserOverrides,
};
const mockGroups: Group[] = [
{
id: "grp-1",
name: "engineering",
display_name: "Engineering",
organization_id: "org-1",
members: [],
total_member_count: 12,
avatar_url: "",
quota_allowance: 0,
source: "user",
organization_name: "default",
organization_display_name: "Default",
},
{
id: "grp-2",
name: "design",
display_name: "Design",
organization_id: "org-1",
members: [],
total_member_count: 5,
avatar_url: "",
quota_allowance: 0,
source: "user",
organization_name: "default",
organization_display_name: "Default",
},
];
const meta: Meta<typeof ConfigureAgentsDialog> = {
title: "pages/AgentsPage/ConfigureAgentsDialog",
component: ConfigureAgentsDialog,
@@ -251,3 +320,29 @@ export const UsageTab: Story = {
spyOn(API, "getChatCostSummary").mockResolvedValue(mockUsageSummary);
},
};
/** Admin sees the Limits tab with global, group, and user override data. */
export const LimitsTab: Story = {
args: {
initialSection: "limits",
canManageChatModelConfigs: true,
},
parameters: {
queries: [
...chatQueries,
{ key: chatUsageLimitConfigKey, data: mockLimitConfig },
{ key: groupsQueryKey, data: mockGroups },
],
},
beforeEach: () => {
spyOn(API, "updateChatUsageLimitConfig").mockResolvedValue(mockLimitConfig);
spyOn(API, "upsertChatUsageLimitGroupOverride").mockResolvedValue(
mockGroupOverrides[0]!,
);
spyOn(API, "deleteChatUsageLimitGroupOverride").mockResolvedValue();
spyOn(API, "upsertChatUsageLimitOverride").mockResolvedValue(
mockUserOverrides[0]!,
);
spyOn(API, "deleteChatUsageLimitOverride").mockResolvedValue();
},
};
@@ -44,6 +44,7 @@ import {
BarChart3Icon,
BoxesIcon,
KeyRoundIcon,
ShieldAlertIcon,
ShieldIcon,
UserIcon,
XIcon,
@@ -60,11 +61,13 @@ import { formatCostMicros, formatTokenCount } from "utils/analytics";
import { cn } from "utils/cn";
import { ChatCostSummaryView } from "./ChatCostSummaryView";
import { ChatModelAdminPanel } from "./ChatModelAdminPanel/ChatModelAdminPanel";
import { LimitsTab } from "./LimitsTab";
import { SectionHeader } from "./SectionHeader";
export type ConfigureAgentsSection =
| "providers"
| "models"
| "limits"
| "behavior"
| "usage";
@@ -437,6 +440,12 @@ export const ConfigureAgentsDialog: FC<ConfigureAgentsDialogProps> = ({
icon: BoxesIcon,
adminOnly: true,
});
options.push({
id: "limits",
label: "Limits",
icon: ShieldAlertIcon,
adminOnly: true,
});
options.push({
id: "usage",
label: "Usage",
@@ -513,9 +522,9 @@ export const ConfigureAgentsDialog: FC<ConfigureAgentsDialogProps> = ({
})}
</nav>
<div className="flex min-h-0 flex-1 flex-col overflow-y-auto px-6 py-5 [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent]">
<div className="flex min-h-0 flex-1 flex-col overflow-hidden">
{activeSection === "behavior" && (
<>
<div className="flex-1 overflow-y-auto px-6 py-5 [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent]">
<SectionHeader
label="Behavior"
description="Custom instructions that shape how the agent responds in your chats."
@@ -616,26 +625,35 @@ export const ConfigureAgentsDialog: FC<ConfigureAgentsDialogProps> = ({
</form>
</>
)}
</>
</div>
)}
{activeSection === "providers" && canManageChatModelConfigs && (
<ChatModelAdminPanel
section="providers"
sectionLabel="Providers"
sectionDescription="Connect third-party LLM services like OpenAI, Anthropic, or Google. Each provider supplies models that users can select for their chats."
sectionBadge={<AdminBadge />}
/>
<div className="flex-1 overflow-y-auto px-6 py-5 [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent]">
<ChatModelAdminPanel
section="providers"
sectionLabel="Providers"
sectionDescription="Connect third-party LLM services like OpenAI, Anthropic, or Google. Each provider supplies models that users can select for their chats."
sectionBadge={<AdminBadge />}
/>
</div>
)}
{activeSection === "models" && canManageChatModelConfigs && (
<ChatModelAdminPanel
section="models"
sectionLabel="Models"
sectionDescription="Choose which models from your configured providers are available for users to select. You can set a default and adjust context limits."
sectionBadge={<AdminBadge />}
/>
<div className="flex-1 overflow-y-auto px-6 py-5 [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent]">
<ChatModelAdminPanel
section="models"
sectionLabel="Models"
sectionDescription="Choose which models from your configured providers are available for users to select. You can set a default and adjust context limits."
sectionBadge={<AdminBadge />}
/>
</div>
)}
{activeSection === "limits" && canManageChatModelConfigs && (
<LimitsTab />
)}
{activeSection === "usage" && canManageChatModelConfigs && (
<UsageContent />
<div className="flex-1 overflow-y-auto px-6 py-5 [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent]">
<UsageContent />
</div>
)}
</div>
</DialogContent>
+699
View File
@@ -0,0 +1,699 @@
import { getErrorMessage } from "api/errors";
import {
chatUsageLimitConfig,
deleteChatUsageLimitGroupOverride,
deleteChatUsageLimitOverride,
updateChatUsageLimitConfig,
upsertChatUsageLimitGroupOverride,
upsertChatUsageLimitOverride,
} from "api/queries/chats";
import { groups } from "api/queries/groups";
import type { ChatUsageLimitPeriod, Group, User } from "api/typesGenerated";
import { Autocomplete } from "components/Autocomplete/Autocomplete";
import { AvatarData } from "components/Avatar/AvatarData";
import { Button } from "components/Button/Button";
import { Input } from "components/Input/Input";
import { Label } from "components/Label/Label";
import {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "components/Select/Select";
import { Spinner } from "components/Spinner/Spinner";
import { Switch } from "components/Switch/Switch";
import {
Table,
TableBody,
TableCell,
TableHead,
TableHeader,
TableRow,
} from "components/Table/Table";
import {
Tooltip,
TooltipContent,
TooltipProvider,
TooltipTrigger,
} from "components/Tooltip/Tooltip";
import { UserAutocomplete } from "components/UserAutocomplete/UserAutocomplete";
import { Check, ShieldIcon, TriangleAlertIcon } from "lucide-react";
import { getGroupSubtitle } from "modules/groups";
import { type FC, useEffect, useMemo, useRef, useState } from "react";
import { useMutation, useQuery, useQueryClient } from "react-query";
import { formatCostMicros } from "utils/analytics";
import { SectionHeader } from "./SectionHeader";
const microsToDollars = (micros: number): number =>
Math.round(micros / 10_000) / 100;
const dollarsToMicros = (dollars: string): number =>
Math.round(Number(dollars) * 1_000_000);
const sectionPanelClassName = "space-y-4 rounded-lg border border-border p-4";
const AdminBadge: FC = () => (
<TooltipProvider delayDuration={0}>
<Tooltip>
<TooltipTrigger asChild>
<span className="inline-flex cursor-default items-center gap-1 rounded bg-surface-tertiary/60 px-1.5 py-px text-[11px] font-medium text-content-secondary">
<ShieldIcon className="h-3 w-3" />
Admin
</span>
</TooltipTrigger>
<TooltipContent side="right">
Only visible to deployment administrators.
</TooltipContent>
</Tooltip>
</TooltipProvider>
);
export const LimitsTab: FC = () => {
const queryClient = useQueryClient();
const configQuery = useQuery(chatUsageLimitConfig());
const updateConfigMutation = useMutation(
updateChatUsageLimitConfig(queryClient),
);
const upsertOverrideMutation = useMutation(
upsertChatUsageLimitOverride(queryClient),
);
const deleteOverrideMutation = useMutation(
deleteChatUsageLimitOverride(queryClient),
);
const groupsQuery = useQuery(groups());
const upsertGroupOverrideMutation = useMutation(
upsertChatUsageLimitGroupOverride(queryClient),
);
const deleteGroupOverrideMutation = useMutation(
deleteChatUsageLimitGroupOverride(queryClient),
);
const [enabled, setEnabled] = useState(false);
const [period, setPeriod] = useState<ChatUsageLimitPeriod>("month");
const [amountDollars, setAmountDollars] = useState("");
const [showGroupForm, setShowGroupForm] = useState(false);
const [selectedGroup, setSelectedGroup] = useState<Group | null>(null);
const [groupAmount, setGroupAmount] = useState("");
const [showUserForm, setShowUserForm] = useState(false);
const [selectedUser, setSelectedUser] = useState<User | null>(null);
const [userOverrideAmount, setUserOverrideAmount] = useState("");
const lastSyncedRef = useRef<string | null>(null);
useEffect(() => {
if (!configQuery.data) {
return;
}
const snapshot = JSON.stringify({
spend_limit_micros: configQuery.data.spend_limit_micros,
period: configQuery.data.period,
});
if (lastSyncedRef.current === snapshot) {
return;
}
lastSyncedRef.current = snapshot;
const spendLimitMicros = configQuery.data.spend_limit_micros;
const hasLimit = spendLimitMicros !== null;
setEnabled(hasLimit);
setPeriod(configQuery.data.period ?? "month");
setAmountDollars(
hasLimit ? microsToDollars(spendLimitMicros).toString() : "",
);
}, [configQuery.data]);
const existingGroupIds = useMemo(
() =>
new Set((configQuery.data?.group_overrides ?? []).map((g) => g.group_id)),
[configQuery.data?.group_overrides],
);
const existingUserIds = useMemo(
() => new Set((configQuery.data?.overrides ?? []).map((o) => o.user_id)),
[configQuery.data?.overrides],
);
const availableGroups = useMemo(
() => (groupsQuery.data ?? []).filter((g) => !existingGroupIds.has(g.id)),
[groupsQuery.data, existingGroupIds],
);
const selectedUserAlreadyOverridden = selectedUser
? existingUserIds.has(selectedUser.id)
: false;
const groupAutocompleteNoOptionsText = groupsQuery.isLoading
? "Loading groups..."
: availableGroups.length === 0
? "All groups already have overrides"
: "No groups available";
const isAmountValid =
!enabled ||
(amountDollars.trim() !== "" &&
!Number.isNaN(Number(amountDollars)) &&
Number(amountDollars) >= 0);
const handleSaveDefault = async () => {
const spendLimitMicros = enabled ? dollarsToMicros(amountDollars) : null;
try {
await updateConfigMutation.mutateAsync({
spend_limit_micros: spendLimitMicros,
period,
updated_at: new Date().toISOString(),
});
lastSyncedRef.current = JSON.stringify({
spend_limit_micros: spendLimitMicros,
period,
});
} catch {
// Keep the current form state so the inline mutation error is visible.
}
};
const handleAddOverride = async () => {
if (!selectedUser) {
return;
}
try {
await upsertOverrideMutation.mutateAsync({
userID: selectedUser.id,
req: { spend_limit_micros: dollarsToMicros(userOverrideAmount) },
});
setSelectedUser(null);
setUserOverrideAmount("");
setShowUserForm(false);
} catch {
// Keep the current form state so the inline mutation error is visible.
}
};
const handleAddGroupOverride = async () => {
if (!selectedGroup) {
return;
}
try {
await upsertGroupOverrideMutation.mutateAsync({
groupID: selectedGroup.id,
req: { spend_limit_micros: dollarsToMicros(groupAmount) },
});
setSelectedGroup(null);
setGroupAmount("");
setShowGroupForm(false);
} catch {
// Keep the current form state so the inline mutation error is visible.
}
};
const handleDeleteGroupOverride = async (groupID: string) => {
try {
await deleteGroupOverrideMutation.mutateAsync(groupID);
} catch {
// Keep the current UI state so the inline mutation error is visible.
}
};
const handleDeleteOverride = async (userID: string) => {
try {
await deleteOverrideMutation.mutateAsync(userID);
} catch {
// Keep the current UI state so the inline mutation error is visible.
}
};
if (configQuery.isLoading) {
return (
<div className="flex min-h-0 flex-1 flex-col overflow-hidden">
<div className="flex flex-1 items-center justify-center px-6 py-5">
<Spinner loading className="h-6 w-6" />
</div>
</div>
);
}
if (configQuery.isError) {
return (
<div className="flex min-h-0 flex-1 flex-col overflow-hidden">
<div className="flex flex-1 items-center justify-center px-6 py-5">
<div className="space-y-4 py-4 text-center">
<p className="text-sm text-content-secondary">
{getErrorMessage(
configQuery.error,
"Failed to load spend limit settings.",
)}
</p>
<Button
variant="outline"
size="sm"
type="button"
onClick={() => void configQuery.refetch()}
>
Retry
</Button>
</div>
</div>
</div>
);
}
const groupOverrides = configQuery.data?.group_overrides ?? [];
const overrides = configQuery.data?.overrides ?? [];
const unpricedModelCount = configQuery.data?.unpriced_model_count ?? 0;
return (
<div className="flex min-h-0 flex-1 flex-col overflow-hidden">
<div className="flex-1 overflow-y-auto px-6 py-5 pb-24 [scrollbar-width:thin] [scrollbar-color:hsl(var(--surface-quaternary))_transparent]">
<div className="space-y-6">
<SectionHeader
label="Default Spend Limit"
description="Set a deployment-wide spend cap that applies to all users by default."
badge={<AdminBadge />}
/>
<div className="space-y-4 rounded-lg border border-border p-4">
<div className="flex items-center justify-between gap-4">
<div>
<p className="m-0 text-sm font-medium text-content-primary">
Enable spend limit
</p>
<p className="m-0 text-xs text-content-secondary">
When disabled, users have unlimited spending.
</p>
</div>
<Switch checked={enabled} onCheckedChange={setEnabled} />
</div>
{enabled && (
<div className="flex flex-col gap-3 md:flex-row md:items-end">
<div className="flex-1 space-y-1">
<Label htmlFor="chat-limit-period">Period</Label>
<Select
value={period}
onValueChange={(value) =>
setPeriod(value as ChatUsageLimitPeriod)
}
>
<SelectTrigger
id="chat-limit-period"
className="h-9 min-w-0 text-[13px]"
>
<SelectValue />
</SelectTrigger>
<SelectContent>
<SelectItem value="day">Day</SelectItem>
<SelectItem value="week">Week</SelectItem>
<SelectItem value="month">Month</SelectItem>
</SelectContent>
</Select>
</div>
<div className="flex-1 space-y-1">
<Label htmlFor="chat-limit-amount">Amount ($)</Label>
<Input
id="chat-limit-amount"
type="number"
step="0.01"
min="0"
className="h-9 min-w-0 text-[13px]"
value={amountDollars}
onChange={(event) => setAmountDollars(event.target.value)}
placeholder="0.00"
/>
</div>
</div>
)}
</div>
{enabled && unpricedModelCount > 0 && (
<div className="flex items-start gap-3 rounded-lg border border-border-warning bg-surface-warning p-4 text-sm text-content-primary">
<TriangleAlertIcon className="h-5 w-5 shrink-0 text-content-warning" />
<div>
{unpricedModelCount === 1
? "1 enabled model does not have pricing configured."
: `${unpricedModelCount} enabled models do not have pricing configured.`}{" "}
Usage of unpriced models cannot be tracked against the spend
limit.
</div>
</div>
)}
<section className="space-y-4">
<SectionHeader
label="Group Limits"
description="Override the default limit for specific groups."
/>
<div className={sectionPanelClassName}>
{groupOverrides.length > 0 ? (
<Table>
<TableHeader>
<TableRow>
<TableHead>Group</TableHead>
<TableHead>Members</TableHead>
<TableHead>Spend Limit</TableHead>
<TableHead className="w-[80px]">Actions</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{groupOverrides.map((override) => (
<TableRow key={override.group_id}>
<TableCell>
<AvatarData
title={
override.group_display_name || override.group_name
}
subtitle={override.group_name}
src={override.group_avatar_url}
imgFallbackText={override.group_name}
/>
</TableCell>
<TableCell>{override.member_count}</TableCell>
<TableCell>
{override.spend_limit_micros !== null
? formatCostMicros(override.spend_limit_micros)
: "Unlimited"}
</TableCell>
<TableCell>
<Button
variant="outline"
size="sm"
type="button"
onClick={() =>
void handleDeleteGroupOverride(override.group_id)
}
disabled={deleteGroupOverrideMutation.isPending}
>
Delete
</Button>
</TableCell>
</TableRow>
))}
</TableBody>
</Table>
) : (
<div className="rounded-lg border border-border bg-surface-secondary px-4 py-6 text-center text-sm text-content-secondary">
No group overrides configured.
</div>
)}
{deleteGroupOverrideMutation.isError && (
<p className="text-xs text-content-destructive">
{getErrorMessage(
deleteGroupOverrideMutation.error,
"Failed to delete group override.",
)}
</p>
)}
{!showGroupForm ? (
<Button
variant="outline"
size="sm"
type="button"
onClick={() => setShowGroupForm(true)}
disabled={
groupsQuery.isLoading || availableGroups.length === 0
}
>
Add Group
</Button>
) : (
<div className="space-y-3 rounded-lg border border-border bg-surface-secondary/40 p-4">
<div className="flex flex-col gap-3 md:flex-row md:items-end">
<div className="flex-1 space-y-1">
<Label htmlFor="group-override-autocomplete">Group</Label>
<Autocomplete
id="group-override-autocomplete"
value={selectedGroup}
onChange={setSelectedGroup}
options={availableGroups}
getOptionValue={(group) => group.id}
getOptionLabel={(group) =>
group.display_name || group.name
}
isOptionEqualToValue={(option, optionValue) =>
option.id === optionValue.id
}
renderOption={(option, isSelected) => (
<div className="flex w-full items-center justify-between gap-2">
<AvatarData
title={option.display_name || option.name}
subtitle={getGroupSubtitle(option)}
src={option.avatar_url}
imgFallbackText={option.name}
/>
{isSelected && (
<Check className="size-4 shrink-0" />
)}
</div>
)}
placeholder="Search groups..."
noOptionsText={groupAutocompleteNoOptionsText}
loading={groupsQuery.isLoading}
disabled={groupsQuery.isLoading}
className="w-full"
/>
</div>
<div className="flex-1 space-y-1">
<Label htmlFor="group-override-amount">
Spend Limit ($)
</Label>
<Input
id="group-override-amount"
type="number"
step="0.01"
min="0"
className="h-9 min-w-0 text-[13px]"
value={groupAmount}
onChange={(event) => setGroupAmount(event.target.value)}
placeholder="0.00"
/>
</div>
<div className="flex gap-2 md:pb-0.5">
<Button
size="sm"
type="button"
onClick={() => void handleAddGroupOverride()}
disabled={
upsertGroupOverrideMutation.isPending ||
selectedGroup === null ||
groupAmount.trim() === "" ||
Number.isNaN(Number(groupAmount)) ||
Number(groupAmount) < 0
}
>
{upsertGroupOverrideMutation.isPending ? (
<Spinner loading className="h-4 w-4" />
) : null}
Add
</Button>
<Button
variant="outline"
size="sm"
type="button"
onClick={() => {
setShowGroupForm(false);
setSelectedGroup(null);
setGroupAmount("");
}}
>
Cancel
</Button>
</div>
</div>
</div>
)}
{upsertGroupOverrideMutation.isError && (
<p className="text-xs text-content-destructive">
{getErrorMessage(
upsertGroupOverrideMutation.error,
"Failed to save group override.",
)}
</p>
)}
{groupsQuery.isError && (
<p className="text-xs text-content-destructive">
{getErrorMessage(groupsQuery.error, "Failed to load groups.")}
</p>
)}
</div>
</section>
<section className="space-y-4">
<SectionHeader
label="Per-User Overrides"
description="Override the deployment default spend limit for specific users."
/>
<div className={sectionPanelClassName}>
{overrides.length > 0 ? (
<Table>
<TableHeader>
<TableRow>
<TableHead>User</TableHead>
<TableHead>Spend Limit</TableHead>
<TableHead className="w-[80px]">Actions</TableHead>
</TableRow>
</TableHeader>
<TableBody>
{overrides.map((override) => (
<TableRow key={override.user_id}>
<TableCell>
<AvatarData
title={override.name || override.username}
subtitle={`@${override.username}`}
src={override.avatar_url}
imgFallbackText={override.username}
/>
</TableCell>
<TableCell>
{override.spend_limit_micros !== null
? formatCostMicros(override.spend_limit_micros)
: "Unlimited"}
</TableCell>
<TableCell>
<Button
variant="outline"
size="sm"
type="button"
onClick={() =>
void handleDeleteOverride(override.user_id)
}
disabled={deleteOverrideMutation.isPending}
>
Delete
</Button>
</TableCell>
</TableRow>
))}
</TableBody>
</Table>
) : (
<div className="rounded-lg border border-border bg-surface-secondary px-4 py-6 text-center text-sm text-content-secondary">
No overrides configured.
</div>
)}
{deleteOverrideMutation.isError && (
<p className="text-xs text-content-destructive">
{getErrorMessage(
deleteOverrideMutation.error,
"Failed to delete override.",
)}
</p>
)}
{!showUserForm ? (
<Button
variant="outline"
size="sm"
type="button"
onClick={() => setShowUserForm(true)}
>
Add User
</Button>
) : (
<div className="space-y-3 rounded-lg border border-border bg-surface-secondary/40 p-4">
<div className="flex flex-col gap-3 md:flex-row md:items-end">
<div className="flex-1">
<UserAutocomplete
value={selectedUser}
onChange={setSelectedUser}
label="User"
/>
</div>
<div className="flex-1 space-y-1">
<Label htmlFor="user-override-amount">
Spend Limit ($)
</Label>
<Input
id="user-override-amount"
type="number"
step="0.01"
min="0"
className="h-9 min-w-0 text-[13px]"
value={userOverrideAmount}
onChange={(event) =>
setUserOverrideAmount(event.target.value)
}
placeholder="0.00"
/>
</div>
<div className="flex gap-2 md:pb-0.5">
<Button
size="sm"
type="button"
onClick={() => void handleAddOverride()}
disabled={
upsertOverrideMutation.isPending ||
!selectedUser ||
selectedUserAlreadyOverridden ||
userOverrideAmount.trim() === "" ||
Number.isNaN(Number(userOverrideAmount)) ||
Number(userOverrideAmount) < 0
}
>
{upsertOverrideMutation.isPending ? (
<Spinner loading className="h-4 w-4" />
) : null}
Add
</Button>
<Button
variant="outline"
size="sm"
type="button"
onClick={() => {
setShowUserForm(false);
setSelectedUser(null);
setUserOverrideAmount("");
}}
>
Cancel
</Button>
</div>
</div>
</div>
)}
{selectedUserAlreadyOverridden && (
<p className="text-xs text-content-warning">
This user already has an override.
</p>
)}
{upsertOverrideMutation.isError && (
<p className="text-xs text-content-destructive">
{getErrorMessage(
upsertOverrideMutation.error,
"Failed to save the override.",
)}
</p>
)}
</div>
</section>
</div>
</div>
<div className="sticky bottom-0 flex shrink-0 flex-col gap-2 border-t border-border bg-surface-primary px-6 py-3 sm:flex-row sm:items-center sm:justify-between">
<div className="min-h-4 text-xs">
{updateConfigMutation.isError && (
<p className="m-0 text-content-destructive">
{getErrorMessage(
updateConfigMutation.error,
"Failed to save the default spend limit.",
)}
</p>
)}
{updateConfigMutation.isSuccess && (
<p className="m-0 text-content-success">Saved!</p>
)}
</div>
<Button
size="sm"
type="button"
onClick={() => void handleSaveDefault()}
disabled={updateConfigMutation.isPending || !isAmountValid}
>
{updateConfigMutation.isPending ? (
<Spinner loading className="h-4 w-4" />
) : null}
Save default limit
</Button>
</div>
</div>
);
};