coder · evgeniy-scherbina · Jun 16, 2026 · Jun 10, 2026 · Jun 12, 2026 · Jun 12, 2026
diff --git a/coderd/aibridge/budget/budget.go b/coderd/aibridge/budget/budget.go
@@ -26,6 +26,13 @@ const (
 	SourceGroup LimitSource = "group"
 )
 
+// Store is the subset of database.Store needed to resolve a user's effective
+// AI budget.
+type Store interface {
+	GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAIBudgetOverride, error)
+	GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error)
+}
+
 // EffectiveBudget is the AI budget that applies to a user after override and
 // policy resolution.
 type EffectiveBudget struct {
@@ -41,7 +48,7 @@ type EffectiveBudget struct {
 // return value is false when no budget is configured for the user. A per-user
 // override wins unconditionally; otherwise the budget is selected from the
 // user's groups according to policy.
-func ResolveUserAIBudget(ctx context.Context, db database.Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, bool, error) {
+func ResolveUserAIBudget(ctx context.Context, db Store, userID uuid.UUID, policy codersdk.AIBudgetPolicy) (EffectiveBudget, bool, error) {
 	// A per-user override always wins.
 	override, err := db.GetUserAIBudgetOverride(ctx, userID)
 	if err == nil {

diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go
@@ -67,6 +67,13 @@ type store interface {
 	UpdateAIBridgeInterceptionEnded(ctx context.Context, intcID database.UpdateAIBridgeInterceptionEndedParams) (database.AIBridgeInterception, error)
 	GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (database.GetAIBridgeInterceptionLineageByToolCallIDRow, error)
 
+	// Cost-attribution queries, used to snapshot price and effective group on
+	// each token usage record.
+	GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (database.AIBridgeInterception, error)
+	GetAIModelPriceByProviderModel(ctx context.Context, arg database.GetAIModelPriceByProviderModelParams) (database.AIModelPrice, error)
+	GetUserAIBudgetOverride(ctx context.Context, userID uuid.UUID) (database.UserAIBudgetOverride, error)
+	GetHighestGroupAIBudgetByUser(ctx context.Context, userID uuid.UUID) (database.GetHighestGroupAIBudgetByUserRow, error)
+
 	// MCPConfigurator-related queries.
 	GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error)
 
@@ -87,6 +94,9 @@ type Server struct {
 	coderMCPConfig    *proto.MCPServerConfig // may be nil if not available
 	structuredLogging bool
 	aiSeatTracker     aiseats.SeatTracker
+	// budgetPolicy selects the effective group when a user belongs to multiple
+	// budgeted groups, used for cost attribution on token usage records.
+	budgetPolicy codersdk.AIBudgetPolicy
 }
 
 func NewServer(lifecycleCtx context.Context, store store, logger slog.Logger, accessURL string,
@@ -110,6 +120,7 @@ func NewServer(lifecycleCtx context.Context, store store, logger slog.Logger, ac
 		externalAuthConfigs: eac,
 		structuredLogging:   bridgeCfg.StructuredLogging.Value(),
 		aiSeatTracker:       aiSeatTracker,
+		budgetPolicy:        codersdk.NewAIBudgetPolicyFromString(bridgeCfg.BudgetPolicy),
 	}
 
 	if bridgeCfg.InjectCoderMCPTools {
@@ -272,6 +283,21 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag
 		s.logger.Warn(ctx, "failed to marshal aibridge metadata from proto to JSON", slog.F("metadata", in), slog.Error(err))
 	}
 
+	// The interception is always recorded before any of its token usages,
+	// so it must exist. It carries the provider, model, and initiator needed
+	// for cost attribution.
+	intc, err := s.store.GetAIBridgeInterceptionByID(ctx, intcID)
+	if err != nil {
+		return nil, xerrors.Errorf("get interception %q: %w", intcID, err)
+	}
+
+	// Snapshot the effective group, per-token prices and compute cost. A
+	// missing price row or unbudgeted user 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,
@@ -282,6 +308,12 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag
 		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)