fix(kernel): pre-dispatch provider budget gate prevents streaming token leak#4801
fix(kernel): pre-dispatch provider budget gate prevents streaming token leak#4801DaBlitzStein wants to merge 5 commits into
Conversation
github-actions
Bot
added
area/kernel
…aste Both send_message_full and send_message_ephemeral checked provider budgets only post-call via check_all_and_record. Streaming responses that had already consumed tokens could not be rejected in time, causing providers to overshoot their hourly token budgets by the full cost of the leaked call. Wire check_provider_budget before the LLM dispatch in both paths. On the main dispatch path, roll back the USD and token reservations on rejection so no phantom holds accumulate in the in-memory ledger.
f65d97f to
93fb234
Compare
|
Closing — two issues surfaced in review that make this PR's claim ("ready, all checks passed locally") not hold: 1. Compile bug (will fail CI's 2. The streaming dispatch path is not modified. The PR description's reproduction is a streaming overshoot ("16,924 tokens leaked through a single streaming call"), but neither of the two sites this PR touches sits on the streaming dispatch chain. All Beyond those two: putting the pre-check before the Suggestion: this is a sign the wrapper-driver layer (#4757's original direction) was probably the right one — there is exactly one chokepoint at the driver call where pre-check belongs, and it covers streaming and non-streaming uniformly. The "minimal patch" approach has to find every dispatch entry and the streaming one was missed. Thanks for picking this up; happy to look at a v2 that either (a) covers all three entries ( Issue #4800 stays open. |
Cover send_message_full_with_upstream, send_message_ephemeral, and send_message_streaming_with_sender_and_opts with check_provider_budget before the LLM round-trip. Uses self.metering.engine (not the subsystem wrapper) and budget_config.load() to avoid cloning. Checks run after the suspended-agent gate to skip SQL for no-ops. Addresses librefang#4800 review feedback on closed librefang#4801.
2 participants
Summary
Closes #4800
check_provider_budget()pre-dispatch insend_message_full_with_upstream— after the existingcheck_quota_and_reservegate, with proper USD + token reservation rollback on rejectioncheck_provider_budget()pre-dispatch insend_message_ephemeral— after the suspended-agent checkBoth paths previously checked provider budgets only post-call via
check_all_and_record(). Streaming responses that had already consumed tokens could not be rejected in time, causing providers to overshoot their hourly token budgets. Observed on a production instance where Moonshot's 250k hourly budget was exceeded to 317k — 16,924 tokens leaked through a single streaming call after the budget was already exhausted.1 file changed, 27 insertions. No new types, no new crates — wires the existing
MeteringEngine::check_provider_budget(which has unit test coverage) into the two dispatch entry points.Supersedes the approach in #4757 (closed:
BudgetGatedDriverwas never instantiated). This PR takes the minimal path: direct calls at the dispatch sites instead of a wrapper type.Test plan
cargo check --workspace --lib— zero errorscargo clippy --workspace --all-targets -- -D warnings— zero warningscargo test -p librefang-kernel-metering— 28/28 passedcargo test -p librefang-kernel— 831/831 passed (1 pre-existing failure unrelated:reload_config_with_invalid_toml_preserves_live_config)