ResolveUsageLimitStatus resolves the current usage-limit status for userID within organizationID. When organizationID is invalid (Valid == false), limits and spend are computed globally across all organizations (legacy behavior). Note: There is a potential race condition where two concurrent messag
(ctx context.Context, db database.Store, userID uuid.UUID, organizationID uuid.NullUUID, now time.Time)
| 64 | // for each period and compare each spend/limit pair Go-side, blocking on |
| 65 | // whichever period is tightest. |
| 66 | func ResolveUsageLimitStatus(ctx context.Context, db database.Store, userID uuid.UUID, organizationID uuid.NullUUID, now time.Time) (*codersdk.ChatUsageLimitStatus, error) { |
| 67 | //nolint:gocritic // AsChatd provides narrowly-scoped daemon access for |
| 68 | // deployment config reads and cross-user chat spend aggregation. |
| 69 | authCtx := dbauthz.AsChatd(ctx) |
| 70 | |
| 71 | config, err := db.GetChatUsageLimitConfig(authCtx) |
| 72 | if err != nil { |
| 73 | if errors.Is(err, sql.ErrNoRows) { |
| 74 | return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits. |
| 75 | } |
| 76 | return nil, err |
| 77 | } |
| 78 | if !config.Enabled { |
| 79 | return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits. |
| 80 | } |
| 81 | |
| 82 | period, ok := mapDBPeriodToSDK(config.Period) |
| 83 | if !ok { |
| 84 | return nil, xerrors.Errorf("invalid chat usage limit period %q", config.Period) |
| 85 | } |
| 86 | |
| 87 | // Resolve effective limit in a single query: |
| 88 | // individual override > group limit > global default. |
| 89 | limitResult, err := db.ResolveUserChatSpendLimit(authCtx, database.ResolveUserChatSpendLimitParams{ |
| 90 | UserID: userID, |
| 91 | OrganizationID: organizationID, |
| 92 | }) |
| 93 | if err != nil { |
| 94 | return nil, err |
| 95 | } |
| 96 | // -1 means limits are disabled (shouldn't happen since we checked |
| 97 | // above, but handle gracefully). |
| 98 | if limitResult.EffectiveLimitMicros < 0 { |
| 99 | return nil, nil //nolint:nilnil // Nil status cleanly signals disabled limits. |
| 100 | } |
| 101 | |
| 102 | start, end := ComputeUsagePeriodBounds(now, period) |
| 103 | |
| 104 | // When the winning limit tier is org-scoped (group), scope spend |
| 105 | // to the same org. When the limit is global (user override or |
| 106 | // deployment default), check spend globally to prevent a user |
| 107 | // from exceeding their limit by spreading spend across orgs. |
| 108 | spendOrgID := organizationID |
| 109 | if limitResult.LimitSource != limitSourceGroup { |
| 110 | spendOrgID = uuid.NullUUID{} |
| 111 | } |
| 112 | |
| 113 | spendTotal, err := db.GetUserChatSpendInPeriod(authCtx, database.GetUserChatSpendInPeriodParams{ |
| 114 | UserID: userID, |
| 115 | OrganizationID: spendOrgID, |
| 116 | StartTime: start, |
| 117 | EndTime: end, |
| 118 | }) |
| 119 | if err != nil { |
| 120 | return nil, err |
| 121 | } |
| 122 | |
| 123 | effectiveLimit := limitResult.EffectiveLimitMicros |