soft-agent: --budget flag + LexHost::with_policy (lex 0.3 #225)

agents/llm_driver.lex

@@ -50,6 +50,17 @@ fn on_decide(
   # client; before v0.2.2 it was a stub. The result is `Result[Str, Str]`.
   # The `[llm_local]` effect on this function's return type is what
   # promotes the call from "type error" to "allowed". Compile-time check.
+  #
+  # Note: lex 0.3 ships a `[budget(N)]` effect (#225). Adding it here
+  # would be documentation-only today: when soft-runner invokes
+  # on_decide via `host.call(...)` (a Rust → Lex entry), the runtime
+  # doesn't charge N for the entry call itself — only Lex-internal
+  # `Op::Call`s deduct. With on_decide making just one stdlib call and
+  # one local helper call, neither of which currently declares a
+  # matching `[budget(N)]`, no deduction would happen. See
+  # `crates/soft-agent/tests/budget_probe.rs` for the pinned semantics.
+  # `--budget N` on soft-run still constrains agents whose handlers
+  # call into matching-budget Lex helpers.
   let answer := unwrap_or(agent.local_complete(state.prompt), "(llm unavailable)")
   [{
     kind: "send_a2a", server: "", tool: "", args_json: "",

crates/soft-agent/src/lex_host.rs

@@ -92,6 +92,21 @@ impl LexHost {
         self
     }
 
+    /// Install a `DefaultHandler` initialized with a non-default
+    /// [`Policy`]. The most common reason to call this is to set
+    /// `policy.budget = Some(N)` so the lex 0.3 `[budget(N)]` effect
+    /// is enforced — calls deduct from the shared pool atomically and
+    /// fail with `"budget exceeded: requested N, used so far M,
+    /// ceiling C"` when exhausted.
+    ///
+    /// Cannot be combined with [`Self::with_handler_factory`] — the
+    /// last call wins.
+    pub fn with_policy(self, policy: Policy) -> Self {
+        self.with_handler_factory(move || -> Box<dyn EffectHandler> {
+            Box::new(DefaultHandler::new(policy.clone()))
+        })
+    }
+
     /// Call a Lex function by name with the given arguments.
     ///
     /// Returns the function's value plus a [`TraceTree`] capturing call

crates/soft-agent/src/lib.rs

@@ -41,6 +41,7 @@ pub use executor::{ActionExecutor, ExecError, MockExecutor};
 pub use gate::{Gate, Verdict};
 pub use lex_dsl::{parse_lex_config, LexAgentSetup, DSL_PREAMBLE};
 pub use lex_host::{HandlerFactory, LexCall, LexHost};
+pub use lex_runtime::Policy;
 pub use mailbox::{A2aMessage, Mailbox, MailboxSender};
 pub use metrics::Metrics;
 pub use router::{InProcessExecutor, InProcessRouter};

crates/soft-agent/tests/budget_probe.rs

@@ -0,0 +1,116 @@
+//! Probe: does soft-agent on lex 0.3 honor `[budget(N)]` on a Lex
+//! function? If yes, we can plumb a `--budget` CLI flag into
+//! soft-runner; if no, this test surfaces the failure mode.
+//!
+//! Tests both the happy path (budget high enough) and the rejection
+//! path (budget exhausted). Setup is deliberately tiny — a single
+//! function annotated with `[budget(N)]` and called via `LexHost`.
+
+use soft_agent::{LexHost, Policy};
+
+const SOURCE: &str = r#"
+fn step() -> [budget(10)] Int { 1 }
+fn run_three() -> [budget(10)] Int {
+  step() + step() + step()
+}
+"#;
+
+fn build_host_with_budget(ceiling: Option<u64>) -> LexHost {
+    let host = LexHost::from_source(SOURCE).expect("compile");
+    let mut policy = Policy::permissive();
+    policy.budget = ceiling;
+    host.with_policy(policy)
+}
+
+#[test]
+fn budget_unbounded_allows_any_calls() {
+    let host = build_host_with_budget(None);
+    let result = host.call("run_three", vec![]).expect("call ok");
+    assert_eq!(result.value.to_json().as_i64(), Some(3));
+}
+
+#[test]
+fn budget_high_enough_allows_call() {
+    // run_three calls step() three times; each [budget(10)] call
+    // deducts cost ≥1 from the pool. Ceiling 1000 should be plenty.
+    let host = build_host_with_budget(Some(1000));
+    let result = host.call("run_three", vec![]).expect("call ok");
+    assert_eq!(result.value.to_json().as_i64(), Some(3));
+}
+
+#[test]
+fn budget_exhausted_surfaces_error() {
+    // Set ceiling to 0 — even the first call should be rejected.
+    let host = build_host_with_budget(Some(0));
+    let result = host.call("run_three", vec![]);
+    let err = result.expect_err("expected budget-exceeded");
+    let msg = format!("{err}");
+    assert!(
+        msg.contains("budget exceeded"),
+        "expected budget-exceeded error, got: {msg}"
+    );
+}
+
+/// The lex 0.3 type checker requires every Lex caller of a
+/// `[budget(N)]` function to declare a compatible budget effect on
+/// its own return type. Pins the contract.
+#[test]
+fn calling_budgeted_callee_from_unbudgeted_lex_caller_is_a_type_error() {
+    let src = r#"
+        fn budgeted() -> [budget(10)] Int { 1 }
+        fn entry() -> Int { budgeted() }
+    "#;
+    let result = LexHost::from_source(src);
+    let err = match result {
+        Err(e) => e,
+        Ok(_) => panic!("expected typecheck failure, but compile succeeded"),
+    };
+    let msg = format!("{err}");
+    assert!(
+        msg.contains("budget(10)") && msg.contains("EffectNotDeclared"),
+        "expected EffectNotDeclared for budget(10), got: {msg}"
+    );
+}
+
+/// Practical consequence for soft-agent: when soft-runner invokes a
+/// `[budget(N)]`-annotated lex handler from Rust (via
+/// `host.call(fn_name, args)`), the runtime currently does NOT charge
+/// N for the entry call itself — only for `Op::Call`s the handler
+/// makes downstream. So `[budget(50)]` on a top-level handler bounds
+/// the *helper-function* + *stdlib-call* depth inside the handler,
+/// not the act of dispatching to it.
+///
+/// This test pins the behavior so we'd notice if it changes upstream.
+#[test]
+fn rust_caller_into_budgeted_handler_with_internal_calls_deducts_correctly() {
+    // lex 0.3 budget effects don't subsume — caller and callee
+    // declare the same `[budget(N)]` value (mirrors the upstream
+    // budget_runtime.rs test pattern).
+    let src = r#"
+        fn helper() -> [budget(10)] Int { 1 }
+        fn handler() -> [budget(10)] Int {
+            helper() + helper() + helper()
+        }
+    "#;
+    let host = LexHost::from_source(src).expect("compile");
+
+    // Ceiling 100 — plenty for 3 × 5 = 15 deductions plus overhead.
+    let mut p = Policy::permissive();
+    p.budget = Some(100);
+    let h = host.with_policy(p);
+    let v = h.call("handler", vec![]).expect("ceiling 100 ok");
+    assert_eq!(v.value.to_json().as_i64(), Some(3));
+
+    // Ceiling 1 — first internal helper() call should fail.
+    let host = LexHost::from_source(src).expect("compile");
+    let mut p = Policy::permissive();
+    p.budget = Some(1);
+    let h = host.with_policy(p);
+    let err = h
+        .call("handler", vec![])
+        .expect_err("ceiling 1 should fail");
+    assert!(
+        format!("{err}").contains("budget exceeded"),
+        "expected budget-exceeded, got: {err}"
+    );
+}

crates/soft-runner/src/main.rs

@@ -57,6 +57,7 @@ struct Args {
     store: Option<PathBuf>,
     llm_cloud_provider: LlmCloudProvider,
     shutdown_token: Option<String>,
+    budget: Option<u64>,
 }
 
 fn usage_exit(msg: &str) -> ! {
@@ -69,6 +70,8 @@ fn usage_exit(msg: &str) -> ! {
             [--bind <addr>]                   (default 127.0.0.1) \\\n  \
             [--store <path>]                  (persist trace on shutdown) \\\n  \
             [--llm-cloud-provider <name>]     (default | anthropic) \\\n  \
+            [--budget <N>]                    (lex 0.3 [budget(N)] ceiling; \
+no flag = unbounded) \\\n  \
             [--shutdown-token <T>]            (or SOFT_SHUTDOWN_TOKEN env; \
 required when --bind is non-loopback)"
     );
@@ -106,6 +109,9 @@ fn parse_args() -> Args {
     let mut store = None;
     let mut llm_cloud_provider = LlmCloudProvider::Default;
     let mut shutdown_token: Option<String> = std::env::var("SOFT_SHUTDOWN_TOKEN").ok();
+    let mut budget: Option<u64> = std::env::var("SOFT_BUDGET")
+        .ok()
+        .and_then(|s| s.parse().ok());
 
     while let Some(flag) = iter.next() {
         match flag.as_str() {
@@ -173,6 +179,15 @@ fn parse_args() -> Args {
                     )),
                 };
             }
+            "--budget" => {
+                let v = iter
+                    .next()
+                    .unwrap_or_else(|| usage_exit("--budget needs a u64 ceiling"));
+                budget = Some(
+                    v.parse()
+                        .unwrap_or_else(|e| usage_exit(&format!("--budget: {e}"))),
+                );
+            }
             other => usage_exit(&format!("unknown flag: `{other}`")),
         }
     }
@@ -193,6 +208,7 @@ to protect POST /shutdown. Use --bind 127.0.0.1 for local-only, or supply a toke
         ticks,
         store,
         llm_cloud_provider,
+        budget,
         shutdown_token,
     }
 }
@@ -227,9 +243,10 @@ fn main() -> ExitCode {
         }
     };
 
+    let policy = make_policy(args.budget);
     let host = match args.llm_cloud_provider {
-        LlmCloudProvider::Default => host,
-        LlmCloudProvider::Anthropic => match install_anthropic_handler(host) {
+        LlmCloudProvider::Default => host.with_policy(policy.clone()),
+        LlmCloudProvider::Anthropic => match install_anthropic_handler(host, policy.clone()) {
             Ok(h) => h,
             Err(msg) => {
                 eprintln!("--llm-cloud-provider anthropic: {msg}");
@@ -343,7 +360,7 @@ fn main() -> ExitCode {
     };
 
     eprintln!(
-        "soft-run: agent `{agent_name}` on http://{listen}  peers={}{}{}{}",
+        "soft-run: agent `{agent_name}` on http://{listen}  peers={}{}{}{}{}",
         format_peers(&args.peers),
         format_ticks(&args.ticks),
         match args.store {
@@ -354,6 +371,10 @@ fn main() -> ExitCode {
             LlmCloudProvider::Default => "",
             LlmCloudProvider::Anthropic => "  llm_cloud=anthropic",
         },
+        match args.budget {
+            Some(n) => format!("  budget={n}"),
+            None => String::new(),
+        },
     );
     let _server_handle = server.spawn();
 
@@ -402,20 +423,29 @@ fn main() -> ExitCode {
     ExitCode::SUCCESS
 }
 
-fn install_anthropic_handler(host: LexHost) -> Result<LexHost, String> {
+fn install_anthropic_handler(host: LexHost, policy: Policy) -> Result<LexHost, String> {
     if std::env::var("ANTHROPIC_API_KEY").is_err() {
         return Err("ANTHROPIC_API_KEY env var not set".into());
     }
-    Ok(host.with_handler_factory(|| {
+    Ok(host.with_handler_factory(move || {
         // `from_env` re-reads env per call; cheap, and gives users
         // hot-reload of model / max-tokens between calls.
         Box::new(
-            AnthropicCloudHandler::from_env(Policy::permissive())
+            AnthropicCloudHandler::from_env(policy.clone())
                 .expect("ANTHROPIC_API_KEY checked at startup"),
         )
     }))
 }
 
+/// Build the lex Policy applied to every handler. Today only the
+/// `[budget(N)]` ceiling is configurable here; future per-call
+/// policy bits (effect allowlist, time limits) will land alongside.
+fn make_policy(budget: Option<u64>) -> Policy {
+    let mut p = Policy::permissive();
+    p.budget = budget;
+    p
+}
+
 fn format_peers(peers: &HashMap<String, String>) -> String {
     if peers.is_empty() {
         return "{}".into();