From 7995bb5f5924882703ad1c20ca96029d4c2c4421 Mon Sep 17 00:00:00 2001 From: Elron Bandel Date: Wed, 17 Jun 2026 15:03:22 +0300 Subject: [PATCH] feat(bifrost): enforce the rule-16 spend cap via a governance budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The per-run EVAL_MODEL_MAX_BUDGET cap (models/RULES.md rule 16) was enforced only by core/litellm; the default gpt-5.4--bifrost path ignored it. Add bifrost governance: a budget on the sk-proxy virtual key (the bearer the agent already sends), priced by a pricing_override at the same $2.50/$10-per-1M rates as the litellm config, behind the governance plugin. The start script renders EVAL_MODEL_MAX_BUDGET into the config (mirrors core/litellm's wrapper). Tests: a static guard locks the wiring; an #[ignore] behavioral test (upstream_bifrost_budget_cap_rejects) proves bifrost rejects once spend crosses the cap — also the empirical check that config.json governance enforces in our single-node setup (the silent no-op maximhq/bifrost#2408 is multinode-only). Lean config: the budget attaches directly to the virtual key (no team object, no client block). Enforcement is validated by the #[ignore] test against the live endpoint. Resolves #185. Signed-off-by: Elron Bandel --- containers/gateways/bifrost/start | 13 ++- .../gpt-5.4--bifrost/config.json.template | 17 ++++ tests/run/gateways/test.rs | 95 +++++++++++++++++++ 3 files changed, 124 insertions(+), 1 deletion(-) diff --git a/containers/gateways/bifrost/start b/containers/gateways/bifrost/start index 18ece0ad..3b4f0858 100644 --- a/containers/gateways/bifrost/start +++ b/containers/gateways/bifrost/start @@ -49,8 +49,19 @@ if [ -f "$TEMPLATE" ]; then OTEL_COLLECTOR_URL="${OTEL_EXPORTER_OTLP_ENDPOINT%/}/v1/traces" export OTEL_COLLECTOR_URL + # Per-run hard spend cap (models/RULES.md rule 16). Default $1; must be a + # number so it renders as valid JSON into the governance budget in the + # template (litellm does the same via core/litellm's entrypoint wrapper). + EVAL_MODEL_MAX_BUDGET="${EVAL_MODEL_MAX_BUDGET:-1}" + case "$EVAL_MODEL_MAX_BUDGET" in + ''|.|*[!0-9.]*|*.*.*) + echo "bifrost: EVAL_MODEL_MAX_BUDGET='$EVAL_MODEL_MAX_BUDGET' is not a number" >&2 + exit 64 ;; + esac + export EVAL_MODEL_MAX_BUDGET + # Explicit var list so envsubst leaves bifrost's own $schema key alone. - envsubst '${PROVIDER} ${MODEL_NAME} ${OPENAI_API_BASE} ${OTEL_COLLECTOR_URL}' \ + envsubst '${PROVIDER} ${MODEL_NAME} ${OPENAI_API_BASE} ${OTEL_COLLECTOR_URL} ${EVAL_MODEL_MAX_BUDGET}' \ < "$TEMPLATE" > "$CONFIG" fi diff --git a/containers/models/gpt-5.4--bifrost/config.json.template b/containers/models/gpt-5.4--bifrost/config.json.template index b955e141..5aeb2a01 100644 --- a/containers/models/gpt-5.4--bifrost/config.json.template +++ b/containers/models/gpt-5.4--bifrost/config.json.template @@ -21,6 +21,22 @@ } }, "governance": { + "budgets": [ + { "id": "eval-budget", "max_limit": ${EVAL_MODEL_MAX_BUDGET}, "reset_duration": "1Y" } + ], + "virtual_keys": [ + { "id": "eval-vk", "value": "sk-proxy", "is_active": true, "budget_id": "eval-budget" } + ], + "pricing_overrides": [ + { + "id": "eval-pricing", + "scope_kind": "global", + "match_type": "wildcard", + "pattern": "*", + "request_types": ["chat_completion"], + "pricing_patch": "{\"input_cost_per_token\":0.0000025,\"output_cost_per_token\":0.000010}" + } + ], "routing_rules": [ { "id": "rewrite-to-upstream", @@ -35,6 +51,7 @@ ] }, "plugins": [ + { "enabled": true, "name": "governance", "config": { "is_vk_mandatory": true } }, { "enabled": true, "name": "otel", diff --git a/tests/run/gateways/test.rs b/tests/run/gateways/test.rs index b391677c..2fae699e 100644 --- a/tests/run/gateways/test.rs +++ b/tests/run/gateways/test.rs @@ -164,8 +164,17 @@ async fn gateway_port(c: &ContainerAsync) -> u16 { } fn http() -> Client { + // Mirror the agent: every gateway request carries the fixed `sk-proxy` bearer. + // For bifrost that token is the governance virtual key, so the per-run budget + // applies; litellm/portkey ignore it (they auth upstream from env vars). + let mut headers = reqwest::header::HeaderMap::new(); + headers.insert( + reqwest::header::AUTHORIZATION, + reqwest::header::HeaderValue::from_static("Bearer sk-proxy"), + ); Client::builder() .timeout(Duration::from_secs(60)) + .default_headers(headers) .build() .expect("build reqwest client") } @@ -479,6 +488,92 @@ async fn upstream_portkey_openai() { assert_openai_200("portkey").await } +// ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ +// Spend cap (models/RULES.md rule 16). litellm enforces it via core/litellm's +// entrypoint; bifrost enforces it via governance — a budget on the `sk-proxy` +// virtual key, behind the governance plugin, priced by a pricing_override. +// ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ + +/// Static: the bifrost path wires the rule-16 cap. The model config must put the +/// governance plugin in front of a budget fed by `${EVAL_MODEL_MAX_BUDGET}` on the +/// `sk-proxy` virtual key, and the gateway start must render that env var. Guards +/// against the cap silently vanishing (the gap that motivated this — the gateways +/// shipped no budget enforcement while only core/litellm did). +#[test] +fn static_bifrost_wires_budget_cap() { + let cfg = std::fs::read_to_string( + test_support::repo_root().join("containers/models/gpt-5.4--bifrost/config.json.template"), + ) + .expect("read gpt-5.4--bifrost config.json.template"); + let start = std::fs::read_to_string( + test_support::repo_root().join("containers/gateways/bifrost/start"), + ) + .expect("read gateways/bifrost/start"); + + assert!( + cfg.contains(r#""name": "governance""#), + "bifrost config must enable the governance plugin (it enforces the virtual-key budget)" + ); + assert!( + cfg.contains("${EVAL_MODEL_MAX_BUDGET}"), + "bifrost config governance budget max_limit must come from ${{EVAL_MODEL_MAX_BUDGET}} (rule 16)" + ); + assert!( + cfg.contains(r#""value": "sk-proxy""#), + "bifrost config must define the `sk-proxy` virtual key the agent presents" + ); + assert!( + start.contains("${EVAL_MODEL_MAX_BUDGET}"), + "gateways/bifrost/start must render EVAL_MODEL_MAX_BUDGET into the config" + ); + eprintln!("✓ bifrost wires the rule-16 budget cap (governance plugin + sk-proxy VK budget)"); +} + +/// Behavioral (#[ignore], real upstream): bifrost MUST stop spend once it crosses +/// `EVAL_MODEL_MAX_BUDGET`. With a ~zero cap the first costed call pushes accumulated +/// spend over the limit, so a subsequent call is rejected. This is also the empirical +/// check that config.json governance enforces in our single-node setup (the silent +/// no-op bug maximhq/bifrost#2408 is multinode-only). `http()` sends the `sk-proxy` +/// VK, so the budget applies. +#[tokio::test] +#[ignore] +async fn upstream_bifrost_budget_cap_rejects() { + let (key, base) = upstream_creds(); + let c = start_gateway( + "bifrost", + &[ + ("OPENAI_API_KEY", &key), + ("OPENAI_API_BASE", &base), + ("EVAL_MODEL_MAX_BUDGET", "0.0000001"), + ], + ) + .await; + let port = gateway_port(&c).await; + let url = format!("http://127.0.0.1:{port}/openai/v1/chat/completions"); + + let mut last = 0u16; + let mut rejected = false; + for _ in 0..4 { + let status = http() + .post(&url) + .json(&body_openai()) + .send() + .await + .expect("post chat completions") + .status(); + last = status.as_u16(); + if status.is_client_error() { + rejected = true; + break; + } + } + assert!( + rejected, + "bifrost did not reject after exceeding EVAL_MODEL_MAX_BUDGET (last status {last}) — \ + rule 16 cap not enforced (config.json governance ignored? cf. maximhq/bifrost#2408)" + ); +} + async fn assert_anthropic_200(flavor: &str) { let c = start_with_real_creds(flavor).await; let port = gateway_port(&c).await;