diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 8b539153e98..88fdc8e0900 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -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 { @@ -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 { + s.logger.Debug(ctx, "skipping spend update", + 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 + // so it matches the token usage row's created_at column. + Day: dbtime.StartOfDay(createdAt.UTC()), + 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) diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 725f4e954af..b886f569615 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -1170,6 +1170,8 @@ func TestRecordTokenUsage(t *testing.T) { "key": mustMarshalAny(t, &structpb.Value{Kind: &structpb.Value_StringValue{StringValue: "value"}}), } metadataJSON = `{"key":"value"}` + // Use fixed dates to keep the test deterministic. + now = time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC) ) testRecordMethod(t, @@ -1178,7 +1180,8 @@ func TestRecordTokenUsage(t *testing.T) { }, []testRecordMethodCase[*proto.RecordTokenUsageRequest]{ { - name: "valid token usage with null cost", + // Budget resolves via group lookup, model is priced. + name: "valid token usage with effective group and cost", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", @@ -1186,63 +1189,7 @@ func TestRecordTokenUsage(t *testing.T) { OutputTokens: 200, CacheReadInputTokens: 50, CacheWriteInputTokens: 10, - Metadata: metadataProto, - CreatedAt: timestamppb.Now(), - }, - setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { - interceptionID, err := uuid.Parse(req.GetInterceptionId()) - assert.NoError(t, err, "parse interception UUID") - - // No budget configured and no price row: tokens recorded - // with NULL cost, prices, and group attribution. - intc := newTestInterception(interceptionID) - expectTokenUsageCostLookups(db, intc, nil, nil, nil) - - db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { - if !assert.NotEqual(t, uuid.Nil, p.ID, "ID") || - !assert.Equal(t, interceptionID, p.InterceptionID, "interception ID") || - !assert.Equal(t, req.GetMsgId(), p.ProviderResponseID, "provider response ID") || - !assert.Equal(t, req.GetInputTokens(), p.InputTokens, "input tokens") || - !assert.Equal(t, req.GetOutputTokens(), p.OutputTokens, "output tokens") || - !assert.Equal(t, req.GetCacheReadInputTokens(), p.CacheReadInputTokens, "cache read input tokens") || - !assert.Equal(t, req.GetCacheWriteInputTokens(), p.CacheWriteInputTokens, "cache write input tokens") || - !assert.JSONEq(t, metadataJSON, string(p.Metadata), "metadata") || - !assert.WithinDuration(t, req.GetCreatedAt().AsTime(), p.CreatedAt, time.Second, "created at") || - !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") || - !assert.False(t, p.InputPriceMicros.Valid, "input price null") || - !assert.False(t, p.OutputPriceMicros.Valid, "output price null") || - !assert.False(t, p.CacheReadPriceMicros.Valid, "cache read price null") || - !assert.False(t, p.CacheWritePriceMicros.Valid, "cache write price null") || - !assert.False(t, p.CostMicros.Valid, "cost null") { - return false - } - return true - })).Return(database.AIBridgeTokenUsage{ - ID: uuid.New(), - InterceptionID: interceptionID, - ProviderResponseID: req.GetMsgId(), - InputTokens: req.GetInputTokens(), - OutputTokens: req.GetOutputTokens(), - CacheReadInputTokens: req.GetCacheReadInputTokens(), - CacheWriteInputTokens: req.GetCacheWriteInputTokens(), - Metadata: pqtype.NullRawMessage{ - RawMessage: json.RawMessage(metadataJSON), - Valid: true, - }, - CreatedAt: req.GetCreatedAt().AsTime(), - }, nil) - }, - }, - { - name: "valid token usage with cost", - request: &proto.RecordTokenUsageRequest{ - InterceptionId: uuid.NewString(), - MsgId: "msg_123", - InputTokens: 100, - OutputTokens: 200, - CacheReadInputTokens: 50, - CacheWriteInputTokens: 10, - CreatedAt: timestamppb.Now(), + CreatedAt: timestamppb.New(now), }, setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { interceptionID, err := uuid.Parse(req.GetInterceptionId()) @@ -1265,6 +1212,10 @@ func TestRecordTokenUsage(t *testing.T) { // input 300 + output 1200 + cache read 15 + cache write 40. const wantCost int64 = 1555 + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") || !assert.Equal(t, price.InputPrice, p.InputPriceMicros, "input price") || @@ -1276,15 +1227,23 @@ func TestRecordTokenUsage(t *testing.T) { } return true })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), database.IncrementUserAIDailySpendParams{ + UserID: intc.InitiatorID, + EffectiveGroupID: groupID, + Day: now.UTC().Truncate(24 * time.Hour), + CostMicros: wantCost, + }).Return(database.AIUserDailySpend{}, nil) }, }, { - name: "valid token usage with user override", + // Budget resolves via user override, model is priced. + name: "valid token usage with user override and cost", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", InputTokens: 100, - CreatedAt: timestamppb.Now(), + CreatedAt: timestamppb.New(now), }, setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { interceptionID, err := uuid.Parse(req.GetInterceptionId()) @@ -1305,18 +1264,33 @@ func TestRecordTokenUsage(t *testing.T) { // No group expectTokenUsageCostLookups(db, intc, override, nil, price) + // input 300. + const wantCost int64 = 300 + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { // Override group wins. if !assert.Equal(t, uuid.NullUUID{UUID: overrideGroupID, Valid: true}, p.EffectiveGroupID, "effective group ID") || - !assert.Equal(t, sql.NullInt64{Int64: 300, Valid: true}, p.CostMicros, "cost") { + !assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") { return false } return true })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), database.IncrementUserAIDailySpendParams{ + UserID: intc.InitiatorID, + EffectiveGroupID: overrideGroupID, + Day: now.UTC().Truncate(24 * time.Hour), + CostMicros: wantCost, + }).Return(database.AIUserDailySpend{}, nil) }, }, { - name: "valid token usage with budget but no price", + // Model has no price row, so cost is NULL. + name: "valid token usage with effective group and no price", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", @@ -1338,6 +1312,10 @@ func TestRecordTokenUsage(t *testing.T) { // return on sql.ErrNoRows, while prices and cost stay NULL. expectTokenUsageCostLookups(db, intc, nil, group, nil) + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") || !assert.False(t, p.InputPriceMicros.Valid, "input price null") || @@ -1349,10 +1327,14 @@ func TestRecordTokenUsage(t *testing.T) { } return true })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + // Spend update is skipped because cost is NULL. + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) }, }, { - name: "valid token usage with price but no budget", + // Price row exists with NULL columns, so cost is 0 (Valid). + name: "valid token usage with effective group and NULL prices", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", @@ -1367,37 +1349,48 @@ func TestRecordTokenUsage(t *testing.T) { assert.NoError(t, err, "parse interception UUID") intc := newTestInterception(interceptionID) + groupID := uuid.New() + group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000} + // The price row exists but every price column is NULL. Each + // category is treated as zero for cost, so the columns are + // recorded as NULL while cost is recorded as 0 (not NULL): + // cost's NULL-ness tracks price row presence, not the price + // values. price := &database.AIModelPrice{ Provider: intc.Provider, Model: intc.Model, - InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, - OutputPrice: sql.NullInt64{Int64: 6_000_000, Valid: true}, - CacheReadPrice: sql.NullInt64{Int64: 300_000, Valid: true}, - CacheWritePrice: sql.NullInt64{Int64: 4_000_000, Valid: true}, + InputPrice: sql.NullInt64{Valid: false}, + OutputPrice: sql.NullInt64{Valid: false}, + CacheReadPrice: sql.NullInt64{Valid: false}, + CacheWritePrice: sql.NullInt64{Valid: false}, } - // No budget configured, but the model is priced: cost is - // computed independently of budget resolution, and the group - // attribution stays NULL. - expectTokenUsageCostLookups(db, intc, nil, nil, price) + expectTokenUsageCostLookups(db, intc, nil, group, price) - // input 300 + output 1200 + cache read 15 + cache write 40. - const wantCost int64 = 1555 + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { - if !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") || - !assert.Equal(t, price.InputPrice, p.InputPriceMicros, "input price") || - !assert.Equal(t, price.OutputPrice, p.OutputPriceMicros, "output price") || - !assert.Equal(t, price.CacheReadPrice, p.CacheReadPriceMicros, "cache read price") || - !assert.Equal(t, price.CacheWritePrice, p.CacheWritePriceMicros, "cache write price") || - !assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") { + if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") || + !assert.False(t, p.InputPriceMicros.Valid, "input price null") || + !assert.False(t, p.OutputPriceMicros.Valid, "output price null") || + !assert.False(t, p.CacheReadPriceMicros.Valid, "cache read price null") || + !assert.False(t, p.CacheWritePriceMicros.Valid, "cache write price null") || + // Cost is recorded as 0 (Valid), not NULL, because the + // price row exists. + !assert.Equal(t, sql.NullInt64{Int64: 0, Valid: true}, p.CostMicros, "cost zero") { return false } return true })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + // Spend update is skipped because cost is 0. + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) }, }, { - name: "valid token usage with zero prices", + // Model is priced at zero, so cost is 0 (Valid). + name: "valid token usage with effective group and zero prices", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", @@ -1412,6 +1405,8 @@ func TestRecordTokenUsage(t *testing.T) { assert.NoError(t, err, "parse interception UUID") intc := newTestInterception(interceptionID) + groupID := uuid.New() + group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: groupID, SpendLimitMicros: 1_000_000_000} // A model priced at zero is distinct from an unpriced model: // the price columns and cost are recorded as 0, not NULL. price := &database.AIModelPrice{ @@ -1422,11 +1417,16 @@ func TestRecordTokenUsage(t *testing.T) { CacheReadPrice: sql.NullInt64{Int64: 0, Valid: true}, CacheWritePrice: sql.NullInt64{Int64: 0, Valid: true}, } - expectTokenUsageCostLookups(db, intc, nil, nil, price) + expectTokenUsageCostLookups(db, intc, nil, group, price) + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { zero := sql.NullInt64{Int64: 0, Valid: true} - if !assert.Equal(t, zero, p.InputPriceMicros, "input price zero") || + if !assert.Equal(t, uuid.NullUUID{UUID: groupID, Valid: true}, p.EffectiveGroupID, "effective group ID") || + !assert.Equal(t, zero, p.InputPriceMicros, "input price zero") || !assert.Equal(t, zero, p.OutputPriceMicros, "output price zero") || !assert.Equal(t, zero, p.CacheReadPriceMicros, "cache read price zero") || !assert.Equal(t, zero, p.CacheWritePriceMicros, "cache write price zero") || @@ -1436,10 +1436,14 @@ func TestRecordTokenUsage(t *testing.T) { } return true })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + // Spend update is skipped because cost is 0. + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) }, }, { - name: "valid token usage with all null prices", + // No budget configured, model is priced: group is NULL but cost is computed. + name: "valid token usage with no budget and cost", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", @@ -1454,33 +1458,104 @@ func TestRecordTokenUsage(t *testing.T) { assert.NoError(t, err, "parse interception UUID") intc := newTestInterception(interceptionID) - // The price row exists but every price column is NULL. Each - // category is treated as zero for cost, so the columns are - // recorded as NULL while cost is recorded as 0 (not NULL): - // cost's NULL-ness tracks price row presence, not the price - // values. price := &database.AIModelPrice{ Provider: intc.Provider, Model: intc.Model, - InputPrice: sql.NullInt64{Valid: false}, - OutputPrice: sql.NullInt64{Valid: false}, - CacheReadPrice: sql.NullInt64{Valid: false}, - CacheWritePrice: sql.NullInt64{Valid: false}, + InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, + OutputPrice: sql.NullInt64{Int64: 6_000_000, Valid: true}, + CacheReadPrice: sql.NullInt64{Int64: 300_000, Valid: true}, + CacheWritePrice: sql.NullInt64{Int64: 4_000_000, Valid: true}, } + // No budget configured, but the model is priced: cost is + // computed independently of budget resolution, and the group + // attribution stays NULL. expectTokenUsageCostLookups(db, intc, nil, nil, price) + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + + // input 300 + output 1200 + cache read 15 + cache write 40. + const wantCost int64 = 1555 + + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { + if !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") || + !assert.Equal(t, price.InputPrice, p.InputPriceMicros, "input price") || + !assert.Equal(t, price.OutputPrice, p.OutputPriceMicros, "output price") || + !assert.Equal(t, price.CacheReadPrice, p.CacheReadPriceMicros, "cache read price") || + !assert.Equal(t, price.CacheWritePrice, p.CacheWritePriceMicros, "cache write price") || + !assert.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, p.CostMicros, "cost") { + return false + } + return true + })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + + // Spend update is skipped because the effective group is NULL. + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) + }, + }, + { + // No budget and no price row: group and cost are NULL. + name: "valid token usage with no budget and no price", + request: &proto.RecordTokenUsageRequest{ + InterceptionId: uuid.NewString(), + MsgId: "msg_123", + InputTokens: 100, + OutputTokens: 200, + CacheReadInputTokens: 50, + CacheWriteInputTokens: 10, + Metadata: metadataProto, + CreatedAt: timestamppb.Now(), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { + interceptionID, err := uuid.Parse(req.GetInterceptionId()) + assert.NoError(t, err, "parse interception UUID") + + // No budget configured and no price row: tokens recorded + // with NULL cost, prices, and group attribution. + intc := newTestInterception(interceptionID) + expectTokenUsageCostLookups(db, intc, nil, nil, nil) + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Cond(func(p database.InsertAIBridgeTokenUsageParams) bool { - if !assert.False(t, p.InputPriceMicros.Valid, "input price null") || + if !assert.NotEqual(t, uuid.Nil, p.ID, "ID") || + !assert.Equal(t, interceptionID, p.InterceptionID, "interception ID") || + !assert.Equal(t, req.GetMsgId(), p.ProviderResponseID, "provider response ID") || + !assert.Equal(t, req.GetInputTokens(), p.InputTokens, "input tokens") || + !assert.Equal(t, req.GetOutputTokens(), p.OutputTokens, "output tokens") || + !assert.Equal(t, req.GetCacheReadInputTokens(), p.CacheReadInputTokens, "cache read input tokens") || + !assert.Equal(t, req.GetCacheWriteInputTokens(), p.CacheWriteInputTokens, "cache write input tokens") || + !assert.JSONEq(t, metadataJSON, string(p.Metadata), "metadata") || + !assert.WithinDuration(t, req.GetCreatedAt().AsTime(), p.CreatedAt, time.Second, "created at") || + !assert.False(t, p.EffectiveGroupID.Valid, "effective group ID null") || + !assert.False(t, p.InputPriceMicros.Valid, "input price null") || !assert.False(t, p.OutputPriceMicros.Valid, "output price null") || !assert.False(t, p.CacheReadPriceMicros.Valid, "cache read price null") || !assert.False(t, p.CacheWritePriceMicros.Valid, "cache write price null") || - // Cost is recorded as 0 (Valid), not NULL, because the - // price row exists. - !assert.Equal(t, sql.NullInt64{Int64: 0, Valid: true}, p.CostMicros, "cost zero") { + !assert.False(t, p.CostMicros.Valid, "cost null") { return false } return true - })).Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + })).Return(database.AIBridgeTokenUsage{ + ID: uuid.New(), + InterceptionID: interceptionID, + ProviderResponseID: req.GetMsgId(), + InputTokens: req.GetInputTokens(), + OutputTokens: req.GetOutputTokens(), + CacheReadInputTokens: req.GetCacheReadInputTokens(), + CacheWriteInputTokens: req.GetCacheWriteInputTokens(), + Metadata: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(metadataJSON), + Valid: true, + }, + CreatedAt: req.GetCreatedAt().AsTime(), + }, nil) + + // Spend update is skipped because the effective group and cost are NULL. + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()).Times(0) }, }, { @@ -1541,7 +1616,7 @@ func TestRecordTokenUsage(t *testing.T) { expectedErr: "resolve token usage cost", }, { - name: "insert error", + name: "insert token usage error", request: &proto.RecordTokenUsageRequest{ InterceptionId: uuid.NewString(), MsgId: "msg_123", @@ -1554,17 +1629,53 @@ func TestRecordTokenUsage(t *testing.T) { assert.NoError(t, err, "parse interception UUID") expectTokenUsageCostLookups(db, newTestInterception(interceptionID), nil, nil, nil) + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeTokenUsage{}, sql.ErrConnDone) }, expectedErr: "insert token usage", }, + { + name: "increment user daily spend error", + request: &proto.RecordTokenUsageRequest{ + InterceptionId: uuid.NewString(), + MsgId: "msg_123", + InputTokens: 100, + CreatedAt: timestamppb.Now(), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { + interceptionID, err := uuid.Parse(req.GetInterceptionId()) + assert.NoError(t, err, "parse interception UUID") + + intc := newTestInterception(interceptionID) + group := &database.GetHighestGroupAIBudgetByUserRow{GroupID: uuid.New(), SpendLimitMicros: 1_000_000_000} + price := &database.AIModelPrice{ + Provider: intc.Provider, + Model: intc.Model, + InputPrice: sql.NullInt64{Int64: 3_000_000, Valid: true}, + } + expectTokenUsageCostLookups(db, intc, nil, group, price) + + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()). + Return(database.AIBridgeTokenUsage{ID: uuid.New(), InterceptionID: interceptionID}, nil) + db.EXPECT().IncrementUserAIDailySpend(gomock.Any(), gomock.Any()). + Return(database.AIUserDailySpend{}, sql.ErrConnDone) + }, + expectedErr: "increment user daily spend", + }, }, ) } // TestRecordTokenUsageAuthorized exercises RecordTokenUsage end-to-end against a // real database through the dbauthz layer as subjectAibridged. This catches missing -// RBAC grants on the aibridged subject and verifies the cost columns round-trip to storage. +// RBAC grants on the aibridged subject and verifies the cost columns round-trip +// to storage along with the daily spend row. func TestRecordTokenUsageAuthorized(t *testing.T) { t.Parallel() @@ -1606,6 +1717,9 @@ func TestRecordTokenUsageAuthorized(t *testing.T) { Model: model, }, nil) + // Use fixed dates to keep the test deterministic. + now := time.Date(2026, 6, 25, 14, 30, 0, 0, time.UTC) + // The server runs every store call as subjectAibridged via the authzDB. srv, err := aibridgedserver.NewServer(ctx, authzDB, logger, "/", codersdk.AIBridgeConfig{}, nil, requiredExperiments, agplaiseats.Noop{}) require.NoError(t, err) @@ -1617,23 +1731,37 @@ func TestRecordTokenUsageAuthorized(t *testing.T) { OutputTokens: 200, CacheReadInputTokens: 50, CacheWriteInputTokens: 10, - CreatedAt: timestamppb.Now(), + CreatedAt: timestamppb.New(now), }) require.NoError(t, err, "record token usage") // Read the persisted row back via the raw store and verify the snapshot. - usages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID) + tokenUsages, err := rawDB.GetAIBridgeTokenUsagesByInterceptionID(ctx, intc.ID) require.NoError(t, err) - require.Len(t, usages, 1) - got := usages[0] - - require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, got.EffectiveGroupID, "effective group") - require.Equal(t, sql.NullInt64{Int64: 3_000_000, Valid: true}, got.InputPriceMicros, "input price") - require.Equal(t, sql.NullInt64{Int64: 6_000_000, Valid: true}, got.OutputPriceMicros, "output price") - require.Equal(t, sql.NullInt64{Int64: 300_000, Valid: true}, got.CacheReadPriceMicros, "cache read price") - require.Equal(t, sql.NullInt64{Int64: 4_000_000, Valid: true}, got.CacheWritePriceMicros, "cache write price") + require.Len(t, tokenUsages, 1) + tokenUsage := tokenUsages[0] + + require.Equal(t, uuid.NullUUID{UUID: group.ID, Valid: true}, tokenUsage.EffectiveGroupID, "effective group") + require.Equal(t, sql.NullInt64{Int64: 3_000_000, Valid: true}, tokenUsage.InputPriceMicros, "input price") + require.Equal(t, sql.NullInt64{Int64: 6_000_000, Valid: true}, tokenUsage.OutputPriceMicros, "output price") + require.Equal(t, sql.NullInt64{Int64: 300_000, Valid: true}, tokenUsage.CacheReadPriceMicros, "cache read price") + require.Equal(t, sql.NullInt64{Int64: 4_000_000, Valid: true}, tokenUsage.CacheWritePriceMicros, "cache write price") // input 300 + output 1200 + cache read 15 + cache write 40. - require.Equal(t, sql.NullInt64{Int64: 1555, Valid: true}, got.CostMicros, "cost") + const wantCost int64 = 1555 + require.Equal(t, sql.NullInt64{Int64: wantCost, Valid: true}, tokenUsage.CostMicros, "cost") + + // The daily spend row was incremented for (user, group, today) by the same cost. + today := now.UTC().Truncate(24 * time.Hour) + spend, err := rawDB.GetUserAISpendSince(ctx, database.GetUserAISpendSinceParams{ + UserID: user.ID, + EffectiveGroupID: group.ID, + PeriodStart: today, + }) + require.NoError(t, err, "get user AI spend since") + require.Equal(t, user.ID, spend.UserID, "user ID") + require.Equal(t, group.ID, spend.EffectiveGroupID, "effective group ID") + require.True(t, today.Equal(spend.PeriodStart), "period start: want %s, got %s", today, spend.PeriodStart) + require.Equal(t, wantCost, spend.SpendMicros, "spend micros") } // newTestInterception returns an interception with a fixed initiator, provider, @@ -2162,6 +2290,9 @@ func TestStructuredLogging(t *testing.T) { structuredLogging: true, setupMocks: func(db *dbmock.MockStore, intcID uuid.UUID) { expectTokenUsageCostLookups(db, newTestInterception(intcID), nil, nil, nil) + db.EXPECT().InTx(gomock.Any(), nil).DoAndReturn( + func(fn func(database.Store) error, _ *database.TxOptions) error { return fn(db) }, + ) db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeTokenUsage{ ID: uuid.New(), InterceptionID: intcID,