Documentation
¶
Overview ¶
Package usage's threshold scheduler implements the per-month spend notification dial wired into Manager.Add (token-cost-telemetry-01KQ8TD7 WP06).
Architecture:
Each Manager.Add call invokes Checker.Check after the per-turn row has been persisted. Check reads the calendar-month total of cost_usd from session_messages (assistant rows only), divides by the user-configured MonthlyCostNotifyUSD threshold, and walks the locked tier list [50, 80, 100, 150, 200].
For each tier whose pct is met by the current month's spend AND which has not yet been recorded in cost_threshold_fired for the current YYYY-MM, Check INSERT-OR-IGNOREs a row and (on a fresh insert) publishes a `cost.threshold.crossed` event through the injected Publisher.
Calendar boundary uses time.Now().Local() so the rollover matches how a human reads a billing month. A turn at 23:59:59 local on April 30 fires under "2026-04"; the next turn at 00:00:01 May 1 queries "2026-05" which has zero spend, so the May tier list starts fresh.
Idempotency: cost_threshold_fired (year_month, pct) is the primary key. INSERT OR IGNORE makes restart + replay safe.
Timezone: time.Now().Local() is captured ONCE per Check call and reused for both the SUM-by-month query window and the year_month key. No drift inside a single Check.
Package usage implements per-session token + cost aggregation for the token-cost-telemetry-01KQ8TD7 mission.
Architecture: UsageTurn rows are written to session_messages (via the four nullable columns added in migration 0314) on each assistant turn completion. GetSession reads those columns back via SQL aggregation. The Manager is wired into the chat runner's driveRun finish path.
Index ¶
Constants ¶
const ThresholdEventTopic = "cost.threshold.crossed"
ThresholdEventTopic is the broker topic the checker publishes on when a new tier crosses. Frontend's CostThresholdToast listens on this topic via the existing event-stream composable.
Variables ¶
var ThresholdTiers = []int{50, 80, 100, 150, 200}
ThresholdTiers is the locked escalating-percent list the spec calls out (plan §2.5). Each tier fires at most once per (year_month, pct) thanks to the cost_threshold_fired primary key.
Functions ¶
Types ¶
type Aggregate ¶
type Aggregate struct {
// PromptTokens is the sum of all prompt_tokens rows for the session.
PromptTokens int
// CompletionTokens is the sum of all completion_tokens rows.
CompletionTokens int
// TotalTokens is PromptTokens + CompletionTokens.
TotalTokens int
// CostUSD is the sum of all non-NULL cost_usd rows (0.0 when no rows
// have cost data yet).
CostUSD float64
// CostSource summarises the cost_source values across all turns:
// "provider" — every turn with a cost had source=provider
// "derived" — every turn with a cost had source=derived
// "mixed" — some turns have source=provider, others derived
// "unknown" — no turns have any cost data
CostSource string
// MessageCount is the number of assistant turns that contributed usage.
MessageCount int
}
Aggregate is the per-session cumulative total returned by GetSession.
type Checker ¶
type Checker struct {
// contains filtered or unexported fields
}
Checker is the threshold scheduler. Built once at API construction and called from Manager.Add's tail.
The checker holds no per-month state: every Check call re-reads the total from MonthlyTotalReader and the threshold from ThresholdReader, so changes to the dial take effect on the very next turn.
func NewChecker ¶
func NewChecker(cfg CheckerConfig) (*Checker, error)
NewChecker constructs a Checker. Returns an error when any required callback is missing — the caller should fail loudly at boot rather than silently shipping a no-op scheduler.
func NewCheckerFromManager ¶
func NewCheckerFromManager(m ManagerWithChecker, threshold ThresholdReader, publisher Publisher) (*Checker, error)
NewCheckerFromManager is the convenience constructor the rpc layer uses at boot. It wires the manager's MonthlyTotalUSD + RecordFired methods straight through and threads the supplied threshold + publisher callbacks. Returns a ready-to-use *Checker that can be installed via Manager.SetThresholdChecker.
func (*Checker) Check ¶
Check evaluates the calendar-month total against the dial and publishes a `cost.threshold.crossed` event for every newly-crossed tier. Returns the slice of tiers actually fired on this call (mostly useful for tests; production callers can ignore it).
Errors from the underlying readers are returned wrapped — the chat runner logs and continues so a transient SQL failure does not break the user's chat turn.
type CheckerConfig ¶
type CheckerConfig struct {
Threshold ThresholdReader
Monthly MonthlyTotalReader
Fired FiredRecorder
Publisher Publisher
// Tiers overrides the default ThresholdTiers. Empty/nil falls back
// to the canonical [50, 80, 100, 150, 200].
Tiers []int
// NowLocal is the time source. nil falls back to time.Now().Local().
// Tests inject a fake clock here.
NowLocal func() time.Time
}
CheckerConfig wires the checker. All fields except Publisher are required; passing a nil Publisher disables event emission but keeps the dedup writes intact.
type FiredRecorder ¶
type FiredRecorder func(ctx context.Context, ym string, pct int, firedAt time.Time) (inserted bool, err error)
FiredRecorder is the dedup primitive the checker uses to attempt recording a (year_month, pct) firing. It MUST INSERT OR IGNORE and return inserted=true ONLY when a fresh row was written. A second caller for the same (year_month, pct) MUST return inserted=false.
type Manager ¶
type Manager interface {
// Add persists usage for one assistant turn. MessageID must be the
// session_messages.id of the assistant row that was just written.
// Returns nil when the write succeeded or when telemetry is disabled.
Add(ctx context.Context, turn UsageTurn) error
// GetSession returns the cumulative aggregate for a session.
GetSession(ctx context.Context, sessionID string) (Aggregate, error)
// SetThresholdChecker wires the WP06 threshold scheduler into
// Add's tail. Safe to call once after construction (the rpc layer
// builds the checker after the broker + settings store exist, so
// this setter avoids a circular dependency at usage.New time).
// Passing nil removes the scheduler; per-turn rows still record.
SetThresholdChecker(c *Checker)
// MonthlyTotalUSD returns the total cost_usd in USD across every
// session for the calendar month containing now (in now's
// location). Drives the WP06 threshold scheduler's
// MonthlyTotalReader callback and is exposed for tests +
// integration tests.
MonthlyTotalUSD(ctx context.Context, now time.Time) (float64, error)
}
Manager is the per-process token + cost aggregation surface.
type ManagerWithChecker ¶
type ManagerWithChecker interface {
MonthlyTotalUSD(ctx context.Context, now time.Time) (float64, error)
RecordFired(ctx context.Context, ym string, pct int, firedAt time.Time) (bool, error)
}
ManagerWithChecker is the subset of Manager the threshold helper needs to wire a Checker. Both the production sqlManager and any test fake satisfying these methods qualify.
type MonthlyTotalReader ¶
MonthlyTotalReader returns the total cost_usd in USD for the calendar month containing now (in now's local timezone). Empty session_messages or NULL cost_usd contribute zero.
type Publisher ¶
Publisher is the broker-emit surface the threshold checker uses to announce a `cost.threshold.crossed` event. The rpc layer adapts its StreamBroker to this interface in api.go; tests can pass a recording fake. A nil Publisher silently disables event emission (the database row still goes in so a later restart with a real publisher does not re-fire the same tier).
type ThresholdCrossedPayload ¶
type ThresholdCrossedPayload struct {
// Pct is the tier that just crossed (50/80/100/150/200).
Pct int `json:"pct"`
// MonthTotalUSD is the calendar-month total spend in USD as of the
// turn that triggered the crossing.
MonthTotalUSD float64 `json:"monthTotalUsd"`
// ThresholdUSD is the user-configured MonthlyCostNotifyUSD value
// (the dial setting) at the moment of the firing. Surfaced so the
// toast can render "you've used 80% of your $25/mo budget."
ThresholdUSD float64 `json:"thresholdUsd"`
// YearMonth is the local-time YYYY-MM key the row was written under.
YearMonth string `json:"yearMonth"`
// FiredAt is the wall clock at which the tier was recorded. Useful
// for downstream UIs that want to show "fired 3 minutes ago."
FiredAt time.Time `json:"firedAt"`
}
ThresholdCrossedPayload is the typed event payload. The frontend types in lib/types.ts mirror these field names exactly (camelCase via the json tags) so the toast can render without any translation shim.
type ThresholdReader ¶
ThresholdReader is the func() (float64, error) callback the checker uses to fetch the current MonthlyCostNotifyUSD setting. We use a callback (not a fixed float64) so the dial takes effect on the next Add without restarting the process — the settings store reads the freshest value from disk on every call.
type UsageTurn ¶
type UsageTurn struct {
// SessionID identifies the session this turn belongs to.
SessionID string
// MessageID is the session_messages.id of the assistant row that
// was just written. When empty, the manager skips the write.
MessageID string
// ProviderKind is the provider kind string ("anthropic", "openrouter", …).
ProviderKind string
// ModelID is the model id string (e.g. "claude-sonnet-4-5").
ModelID string
// PromptTokens is the provider-reported input token count.
PromptTokens int
// CompletionTokens is the provider-reported output/completion token count.
CompletionTokens int
// CostUSD is the derived-or-provider-reported cost in USD. Nil means
// the cost is unknown (no pricing entry + no provider cost).
CostUSD *float64
// CostSource is one of "provider", "derived", "unknown".
CostSource string
}
UsageTurn is the per-turn accounting snapshot passed to Manager.Add.
Source Files