Compare commits
24 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 713facdad5 | |||
| 9a70780820 | |||
| da3aca6d52 | |||
| 8c15ad5950 | |||
| b3ad420db2 | |||
| 6369ba1cf8 | |||
| 56e369b0d6 | |||
| 6fbbd155e7 | |||
| 71dca287f9 | |||
| 8156ef36c9 | |||
| 796f4ac65d | |||
| c3ae8228ef | |||
| 095d18dd10 | |||
| f470475e45 | |||
| 01010fbafc | |||
| fb385c60d0 | |||
| 2bfc2eb6c7 | |||
| efcf39e98c | |||
| 48de0677a4 | |||
| f67c20b71b | |||
| 9ef9e78bf2 | |||
| e1c1226aee | |||
| 960bafddf2 | |||
| 596752c4f2 |
@@ -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(
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
|
||||
Generated
+85
@@ -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');
|
||||
@@ -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"`
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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";
|
||||
|
||||
@@ -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,
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
@@ -10,7 +10,8 @@ type GroupSortOrder = "asc" | "desc";
|
||||
|
||||
export const groupsQueryKey = ["groups"];
|
||||
|
||||
const groups = () => {
|
||||
/** @public */
|
||||
export const groups = () => {
|
||||
return {
|
||||
queryKey: groupsQueryKey,
|
||||
queryFn: () => API.getGroups(),
|
||||
|
||||
Generated
+92
@@ -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;
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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>
|
||||
);
|
||||
};
|
||||
Reference in New Issue
Block a user