Skip to content

fix(metering): gate provider hourly token budget pre-call so exhaustion flags the fallback chain (#5980)#5988

Merged
houko merged 8 commits into
librefang:mainfrom
DaBlitzStein:fix/5980-clean
Jun 23, 2026
Merged

fix(metering): gate provider hourly token budget pre-call so exhaustion flags the fallback chain (#5980)#5988
houko merged 8 commits into
librefang:mainfrom
DaBlitzStein:fix/5980-clean

Conversation

@DaBlitzStein

@DaBlitzStein DaBlitzStein commented Jun 2, 2026

Copy link
Copy Markdown
Contributor

Summary

Restores the #4807 pre-dispatch per-provider budget design and makes the flag load-bearing, closing the no-op that @houko flagged in CHANGES_REQUESTED.

The original commit flagged an over-cap provider in the shared ProviderExhaustionStore but the flag was never consulted in the reported single-Moonshot-agent scenario, so the agent kept running ~4x past its hourly cap.

Changes

Commit 1 (ec6c7b4) — the pre-dispatch gate: flag_provider_budget_if_exhausted runs MeteringEngine::check_provider_budget before dispatch so a breach flags the provider in the shared store. Called from the three dispatch sites (agent_execution.rs, messaging.rs ephemeral + streaming).

Commit 2 (4b1d83f) — wires the flag to the two paths that ignored it:

  • Multi-provider agent path (llm_drivers.rs resolve_driver): the agent-fallback chain now carries each slot's provider name and builds the FallbackDriver via with_models_and_providers(...) with the same MeteringEngine::exhaustion_store() handle the gate flags (mirrors boot.rs:698-714). Previously it used with_models(chain) — no store, empty provider names — so is_slot_exhausted short-circuited to false and re-dispatched the exhausted provider. When the store is unwired the provider-less builder is kept (no regression).
  • Single-provider agent path: there is no slot to skip to, so flag_provider_budget_if_exhausted now returns whether the provider is over cap, and the three dispatch sites hard-return QuotaExceeded when resolve_effective_fallbacks(...) is empty. A new pure helper provider_exhausted_blocks_call(exhausted, manifest, cfg) encapsulates that decision. When a fallback chain exists the call sites do not hard-return — they defer to the store-aware FallbackDriver, preserving the feat(kernel): budget-aware fallback chain — skip exhausted providers, try next in chain #4807 philosophy.

Verification

  • cargo check --workspace --lib — clean.
  • cargo clippy --workspace --all-targets -- -D warnings — 0 warnings.
  • New librefang-kernel tests kernel::tests::provider_budget_gate_5980 (3 passed):
    • single_provider_over_cap_is_hard_blocked — over-cap single-provider agent is blocked, store reports BudgetExceeded, and the gate stays blocking on re-check (counter never drains via a leaked dispatch).
    • multi_provider_over_cap_defers_to_fallback_skip — over-cap multi-provider agent is NOT hard-blocked; the FallbackDriver built over the same store handle pre-skips the exhausted moonshot slot and serves from healthy groq (actual_provider == Some("groq")).
    • no_budget_configured_gate_is_noop — no provider budget configured ⇒ gate returns false, store never flagged, never blocks.

Rebased onto current origin/main to clear the stale base.

Closes #5980

@github-actions github-actions Bot added size/M 50-249 lines changed area/docs Documentation and guides area/kernel Core kernel (scheduling, RBAC, workflows) labels Jun 2, 2026
houko
houko previously requested changes Jun 3, 2026

@houko houko left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for tackling this — the diagnosis in #5980 is solid and a pre-dispatch gate is the right shape. But as written the fix flags a store that the per-agent execution driver never reads, so it's a no-op for the reported scenario. Requesting changes.

The flag is written, but nothing on the agent path consumes it

flag_provider_budget_if_exhaustedcheck_provider_budgetstore.mark_exhausted(provider, BudgetExceeded, …). That store is only honored by FallbackDriver/FallbackChain, and only when (a) a store is attached via with_exhaustion_store and (b) the slot has a non-empty provider_name. The driver the agent loop actually runs comes from resolve_driver:

  • With fallbacks configured, crates/librefang-kernel/src/kernel/llm_drivers.rs:326-328 builds FallbackDriver::with_models(chain) — that constructor sets exhaustion_store: None (fallback.rs:120) and leaves every provider_name empty. Both is_slot_exhausted and record_exhaustion early-return on missing store / empty name (fallback.rs:269-292), so the flag is ignored.
  • With no fallbacks, llm_drivers.rs:331 returns the bare primary driver, which has no store awareness at all — exactly #5980's case (a single Moonshot agent).

The new comment ("resolve_driver reads the same store and skips the exhausted slot") describes the intended behavior, but the diff doesn't change resolve_driver to attach the store or set provider names, so it isn't true yet. To close the loop, resolve_driver needs to use with_models_and_providers (or set_provider_names) so slots carry provider names, and attach the shared ProviderExhaustionStore via with_exhaustion_store — including a single-slot wrapper for the no-fallback case so a solo over-cap provider surfaces a clean AllProvidersExhausted instead of overspending.

Issue's proposed-fix item #2 not addressed

crates/librefang-kernel/src/kernel/agent_execution.rs:1303 still does let _ = self.metering.engine.record(&usage_record) unconditionally after a post-call QuotaExceeded, so the rolling counter keeps climbing (the 4.8M-over-1.2M symptom). Either drop the unconditional record on breach, or rely on the pre-call gate — but the gate has to be wired into a driver that reads it first.

Test coverage

No test added. The metering unit tests only prove the flag is set, not that the execution-path driver honors it — which is the whole bug. Per CLAUDE.md / #3721 this needs a test showing: provider at/over cap → next dispatch is skipped/failed-over before the HTTP send, and the counter doesn't grow unbounded across consecutive ticks.

Happy to re-review once resolve_driver carries the store + provider names through.

@github-actions github-actions Bot added the needs-changes Changes requested by reviewer label Jun 3, 2026
@DaBlitzStein
DaBlitzStein force-pushed the fix/5980-clean branch 3 times, most recently from 0ffc07e to 04f70f0 Compare June 5, 2026 16:13
@houko

houko commented Jun 6, 2026

Copy link
Copy Markdown
Contributor

The pre-dispatch budget gate looks like a no-op on the normal configured-fallback path. The metering side flags exhaustion correctly, but the FallbackDriver that actually serves the request on that path is built without the store handle (and with empty per-slot provider names), so it can never read the flag and never skips the exhausted slot.

Where the flag is set vs. where the driver is built

crates/librefang-kernel/src/kernel/agent_execution.rs:954
    self.flag_provider_budget_if_exhausted(&manifest.model.provider);
crates/librefang-kernel/src/kernel/agent_execution.rs:957
    let driver = self.resolve_driver(&manifest)?;

crates/librefang-kernel/src/kernel/llm_drivers.rs:326
    return Ok(Arc::new(librefang_runtime::drivers::fallback::FallbackDriver::with_models(chain),));

with_models leaves the driver unwired, and is_slot_exhausted short-circuits to false on both guards:

crates/librefang-llm-drivers/src/drivers/fallback.rs:120   exhaustion_store: None,
crates/librefang-llm-drivers/src/drivers/fallback.rs:114   provider_name: String::new(),
crates/librefang-llm-drivers/src/drivers/fallback.rs:285   let Some(store) = &self.exhaustion_store else { return false; };
crates/librefang-llm-drivers/src/drivers/fallback.rs:288   if entry.provider_name.is_empty() { return false; }

The boot-time default_driver (boot.rs:705-714) is wired correctly, but resolve_driver only returns it in the narrow claude_code + cli_profile_dirs branch (llm_drivers.rs:151) or on a driver-creation-failure fallback (line 249) — never on the normal configured-fallback path this PR targets.

Scenario: operator sets [budget.providers.groq] max_cost_per_hour_usd with a groq -> openai fallback chain. groq usage crosses the cap. Next turn: flag_provider_budget_if_exhausted("groq") marks groq exhausted in the metering engine's store; resolve_driver returns a fresh FallbackDriver::with_models(chain) whose exhaustion_store == None and whose slots all have empty provider_name; on dispatch is_slot_exhausted returns false for groq on both guards, so health_order tries groq first and the over-cap provider is dispatched again — the exact overspend #5980 set out to fix. The store the gate flags is a different object from the (absent) store this driver reads.

Suggested fix: in resolve_driver, build the chain as (driver, model, provider) triples and construct via FallbackDriver::with_models_and_providers(...).with_exhaustion_store(<shared ProviderExhaustionStore handle>), mirroring boot.rs:705-714 (the store is Arc-backed, so cloning preserves shared state). Both the store handle and non-empty per-slot provider names are required, or is_slot_exhausted short-circuits. The existing check_provider_budget_flags_exhaustion_store test only checks the metering engine's store, not the driver returned by resolve_driver, so it can't catch this — worth an integration test that asserts end-to-end skip behavior through the resolved driver.

Reviewed at 04f70f0.

@DaBlitzStein

Copy link
Copy Markdown
Contributor Author

Rebased onto current origin/main to clear the stale base. On test coverage: the substantive budget-exhaustion logic exercised by this change lives in MeteringEngine::check_provider_budget, which already has unit coverage in librefang-kernel-metering (test_check_provider_budget_under_limit, _over_limit, _zero_limit_skipped, _separate_providers_isolated, _tokens_per_hour). The new flag_provider_budget_if_exhausted is a thin forwarding wrapper to that tested function (empty-provider guard + lookup + side-effecting call), and the call-site is pre-call wiring in execute_llm_agent/messaging.rs. Happy to add a kernel-level wrapper test if you'd prefer explicit coverage there.

@houko

houko commented Jun 8, 2026

Copy link
Copy Markdown
Contributor

Thanks for picking this up — the direction is right (restore the #4807 design and call check_provider_budget pre-call so it can flag exhaustion). But as written the gate is a no-op in the single-provider case, which is exactly the scenario #5980 reports (one Moonshot agent running ~4× over the hourly cap), and there are no tests. Two concrete gaps:

1. The flag is written to a store nothing reads when there's no fallback chain.
flag_provider_budget_if_exhaustedcheck_provider_budgetflag_provider_budget_exhausted writes to MeteringEngine.exhaustion. For that flag to actually skip a call, a FallbackDriver must read the same store via is_slot_exhausted. But boot only builds a store-aware FallbackDriver when the chain has >1 entry (boot.rs:701-714); the single-driver branch (boot.rs:716-717) returns a bare driver with no exhaustion store and no slot to skip to. So for a single-provider agent the flag lands in a store no driver consults, and there is nowhere to fall back even if it did — the call still goes out and the counter keeps climbing.

2. Agent-level fallback_models builds a FallbackDriver without the store.
resolve_driver constructs FallbackDriver::with_models(chain) (llm_drivers.rs:326-329) without .with_exhaustion_store(...), so is_slot_exhausted short-circuits to false (fallback.rs:285-287); the slots also have empty provider_name, which the same function treats as not-exhausted. The flag is ignored on this path too.

The only path where the flag currently does anything is: default provider + a global fallback_providers chain of ≥2, going through boot's len > 1 branch — which does not include the reported single-provider case.

Suggested fix: to stop the bleed even with no fallback, the flag-only approach isn't sufficient. When the gate hits and the resolved driver isn't a store-aware FallbackDriver with a healthy slot, the current call should hard-return QuotaExceeded pre-call rather than relying on a fallback slot that may not exist. The minimal form is: gate hit → return Err(QuotaExceeded); treat "let the fallback chain take over on multi-provider" as a separate enhancement on top of that. (If you want to keep the flag-only path for multi-provider, also pass .with_exhaustion_store(...) + with_models_and_providers(...) in the resolve_driver with_models branch.)

Tests (required, per the repo's integration-testing policy and #5980's acceptance items): a kernel→driver e2e test, not just the metering-layer unit coverage, asserting (a) a single-provider agent over the hourly cap has its next dispatch blocked and the counter stops; (b) a multi-provider agent over the cap falls back to a healthy provider; (c) no budget configured → behaviour unchanged.

@DaBlitzStein
DaBlitzStein force-pushed the fix/5980-clean branch 2 times, most recently from ce66d21 to 4b1d83f Compare June 10, 2026 18:42
@github-actions github-actions Bot added size/L 250-999 lines changed and removed size/M 50-249 lines changed labels Jun 10, 2026
@DaBlitzStein

Copy link
Copy Markdown
Contributor Author

Thanks @houko — addressed both defects from your review. Pushed 4b1d83f (rebased onto current origin/main).

1. Multi-provider path now store-aware. resolve_driver (crates/librefang-kernel/src/kernel/llm_drivers.rs) builds the agent-fallback chain as (driver, model, provider) triples via with_models_and_providers(...) and attaches the shared handle from self.metering.engine.exhaustion_store() — the same store the gate flags — mirroring boot.rs:698-714. So a budget-exhausted slot is now record_skip-ed on dispatch instead of re-invoked. When the accessor returns None the provider-less builder is kept, so the no-store case is unchanged.

2. Single-provider path now hard-returns. As you noted, the flag-only approach can't help when there is no slot to skip to. flag_provider_budget_if_exhausted now returns bool (over-cap?), and the three dispatch sites (agent_execution.rs:~974, messaging.rs:~581 ephemeral, messaging.rs:~2280 streaming) hard-return QuotaExceeded when resolve_effective_fallbacks(...) is empty, via the new pure helper provider_exhausted_blocks_call. When a fallback chain exists they defer to the now-store-aware FallbackDriver rather than hard-blocking, keeping the #4807 behaviour you wanted to protect.

3. Kernel→driver e2e tests (kernel::tests::provider_budget_gate_5980), covering your three required cases: single-provider over-cap blocked + counter stops; multi-provider over-cap skips the exhausted slot through the resolved driver and serves from the healthy one; no budget configured ⇒ unchanged. cargo clippy --workspace --all-targets -- -D warnings is clean.

One thing I want to flag for a possible follow-up (pre-existing, not introduced here): the gate keys on the raw manifest.model.provider, so an agent configured with provider = "default" (resolved to a concrete provider only inside resolve_driver) isn't matched against [budget.providers.<id>]. The reported case uses an explicit provider so it's covered; happy to thread the resolved provider through the gate in a separate change if you think it's worth it.

@houko

houko commented Jun 11, 2026

Copy link
Copy Markdown
Contributor

Re-reviewed the two follow-up commits — this resolves my earlier CHANGES_REQUESTED. The flag is now actually consumed on the agent path, both directions:

  • Multi-provider agents: resolve_driver builds the chain via with_models_and_providers(...) (carrying each slot's provider name) and attaches the same handle from MeteringEngine::exhaustion_store() via with_exhaustion_store(store) — so the slot the gate flagged is now non-empty and is_slot_exhausted actually skips it, instead of the old with_models path that left the store unset and short-circuited to false. That was the exact no-op I called out.
  • Single-provider agents (the reported single-Moonshot case, no slot to fail over to): provider_exhausted_blocks_call returns exhausted && fallback-chain-is-empty, and the three dispatch sites hard-return QuotaExceeded instead of dispatching past the cap.

The split is correct: block only when exhausted and there's no viable fallback; otherwise let the now-store-aware FallbackDriver do the skipping. The provider_budget_gate_5980 tests cover both. The match exhaustion_store() { None => fb } arm also degrades safely when metering is off (the gate writes to the same store, so they stay consistent).

LGTM from my side — clearing my changes-requested. Nice fix.

@houko houko left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two issues found during daily review pass:

1. Misplaced doc comment and #[allow] attribute in agent_execution.rs

The new functions are inserted after the existing context lines /// Execute the default LLM-based agent loop. and #[allow(clippy::too_many_arguments)], so in the resulting file those two items are syntactically attached to flag_provider_budget_if_exhausted (not execute_llm_agent):

/// Execute the default LLM-based agent loop.          ← wrong: was execute_llm_agent's doc
#[allow(clippy::too_many_arguments)]                   ← wrong: was execute_llm_agent's allow
/// Pre-dispatch per-provider budget gate (#5980…)     ← flag_provider_budget_if_exhausted's doc
pub(crate) fn flag_provider_budget_if_exhausted() {}

pub(crate) fn provider_exhausted_blocks_call() {}

#[allow(clippy::too_many_arguments)]   ← correctly precedes execute_llm_agent
#[instrument()]
pub(super) fn execute_llm_agent() {}   ← no doc comment

execute_llm_agent loses its /// Execute the default LLM-based agent loop. doc. flag_provider_budget_if_exhausted gains it (merged with the correct doc as a two-paragraph doc comment) plus an unnecessary #[allow].

Fix: move the two new methods to before the /// Execute the default LLM-based agent loop. line (i.e. earlier in the impl block), or extract them to a dedicated impl block so the existing doc + attribute stay cleanly above execute_llm_agent.

2. Duplicate CHANGELOG entry for #5980

The diff adds two bullets that both reference (#5980): a detailed one near the top of the changes in the hunk and a short one-liner further down. Remove the shorter duplicate.


Generated by Claude Code

@github-actions github-actions Bot added the has-conflicts PR has merge conflicts that need resolution label Jun 15, 2026
@github-actions github-actions Bot added area/runtime Agent loop, LLM drivers, WASM sandbox and removed has-conflicts PR has merge conflicts that need resolution labels Jun 15, 2026
@github-actions github-actions Bot removed the area/runtime Agent loop, LLM drivers, WASM sandbox label Jun 15, 2026

@houko houko left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One CHANGELOG issue to fix before merge: duplicate entry for #5980 (see inline comment on CHANGELOG.md).


Generated by Claude Code

Comment thread CHANGELOG.md

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are two [Unreleased] entries for the same fix (#5980):

  1. The detailed one near the top of the hunk — **metering(budget): the pre-dispatch per-provider budget gate now actually enforces the hourly cap…** — describes the full root cause, two fixes, and the three regression test names.
  2. A shorter duplicate further down — **metering(budget): enforce per-provider hourly token/cost budgets pre-dispatch…** — covers the same change in one line.

The shorter duplicate should be removed; the detailed entry is the correct one to keep. Having two entries for the same issue will also trip the pre-commit duplicate-[Unreleased] guard.


Generated by Claude Code

DaBlitzStein added a commit to DaBlitzStein/librefang that referenced this pull request Jun 17, 2026
DaBlitzStein added a commit to DaBlitzStein/librefang that referenced this pull request Jun 17, 2026
DaBlitzStein added a commit to DaBlitzStein/librefang that referenced this pull request Jun 17, 2026
DaBlitzStein added a commit to DaBlitzStein/librefang that referenced this pull request Jun 17, 2026
DaBlitzStein added a commit to DaBlitzStein/librefang that referenced this pull request Jun 17, 2026
DaBlitzStein added a commit to DaBlitzStein/librefang that referenced this pull request Jun 17, 2026
…block single-provider over-cap (librefang#5980)

The librefang#5980 pre-dispatch budget gate flagged an over-cap provider in the
shared ProviderExhaustionStore but never had its flag consulted in the
reported scenario.
The agent-fallback path built its FallbackDriver via with_models(chain),
which leaves the exhaustion store unset and every slot's provider name
empty, so is_slot_exhausted short-circuited to false and the over-cap
provider was dispatched again.
A single-provider agent had no fallback at all, so the flag landed in a
store nothing read and the call still went out.

resolve_driver now builds the multi-provider chain via
with_models_and_providers(...) carrying each slot's provider name and
attaches the same exhaustion-store handle the gate flags (from
MeteringEngine::exhaustion_store()), mirroring the boot-time path, so a
budget-exhausted slot is pre-skipped and the call fails over to a healthy
provider.
flag_provider_budget_if_exhausted now returns whether the provider is
over cap, and the three dispatch call sites hard-return QuotaExceeded
when the agent has no viable fallback chain instead of dispatching past
the cap.

Regression tests in librefang-kernel cover all three cases: a
single-provider agent over the hourly cap is blocked, a multi-provider
agent defers to the store-aware FallbackDriver which skips the exhausted
slot and serves from a healthy one, and an agent with no configured
budget dispatches unchanged.
…t gate

The 'Execute the default LLM-based agent loop' doc comment and its
#[allow(clippy::too_many_arguments)] were duplicated above the new
flag_provider_budget_if_exhausted — they already sit above
execute_llm_agent where they belong.
… exhaustion is visible to FallbackDriver (librefang#5980)

The primary slot's provider_name was String::new() (empty), which made
is_slot_exhausted always return false for the first slot — the budget
gate flagged exhaustion correctly but the FallbackDriver never skipped
the exhausted primary because it couldn't match the empty name against
the store entry.
DaBlitzStein added a commit to DaBlitzStein/librefang that referenced this pull request Jun 20, 2026
houko
houko previously requested changes Jun 23, 2026

@houko houko left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The metering logic looks good and the earlier review concerns are addressed — one small thing still blocks merge.

CHANGELOG.md: duplicate [Unreleased] entry for #5980. There are two bullets for this change:

  • the detailed one (the pre-dispatch per-provider budget gate now actually enforces…) — keep this one
  • a redundant one-liner further down (enforce per-provider hourly token/cost budgets pre-dispatch…) — please remove this duplicate

This is the item from the prior review that's still outstanding.

Minor, while you're in there: execute_llm_agent lost its /// Execute the default LLM-based agent loop. doc comment (it's present on main). Please restore it.

Once the duplicate CHANGELOG bullet is removed, this is good to merge.

…ent doc

Removes the redundant one-line [Unreleased] bullet for librefang#5980 (the detailed entry above is the keeper) and restores the `/// Execute the default LLM-based agent loop.` doc comment that was dropped from execute_llm_agent. Addresses the review feedback on this PR.
@houko
houko dismissed stale reviews from themself June 23, 2026 05:10

The duplicate [Unreleased] CHANGELOG bullet is removed and the execute_llm_agent doc comment restored (pushed to the branch); budget logic was already verified. Clearing for merge.

@houko houko removed the needs-changes Changes requested by reviewer label Jun 23, 2026
@houko
houko enabled auto-merge (squash) June 23, 2026 05:11
@github-actions github-actions Bot added the has-conflicts PR has merge conflicts that need resolution label Jun 23, 2026
houko added a commit that referenced this pull request Jun 23, 2026
…top (#6285)

CHANGELOG.md had two `## [Unreleased]` sections: one buried between released version sections (it predates #6281, which then added a duplicate at the top). The release tooling reads only the first [Unreleased] and silently drops the rest, so ~150 stranded entries were never reaching release notes; the duplicate header also tripped the `## [Unreleased]` guard, blocking every CHANGELOG-touching commit (e.g. #5988).

Consolidate both blocks into a single [Unreleased] at the top, merging same-named `###` subsections (collapsed 10 redundant subsection headers). Every bullet is preserved verbatim — verified that the multiset of non-header content lines is byte-identical before/after. Drop the two entries already published in released sections: #6272 (in [2026.6.22]) and #6171 (in [2026.6.17]).

Co-authored-by: Evan <[email protected]>
# Conflicts:
#	.secrets.baseline
@github-actions github-actions Bot added ready-for-review PR is ready for maintainer review and removed has-conflicts PR has merge conflicts that need resolution labels Jun 23, 2026
@houko
houko merged commit 67e9fe2 into librefang:main Jun 23, 2026
32 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/docs Documentation and guides area/kernel Core kernel (scheduling, RBAC, workflows) ready-for-review PR is ready for maintainer review size/L 250-999 lines changed

Projects

None yet

Development

Successfully merging this pull request may close these issues.

metering: provider hourly token budget is enforced post-call only; exhaustion-flag path (check_provider_budget) is never called

2 participants