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

Skip to content

feat(coderd): accumulate user daily AI spend on token usage #26741

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 3 commits into from
Jul 2, 2026
Merged

feat(coderd): accumulate user daily AI spend on token usage #26741

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a6c50a4
feat(coderd): accumulate user daily AI spend on token usage
ssncferreira Jun 26, 2026
3d5c156
chore: address comments
ssncferreira Jun 29, 2026
4aa77ba
chore: replace 'unbudgeted' term with 'no effective group'
ssncferreira Jun 29, 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
78 changes: 57 additions & 21 deletions coderd/aibridgedserver/aibridgedserver.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,9 +86,10 @@ type store interface {
// ProviderConfigurator-related queries. InTx wraps the provider and key
// reads in a single read-only transaction; AcquireLock serializes against
// any in-flight env seed holding LockIDAIProvidersEnvSeed.
InTx(func(database.Store) error, *database.TxOptions) error
GetAIProviders(ctx context.Context, arg database.GetAIProvidersParams) ([]database.AIProvider, error)
GetAIProviderKeysByProviderIDs(ctx context.Context, providerIDs []uuid.UUID) ([]database.AIProviderKey, error)

InTx(func(database.Store) error, *database.TxOptions) error
}

type Server struct {
Expand Down Expand Up @@ -301,36 +302,71 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag
}

// Snapshot the effective group, per-token prices and compute cost. A
// missing price row or unbudgeted user yields NULL columns.
// missing price row or no effective group yields NULL columns.
cost, err := s.resolveTokenUsageCost(ctx, intc, in)
if err != nil {
return nil, xerrors.Errorf("resolve token usage cost: %w", err)
}

_, err = s.store.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{
ID: uuid.New(),
InterceptionID: intcID,
ProviderResponseID: in.GetMsgId(),
InputTokens: in.GetInputTokens(),
OutputTokens: in.GetOutputTokens(),
CacheReadInputTokens: in.GetCacheReadInputTokens(),
CacheWriteInputTokens: in.GetCacheWriteInputTokens(),
Metadata: out,
CreatedAt: in.GetCreatedAt().AsTime(),
EffectiveGroupID: cost.effectiveGroupID,
InputPriceMicros: cost.inputPriceMicros,
OutputPriceMicros: cost.outputPriceMicros,
CacheReadPriceMicros: cost.cacheReadPriceMicros,
CacheWritePriceMicros: cost.cacheWritePriceMicros,
CostMicros: cost.costMicros,
})
if err != nil {
return nil, xerrors.Errorf("insert token usage: %w", err)
if err := s.recordTokenUsageAndSpend(ctx, intc, cost, in, out); err != nil {
return nil, xerrors.Errorf("record token usage and spend: %w", err)
}

return &proto.RecordTokenUsageResponse{}, nil
}

// recordTokenUsageAndSpend atomically records the token usage (including the
// interception's cost) and, when the user is budgeted and the computed cost is
// positive, accumulates that cost into the user's daily spend.
func (s *Server) recordTokenUsageAndSpend(ctx context.Context, intc database.AIBridgeInterception, cost tokenUsageCost, in *proto.RecordTokenUsageRequest, metadataJSON []byte) error {
createdAt := in.GetCreatedAt().AsTime()
return s.store.InTx(func(tx database.Store) error {
if _, err := tx.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{
ID: uuid.New(),
InterceptionID: intc.ID,
ProviderResponseID: in.GetMsgId(),
InputTokens: in.GetInputTokens(),
OutputTokens: in.GetOutputTokens(),
CacheReadInputTokens: in.GetCacheReadInputTokens(),
CacheWriteInputTokens: in.GetCacheWriteInputTokens(),
Metadata: metadataJSON,
CreatedAt: createdAt,
EffectiveGroupID: cost.effectiveGroupID,
InputPriceMicros: cost.inputPriceMicros,
OutputPriceMicros: cost.outputPriceMicros,
CacheReadPriceMicros: cost.cacheReadPriceMicros,
CacheWritePriceMicros: cost.cacheWritePriceMicros,
CostMicros: cost.costMicros,
}); err != nil {
return xerrors.Errorf("insert token usage: %w", err)
}

// Skip the spend update when there is no effective group or the interception has no cost.
if !cost.effectiveGroupID.Valid || !cost.costMicros.Valid || cost.costMicros.Int64 <= 0 {
Comment thread
ssncferreira marked this conversation as resolved.
Comment thread
ssncferreira marked this conversation as resolved.
s.logger.Debug(ctx, "skipping spend update",
Comment thread
ssncferreira marked this conversation as resolved.
slog.F("interception_id", intc.ID),
slog.F("initiator_id", intc.InitiatorID),
slog.F("has_effective_group", cost.effectiveGroupID.Valid),
slog.F("has_cost", cost.costMicros.Valid),
slog.F("cost_micros", cost.costMicros.Int64),
)
return nil
}

if _, err := tx.IncrementUserAIDailySpend(ctx, database.IncrementUserAIDailySpendParams{
UserID: intc.InitiatorID,
EffectiveGroupID: cost.effectiveGroupID.UUID,
// Day is derived from the record usage request CreatedAt
Comment thread
ssncferreira marked this conversation as resolved.
// so it matches the token usage row's created_at column.
Day: dbtime.StartOfDay(createdAt.UTC()),

@ssncferreira ssncferreira Jun 29, 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.

Initially, I was using a new clock for Day, but then realized the request's CreatedAt already carries the time the token usage was created on the aibridge side. To keep the database consistent and the spend Day deterministic from the aibridge_token_usages table (source of truth), I switched to deriving it from CreatedAt. This handles the edge case where token usage is created on the aibridge side on day X, but the request reaches coderd after midnight, so time.Now() would return day X+1. Let me know if you see any issues.

CostMicros: cost.costMicros.Int64,
}); err != nil {
return xerrors.Errorf("increment user daily spend: %w", err)
}
return nil
}, nil)
}

func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
Expand Down
Loading
Loading