Thanks to visit codestin.com
Credit goes to github.com

Skip to content

feat: add user AI spend endpoint #26978

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ssncferreira merged 4 commits into from
Jul 8, 2026
Merged

feat: add user AI spend endpoint #26978

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
a6a3d3a
feat: add user AI spend endpoint
ssncferreira Jul 2, 2026
6edb59b
chore: address comments
ssncferreira Jul 6, 2026
f4d8c4f
chore: minor improvements
ssncferreira Jul 6, 2026
9f18e80
chore: improve log messages
ssncferreira Jul 8, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 3 additions & 14 deletions coderd/aibridge/budget/budget.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,17 +15,6 @@ import (
"github.com/coder/coder/v2/codersdk"
)

// LimitSource identifies which tier produced an EffectiveBudget.
type LimitSource string

const (
// SourceUserOverride indicates the budget came from a per-user override.
SourceUserOverride LimitSource = "user_override"
// SourceGroup indicates the budget came from a group budget selected by the
// deployment policy.
SourceGroup LimitSource = "group"
)

// Store is the subset of database.Store needed to resolve a user's effective
// AI budget.
type Store interface {
Expand All @@ -41,7 +30,7 @@ type EffectiveBudget struct {
// SpendLimitMicros is the effective spend limit in micro-units
// (1 unit = 1,000,000).
SpendLimitMicros int64
Source LimitSource
Source codersdk.AIBudgetLimitSource
}

// ResolveUserAIBudget returns the effective AI budget for userID. The second
Expand All @@ -55,7 +44,7 @@ func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy
return EffectiveBudget{
GroupID: override.GroupID,
SpendLimitMicros: override.SpendLimitMicros,
Source: SourceUserOverride,
Source: codersdk.AIBudgetLimitSourceUserOverride,
}, true, nil
}
if !errors.Is(err, sql.ErrNoRows) {
Expand All @@ -75,7 +64,7 @@ func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy
return EffectiveBudget{
GroupID: row.GroupID,
SpendLimitMicros: row.SpendLimitMicros,
Source: SourceGroup,
Source: codersdk.AIBudgetLimitSourceGroup,
}, true, nil
default:
return EffectiveBudget{}, false, xerrors.Errorf("unsupported AI budget policy: %q", policy)
Expand Down
14 changes: 7 additions & 7 deletions coderd/aibridge/budget/budget_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ func TestResolveUserAIBudget(t *testing.T) {
SpendLimitMicros: 1_000_000,
})
require.NoError(t, err)
return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: budget.SourceUserOverride}, true
return user.ID, budget.EffectiveBudget{GroupID: og.ID, SpendLimitMicros: 1_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true
},
},
{
Expand All @@ -82,7 +82,7 @@ func TestResolveUserAIBudget(t *testing.T) {
org := dbgen.Organization(t, db, database.Organization{})
user := dbgen.User(t, db, database.User{})
gid := budgetedGroup(t, ctx, db, org.ID, user.ID, "only", 8_000_000)
return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: budget.SourceGroup}, true
return user.ID, budget.EffectiveBudget{GroupID: gid, SpendLimitMicros: 8_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
},
},
{
Expand All @@ -94,7 +94,7 @@ func TestResolveUserAIBudget(t *testing.T) {
budgetedGroup(t, ctx, db, org.ID, user.ID, "low", 5_000_000)
budgetedGroup(t, ctx, db, org.ID, user.ID, "mid", 20_000_000)
high := budgetedGroup(t, ctx, db, org.ID, user.ID, "high", 50_000_000)
return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: budget.SourceGroup}, true
return user.ID, budget.EffectiveBudget{GroupID: high, SpendLimitMicros: 50_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
},
},
{
Expand All @@ -106,7 +106,7 @@ func TestResolveUserAIBudget(t *testing.T) {
// Equal limits; "alpha" must win over "beta" by name ascending.
alpha := budgetedGroup(t, ctx, db, org.ID, user.ID, "alpha", 10_000_000)
budgetedGroup(t, ctx, db, org.ID, user.ID, "beta", 10_000_000)
return user.ID, budget.EffectiveBudget{GroupID: alpha, SpendLimitMicros: 10_000_000, Source: budget.SourceGroup}, true
return user.ID, budget.EffectiveBudget{GroupID: alpha, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
},
},
{
Expand All @@ -124,7 +124,7 @@ func TestResolveUserAIBudget(t *testing.T) {
if bytes.Compare(g2[:], g1[:]) < 0 {
winner = g2
}
return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: budget.SourceGroup}, true
return user.ID, budget.EffectiveBudget{GroupID: winner, SpendLimitMicros: 10_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
},
},
{
Expand All @@ -147,7 +147,7 @@ func TestResolveUserAIBudget(t *testing.T) {
// Membership is via organization_members only (no group_members row),
// exercising the org-members half of group_members_expanded.
everyoneID := budgetedEveryoneGroup(t, ctx, db, org.ID, user.ID, 7_000_000)
return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 7_000_000, Source: budget.SourceGroup}, true
return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 7_000_000, Source: codersdk.AIBudgetLimitSourceGroup}, true
},
},
{
Expand All @@ -165,7 +165,7 @@ func TestResolveUserAIBudget(t *testing.T) {
SpendLimitMicros: 2_000_000,
})
require.NoError(t, err)
return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: budget.SourceUserOverride}, true
return user.ID, budget.EffectiveBudget{GroupID: everyoneID, SpendLimitMicros: 2_000_000, Source: codersdk.AIBudgetLimitSourceUserOverride}, true
},
},
{
Expand Down
85 changes: 85 additions & 0 deletions coderd/apidoc/docs.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

78 changes: 78 additions & 0 deletions coderd/apidoc/swagger.json
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

62 changes: 62 additions & 0 deletions codersdk/aibridge.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,49 @@ import (
"golang.org/x/xerrors"
)

// AIBudgetLimitSource identifies which tier produced the user's
// effective budget limit.
type AIBudgetLimitSource string

const (
// AIBudgetLimitSourceUserOverride indicates the limit came from a
// per-user override.
AIBudgetLimitSourceUserOverride AIBudgetLimitSource = "user_override"

@coder-agents-review coder-agents-review Bot Jul 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note [CRF-8] The pre-existing manual frontend type UserAISpend in site/src/api/api.ts:231 uses "override" for the user-override source, while this constant is "user_override". Currently harmless (the frontend never inspects limit_source), but the migration from the manual type to the generated UserAISpendStatus will change the wire value. The AIGOV-473 ticket should capture this discrepancy. (Pariston)

🤖

@ssncferreira ssncferreira Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@EhabY FYI

// AIBudgetLimitSourceGroup indicates the limit came from a group
// budget selected by the deployment budget policy.
AIBudgetLimitSourceGroup AIBudgetLimitSource = "group"
)

// UserAIBudgetSummary is the effective AI budget for a user. When no
// budget applies, all fields except UserID are null.
type UserAIBudgetSummary struct {
UserID uuid.UUID `json:"user_id" format:"uuid"`
// EffectiveGroupID is the group the spend is attributed to. Null when
// no budget applies.
EffectiveGroupID *uuid.UUID `json:"effective_group_id" format:"uuid"`
// SpendLimitMicros is the effective spend limit in micro-units.
// Null when no budget applies to the user (unlimited).
SpendLimitMicros *int64 `json:"spend_limit_micros"`
// LimitSource identifies which tier produced the limit. Null when no
// budget applies.
LimitSource *AIBudgetLimitSource `json:"limit_source"`
}

// UserAISpendStatus is the current AI spend snapshot for a user within
// the active budget period.
type UserAISpendStatus struct {
UserAIBudgetSummary
// CurrentSpendMicros is the user's spend on their effective group over
// the current budget period.
CurrentSpendMicros int64 `json:"current_spend_micros"`
Comment thread
ssncferreira marked this conversation as resolved.
// PeriodStart is the inclusive lower bound of the current budget
// period.
PeriodStart time.Time `json:"period_start" format:"date-time"`
// PeriodEnd is the exclusive upper bound of the current budget
// period.
PeriodEnd time.Time `json:"period_end" format:"date-time"`
}

type AIBridgeSession struct {
ID string `json:"id"`
Initiator MinimalUser `json:"initiator"`
Expand Down Expand Up @@ -375,3 +418,22 @@ func (c *Client) DeleteUserAIBudgetOverride(ctx context.Context, user uuid.UUID)
}
return nil
}

// UserAISpendStatus returns the current AI spend snapshot for the given user
// within the active budget period.
func (c *Client) UserAISpendStatus(ctx context.Context, user uuid.UUID) (UserAISpendStatus, error) {
res, err := c.Request(ctx, http.MethodGet,
fmt.Sprintf("/api/v2/users/%s/ai/spend", user.String()),
nil,
)
if err != nil {
return UserAISpendStatus{}, xerrors.Errorf("make request: %w", err)
}
defer res.Body.Close()

if res.StatusCode != http.StatusOK {
return UserAISpendStatus{}, ReadBodyAsError(res)
}
var resp UserAISpendStatus
return resp, json.NewDecoder(res.Body).Decode(&resp)
}
Loading
Loading