From b32164b08659c4bc7a044ef953e5171dbbf9b98e Mon Sep 17 00:00:00 2001 From: Bryan Helmkamp Date: Thu, 6 Aug 2026 14:22:10 -0400 Subject: [PATCH] Add [run.limits] run budgets: halt runs on token/cost spend A run can now carry an optional spend budget. New [run.limits] settings table with max_cost (USD) and max_tokens; when accumulated LLM spend exceeds either limit, the run halts immediately and fails with the existing BudgetExhausted failure reason. - RunBudget (fabro-model): thread-safe accumulator shared across the engine, agent sessions, hooks, and the prompt path. Charges per LLM response; queues one-shot 80% warnings drained into run notices. - Enforcement is inline at every usage-recording site: the agent turn loop (interrupting like the wall-clock timeout via a new InterruptReason::BudgetExhausted), subagent sessions, the one-shot prompt path with schema-repair retries, context compaction, web_fetch summarization, and prompt/agent hooks (which skip once the budget is spent, matching their fail-open posture). - Budget exhaustion terminates the run as an engine-level error (fabro_core::Error::BudgetExhausted, following the StallTimeout pattern) so fail edges are not followed, plus a between-stage gate for non-LLM stages and a run-start gate that makes a resumed budget-failed run fail fast until the limit is raised (spend is seeded from the billing projection). - Exposed on the RunNamespace OpenAPI schema (RunLimitsSettings); Rust and TypeScript clients regenerated; docs updated. Known gaps, by design for now: ACP external agents report no usable per-response usage (cumulative UsageUpdate capture is a follow-up), and post-run PR content generation is outside the engine. Co-Authored-By: Claude Fable 5 --- docs/public/api-reference/fabro-api.yaml | 23 ++ docs/public/execution/run-configuration.mdx | 24 ++ docs/public/reference/user-configuration.mdx | 15 + lib/apps/fabro-cli/tests/it/cmd/attach.rs | 1 + lib/apps/fabro-cli/tests/it/cmd/inspect.rs | 1 + lib/components/fabro-agent/src/cli.rs | 5 +- lib/components/fabro-agent/src/compaction.rs | 12 + lib/components/fabro-agent/src/config.rs | 8 +- lib/components/fabro-agent/src/error.rs | 3 + lib/components/fabro-agent/src/session.rs | 34 ++ lib/components/fabro-agent/src/tools.rs | 19 +- .../fabro-agent/tests/it/parity_matrix.rs | 5 +- lib/components/fabro-hooks/src/executor.rs | 143 ++++++- lib/components/fabro-hooks/src/runner.rs | 12 +- lib/components/fabro-workflow/src/error.rs | 11 +- .../fabro-workflow/src/handler/llm/api.rs | 382 +++++++++++++++++- .../fabro-workflow/src/node_handler.rs | 16 + .../fabro-workflow/src/pipeline/execute.rs | 33 +- .../fabro-workflow/src/pipeline/initialize.rs | 57 ++- lib/components/fabro-workflow/src/services.rs | 22 +- .../fabro-workflow/tests/it/integration.rs | 117 ++++++ .../tests/workflow_settings_round_trip.rs | 28 ++ lib/foundation/fabro-config/src/layers/mod.rs | 6 +- lib/foundation/fabro-config/src/layers/run.rs | 28 ++ lib/foundation/fabro-config/src/lib.rs | 13 +- .../fabro-config/src/resolve/run.rs | 46 ++- .../fabro-config/src/tests/combine.rs | 22 + .../fabro-config/src/tests/resolve_run.rs | 75 ++++ lib/foundation/fabro-core/src/error.rs | 5 + .../src/commands/docs_options_reference.rs | 6 + lib/foundation/fabro-model/src/budget.rs | 348 ++++++++++++++++ lib/foundation/fabro-model/src/lib.rs | 2 + .../fabro-types/src/run_event/infra.rs | 1 + .../fabro-types/src/settings/mod.rs | 4 +- .../fabro-types/src/settings/run.rs | 31 ++ .../src/.openapi-generator/FILES | 1 + .../fabro-api-client/src/models/index.ts | 1 + .../src/models/run-limits-settings.ts | 29 ++ .../src/models/run-namespace.ts | 4 + 39 files changed, 1522 insertions(+), 71 deletions(-) create mode 100644 lib/foundation/fabro-model/src/budget.rs create mode 100644 lib/packages/fabro-api-client/src/models/run-limits-settings.ts diff --git a/docs/public/api-reference/fabro-api.yaml b/docs/public/api-reference/fabro-api.yaml index 261badbbb..d5815ff54 100644 --- a/docs/public/api-reference/fabro-api.yaml +++ b/docs/public/api-reference/fabro-api.yaml @@ -14210,6 +14210,7 @@ components: - git - prepare - execution + - limits - checkpoint - clone - run_branch @@ -14246,6 +14247,8 @@ components: $ref: "#/components/schemas/RunPrepareSettings" execution: $ref: "#/components/schemas/RunExecutionSettings" + limits: + $ref: "#/components/schemas/RunLimitsSettings" checkpoint: $ref: "#/components/schemas/RunCheckpointSettings" clone: @@ -14458,6 +14461,26 @@ components: type: string enum: [prompt, auto] + RunLimitsSettings: + description: | + Resolved `[run.limits]` spend budget for the run. An absent field + means that limit is not configured. When accumulated LLM spend + exceeds either limit, the run halts and fails with the + `budget_exhausted` failure reason. + type: object + properties: + max_cost_usd_micros: + type: integer + format: int64 + description: Maximum total LLM cost for the run, in USD micros. + max_tokens: + type: integer + format: int64 + minimum: 0 + description: | + Maximum total LLM tokens for the run, counting all token + buckets (input, output, reasoning, cache reads, cache writes). + RunCheckpointSettings: type: object required: [exclude_globs, skip_git_hooks] diff --git a/docs/public/execution/run-configuration.mdx b/docs/public/execution/run-configuration.mdx index a8f1bc927..30dafb11d 100644 --- a/docs/public/execution/run-configuration.mdx +++ b/docs/public/execution/run-configuration.mdx @@ -409,6 +409,30 @@ commit_timeout = "30s" `exclude_globs` replaces across layers — the higher-precedence layer wins wholesale. `skip_git_hooks` and `commit_timeout` use normal override semantics: the highest layer that sets the field wins. +### `[run.limits]` + +Cap a run's total LLM spend. When accumulated spend exceeds either limit, the run halts immediately and fails with the `budget_exhausted` failure reason. Both limits are optional; an absent limit means unlimited. + +```toml title="run.toml" +[run.limits] +max_cost = 25.00 +max_tokens = 5000000 +``` + +| Field | Description | +|---|---| +| `max_cost` | Maximum total LLM cost for the run, in USD. Uses provider-reported cost when available (for example, OpenRouter reports authoritative cost on every response) and model catalog prices otherwise. | +| `max_tokens` | Maximum total LLM tokens for the run, counting all token buckets: input, output, reasoning, cache reads, and cache writes. | + +Notes: + +- Fabro checks the budget as each LLM response reports its usage, so a run can exceed the budget by the requests already in flight before it halts. +- When a response reports no cost and the model has no catalog price, `max_cost` cannot observe that spend. Set `max_tokens` as a backstop when you use unpriced models. +- The budget covers agent stages (including subagents, context compaction, and `web_fetch` summarization), prompt stages, and `[[run.hooks]]` prompt/agent hooks. Once the budget is exhausted, prompt/agent hooks stop firing and proceed without evaluating. It cannot observe spend from external agents driven over ACP, because those agents do not report per-response usage to Fabro, or from post-run pull-request content generation. +- Fabro emits a warning notice when spend crosses 80% of either limit. +- A budget-exhausted run keeps its last checkpoint. Resuming it fails immediately until you raise or remove the limit, because the spent totals carry forward from the run's billing history. +- Child runs are budgeted independently: a parent's budget does not count usage from runs it spawns. + ### `[run.inputs]` Define inputs that are rendered into final workflow string attributes. See [Variables](/workflows/variables) for the full reference. diff --git a/docs/public/reference/user-configuration.mdx b/docs/public/reference/user-configuration.mdx index cf60ea2a4..a4f3d3207 100644 --- a/docs/public/reference/user-configuration.mdx +++ b/docs/public/reference/user-configuration.mdx @@ -434,6 +434,21 @@ enabled = true | `enabled` | boolean | false | Automatically create a PR after successful runs. | | `merge_strategy` | "merge" \| "squash" \| "rebase" | "squash" | Merge method to configure for the pull request. | +## `[run.limits]` + +`[run.limits]` — run-scoped LLM spend budgets. Absent means unlimited + +```toml title="settings.toml" +[run.limits] +max_cost = 25.00 +max_tokens = 5000000 +``` + +| Key | Type / values | Default | Description | +|---|---|---|---| +| `max_cost` | float | None | Maximum total LLM cost for the run, in USD. When accumulated cost
exceeds this value, the run halts and fails. | +| `max_tokens` | integer | None | Maximum total LLM tokens for the run, counting all token buckets
(input, output, reasoning, cache reads, cache writes). When
accumulated tokens exceed this value, the run halts and fails. | + ## `[run.agent]` `[run.agent]` — agent knobs only (Fabro tools and MCPs) diff --git a/lib/apps/fabro-cli/tests/it/cmd/attach.rs b/lib/apps/fabro-cli/tests/it/cmd/attach.rs index 813bf7424..a33501999 100644 --- a/lib/apps/fabro-cli/tests/it/cmd/attach.rs +++ b/lib/apps/fabro-cli/tests/it/cmd/attach.rs @@ -984,6 +984,7 @@ fn attach_json_errors_without_prompting_for_human_input() { "provider": null, "slack": null }, + "limits": {}, "meta_branch": { "enabled": true, "push": true diff --git a/lib/apps/fabro-cli/tests/it/cmd/inspect.rs b/lib/apps/fabro-cli/tests/it/cmd/inspect.rs index d01e477e1..279050278 100644 --- a/lib/apps/fabro-cli/tests/it/cmd/inspect.rs +++ b/lib/apps/fabro-cli/tests/it/cmd/inspect.rs @@ -146,6 +146,7 @@ fn inspect_resolves_selector_via_server_endpoint() { "mode": "normal", "approval": "prompt" }, + "limits": {}, "checkpoint": { "exclude_globs": [], "skip_git_hooks": false, diff --git a/lib/components/fabro-agent/src/cli.rs b/lib/components/fabro-agent/src/cli.rs index 720b8b0af..5164dd7b7 100644 --- a/lib/components/fabro-agent/src/cli.rs +++ b/lib/components/fabro-agent/src/cli.rs @@ -218,8 +218,9 @@ fn build_summarizer( llm_client: Client, ) -> WebFetchSummarizer { WebFetchSummarizer { - client: llm_client, - model_id: summarizer_model_id(provider_id, catalog, model), + client: llm_client, + model_id: summarizer_model_id(provider_id, catalog, model), + run_budget: None, } } diff --git a/lib/components/fabro-agent/src/compaction.rs b/lib/components/fabro-agent/src/compaction.rs index ad634955d..1a9bb7853 100644 --- a/lib/components/fabro-agent/src/compaction.rs +++ b/lib/components/fabro-agent/src/compaction.rs @@ -2,6 +2,7 @@ use std::fmt::Write; use fabro_llm::client::Client; use fabro_llm::types::{Message as LlmMessage, Request}; +use fabro_model::{RunBudget, UsdMicros}; use tracing::debug; use crate::agent_profile::AgentProfile; @@ -85,6 +86,7 @@ pub(crate) async fn compact_context( estimate: ContextEstimate, emitter: &Emitter, session_id: &str, + run_budget: Option<&RunBudget>, ) -> Result<(), Error> { let original_turn_count = history.turns().len(); let preserve_start = history.compact_preserve_start(preserve_count); @@ -165,6 +167,15 @@ function names, error messages, and exact values. Omit pleasantries and conversa .await .map_err(CompactionError::Llm)?; + // Summarization spend counts against the run budget. A trip here is only + // charged; the session's next per-turn check halts it. + if let Some(budget) = run_budget { + let _ = budget.charge( + response.usage.total_tokens(), + response.cost_usd.map(UsdMicros::from_usd), + ); + } + let response_text = response.text(); let summary_text = response_text.trim(); @@ -759,6 +770,7 @@ mod tests { }, &emitter, "sess", + None, ) .await; diff --git a/lib/components/fabro-agent/src/config.rs b/lib/components/fabro-agent/src/config.rs index a6c8b5c39..8e9f48ad5 100644 --- a/lib/components/fabro-agent/src/config.rs +++ b/lib/components/fabro-agent/src/config.rs @@ -4,7 +4,7 @@ use std::time::Duration; use fabro_llm::types::{ReasoningEffort, Speed}; use fabro_mcp::config::McpServerSettings; -use fabro_model::AgentProfileKind; +use fabro_model::{AgentProfileKind, RunBudget}; use fabro_types::PermissionLevel; /// Callback invoked before each tool execution. Return `Ok(())` to allow, @@ -192,6 +192,10 @@ pub struct SessionOptions { /// Wall-clock timeout for the entire `process_input` call. /// When set, the session's cancel token is triggered after this duration. pub wall_clock_timeout: Option, + /// Shared run budget. When set, every LLM response's usage is charged + /// here, and the session interrupts with + /// `InterruptReason::BudgetExhausted` once a limit is exceeded. + pub run_budget: Option>, } impl std::fmt::Debug for SessionOptions { @@ -226,6 +230,7 @@ impl std::fmt::Debug for SessionOptions { .field("skill_dirs", &self.skill_dirs) .field("mcp_servers", &self.mcp_servers.len()) .field("wall_clock_timeout", &self.wall_clock_timeout) + .field("run_budget", &self.run_budget.is_some()) .finish() } } @@ -253,6 +258,7 @@ impl Default for SessionOptions { skill_dirs: None, mcp_servers: Vec::new(), wall_clock_timeout: None, + run_budget: None, } } } diff --git a/lib/components/fabro-agent/src/error.rs b/lib/components/fabro-agent/src/error.rs index b274bd4b7..bfe6cca32 100644 --- a/lib/components/fabro-agent/src/error.rs +++ b/lib/components/fabro-agent/src/error.rs @@ -6,6 +6,8 @@ use fabro_llm::Error as LlmError; pub enum InterruptReason { WallClockTimeout, Cancelled, + /// The run's token/cost budget was exhausted (`[run.limits]`). + BudgetExhausted, } impl std::fmt::Display for InterruptReason { @@ -13,6 +15,7 @@ impl std::fmt::Display for InterruptReason { match self { Self::WallClockTimeout => write!(f, "wall clock timeout"), Self::Cancelled => write!(f, "cancelled"), + Self::BudgetExhausted => write!(f, "run budget exhausted"), } } } diff --git a/lib/components/fabro-agent/src/session.rs b/lib/components/fabro-agent/src/session.rs index 3d0cf6ef0..d731b7652 100644 --- a/lib/components/fabro-agent/src/session.rs +++ b/lib/components/fabro-agent/src/session.rs @@ -1097,6 +1097,21 @@ impl Session { self.interrupt_reason.clone() } + /// True when the shared run budget is exhausted. On the first + /// observation, marks the session interrupted with `BudgetExhausted` and + /// cancels its token so in-flight work stops too. + fn budget_exceeded(&self) -> bool { + let Some(budget) = &self.config.run_budget else { + return false; + }; + if budget.exceeded().is_none() { + return false; + } + self.set_interrupt_reason(InterruptReason::BudgetExhausted); + self.cancel_token.cancel(); + true + } + fn set_interrupt_reason(&self, reason: InterruptReason) { let mut guard = self .interrupt_reason @@ -1491,6 +1506,14 @@ impl Session { return Err(self.interrupted_error()); } + // A budget exhausted outside this session (a parallel stage, a + // prior input, seeded resume spend) stops it before the next + // request. + if self.budget_exceeded() { + self.shutdown(SessionShutdownReason::Cancelled).await; + return Err(self.interrupted_error()); + } + if round_was_interrupted { let generations = { let mut control = self @@ -1889,6 +1912,16 @@ impl Session { )); *usage_accumulator += usage.clone(); UsdMicros::accumulate(cost_accumulator, response.cost_usd.map(UsdMicros::from_usd)); + if let Some(budget) = &self.config.run_budget { + let response_cost = response.cost_usd.map(UsdMicros::from_usd); + if budget.charge(usage.total_tokens(), response_cost).is_some() { + // Halt like the wall-clock timeout: the cancelled token + // stops tool execution and the next turn, and the reason + // makes `interrupted_error()` report the budget. + self.set_interrupt_reason(InterruptReason::BudgetExhausted); + self.cancel_token.cancel(); + } + } if let Some(reminder) = pending_task_reminder { self.history.push(reminder); @@ -2057,6 +2090,7 @@ impl Session { estimate, &self.event_emitter, &self.id, + self.config.run_budget.as_deref(), ) .await { diff --git a/lib/components/fabro-agent/src/tools.rs b/lib/components/fabro-agent/src/tools.rs index f892c971c..557a62482 100644 --- a/lib/components/fabro-agent/src/tools.rs +++ b/lib/components/fabro-agent/src/tools.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use fabro_llm::client::Client; use fabro_llm::types::{Message, Request, ToolDefinition}; -use fabro_model::ModelHandle; +use fabro_model::{ModelHandle, RunBudget, UsdMicros}; #[cfg(test)] use fabro_static::EnvVars; use futures::{StreamExt, stream}; @@ -22,8 +22,11 @@ pub(crate) const DEFAULT_READ_LINES: usize = 2000; /// Configuration for the optional LLM-based summarizer used by `web_fetch`. #[derive(Clone)] pub struct WebFetchSummarizer { - pub client: Client, - pub model_id: ModelHandle, + pub client: Client, + pub model_id: ModelHandle, + /// Shared run budget charged for summarization spend; `None` outside + /// budgeted workflow runs. + pub run_budget: Option>, } /// Returns true if the input looks like it contains HTML markup. @@ -794,6 +797,14 @@ pub(crate) fn make_web_fetch_tool(summarizer: Option) -> Reg let response = s.client.complete(&request).await.map_err(|e| { format!("web_fetch summarization (model={}) failed: {e}", s.model_id.model_id()) })?; + // Summarization spend counts against the run budget; + // the session's next per-turn check acts on a trip. + if let Some(budget) = &s.run_budget { + let _ = budget.charge( + response.usage.total_tokens(), + response.cost_usd.map(UsdMicros::from_usd), + ); + } Ok(response.text()) } (Some(_), None) => { @@ -2077,6 +2088,7 @@ mod tests { provider: ProviderId::anthropic(), model: "mock-model".to_string(), }, + run_budget: None, }; let tool = make_web_fetch_tool(Some(summarizer)); @@ -2186,6 +2198,7 @@ mod tests { provider: ProviderId::anthropic(), model: "target-model".to_string(), }, + run_budget: None, }; let tool = make_web_fetch_tool(Some(summarizer)); diff --git a/lib/components/fabro-agent/tests/it/parity_matrix.rs b/lib/components/fabro-agent/tests/it/parity_matrix.rs index 7f64a29fb..9f822ff00 100644 --- a/lib/components/fabro-agent/tests/it/parity_matrix.rs +++ b/lib/components/fabro-agent/tests/it/parity_matrix.rs @@ -49,8 +49,9 @@ fn summarizer_model_id(provider: &Provider) -> ModelHandle { fn build_summarizer(provider: &Provider, client: &Client) -> WebFetchSummarizer { WebFetchSummarizer { - client: client.clone(), - model_id: summarizer_model_id(provider), + client: client.clone(), + model_id: summarizer_model_id(provider), + run_budget: None, } } diff --git a/lib/components/fabro-hooks/src/executor.rs b/lib/components/fabro-hooks/src/executor.rs index 4aa87d168..2086d36cf 100644 --- a/lib/components/fabro-hooks/src/executor.rs +++ b/lib/components/fabro-hooks/src/executor.rs @@ -10,7 +10,7 @@ use fabro_auth::CredentialSource; use fabro_llm::client::Client as LlmClient; use fabro_llm::generate::{GenerateParams, generate_object}; use fabro_llm::types::{Message, Request, ToolResult}; -use fabro_model::Catalog; +use fabro_model::{Catalog, RunBudget, UsdMicros}; use fabro_redact::redacted_url_for_log; use fabro_types::settings::{InterpString, ResolveCtx, ResolveError}; use tokio::process::Command as TokioCommand; @@ -83,9 +83,20 @@ fn safe_url_source_for_log(url: &InterpString) -> String { } /// Executes hooks via shell commands or HTTP POST. -pub struct HookExecutorImpl; +#[derive(Default)] +pub struct HookExecutorImpl { + /// Shared `[run.limits]` spend budget. Prompt/agent hook LLM calls are + /// charged against it, and hooks stop firing once it is exhausted. + /// `None` outside budgeted runs. + run_budget: Option>, +} impl HookExecutorImpl { + #[must_use] + pub fn with_run_budget(run_budget: Option>) -> Self { + Self { run_budget } + } + /// Parse a hook decision from JSON stdout and exit code. fn parse_decision(exit_code: i32, stdout: &str) -> HookDecision { if exit_code == 0 { @@ -284,7 +295,13 @@ impl HookExecutorImpl { context: &HookContext, llm_source: &dyn CredentialSource, catalog: Arc, + run_budget: Option<&RunBudget>, ) -> HookDecision { + if run_budget.is_some_and(|budget| budget.exceeded().is_some()) { + tracing::warn!("prompt hook skipped: run budget exhausted"); + return HookDecision::Proceed; + } + let (prompt, model) = match Self::resolve_prompt_and_model(prompt, model) { Ok(resolved) => resolved, Err(error) => { @@ -313,19 +330,29 @@ impl HookExecutorImpl { .max_tokens(1024); match generate_object(params, HOOK_RESPONSE_SCHEMA.clone()).await { - Ok(result) => if let Some(obj) = result.output { match serde_json::from_value::(obj) { - Ok(resp) if resp.ok => HookDecision::Proceed, - Ok(resp) => HookDecision::Block { - reason: resp.reason, - }, - Err(e) => { - tracing::warn!(error = %e, "prompt hook response deserialize failed, proceeding"); + Ok(result) => { + // Hook spend counts against the run budget; a trip here + // stops later hooks and stages, not this decision. + if let Some(budget) = run_budget { + let _ = budget.charge( + result.total_usage.total_tokens(), + result.response.cost_usd.map(UsdMicros::from_usd), + ); + } + if let Some(obj) = result.output { match serde_json::from_value::(obj) { + Ok(resp) if resp.ok => HookDecision::Proceed, + Ok(resp) => HookDecision::Block { + reason: resp.reason, + }, + Err(e) => { + tracing::warn!(error = %e, "prompt hook response deserialize failed, proceeding"); + HookDecision::Proceed + } + } } else { + tracing::warn!("prompt hook returned no structured output, proceeding"); HookDecision::Proceed } - } } else { - tracing::warn!("prompt hook returned no structured output, proceeding"); - HookDecision::Proceed - }, + } Err(e) => { tracing::warn!(error = %e, "prompt hook LLM call failed, proceeding"); HookDecision::Proceed @@ -349,7 +376,13 @@ impl HookExecutorImpl { sandbox: Arc, llm_source: &dyn CredentialSource, catalog: Arc, + run_budget: Option<&RunBudget>, ) -> HookDecision { + if run_budget.is_some_and(|budget| budget.exceeded().is_some()) { + tracing::warn!("agent hook skipped: run budget exhausted"); + return HookDecision::Proceed; + } + let (prompt, model) = match Self::resolve_prompt_and_model(prompt, model) { Ok(resolved) => resolved, Err(error) => { @@ -386,6 +419,13 @@ impl HookExecutorImpl { let cancel = CancellationToken::new(); for _ in 0..rounds { + // The budget can be exhausted before this hook fired or by + // another concurrent charge; stop before the next request. + if run_budget.is_some_and(|budget| budget.exceeded().is_some()) { + tracing::warn!("agent hook stopped: run budget exhausted"); + return HookDecision::Proceed; + } + let request = Request { model: resolved_model.clone(), messages: messages.clone(), @@ -411,10 +451,25 @@ impl HookExecutorImpl { } }; + let budget_tripped = run_budget.is_some_and(|budget| { + budget + .charge( + response.usage.total_tokens(), + response.cost_usd.map(UsdMicros::from_usd), + ) + .is_some() + }); + let tool_calls = response.tool_calls(); if tool_calls.is_empty() { + // The hook finished its work; parse the decision even if + // this response tripped the budget. return Self::parse_prompt_response(&response.text()); } + if budget_tripped { + tracing::warn!("agent hook stopped mid-loop: run budget exhausted"); + return HookDecision::Proceed; + } messages.push(response.message.clone()); @@ -685,6 +740,7 @@ impl HookExecutor for HookExecutorImpl { context, llm_source, Arc::clone(&catalog), + self.run_budget.as_deref(), ) .await } @@ -709,6 +765,7 @@ impl HookExecutor for HookExecutorImpl { sandbox, llm_source, Arc::clone(&catalog), + self.run_budget.as_deref(), ) .await } @@ -830,7 +887,7 @@ mod tests { #[tokio::test] async fn command_executor_host_success() { - let executor = HookExecutorImpl; + let executor = HookExecutorImpl::default(); let def = make_definition("exit 0"); let ctx = make_context(); let sandbox = make_sandbox(); @@ -851,7 +908,7 @@ mod tests { #[tokio::test] async fn command_executor_host_failure() { - let executor = HookExecutorImpl; + let executor = HookExecutorImpl::default(); let def = make_definition("exit 1"); let ctx = make_context(); let sandbox = make_sandbox(); @@ -871,7 +928,7 @@ mod tests { #[tokio::test] async fn command_executor_host_skip_via_exit_2() { - let executor = HookExecutorImpl; + let executor = HookExecutorImpl::default(); let def = make_definition("exit 2"); let ctx = make_context(); let sandbox = make_sandbox(); @@ -891,7 +948,7 @@ mod tests { #[tokio::test] async fn command_executor_host_json_decision() { - let executor = HookExecutorImpl; + let executor = HookExecutorImpl::default(); let def = make_definition(r#"echo '{"decision": "skip", "reason": "test skip"}'"#); let ctx = make_context(); let sandbox = make_sandbox(); @@ -913,7 +970,7 @@ mod tests { #[tokio::test] async fn command_executor_env_vars_set() { - let executor = HookExecutorImpl; + let executor = HookExecutorImpl::default(); // Print env vars to stdout for verification let def = make_definition("echo $ARC_EVENT:$ARC_RUN_ID:$ARC_WORKFLOW"); let mut ctx = make_context(); @@ -935,7 +992,7 @@ mod tests { #[tokio::test] async fn no_hook_type_blocks() { - let executor = HookExecutorImpl; + let executor = HookExecutorImpl::default(); let def = HookDefinition { name: None, event: HookEvent::StageStart, @@ -1417,7 +1474,7 @@ mod tests { }) .await; - let executor = HookExecutorImpl; + let executor = HookExecutorImpl::default(); let def = HookDefinition { name: Some("http-test".into()), event: HookEvent::StageStart, @@ -1477,6 +1534,7 @@ mod tests { &make_context(), test_llm_source().as_ref(), test_catalog(), + None, ) .await; @@ -1493,6 +1551,50 @@ mod tests { } } + // A spent run budget short-circuits the hook before interpolation or any + // LLM call: the same prompt that Blocks above now skips with Proceed. + #[tokio::test] + async fn prompt_hook_skips_when_run_budget_exhausted() { + let budget = RunBudget::new(None, Some(10)); + assert!(budget.charge(50, None).is_some()); + + let decision = HookExecutorImpl::execute_prompt( + &make_definition("unused"), + &interp("{{ env.MISSING_HOOK_VALUE }}"), + None, + &make_context(), + test_llm_source().as_ref(), + test_catalog(), + Some(&budget), + ) + .await; + + assert!(matches!(decision, HookDecision::Proceed)); + } + + #[tokio::test] + async fn agent_hook_stops_when_run_budget_exhausted() { + let budget = RunBudget::new(None, Some(10)); + assert!(budget.charge(50, None).is_some()); + + let decision = HookExecutorImpl::execute_agent( + &make_definition("unused"), + &interp("{{ env.MISSING_HOOK_VALUE }}"), + None, + Some(1), + &make_context(), + make_sandbox(), + test_llm_source().as_ref(), + test_catalog(), + Some(&budget), + ) + .await; + + // The budget check fires before interpolation, so the prompt that + // would otherwise Block skips with Proceed instead. + assert!(matches!(decision, HookDecision::Proceed)); + } + // Fail-closed: an agent hook with a missing token blocks instead of firing. #[tokio::test] async fn agent_hook_missing_env_blocks() { @@ -1505,6 +1607,7 @@ mod tests { make_sandbox(), test_llm_source().as_ref(), test_catalog(), + None, ) .await; diff --git a/lib/components/fabro-hooks/src/runner.rs b/lib/components/fabro-hooks/src/runner.rs index c3ec53987..f63479d26 100644 --- a/lib/components/fabro-hooks/src/runner.rs +++ b/lib/components/fabro-hooks/src/runner.rs @@ -5,7 +5,7 @@ use fabro_agent::Sandbox; use fabro_auth::CredentialSource; #[cfg(test)] use fabro_auth::test_support; -use fabro_model::Catalog; +use fabro_model::{Catalog, RunBudget}; use crate::config::{HookDefinition, HookSettings}; use crate::executor::{HookExecutor, HookExecutorImpl}; @@ -32,13 +32,21 @@ impl HookRunner { let compiled_matchers = Self::compile_matchers(&config); Self { config, - executor: Arc::new(HookExecutorImpl), + executor: Arc::new(HookExecutorImpl::default()), llm_source, catalog, compiled_matchers, } } + /// Attach the shared `[run.limits]` spend budget. Prompt/agent hook LLM + /// calls charge it, and those hooks stop firing once it is exhausted. + #[must_use] + pub fn with_run_budget(mut self, run_budget: Option>) -> Self { + self.executor = Arc::new(HookExecutorImpl::with_run_budget(run_budget)); + self + } + /// Create a HookRunner with a custom executor (for testing). #[cfg(test)] pub fn with_executor(config: HookSettings, executor: Arc) -> Self { diff --git a/lib/components/fabro-workflow/src/error.rs b/lib/components/fabro-workflow/src/error.rs index f8c08ca15..721b2a368 100644 --- a/lib/components/fabro-workflow/src/error.rs +++ b/lib/components/fabro-workflow/src/error.rs @@ -337,6 +337,12 @@ pub enum Error { #[error("Pipeline cancelled")] Cancelled, + + /// The run's `[run.limits]` spend budget was exhausted. Terminal and + /// non-retryable: the run halts immediately and fails with + /// `FailureReason::BudgetExhausted`. + #[error("Run budget exhausted: {0}")] + BudgetExhausted(String), } impl Error { @@ -497,7 +503,8 @@ impl Error { | Self::RunNotFound(_) | Self::Unsupported(_) | Self::OutputSchemaValidation(_) - | Self::Cancelled => false, + | Self::Cancelled + | Self::BudgetExhausted(_) => false, } } @@ -506,6 +513,7 @@ impl Error { pub fn failure_category(&self) -> FailureCategory { match self { Self::Cancelled => FailureCategory::Canceled, + Self::BudgetExhausted(_) => FailureCategory::BudgetExhausted, Self::Llm(sdk_err) => classify_sdk_error(sdk_err), Self::Io(_) => FailureCategory::TransientInfra, Self::Parse(_) @@ -529,6 +537,7 @@ impl Error { pub fn failure_reason(&self) -> FailureReason { match self { Self::Cancelled => FailureReason::Cancelled, + Self::BudgetExhausted(_) => FailureReason::BudgetExhausted, Self::Stage { stage: ErrorStage::Publish, .. diff --git a/lib/components/fabro-workflow/src/handler/llm/api.rs b/lib/components/fabro-workflow/src/handler/llm/api.rs index 02be6fa2a..8df828d59 100644 --- a/lib/components/fabro-workflow/src/handler/llm/api.rs +++ b/lib/components/fabro-workflow/src/handler/llm/api.rs @@ -21,7 +21,8 @@ use fabro_mcp::config::McpServerSettings; #[cfg(test)] use fabro_model::catalog::LlmCatalogSettings; use fabro_model::{ - AgentProfileKind, Catalog, FallbackTarget, ModelHandle, ModelRef, ProviderId, UsdMicros, + AgentProfileKind, Catalog, FallbackTarget, ModelHandle, ModelRef, ProviderId, RunBudget, + UsdMicros, }; use fabro_types::settings::run::RunModelControls; use fabro_types::{FailoverProps, PermissionLevel, RunId, SessionCapability, StageId, StageTiming}; @@ -41,7 +42,7 @@ use super::routing::ProviderContext; use crate::context::WorkflowContext; use crate::context::keys::Fidelity; use crate::error::Error; -use crate::event::{Emitter, Event, StageScope}; +use crate::event::{Emitter, Event, RunNoticeCode, RunNoticeLevel, StageScope}; use crate::model_fallback::{ModelFallbackNotice, ModelFallbackPolicy}; use crate::outcome::billed_model_usage_from_llm; use crate::services::FabroRunToolServices; @@ -123,7 +124,11 @@ pub struct EffectiveRequestControls { pub(crate) speed: Option, } -fn classify_agent_error(err: fabro_agent::Error, allow_failover: bool) -> AgentApiErrorDisposition { +fn classify_agent_error( + err: fabro_agent::Error, + allow_failover: bool, + budget: Option<&RunBudget>, +) -> AgentApiErrorDisposition { match err { fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::Cancelled) => { AgentApiErrorDisposition::Cancelled @@ -133,6 +138,13 @@ fn classify_agent_error(err: fabro_agent::Error, allow_failover: bool) -> AgentA "Agent session hit its wall-clock timeout".to_string(), )) } + fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::BudgetExhausted) => { + let message = budget.and_then(RunBudget::exceeded).map_or_else( + || "run budget exhausted".to_string(), + |usage| usage.exceeded_message(), + ); + AgentApiErrorDisposition::Terminal(Error::BudgetExhausted(message)) + } fabro_agent::Error::Llm(err) if allow_failover && err.failover_eligible() => { AgentApiErrorDisposition::FailoverEligible(err) } @@ -565,6 +577,7 @@ fn spawn_event_forwarder( scope: StageScope, emitter: Arc, file_tracking: Arc>, + run_budget: Option>, ) -> EventForwarder { let mut rx = session.subscribe(); let root_session_id = session.id().to_string(); @@ -582,6 +595,22 @@ fn spawn_event_forwarder( // Reset watchdog on every event, including streaming deltas emitter.touch(); + // The session charges the budget as each response's usage + // arrives; this task is the closest site with a run emitter, so + // it drains the one-shot warning notices. + if matches!(&event.event, AgentEvent::AssistantMessage { .. }) { + if let Some(budget) = &run_budget { + for warning in budget.take_pending_warnings() { + emitter.notice_scoped( + RunNoticeLevel::Warn, + RunNoticeCode::BudgetNearlyExhausted, + warning.warning_message(), + &scope, + ); + } + } + } + // Track file changes from tool calls (including sub-agent events) track_file_event( &event.event, @@ -639,6 +668,9 @@ pub struct AgentApiBackend { /// Messages of fallback-plan notices already emitted for this run, so the /// same configuration warning is not repeated on every LLM call. emitted_plan_notices: Mutex>, + /// Shared run budget: charged by agent sessions and the one-shot prompt + /// path; `None` means unlimited. + run_budget: Option>, tool_env: Option>, mcp_servers: Vec, tool_secrets: ToolSecrets, @@ -730,6 +762,7 @@ struct LiveAgentInvocation { lease: Option>, event_forwarder: EventForwarder, file_tracking: Arc>, + run_budget: Option>, total_usage: TokenCounts, total_cost: Option, inference_duration: Duration, @@ -761,7 +794,7 @@ impl LiveAgentInvocation { allow_failover: bool, emitter: &Arc, ) -> Result { - let disposition = classify_agent_error(error, allow_failover); + let disposition = classify_agent_error(error, allow_failover, self.run_budget.as_deref()); self.abort_and_discard(emitter).await; match disposition { AgentApiErrorDisposition::Cancelled => Err(Error::Cancelled), @@ -818,6 +851,7 @@ impl AgentApiBackend { fallbacks, sessions: Mutex::new(HashMap::new()), emitted_plan_notices: Mutex::new(HashSet::new()), + run_budget: None, tool_env: None, mcp_servers: Vec::new(), tool_secrets: ToolSecrets::default(), @@ -865,6 +899,12 @@ impl AgentApiBackend { self } + #[must_use] + pub fn with_run_budget(mut self, budget: Option>) -> Self { + self.run_budget = budget; + self + } + fn resolve_effective_request_controls( &self, node: &Node, @@ -1028,6 +1068,7 @@ impl AgentApiBackend { self.mcp_servers.clone(), self.tool_secrets.clone(), self.fabro_run_tools.clone(), + self.run_budget.clone(), ) .await?; Ok(( @@ -1052,6 +1093,7 @@ impl AgentApiBackend { mcp_servers: Vec, tool_secrets: ToolSecrets, fabro_run_tools: Option, + run_budget: Option>, ) -> Result { let client = Client::from_source(source, Arc::clone(&catalog)) .await @@ -1066,11 +1108,12 @@ impl AgentApiBackend { .with_tool_secrets(tool_secrets); let profile_builder = if provider.profile_kind == AgentProfileKind::Claude5 { profile_builder.with_web_fetch_summarizer(Some(WebFetchSummarizer { - client: client.clone(), - model_id: ModelHandle::ByName { + client: client.clone(), + model_id: ModelHandle::ByName { provider: provider.provider_id.clone(), model: model.to_string(), }, + run_budget: run_budget.clone(), })) } else { profile_builder @@ -1083,6 +1126,7 @@ impl AgentApiBackend { speed: controls.speed, tool_hooks, mcp_servers, + run_budget: run_budget.clone(), // Workflow agents run with no `tool_access_policy`, which exposes // the entire tool registry (read, write, shell, subagent, MCP) and // skips approval gating. Report that truthfully so the UI doesn't @@ -1106,6 +1150,7 @@ impl AgentApiBackend { let factory_fabro_run_tools = fabro_run_tools.clone(); let factory_permission_level = config.permission_level; let factory_tool_hooks = config.tool_hooks.clone(); + let factory_run_budget = run_budget; let factory: SessionFactory = Arc::new(move || { let mut child_profile = factory_profile_builder.build(); if let Some(services) = factory_fabro_run_tools.clone() { @@ -1121,6 +1166,9 @@ impl AgentApiBackend { speed: controls.speed, tool_hooks: factory_tool_hooks.clone(), permission_level: factory_permission_level, + // Subagents burn real spend: charge them against the + // same run budget as the parent session. + run_budget: factory_run_budget.clone(), ..SessionOptions::default() }, None, @@ -1243,6 +1291,7 @@ impl AgentApiBackend { self.mcp_servers.clone(), self.tool_secrets.clone(), self.fabro_run_tools.clone(), + self.run_budget.clone(), ) .await; if request.cancel_token.is_cancelled() { @@ -1263,6 +1312,7 @@ impl AgentApiBackend { stage_scope.clone(), Arc::clone(emitter), Arc::clone(&live.file_tracking), + self.run_budget.clone(), ); begin_session_lifecycle(&live.session, emitter, None); @@ -1486,6 +1536,12 @@ impl CodergenBackend for AgentApiBackend { let mut inference_duration = Duration::ZERO; loop { + if let Some(budget) = &self.run_budget { + if let Some(exceeded) = budget.exceeded() { + return Err(Error::BudgetExhausted(exceeded.exceeded_message())); + } + } + let request = self.route_request( node, fallback_plan.current(), @@ -1511,6 +1567,23 @@ impl CodergenBackend for AgentApiBackend { &mut total_cost, completion.response.cost_usd.map(UsdMicros::from_usd), ); + if let Some(budget) = &self.run_budget { + let exceeded = budget.charge( + completion.response.usage.total_tokens(), + completion.response.cost_usd.map(UsdMicros::from_usd), + ); + for warning in budget.take_pending_warnings() { + emitter.notice_scoped( + RunNoticeLevel::Warn, + RunNoticeCode::BudgetNearlyExhausted, + warning.warning_message(), + stage_scope, + ); + } + if let Some(exceeded) = exceeded { + return Err(Error::BudgetExhausted(exceeded.exceeded_message())); + } + } let response_text = completion.response.text(); let validation_error = if let Some(schema) = &output_schema { @@ -1626,6 +1699,7 @@ impl CodergenBackend for AgentApiBackend { stage_scope.clone(), Arc::clone(emitter), Arc::clone(&file_tracking), + self.run_budget.clone(), ); // Activate with the steering hub after initialization so HTTP @@ -1639,6 +1713,7 @@ impl CodergenBackend for AgentApiBackend { lease: None, event_forwarder, file_tracking, + run_budget: self.run_budget.clone(), total_usage: TokenCounts::default(), total_cost: None, inference_duration: Duration::ZERO, @@ -3690,6 +3765,280 @@ enabled = true assert_eq!(usage.total_usd_micros, Some(100_000)); } + fn token_budget(max_tokens: u64) -> Arc { + Arc::new(RunBudget::new(None, Some(max_tokens))) + } + + fn collect_run_events(emitter: &Emitter) -> Arc>> { + let seen = Arc::new(Mutex::new(Vec::new())); + emitter.on_event({ + let seen = Arc::clone(&seen); + move |event| seen.lock().unwrap().push(event.clone()) + }); + seen + } + + fn budget_notice_messages(events: &[fabro_types::RunEvent]) -> Vec { + events + .iter() + .filter_map(|event| match &event.body { + fabro_types::EventBody::RunNotice(props) + if props.code == "budget_nearly_exhausted" => + { + Some(props.message.clone()) + } + _ => None, + }) + .collect() + } + + #[tokio::test] + async fn one_shot_fails_when_budget_exceeded_by_response() { + let server = MockServer::start(); + let completion = server.mock(|when, then| { + when.method(POST).path("/chat/completions"); + then.status(200) + .header("content-type", "application/json") + .json_body(chat_completion_response("done", 40, 10)); + }); + let backend = mock_api_backend(&server).with_run_budget(Some(token_budget(30))); + let node = Node::new("audit"); + let context = Context::new(); + let stage_scope = StageScope::for_handler(&context, &node.id); + let emitter = Arc::new(Emitter::new(fabro_types::RunId::new())); + let workspace = tempfile::tempdir().unwrap(); + let sandbox: Arc = + Arc::new(LocalSandbox::new(workspace.path().to_path_buf())); + + let result = backend + .one_shot(OneShotRequest { + node: &node, + prompt: "Audit the result", + system_prompt: None, + emitter: &emitter, + stage_scope: &stage_scope, + sandbox: &sandbox, + cancel_token: CancellationToken::new(), + }) + .await; + let Err(error) = result else { + panic!("expected the budget to fail the call") + }; + + completion.assert_calls(1); + assert!( + matches!(error, Error::BudgetExhausted(_)), + "expected budget exhausted, got {error:?}" + ); + assert!( + error.to_string().contains("run.limits.max_tokens"), + "message should name the limit: {error}" + ); + } + + #[tokio::test] + async fn one_shot_skips_the_request_when_budget_is_already_exhausted() { + let server = MockServer::start(); + let completion = server.mock(|when, then| { + when.method(POST).path("/chat/completions"); + then.status(200) + .header("content-type", "application/json") + .json_body(chat_completion_response("done", 1, 1)); + }); + let budget = token_budget(10); + assert!(budget.charge(50, None).is_some()); + let backend = mock_api_backend(&server).with_run_budget(Some(Arc::clone(&budget))); + let node = Node::new("audit"); + let context = Context::new(); + let stage_scope = StageScope::for_handler(&context, &node.id); + let emitter = Arc::new(Emitter::new(fabro_types::RunId::new())); + let workspace = tempfile::tempdir().unwrap(); + let sandbox: Arc = + Arc::new(LocalSandbox::new(workspace.path().to_path_buf())); + + let result = backend + .one_shot(OneShotRequest { + node: &node, + prompt: "Audit the result", + system_prompt: None, + emitter: &emitter, + stage_scope: &stage_scope, + sandbox: &sandbox, + cancel_token: CancellationToken::new(), + }) + .await; + let Err(error) = result else { + panic!("expected the budget to fail the call") + }; + + completion.assert_calls(0); + assert!( + matches!(error, Error::BudgetExhausted(_)), + "expected budget exhausted, got {error:?}" + ); + } + + #[tokio::test] + async fn one_shot_emits_the_eighty_percent_budget_notice_once() { + let server = MockServer::start(); + let completion = server.mock(|when, then| { + when.method(POST).path("/chat/completions"); + then.status(200) + .header("content-type", "application/json") + .json_body(chat_completion_response("done", 75, 10)); + }); + // 85 of 100 tokens: past the 80% warning threshold, under the limit. + let backend = mock_api_backend(&server).with_run_budget(Some(token_budget(100))); + let node = Node::new("audit"); + let context = Context::new(); + let stage_scope = StageScope::for_handler(&context, &node.id); + let emitter = Arc::new(Emitter::new(fabro_types::RunId::new())); + let events = collect_run_events(&emitter); + let workspace = tempfile::tempdir().unwrap(); + let sandbox: Arc = + Arc::new(LocalSandbox::new(workspace.path().to_path_buf())); + + backend + .one_shot(OneShotRequest { + node: &node, + prompt: "Audit the result", + system_prompt: None, + emitter: &emitter, + stage_scope: &stage_scope, + sandbox: &sandbox, + cancel_token: CancellationToken::new(), + }) + .await + .unwrap(); + + completion.assert_calls(1); + let notices = budget_notice_messages(&events.lock().unwrap()); + assert_eq!(notices.len(), 1, "exactly one warning notice: {notices:?}"); + assert!( + notices[0].contains("80% of run.limits.max_tokens"), + "got {notices:?}" + ); + } + + #[tokio::test] + async fn agent_run_fails_fast_when_budget_is_already_exhausted() { + let server = MockServer::start(); + let completion = server.mock(|when, then| { + when.method(POST).path("/chat/completions"); + then.status(200) + .header("content-type", "text/event-stream") + .body(chat_completion_stream("done", 1, 1)); + }); + let budget = token_budget(10); + assert!(budget.charge(50, None).is_some()); + let backend = mock_api_backend(&server).with_run_budget(Some(budget)); + let node = Node::new("code"); + let context = Context::new(); + let emitter = Arc::new(Emitter::new(fabro_types::RunId::new())); + let workspace = tempfile::tempdir().unwrap(); + let sandbox: Arc = + Arc::new(LocalSandbox::new(workspace.path().to_path_buf())); + + let result = backend + .run(CodergenRunRequest { + node: &node, + prompt: "Implement the feature", + context: &context, + thread_id: None, + emitter: &emitter, + sandbox: &sandbox, + tool_hooks: None, + cancel_token: CancellationToken::new(), + agent_tool_runtime: fabro_agent::AgentToolRuntime::default(), + }) + .await; + let Err(error) = result else { + panic!("expected the budget to fail the call") + }; + + completion.assert_calls(0); + assert!( + matches!(error, Error::BudgetExhausted(_)), + "expected budget exhausted, got {error:?}" + ); + } + + #[tokio::test] + async fn agent_run_halts_mid_session_when_a_response_trips_the_budget() { + let server = MockServer::start(); + // First turn: a tool call whose usage blows the budget. The session + // must interrupt without sending the follow-up turn. + let usage_chunk = serde_json::json!({ + "id": uuid::Uuid::new_v4().to_string(), + "model": "mock-model", + "choices": [], + "usage": { "prompt_tokens": 90, "completion_tokens": 20 } + }); + let first_body = chat_completion_tool_call_stream( + "read_file", + "call_budget", + r#"{"file_path":"data.txt"}"#, + ) + .replace( + "data: [DONE]", + &format!("data: {usage_chunk}\n\ndata: [DONE]"), + ); + let first = server.mock(|when, then| { + when.method(POST) + .path("/chat/completions") + .body_excludes(r#""role":"tool""#); + then.status(200) + .header("content-type", "text/event-stream") + .body(first_body); + }); + let followup = server.mock(|when, then| { + when.method(POST) + .path("/chat/completions") + .body_includes(r#""role":"tool""#); + then.status(200) + .header("content-type", "text/event-stream") + .body(chat_completion_stream("done", 1, 1)); + }); + let backend = mock_api_backend(&server).with_run_budget(Some(token_budget(100))); + let node = Node::new("code"); + let context = Context::new(); + let emitter = Arc::new(Emitter::new(fabro_types::RunId::new())); + let workspace = tempfile::tempdir().unwrap(); + tokio::fs::write(workspace.path().join("data.txt"), "hello\n") + .await + .unwrap(); + let sandbox: Arc = + Arc::new(LocalSandbox::new(workspace.path().to_path_buf())); + + let result = backend + .run(CodergenRunRequest { + node: &node, + prompt: "Implement the feature", + context: &context, + thread_id: None, + emitter: &emitter, + sandbox: &sandbox, + tool_hooks: None, + cancel_token: CancellationToken::new(), + agent_tool_runtime: fabro_agent::AgentToolRuntime::default(), + }) + .await; + let Err(error) = result else { + panic!("expected the budget to fail the call") + }; + + first.assert_calls(1); + followup.assert_calls(0); + assert!( + matches!(error, Error::BudgetExhausted(_)), + "expected budget exhausted, got {error:?}" + ); + assert!( + error.to_string().contains("run.limits.max_tokens"), + "message should name the limit: {error}" + ); + } + #[tokio::test] async fn agent_run_repairs_custom_output_schema_in_same_session() { let server = MockServer::start(); @@ -4075,6 +4424,7 @@ enabled = true scope, Arc::clone(&emitter), file_tracking, + None, ); session.sub_agent_event_callback()( @@ -4277,7 +4627,7 @@ enabled = true fn classify_interrupted_cancelled_is_cancelled() { let err = fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::Cancelled); assert!(matches!( - classify_agent_error(err, true), + classify_agent_error(err, true, None), AgentApiErrorDisposition::Cancelled )); } @@ -4285,7 +4635,7 @@ enabled = true #[test] fn classify_interrupted_wall_clock_is_terminal_precondition() { let err = fabro_agent::Error::Interrupted(fabro_agent::InterruptReason::WallClockTimeout); - match classify_agent_error(err, true) { + match classify_agent_error(err, true, None) { AgentApiErrorDisposition::Terminal(Error::Precondition(msg)) => { assert!(msg.contains("wall-clock")); } @@ -4297,7 +4647,7 @@ enabled = true fn classify_failover_eligible_llm_returns_failover_when_allowed() { let err = fabro_agent::Error::Llm(failover_eligible_llm_error()); assert!(matches!( - classify_agent_error(err, true), + classify_agent_error(err, true, None), AgentApiErrorDisposition::FailoverEligible(_) )); } @@ -4305,7 +4655,7 @@ enabled = true #[test] fn classify_failover_eligible_llm_returns_terminal_when_not_allowed() { let err = fabro_agent::Error::Llm(failover_eligible_llm_error()); - match classify_agent_error(err, false) { + match classify_agent_error(err, false, None) { AgentApiErrorDisposition::Terminal(Error::Llm(_)) => {} _ => panic!("expected Terminal(Error::Llm) when failover disallowed"), } @@ -4314,7 +4664,7 @@ enabled = true #[test] fn classify_non_failover_eligible_llm_is_terminal_llm() { let err = fabro_agent::Error::Llm(non_failover_llm_error()); - match classify_agent_error(err, true) { + match classify_agent_error(err, true, None) { AgentApiErrorDisposition::Terminal(Error::Llm(_)) => {} _ => panic!("expected Terminal(Error::Llm) for non-failover-eligible LLM error"), } @@ -4324,7 +4674,7 @@ enabled = true fn classify_refusal_llm_returns_failover_when_allowed() { let err = fabro_agent::Error::Llm(refusal_llm_error()); assert!(matches!( - classify_agent_error(err, true), + classify_agent_error(err, true, None), AgentApiErrorDisposition::FailoverEligible(_) )); } @@ -4332,7 +4682,7 @@ enabled = true #[test] fn classify_refusal_llm_returns_terminal_when_not_allowed() { let err = fabro_agent::Error::Llm(refusal_llm_error()); - match classify_agent_error(err, false) { + match classify_agent_error(err, false, None) { AgentApiErrorDisposition::Terminal(Error::Llm(llm_err)) => { assert!(llm_err.to_string().contains("claude-fable-5 refused")); } @@ -4343,7 +4693,7 @@ enabled = true #[test] fn classify_session_closed_is_terminal_precondition() { let err = fabro_agent::Error::SessionClosed; - match classify_agent_error(err, true) { + match classify_agent_error(err, true, None) { AgentApiErrorDisposition::Terminal(Error::Precondition(message)) => { assert!(message.contains("Agent session failed")); } @@ -4354,7 +4704,7 @@ enabled = true #[test] fn classify_invalid_state_is_terminal_precondition() { let err = fabro_agent::Error::InvalidState("oops".into()); - match classify_agent_error(err, true) { + match classify_agent_error(err, true, None) { AgentApiErrorDisposition::Terminal(Error::Precondition(message)) => { assert!(message.contains("Agent session failed")); } @@ -4365,7 +4715,7 @@ enabled = true #[test] fn classify_tool_execution_is_terminal_precondition() { let err = fabro_agent::Error::ToolExecution("tool blew up".into()); - match classify_agent_error(err, true) { + match classify_agent_error(err, true, None) { AgentApiErrorDisposition::Terminal(Error::Precondition(message)) => { assert!(message.contains("Agent session failed")); } diff --git a/lib/components/fabro-workflow/src/node_handler.rs b/lib/components/fabro-workflow/src/node_handler.rs index 70a35a34b..db29a8313 100644 --- a/lib/components/fabro-workflow/src/node_handler.rs +++ b/lib/components/fabro-workflow/src/node_handler.rs @@ -89,6 +89,19 @@ pub(crate) async fn execute_single_attempt( run_dir: &Path, services: &EngineServices, ) -> CoreResult { + // Between-stage budget gate. LLM stages also charge and check per + // response; this stops every stage type (command, human, parallel) from + // starting once the run's `[run.limits]` budget is spent. + if let Some(exceeded) = services + .run + .run_budget() + .and_then(|budget| budget.exceeded()) + { + return Err(CoreError::BudgetExhausted { + message: exceeded.exceeded_message(), + }); + } + let handler = services.registry.resolve(node); let wf_context = artifact::resolve_context_for_execution( @@ -149,6 +162,9 @@ pub(crate) async fn execute_single_attempt( match timed_result { Ok(Ok(wf_outcome)) => Ok(wf_outcome), Ok(Err(Error::Cancelled)) => Err(CoreError::Cancelled), + // Budget exhaustion terminates the run: it must not become a fail + // outcome, which would follow fail edges and keep executing stages. + Ok(Err(Error::BudgetExhausted(message))) => Err(CoreError::BudgetExhausted { message }), Ok(Err(fabro_err)) => { let retryable = handler.should_retry(&fabro_err); Err(CoreError::handler(HandlerErrorDetail { diff --git a/lib/components/fabro-workflow/src/pipeline/execute.rs b/lib/components/fabro-workflow/src/pipeline/execute.rs index 3b0fb81f9..2650ba064 100644 --- a/lib/components/fabro-workflow/src/pipeline/execute.rs +++ b/lib/components/fabro-workflow/src/pipeline/execute.rs @@ -13,7 +13,7 @@ use super::types::{Executed, Initialized}; use crate::artifact; use crate::context::{self, Context}; use crate::error::Error; -use crate::event::{Emitter, Event}; +use crate::event::{Emitter, Event, RunNoticeCode, RunNoticeLevel}; use crate::graph::WorkflowGraph; use crate::interview_runtime::InterviewBlockState; use crate::lifecycle::WorkflowLifecycle; @@ -139,6 +139,34 @@ pub async fn execute(init: Initialized) -> Executed { } let start = Instant::now(); + + // Budget gate: a run whose seeded spend (a resume after a budget failure) + // already exceeds `[run.limits]` fails before the first node. Seeded 80% + // warnings drain here so they surface once at start. + if let Some(budget) = engine.run.run_budget() { + for warning in budget.take_pending_warnings() { + engine.run.emitter.notice( + RunNoticeLevel::Warn, + RunNoticeCode::BudgetNearlyExhausted, + warning.warning_message(), + ); + } + if let Some(exceeded) = budget.exceeded() { + return Executed { + graph, + outcome: Err(Error::BudgetExhausted(format!( + "{} — raise [run.limits] to resume", + exceeded.exceeded_message() + ))), + run_options, + wall_time_ms: crate::millis_u64(start.elapsed()), + final_context: seed_context_from_checkpoint(checkpoint.as_ref()), + engine, + model, + }; + } + } + let graph_arc = Arc::new(graph.clone()); let wf_graph = WorkflowGraph(Arc::clone(&graph_arc)); @@ -326,6 +354,9 @@ pub async fn execute(init: Initialized) -> Executed { ) } Err(fabro_core::Error::Cancelled) => (Err(Error::Cancelled), initial_context), + Err(fabro_core::Error::BudgetExhausted { message }) => { + (Err(Error::BudgetExhausted(message)), initial_context) + } Err(fabro_core::Error::Blocked { message }) => { (Err(Error::engine(message)), initial_context) } diff --git a/lib/components/fabro-workflow/src/pipeline/initialize.rs b/lib/components/fabro-workflow/src/pipeline/initialize.rs index 5a7dd6069..a630d13b4 100644 --- a/lib/components/fabro-workflow/src/pipeline/initialize.rs +++ b/lib/components/fabro-workflow/src/pipeline/initialize.rs @@ -9,12 +9,13 @@ use fabro_auth::{ }; use fabro_graphviz::graph; use fabro_hooks::{HookContext, HookDecision, HookEvent, HookExecutionContext, HookRunner}; -use fabro_model::Catalog; +use fabro_model::{Catalog, RunBudget, UsdMicros}; use fabro_sandbox::{ GitSetupIntent, SandboxEventCallback, SandboxSpec, reconnect_for_run_with_callback, shell_quote, }; use fabro_static::EnvVars; use fabro_types::RunSandboxKind; +use fabro_types::settings::run::RunLimitsSettings; use fabro_vault::Vault; use tokio::runtime::Handle; use tokio::sync::RwLock as AsyncRwLock; @@ -30,6 +31,7 @@ use crate::handler::{HandlerRegistry, default_registry}; use crate::model_fallback::ModelFallbackPolicy; use crate::run_metadata::{RunMetadataRuntime, build_metadata_writer, metadata_branch_name}; use crate::run_options::{GitCheckpointOptions, RunOptions}; +use crate::runtime_store::RunStoreHandle; use crate::sandbox_git_runtime::SandboxGitRuntime; use crate::services::{ EngineServices, FabroRunToolServices, RunLocations, RunServices, WorkflowToolEnvProvider, @@ -133,6 +135,29 @@ fn build_sandbox_env( Ok((env, source)) } +/// Build the shared run budget from resolved `[run.limits]`, seeded with the +/// run's prior spend from the billing projection so a resumed run carries its +/// spend forward. Returns `None` when no limit is configured. +async fn build_run_budget( + limits: &RunLimitsSettings, + run_store: &RunStoreHandle, + catalog: &Arc, +) -> Option> { + if limits.is_unlimited() { + return None; + } + let budget = RunBudget::new(limits.max_cost, limits.max_tokens); + if let Ok(projection) = run_store.state().await { + let totals = crate::billing_rollup_from_projection(&projection, Some(catalog)).totals; + if totals.total_tokens != 0 || totals.total_usd_micros.is_some() { + // The seed charge queues an 80% warning and arms the exceeded + // state; `execute` drains and enforces them at run start. + let _ = budget.charge(totals.total_tokens, totals.total_usd_micros.map(UsdMicros)); + } + } + Some(Arc::new(budget)) +} + async fn build_registry( spec: &LlmSpec, interviewer: Arc, @@ -144,6 +169,7 @@ async fn build_registry( catalog: Arc, tool_secrets: ToolSecrets, fabro_run_tools: Option, + run_budget: Option>, ) -> Result<(Arc, bool), Error> { let no_backend_interviewer = Arc::clone(&interviewer); let build_no_backend = move || { @@ -178,6 +204,7 @@ async fn build_registry( let steering_hub_for_api = Arc::clone(&steering_hub); let tool_env_provider_for_backend = Arc::clone(&tool_env_provider); let fabro_run_tools_for_api = fabro_run_tools.clone(); + let run_budget_for_api = run_budget.clone(); Arc::new(default_registry(interviewer, move || { let tool_env_provider = Arc::clone(&tool_env_provider_for_backend); let mut api = AgentApiBackend::new_with_catalog( @@ -191,7 +218,8 @@ async fn build_registry( .with_run_model_controls(model_controls.clone()) .with_tool_env_provider(tool_env_provider.clone()) .with_tool_secrets(tool_secrets_for_api.clone()) - .with_mcp_servers(mcp_servers.clone()); + .with_mcp_servers(mcp_servers.clone()) + .with_run_budget(run_budget_for_api.clone()); if let Some(services) = fabro_run_tools_for_api.clone() { api = api.with_fabro_run_tools(services); } @@ -292,14 +320,24 @@ pub async fn initialize( let sandbox_git = Arc::new(SandboxGitRuntime::new()); let metadata_runtime = Arc::new(RunMetadataRuntime::new()); + let run_budget = build_run_budget( + &options.run_options.settings.run.limits, + &options.run_store, + &catalog, + ) + .await; + let hook_runner = if options.hooks.hooks.is_empty() { None } else { - Some(Arc::new(HookRunner::new( - options.hooks.clone(), - Arc::clone(&llm_source), - Arc::clone(&catalog), - ))) + Some(Arc::new( + HookRunner::new( + options.hooks.clone(), + Arc::clone(&llm_source), + Arc::clone(&catalog), + ) + .with_run_budget(run_budget.clone()), + )) }; let is_resume = checkpoint.is_some(); @@ -474,6 +512,7 @@ pub async fn initialize( Arc::clone(&catalog), tool_secrets.clone(), options.fabro_run_tools.clone(), + run_budget.clone(), ) .await? }; @@ -630,7 +669,8 @@ pub async fn initialize( metadata_runtime, metadata_writer, StageExecutionTracker::seeded(stage_executions), - ); + ) + .with_run_budget(run_budget); let engine = Arc::new(EngineServices { run: Arc::clone(&run_services), registry, @@ -1145,6 +1185,7 @@ mod tests { test_catalog(), ToolSecrets::default(), None, + None, ) .await .unwrap(); diff --git a/lib/components/fabro-workflow/src/services.rs b/lib/components/fabro-workflow/src/services.rs index 2af76a9c7..67b066d4b 100644 --- a/lib/components/fabro-workflow/src/services.rs +++ b/lib/components/fabro-workflow/src/services.rs @@ -10,7 +10,7 @@ use fabro_auth::CredentialSource; use fabro_auth::ResolvedCredentials; use fabro_hooks::{HookContext, HookDecision, HookExecutionContext, HookRunner}; use fabro_interview::Interviewer; -use fabro_model::{Catalog, ProviderId}; +use fabro_model::{Catalog, ProviderId, RunBudget}; use fabro_types::{ManifestPath, RunId}; use tokio_util::sync::CancellationToken; @@ -110,6 +110,8 @@ pub struct RunServices { /// Run-scoped stage execution allocator, shared between the core /// lifecycle and direct-dispatch handlers such as parallel branches. pub(crate) stage_executions: StageExecutionTracker, + /// Shared `[run.limits]` spend budget; `None` means unlimited. + pub(crate) run_budget: Option>, } impl RunServices { @@ -146,9 +148,16 @@ impl RunServices { metadata_writer, interview_blocker: Arc::new(RunInterviewBlocker::new()), stage_executions, + run_budget: None, }) } + /// The shared `[run.limits]` spend budget, when one is configured. + #[must_use] + pub fn run_budget(&self) -> Option<&Arc> { + self.run_budget.as_ref() + } + /// The run-level cancellation token. Cancel this to terminate the run. /// Derive child tokens via `cancel_token().child_token()` for sandbox /// command invocations. @@ -171,6 +180,17 @@ impl RunServices { .await } + #[must_use] + pub(crate) fn with_run_budget( + self: &Arc, + run_budget: Option>, + ) -> Arc { + Arc::new(Self { + run_budget, + ..self.as_ref().clone() + }) + } + #[must_use] pub fn with_run_store(self: &Arc, run_store: RunStoreHandle) -> Arc { Arc::new(Self { diff --git a/lib/components/fabro-workflow/tests/it/integration.rs b/lib/components/fabro-workflow/tests/it/integration.rs index 7ddb72647..29baf971c 100644 --- a/lib/components/fabro-workflow/tests/it/integration.rs +++ b/lib/components/fabro-workflow/tests/it/integration.rs @@ -1124,6 +1124,123 @@ impl Handler for AlwaysFailHandler { } } +/// A handler that reports the run budget as exhausted, as the LLM paths do +/// when accumulated usage passes `[run.limits]`. +struct BudgetExhaustedHandler; + +#[async_trait::async_trait] +impl Handler for BudgetExhaustedHandler { + async fn execute( + &self, + _node: &Node, + _context: &fabro_workflow::context::Context, + _graph: &Graph, + _run_dir: &Path, + _services: &fabro_workflow::handler::EngineServices, + ) -> Result { + Err(fabro_workflow::error::Error::BudgetExhausted( + "run used 150 tokens, exceeding run.limits.max_tokens 100".to_string(), + )) + } +} + +#[tokio::test] +async fn budget_exhausted_ends_the_run_without_following_fail_edges() { + // burn has an explicit fail edge to recover. A normal stage failure would + // route there; budget exhaustion must terminate the run instead, and the + // terminal event must carry the budget_exhausted failure reason. + let mut graph = Graph::new("BudgetHalt"); + + let mut start = Node::new("start"); + start.attrs.insert( + "shape".to_string(), + AttrValue::String("Mdiamond".to_string()), + ); + graph.nodes.insert("start".to_string(), start); + + let mut burn = Node::new("burn"); + burn.attrs.insert( + "type".to_string(), + AttrValue::String("budget_exhausted".to_string()), + ); + graph.nodes.insert("burn".to_string(), burn); + + graph + .nodes + .insert("recover".to_string(), Node::new("recover")); + + let mut exit = Node::new("exit"); + exit.attrs.insert( + "shape".to_string(), + AttrValue::String("Msquare".to_string()), + ); + graph.nodes.insert("exit".to_string(), exit); + + graph.edges.push(Edge::new("start", "burn")); + graph.edges.push(Edge::new("burn", "exit")); + let mut fail_edge = Edge::new("burn", "recover"); + fail_edge.attrs.insert( + "condition".to_string(), + AttrValue::String("outcome=failed".to_string()), + ); + graph.edges.push(fail_edge); + graph.edges.push(Edge::new("recover", "exit")); + + let dir = tempfile::tempdir().unwrap(); + let mut registry = HandlerRegistry::new(Box::new(StartHandler)); + registry.register("start", Box::new(StartHandler)); + registry.register("exit", Box::new(ExitHandler)); + registry.register("budget_exhausted", Box::new(BudgetExhaustedHandler)); + registry.register("recover", Box::new(ContextSetterHandler)); + + let emitter = Arc::new(Emitter::default()); + let events = collect_events(&emitter); + let engine = WorkflowRunner::new(registry, Arc::clone(&emitter), local_env()); + let run_options = RunOptions { + settings: WorkflowSettings::default(), + run_dir: dir.path().to_path_buf(), + cancel_token: CancellationToken::new(), + run_id: test_run_id("budget-halt"), + labels: std::collections::HashMap::new(), + workflow_slug: None, + github_app: None, + base_branch: None, + display_base_sha: None, + pre_run_git: None, + fork_source_ref: None, + git: None, + }; + + let error = engine + .run(&graph, &run_options) + .await + .expect_err("budget exhaustion should terminate the run as an error"); + assert!( + matches!(error, fabro_workflow::error::Error::BudgetExhausted(_)), + "expected budget exhausted, got {error:?}" + ); + + let events = events.lock().unwrap(); + let run_failed_reasons = events + .iter() + .filter_map(|event| match &event.body { + EventBody::RunFailed(props) => Some(props.failure.reason), + _ => None, + }) + .collect::>(); + assert_eq!( + run_failed_reasons, + vec![fabro_types::FailureReason::BudgetExhausted], + "the terminal event should carry the budget_exhausted reason" + ); + assert!( + !events + .iter() + .any(|event| event.node_id.as_deref() == Some("recover")), + "the fail edge must not be followed after budget exhaustion" + ); +} + #[tokio::test] async fn goal_gate_routes_to_retry_target_on_failure() { // Pipeline: diff --git a/lib/foundation/fabro-api/tests/workflow_settings_round_trip.rs b/lib/foundation/fabro-api/tests/workflow_settings_round_trip.rs index c253f217b..fd8a84cbb 100644 --- a/lib/foundation/fabro-api/tests/workflow_settings_round_trip.rs +++ b/lib/foundation/fabro-api/tests/workflow_settings_round_trip.rs @@ -105,6 +105,34 @@ fn workflow_settings_default_run_checkpoint_skip_git_hooks_is_false() { assert_eq!(json["run"]["checkpoint"]["skip_git_hooks"], false); } +#[test] +fn workflow_settings_json_includes_run_limits() { + let settings = workflow_settings_from_toml( + r#" +_version = 1 + +[run.limits] +max_cost = 25.50 +max_tokens = 5000000 +"#, + ); + + let json = serde_json::to_value(&settings).expect("workflow settings should serialize"); + assert_eq!(json["run"]["limits"]["max_cost_usd_micros"], 25_500_000); + assert_eq!(json["run"]["limits"]["max_tokens"], 5_000_000); + + let round_trip: ApiWorkflowSettings = + serde_json::from_value(json).expect("workflow settings should deserialize"); + assert_eq!(round_trip, settings); +} + +#[test] +fn workflow_settings_default_run_limits_are_absent() { + let settings = workflow_settings_from_toml("_version = 1\n"); + let json = serde_json::to_value(&settings).expect("workflow settings should serialize"); + assert_eq!(json["run"]["limits"], serde_json::json!({})); +} + fn assert_same_type() { assert_eq!( TypeId::of::(), diff --git a/lib/foundation/fabro-config/src/layers/mod.rs b/lib/foundation/fabro-config/src/layers/mod.rs index c3fa1632c..7decef054 100644 --- a/lib/foundation/fabro-config/src/layers/mod.rs +++ b/lib/foundation/fabro-config/src/layers/mod.rs @@ -33,9 +33,9 @@ pub use run::{ InterviewsLayer, McpEntryLayer, ModelRefOrSplice, NotificationProviderLayer, NotificationRouteLayer, PrepareStep, RunAgentLayer, RunArtifactsLayer, RunCheckpointLayer, RunCloneLayer, RunExecutionLayer, RunGitLayer, RunGoalLayer, RunIntegrationsGithubLayer, - RunIntegrationsLayer, RunLayer, RunMetaBranchLayer, RunModelControlsLayer, RunModelLayer, - RunPrepareLayer, RunPullRequestLayer, RunRunBranchLayer, RunScmLayer, ScmGitHubLayer, - StringOrSplice, + RunIntegrationsLayer, RunLayer, RunLimitsLayer, RunMetaBranchLayer, RunModelControlsLayer, + RunModelLayer, RunPrepareLayer, RunPullRequestLayer, RunRunBranchLayer, RunScmLayer, + ScmGitHubLayer, StringOrSplice, }; pub use server::{ GithubIntegrationLayer, IntegrationWebhooksLayer, ObjectStoreLocalLayer, ObjectStoreS3Layer, diff --git a/lib/foundation/fabro-config/src/layers/run.rs b/lib/foundation/fabro-config/src/layers/run.rs index f81b4e0f6..812ec3101 100644 --- a/lib/foundation/fabro-config/src/layers/run.rs +++ b/lib/foundation/fabro-config/src/layers/run.rs @@ -36,6 +36,8 @@ pub struct RunLayer { #[serde(default, skip_serializing_if = "Option::is_none")] pub execution: Option, #[serde(default, skip_serializing_if = "Option::is_none")] + pub limits: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] pub checkpoint: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub clone: Option, @@ -286,6 +288,32 @@ pub struct RunExecutionLayer { pub approval: Option, } +/// `[run.limits]` — run-scoped LLM spend budgets. Absent means unlimited. +#[derive( + Debug, + Clone, + Default, + PartialEq, + Serialize, + Deserialize, + fabro_macros::Combine, + fabro_macros::OptionsMetadata, +)] +#[serde(deny_unknown_fields)] +pub struct RunLimitsLayer { + /// Maximum total LLM cost for the run, in USD. When accumulated cost + /// exceeds this value, the run halts and fails. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[option(value_type = "float")] + pub max_cost: Option, + /// Maximum total LLM tokens for the run, counting all token buckets + /// (input, output, reasoning, cache reads, cache writes). When + /// accumulated tokens exceed this value, the run halts and fails. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[option(value_type = "integer")] + pub max_tokens: Option, +} + /// `[run.checkpoint]` — checkpoint policy. #[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] diff --git a/lib/foundation/fabro-config/src/lib.rs b/lib/foundation/fabro-config/src/lib.rs index f6097b9be..2b07d2e47 100644 --- a/lib/foundation/fabro-config/src/lib.rs +++ b/lib/foundation/fabro-config/src/lib.rs @@ -52,12 +52,13 @@ pub use layers::{ PrepareStep, ProjectLayer, ProviderSettings, ReasoningEffortFeature, ReplaceMap, RunAgentLayer, RunArtifactsLayer, RunCheckpointLayer, RunCloneLayer, RunEnvironmentLayer, RunExecutionLayer, RunGitLayer, RunGoalLayer, RunIntegrationsGithubLayer, RunIntegrationsLayer, RunLayer, - RunMetaBranchLayer, RunModelControlsLayer, RunModelLayer, RunPrepareLayer, RunPullRequestLayer, - RunRunBranchLayer, RunScmLayer, ScmGitHubLayer, ServerApiLayer, ServerArtifactsLayer, - ServerAuthGithubLayer, ServerAuthLayer, ServerIntegrationsLayer, ServerLayer, - ServerListenLayer, ServerLoggingLayer, ServerSandboxLayer, ServerSandboxProviderLayer, - ServerSandboxProvidersLayer, ServerSchedulerLayer, ServerSlateDbLayer, ServerStorageLayer, - ServerWebLayer, SettingsLayer, SlackIntegrationLayer, StickyMap, StringOrSplice, WorkflowLayer, + RunLimitsLayer, RunMetaBranchLayer, RunModelControlsLayer, RunModelLayer, RunPrepareLayer, + RunPullRequestLayer, RunRunBranchLayer, RunScmLayer, ScmGitHubLayer, ServerApiLayer, + ServerArtifactsLayer, ServerAuthGithubLayer, ServerAuthLayer, ServerIntegrationsLayer, + ServerLayer, ServerListenLayer, ServerLoggingLayer, ServerSandboxLayer, + ServerSandboxProviderLayer, ServerSandboxProvidersLayer, ServerSchedulerLayer, + ServerSlateDbLayer, ServerStorageLayer, ServerWebLayer, SettingsLayer, SlackIntegrationLayer, + StickyMap, StringOrSplice, WorkflowLayer, }; pub use logging::{resolve_log_destination, resolve_log_destination_with_env}; pub use parse::ParseError; diff --git a/lib/foundation/fabro-config/src/resolve/run.rs b/lib/foundation/fabro-config/src/resolve/run.rs index 9010f25f8..9ba8f74f1 100644 --- a/lib/foundation/fabro-config/src/resolve/run.rs +++ b/lib/foundation/fabro-config/src/resolve/run.rs @@ -1,5 +1,6 @@ use std::collections::{BTreeMap, HashMap}; +use fabro_types::billing::UsdMicros; use fabro_types::settings::InterpString; use fabro_types::settings::run::{ ArtifactsSettings, GitAuthorSettings, HookDefinition, HookType, InterviewProviderSettings, @@ -7,8 +8,9 @@ use fabro_types::settings::run::{ NotificationRouteSettings, PreparedStep, PreparedStepRun, PullRequestSettings, ResolvedMcpEntry, RunAgentSettings, RunBranchSettings, RunCheckpointSettings, RunCloneSettings, RunExecutionSettings, RunGitSettings, RunGoal, RunIntegrationsGithubSettings, - RunIntegrationsSettings, RunInterviewsSettings, RunMetaBranchSettings, RunModelControls, - RunModelSettings, RunNamespace, RunPrepareSettings, RunScmSettings, ScmGitHubSettings, TlsMode, + RunIntegrationsSettings, RunInterviewsSettings, RunLimitsSettings, RunMetaBranchSettings, + RunModelControls, RunModelSettings, RunNamespace, RunPrepareSettings, RunScmSettings, + ScmGitHubSettings, TlsMode, }; use fabro_util::workspace_glob::WorkspaceGlob; @@ -17,7 +19,7 @@ use crate::{ EnvironmentLayer, HookAgentMarker, HookEntry, HookTlsMode, InterviewProviderLayer, InterviewsLayer, McpEntryLayer, MergeMap, ModelRefOrSplice, NotificationProviderLayer, NotificationRouteLayer, RunAgentLayer, RunArtifactsLayer, RunCheckpointLayer, RunCloneLayer, - RunExecutionLayer, RunGitLayer, RunGoalLayer, RunIntegrationsLayer, RunLayer, + RunExecutionLayer, RunGitLayer, RunGoalLayer, RunIntegrationsLayer, RunLayer, RunLimitsLayer, RunMetaBranchLayer, RunModelLayer, RunPrepareLayer, RunPullRequestLayer, RunRunBranchLayer, RunScmLayer, StickyMap, StringOrSplice, }; @@ -64,6 +66,7 @@ pub fn resolve_run( git: resolve_git(layer.git.as_ref()), prepare: resolve_prepare(layer.prepare.as_ref(), errors), execution: resolve_execution(layer.execution.as_ref()), + limits: resolve_limits(layer.limits.as_ref(), errors), checkpoint: resolve_checkpoint(layer.checkpoint.as_ref()), clone, run_branch, @@ -227,6 +230,43 @@ fn resolve_execution(execution: Option<&RunExecutionLayer>) -> RunExecutionSetti } } +fn resolve_limits( + limits: Option<&RunLimitsLayer>, + errors: &mut Vec, +) -> RunLimitsSettings { + let Some(limits) = limits else { + return RunLimitsSettings::default(); + }; + + let max_cost = match limits.max_cost { + Some(cost) if !cost.is_finite() || cost <= 0.0 => { + errors.push(ResolveError::Invalid { + path: "run.limits.max_cost".to_string(), + reason: "must be a positive amount in USD".to_string(), + }); + None + } + Some(cost) => Some(UsdMicros::from_usd(cost)), + None => None, + }; + + let max_tokens = match limits.max_tokens { + Some(0) => { + errors.push(ResolveError::Invalid { + path: "run.limits.max_tokens".to_string(), + reason: "must be a positive token count".to_string(), + }); + None + } + other => other, + }; + + RunLimitsSettings { + max_cost, + max_tokens, + } +} + fn resolve_checkpoint(checkpoint: Option<&RunCheckpointLayer>) -> RunCheckpointSettings { RunCheckpointSettings { exclude_globs: checkpoint diff --git a/lib/foundation/fabro-config/src/tests/combine.rs b/lib/foundation/fabro-config/src/tests/combine.rs index ef296087b..bb6ee7eb4 100644 --- a/lib/foundation/fabro-config/src/tests/combine.rs +++ b/lib/foundation/fabro-config/src/tests/combine.rs @@ -346,3 +346,25 @@ bucket = "higher-bucket" assert_eq!(s3.bucket, Some("higher-bucket".to_string())); assert_eq!(s3.region, None); } + +#[test] +fn run_limits_fields_merge_independently() { + let lower = parse( + r" +[run.limits] +max_cost = 100.0 +max_tokens = 9000000 +", + ); + let higher = parse( + r" +[run.limits] +max_cost = 25.0 +", + ); + + let merged = higher.combine(lower); + let limits = merged.run.unwrap().limits.unwrap(); + assert_eq!(limits.max_cost, Some(25.0)); + assert_eq!(limits.max_tokens, Some(9_000_000)); +} diff --git a/lib/foundation/fabro-config/src/tests/resolve_run.rs b/lib/foundation/fabro-config/src/tests/resolve_run.rs index d416eb651..90ff428ee 100644 --- a/lib/foundation/fabro-config/src/tests/resolve_run.rs +++ b/lib/foundation/fabro-config/src/tests/resolve_run.rs @@ -1,3 +1,4 @@ +use fabro_types::billing::UsdMicros; use fabro_types::settings::InterpString; use fabro_types::settings::run::{ ApprovalMode, EnvironmentNetworkMode, EnvironmentProvider, RunGoal, RunMode, @@ -1481,3 +1482,77 @@ command = ["fs-server"] ); } } + +#[test] +fn run_limits_resolve_to_usd_micros() { + let settings = workflow_settings_from_toml( + r" +_version = 1 + +[run.limits] +max_cost = 25.50 +max_tokens = 5000000 +", + ) + .expect("[run.limits] should resolve") + .run; + + assert_eq!(settings.limits.max_cost, Some(UsdMicros(25_500_000))); + assert_eq!(settings.limits.max_tokens, Some(5_000_000)); + assert!(!settings.limits.is_unlimited()); +} + +#[test] +fn run_limits_accept_integer_max_cost() { + let settings = workflow_settings_from_toml( + r" +_version = 1 + +[run.limits] +max_cost = 25 +", + ) + .expect("an integer max_cost should resolve") + .run; + + assert_eq!(settings.limits.max_cost, Some(UsdMicros(25_000_000))); + assert_eq!(settings.limits.max_tokens, None); +} + +#[test] +fn run_limits_default_to_unlimited() { + let settings = workflow_settings_from_layer(SettingsLayer::default()) + .expect("empty settings should resolve") + .run; + + assert!(settings.limits.max_cost.is_none()); + assert!(settings.limits.max_tokens.is_none()); + assert!(settings.limits.is_unlimited()); +} + +#[test] +fn run_limits_reject_non_positive_values() { + let error = workflow_settings_from_toml( + r" +_version = 1 + +[run.limits] +max_cost = -1.0 +max_tokens = 0 +", + ) + .expect_err("non-positive limits should not resolve"); + + let errors = match error { + crate::Error::Resolve { errors, .. } => errors, + other => panic!("expected structured resolve errors, got {other:#}"), + }; + let paths = errors + .into_iter() + .map(|error| match error { + crate::ResolveError::Invalid { path, .. } => path, + other => panic!("expected invalid-value error, got {other}"), + }) + .collect::>(); + assert_eq!(paths, vec!["run.limits.max_cost", "run.limits.max_tokens"]); +} diff --git a/lib/foundation/fabro-core/src/error.rs b/lib/foundation/fabro-core/src/error.rs index fc4bfbdeb..a5525c2e8 100644 --- a/lib/foundation/fabro-core/src/error.rs +++ b/lib/foundation/fabro-core/src/error.rs @@ -53,6 +53,11 @@ pub enum Error { }, #[error("stall timeout on node \"{node_id}\"")] StallTimeout { node_id: String }, + /// The run's spend budget was exhausted. Terminal: unlike handler + /// failures, this does not convert to a fail outcome, so fail edges are + /// not followed and the run ends immediately. + #[error("run budget exhausted: {message}")] + BudgetExhausted { message: String }, #[error("{detail}")] Handler { detail: Box }, #[error("{message}")] diff --git a/lib/foundation/fabro-dev/src/commands/docs_options_reference.rs b/lib/foundation/fabro-dev/src/commands/docs_options_reference.rs index 406ec0b49..61f5ade5c 100644 --- a/lib/foundation/fabro-dev/src/commands/docs_options_reference.rs +++ b/lib/foundation/fabro-dev/src/commands/docs_options_reference.rs @@ -128,6 +128,12 @@ email = "fabro-bot@company.com""#, "[run.pull_request]", r"[run.pull_request] enabled = true", + ), + Section::of::( + "[run.limits]", + r"[run.limits] +max_cost = 25.00 +max_tokens = 5000000", ), Section::of::( "[run.agent]", diff --git a/lib/foundation/fabro-model/src/budget.rs b/lib/foundation/fabro-model/src/budget.rs new file mode 100644 index 000000000..365dc8298 --- /dev/null +++ b/lib/foundation/fabro-model/src/budget.rs @@ -0,0 +1,348 @@ +//! Run-level LLM spend budget tracking. +//! +//! A [`RunBudget`] is shared (via `Arc`) between the workflow engine and the +//! agent session. Every site that records LLM usage charges it here and acts +//! on the returned [`BudgetCharge`]: a tripped limit halts the run with +//! `BudgetExhausted`, and a newly crossed warning threshold emits a one-shot +//! run notice. + +use std::fmt; +use std::sync::Mutex; + +use crate::billing::UsdMicros; + +/// Numerator/denominator of the warning threshold: warn when spend reaches +/// 4/5 (80%) of a limit. +const WARN_NUMERATOR: i64 = 4; +const WARN_DENOMINATOR: i64 = 5; + +/// A budget dimension's limit and the spend measured against it. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum BudgetUsage { + Cost { limit: UsdMicros, spent: UsdMicros }, + Tokens { limit: u64, spent: u64 }, +} + +impl BudgetUsage { + /// Message for a run halted by this limit. + #[must_use] + pub fn exceeded_message(&self) -> String { + match self { + Self::Cost { limit, spent } => format!( + "run cost {} exceeded run.limits.max_cost {}", + format_usd(*spent), + format_usd(*limit) + ), + Self::Tokens { limit, spent } => { + format!("run used {spent} tokens, exceeding run.limits.max_tokens {limit}") + } + } + } + + /// Message for the one-shot 80% warning notice. + #[must_use] + pub fn warning_message(&self) -> String { + match self { + Self::Cost { limit, spent } => format!( + "run cost {} reached 80% of run.limits.max_cost {}", + format_usd(*spent), + format_usd(*limit) + ), + Self::Tokens { limit, spent } => { + format!("run used {spent} tokens, 80% of run.limits.max_tokens {limit}") + } + } + } +} + +impl fmt::Display for BudgetUsage { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Cost { limit, spent } => { + write!(f, "cost {} of {}", format_usd(*spent), format_usd(*limit)) + } + Self::Tokens { limit, spent } => write!(f, "{spent} of {limit} tokens"), + } + } +} + +fn format_usd(value: UsdMicros) -> String { + format!("${:.2}", value.0 as f64 / 1_000_000.0) +} + +/// Thread-safe accumulator for a run's LLM spend against optional limits. +/// +/// Token spend compares against `TokenCounts::total_tokens()` (all five +/// buckets). Cost spend stays `None` until a cost is observed — a response +/// with no provider-reported cost and no catalog price does not advance it, +/// which is why `max_tokens` is the backstop for unpriced models. +/// +/// Charging and reporting are split so the charging site does not need +/// access to the run's event stream: [`RunBudget::charge`] returns only the +/// exceeded limit (the caller must halt), while newly crossed warning +/// thresholds queue internally until a site that can emit run notices drains +/// them with [`RunBudget::take_pending_warnings`]. +#[derive(Debug)] +pub struct RunBudget { + max_cost: Option, + max_tokens: Option, + state: Mutex, +} + +#[derive(Debug, Default)] +struct BudgetState { + tokens_spent: u64, + cost_spent: Option, + cost_warned: bool, + tokens_warned: bool, + pending_warnings: Vec, +} + +impl RunBudget { + #[must_use] + pub fn new(max_cost: Option, max_tokens: Option) -> Self { + Self { + max_cost, + max_tokens, + state: Mutex::new(BudgetState::default()), + } + } + + /// Whether no limit is configured. An unlimited budget never trips. + #[must_use] + pub fn is_unlimited(&self) -> bool { + self.max_cost.is_none() && self.max_tokens.is_none() + } + + /// Charge usage totals to the budget. Returns the limit the accumulated + /// spend now exceeds, if any — the caller must halt the run. When both + /// limits are exceeded, the cost limit is reported. + /// + /// `tokens` is a `total_tokens()`-style sum; `cost` is `None` when the + /// response carried no cost data. Newly crossed warning thresholds are + /// queued for [`Self::take_pending_warnings`]. + pub fn charge(&self, tokens: i64, cost: Option) -> Option { + let mut state = self.state.lock().expect("run budget lock poisoned"); + state.tokens_spent = state + .tokens_spent + .saturating_add(u64::try_from(tokens).unwrap_or(0)); + UsdMicros::accumulate(&mut state.cost_spent, cost); + + let exceeded = exceeded_locked(self, &state); + + if let (Some(limit), Some(spent)) = (self.max_cost, state.cost_spent) { + if !state.cost_warned && spent <= limit && crossed_warn_threshold(spent.0, limit.0) { + state.cost_warned = true; + state + .pending_warnings + .push(BudgetUsage::Cost { limit, spent }); + } + } + if let Some(limit) = self.max_tokens { + let crossed = crossed_warn_threshold( + i64::try_from(state.tokens_spent).unwrap_or(i64::MAX), + i64::try_from(limit).unwrap_or(i64::MAX), + ); + if !state.tokens_warned && state.tokens_spent <= limit && crossed { + state.tokens_warned = true; + let spent = state.tokens_spent; + state + .pending_warnings + .push(BudgetUsage::Tokens { limit, spent }); + } + } + + exceeded + } + + /// Drain warning thresholds crossed since the last call. Each warning is + /// returned exactly once for the budget's lifetime; the caller emits + /// them as run notices. + #[must_use] + pub fn take_pending_warnings(&self) -> Vec { + std::mem::take( + &mut self + .state + .lock() + .expect("run budget lock poisoned") + .pending_warnings, + ) + } + + /// The currently exceeded limit, if any. Pre-flight check: lets a caller + /// skip an LLM request when the budget is already spent. + #[must_use] + pub fn exceeded(&self) -> Option { + let state = self.state.lock().expect("run budget lock poisoned"); + exceeded_locked(self, &state) + } + + #[must_use] + pub fn tokens_spent(&self) -> u64 { + self.state + .lock() + .expect("run budget lock poisoned") + .tokens_spent + } + + #[must_use] + pub fn cost_spent(&self) -> Option { + self.state + .lock() + .expect("run budget lock poisoned") + .cost_spent + } +} + +fn exceeded_locked(budget: &RunBudget, state: &BudgetState) -> Option { + if let (Some(limit), Some(spent)) = (budget.max_cost, state.cost_spent) { + if spent > limit { + return Some(BudgetUsage::Cost { limit, spent }); + } + } + if let Some(limit) = budget.max_tokens { + if state.tokens_spent > limit { + return Some(BudgetUsage::Tokens { + limit, + spent: state.tokens_spent, + }); + } + } + None +} + +fn crossed_warn_threshold(spent: i64, limit: i64) -> bool { + spent.saturating_mul(WARN_DENOMINATOR) >= limit.saturating_mul(WARN_NUMERATOR) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn usd(value: f64) -> UsdMicros { + UsdMicros::from_usd(value) + } + + #[test] + fn unlimited_budget_never_trips() { + let budget = RunBudget::new(None, None); + assert!(budget.is_unlimited()); + assert_eq!(budget.charge(1_000_000_000, Some(usd(10_000.0))), None); + assert!(budget.take_pending_warnings().is_empty()); + assert!(budget.exceeded().is_none()); + } + + #[test] + fn cost_limit_trips_when_spend_passes_it() { + let budget = RunBudget::new(Some(usd(25.0)), None); + assert!(budget.charge(100, Some(usd(24.0))).is_none()); + let exceeded = budget.charge(100, Some(usd(1.13))); + assert_eq!( + exceeded, + Some(BudgetUsage::Cost { + limit: usd(25.0), + spent: usd(25.13), + }) + ); + assert!(budget.exceeded().is_some()); + } + + #[test] + fn spend_equal_to_the_limit_does_not_trip() { + let budget = RunBudget::new(Some(usd(25.0)), Some(1_000)); + assert!(budget.charge(1_000, Some(usd(25.0))).is_none()); + assert!(budget.exceeded().is_none()); + } + + #[test] + fn token_limit_trips_without_cost_data() { + let budget = RunBudget::new(Some(usd(25.0)), Some(1_000)); + let exceeded = budget.charge(1_500, None); + assert_eq!( + exceeded, + Some(BudgetUsage::Tokens { + limit: 1_000, + spent: 1_500, + }) + ); + } + + #[test] + fn cost_exceeded_wins_over_tokens_exceeded() { + let budget = RunBudget::new(Some(usd(1.0)), Some(100)); + let exceeded = budget.charge(200, Some(usd(2.0))); + assert!(matches!(exceeded, Some(BudgetUsage::Cost { .. }))); + } + + #[test] + fn warnings_fire_once_per_dimension() { + let budget = RunBudget::new(Some(usd(10.0)), Some(1_000)); + + assert!(budget.charge(100, Some(usd(1.0))).is_none()); + assert!(budget.take_pending_warnings().is_empty()); + + assert!(budget.charge(750, Some(usd(7.5))).is_none()); + assert_eq!(budget.take_pending_warnings(), vec![ + BudgetUsage::Cost { + limit: usd(10.0), + spent: usd(8.5), + }, + BudgetUsage::Tokens { + limit: 1_000, + spent: 850, + }, + ]); + + assert!(budget.charge(50, Some(usd(0.5))).is_none()); + assert!(budget.take_pending_warnings().is_empty()); + } + + #[test] + fn a_charge_that_jumps_past_the_limit_reports_exceeded_not_warning() { + let budget = RunBudget::new(Some(usd(10.0)), None); + let exceeded = budget.charge(100, Some(usd(11.0))); + assert!(matches!(exceeded, Some(BudgetUsage::Cost { .. }))); + assert!(budget.take_pending_warnings().is_empty()); + } + + #[test] + fn missing_cost_does_not_advance_cost_spend() { + let budget = RunBudget::new(Some(usd(1.0)), None); + assert!(budget.charge(1_000_000, None).is_none()); + assert_eq!(budget.cost_spent(), None); + } + + #[test] + fn negative_token_totals_are_ignored() { + let budget = RunBudget::new(None, Some(1_000)); + assert!(budget.charge(-500, None).is_none()); + assert_eq!(budget.tokens_spent(), 0); + } + + #[test] + fn seeding_past_the_limit_trips_immediately() { + let budget = RunBudget::new(Some(usd(25.0)), None); + assert!(budget.charge(5_000, Some(usd(26.0))).is_some()); + assert!(budget.exceeded().is_some()); + } + + #[test] + fn messages_name_the_config_keys() { + let cost = BudgetUsage::Cost { + limit: usd(25.0), + spent: usd(25.13), + }; + assert_eq!( + cost.exceeded_message(), + "run cost $25.13 exceeded run.limits.max_cost $25.00" + ); + let tokens = BudgetUsage::Tokens { + limit: 1_000, + spent: 850, + }; + assert_eq!( + tokens.warning_message(), + "run used 850 tokens, 80% of run.limits.max_tokens 1000" + ); + } +} diff --git a/lib/foundation/fabro-model/src/lib.rs b/lib/foundation/fabro-model/src/lib.rs index fc43c95a7..0855f67e9 100644 --- a/lib/foundation/fabro-model/src/lib.rs +++ b/lib/foundation/fabro-model/src/lib.rs @@ -1,6 +1,7 @@ pub mod adapter; pub mod billing; pub mod bootstrap_catalog; +pub mod budget; pub mod catalog; pub mod codec; pub mod ids; @@ -17,6 +18,7 @@ pub use billing::{ ModelBillingFacts, ModelBillingInput, ModelPricing, ModelPricingPolicy, ModelRef, ModelUsage, OpenAiBillingFacts, OpenAiModelPricing, PricePerMTok, Speed, TokenCounts, UsdMicros, }; +pub use budget::{BudgetUsage, RunBudget}; pub use catalog::{ ApiKeyHeaderPolicy, BillingPolicy, Catalog, CredentialRef, CredentialRefParseError, FallbackTarget, ModelSelectionError, ProviderAuthConfig, SelectedModel, diff --git a/lib/foundation/fabro-types/src/run_event/infra.rs b/lib/foundation/fabro-types/src/run_event/infra.rs index 2230e950f..5001dc340 100644 --- a/lib/foundation/fabro-types/src/run_event/infra.rs +++ b/lib/foundation/fabro-types/src/run_event/infra.rs @@ -22,6 +22,7 @@ pub enum RunNoticeCode { ArtifactOffloadFailed, ArtifactSyncFailed, ArtifactUploadFailed, + BudgetNearlyExhausted, CheckpointMetadataDegraded, CheckpointMetadataPushFailed, CheckpointMetadataWriteFailed, diff --git a/lib/foundation/fabro-types/src/settings/mod.rs b/lib/foundation/fabro-types/src/settings/mod.rs index 5bf91e4d9..f915123ac 100644 --- a/lib/foundation/fabro-types/src/settings/mod.rs +++ b/lib/foundation/fabro-types/src/settings/mod.rs @@ -41,8 +41,8 @@ pub use run::{ NotificationProviderSettings, NotificationRouteSettings, PreparedStep, PullRequestSettings, ResolvedMcpEntry, RunAgentSettings, RunCheckpointSettings, RunEnvironmentSettings, RunExecutionSettings, RunGitSettings, RunGoal, RunIntegrationsGithubSettings, - RunIntegrationsSettings, RunInterviewsSettings, RunModelControls, RunModelSettings, - RunNamespace, RunPrepareSettings, RunScmSettings, ScmGitHubSettings, TlsMode, + RunIntegrationsSettings, RunInterviewsSettings, RunLimitsSettings, RunModelControls, + RunModelSettings, RunNamespace, RunPrepareSettings, RunScmSettings, ScmGitHubSettings, TlsMode, }; pub use server::{ GithubIntegrationSettings, IntegrationWebhooksSettings, LogDestination, ObjectStoreSettings, diff --git a/lib/foundation/fabro-types/src/settings/run.rs b/lib/foundation/fabro-types/src/settings/run.rs index 8e589a57a..c8db4cb4b 100644 --- a/lib/foundation/fabro-types/src/settings/run.rs +++ b/lib/foundation/fabro-types/src/settings/run.rs @@ -10,6 +10,7 @@ use std::collections::{BTreeMap, HashMap}; use std::path::PathBuf; use std::time::Duration as StdDuration; +use fabro_model::UsdMicros; use fabro_util::shell; use serde::de::{self, Deserializer}; use serde::ser::SerializeStruct; @@ -31,6 +32,8 @@ pub struct RunNamespace { pub git: RunGitSettings, pub prepare: RunPrepareSettings, pub execution: RunExecutionSettings, + #[serde(default)] + pub limits: RunLimitsSettings, pub checkpoint: RunCheckpointSettings, pub clone: RunCloneSettings, pub run_branch: RunBranchSettings, @@ -61,6 +64,7 @@ impl Default for RunNamespace { git: RunGitSettings::default(), prepare: RunPrepareSettings::default(), execution: RunExecutionSettings::default(), + limits: RunLimitsSettings::default(), checkpoint: RunCheckpointSettings::default(), clone: RunCloneSettings::default(), run_branch: RunBranchSettings::default(), @@ -925,6 +929,33 @@ impl Default for RunExecutionSettings { } } +/// `[run.limits]` — resolved run budget limits. An absent value means +/// unlimited. When a run's accumulated LLM spend exceeds either limit, the +/// run halts and fails with `FailureReason::BudgetExhausted`. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +pub struct RunLimitsSettings { + /// Maximum total LLM cost for the run. Serialized as USD micros. + #[serde( + default, + rename = "max_cost_usd_micros", + skip_serializing_if = "Option::is_none" + )] + pub max_cost: Option, + /// Maximum total LLM tokens for the run, compared against + /// `TokenCounts::total_tokens()` (all five buckets, including cache + /// reads and writes). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, +} + +impl RunLimitsSettings { + /// Whether any budget limit is configured. + #[must_use] + pub fn is_unlimited(&self) -> bool { + self.max_cost.is_none() && self.max_tokens.is_none() + } +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct RunCheckpointSettings { pub exclude_globs: Vec, diff --git a/lib/packages/fabro-api-client/src/.openapi-generator/FILES b/lib/packages/fabro-api-client/src/.openapi-generator/FILES index 4f99a6d3d..f71722bb7 100644 --- a/lib/packages/fabro-api-client/src/.openapi-generator/FILES +++ b/lib/packages/fabro-api-client/src/.openapi-generator/FILES @@ -373,6 +373,7 @@ models/run-integrations-github-settings.ts models/run-integrations-settings.ts models/run-interviews-settings.ts models/run-lifecycle.ts +models/run-limits-settings.ts models/run-links.ts models/run-manifest.ts models/run-meta-branch-settings.ts diff --git a/lib/packages/fabro-api-client/src/models/index.ts b/lib/packages/fabro-api-client/src/models/index.ts index 2fdd2e9fe..c892b1b84 100644 --- a/lib/packages/fabro-api-client/src/models/index.ts +++ b/lib/packages/fabro-api-client/src/models/index.ts @@ -344,6 +344,7 @@ export * from './run-integrations-github-settings'; export * from './run-integrations-settings'; export * from './run-interviews-settings'; export * from './run-lifecycle'; +export * from './run-limits-settings'; export * from './run-links'; export * from './run-manifest'; export * from './run-meta-branch-settings'; diff --git a/lib/packages/fabro-api-client/src/models/run-limits-settings.ts b/lib/packages/fabro-api-client/src/models/run-limits-settings.ts new file mode 100644 index 000000000..5a7536354 --- /dev/null +++ b/lib/packages/fabro-api-client/src/models/run-limits-settings.ts @@ -0,0 +1,29 @@ +/* tslint:disable */ +/* eslint-disable */ +/** + * Fabro Run API + * HTTP API for managing Fabro workflow run executions. + * + * The version of the OpenAPI document: 0.1.0 + * + * + * NOTE: This class is auto generated by OpenAPI Generator (https://openapi-generator.tech). + * https://openapi-generator.tech + * Do not edit the class manually. + */ + + + +/** + * Resolved `[run.limits]` spend budget for the run. An absent field means that limit is not configured. When accumulated LLM spend exceeds either limit, the run halts and fails with the `budget_exhausted` failure reason. + */ +export interface RunLimitsSettings { + /** + * Maximum total LLM cost for the run, in USD micros. + */ + 'max_cost_usd_micros'?: number; + /** + * Maximum total LLM tokens for the run, counting all token buckets (input, output, reasoning, cache reads, cache writes). + */ + 'max_tokens'?: number; +} diff --git a/lib/packages/fabro-api-client/src/models/run-namespace.ts b/lib/packages/fabro-api-client/src/models/run-namespace.ts index cea567a72..51f9eff1a 100644 --- a/lib/packages/fabro-api-client/src/models/run-namespace.ts +++ b/lib/packages/fabro-api-client/src/models/run-namespace.ts @@ -57,6 +57,9 @@ import type { RunIntegrationsSettings } from './run-integrations-settings'; import type { RunInterviewsSettings } from './run-interviews-settings'; // May contain unused imports in some cases // @ts-ignore +import type { RunLimitsSettings } from './run-limits-settings'; +// May contain unused imports in some cases +// @ts-ignore import type { RunMetaBranchSettings } from './run-meta-branch-settings'; // May contain unused imports in some cases // @ts-ignore @@ -80,6 +83,7 @@ export interface RunNamespace { 'git': RunGitSettings; 'prepare': RunPrepareSettings; 'execution': RunExecutionSettings; + 'limits': RunLimitsSettings; 'checkpoint': RunCheckpointSettings; 'clone': RunCloneSettings; 'run_branch': RunBranchSettings;