diff --git a/CHANGELOG.md b/CHANGELOG.md index 300329e7..7147643e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -356,14 +356,47 @@ JamJet uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html). whoever registers workflows. Attribution does not carry into a subgraph's child workflow. In the Python SDK, `Agent(agent_ref=...)` and YAML `agent_ref:` (workflow) / `acts_for:` (node) set it. - - The docs now say what the rest of the card does, because several pages claimed more - than the runtime does: `token_budget` and `cost_budget_usd` are checked only for a - `bounded_autonomous` card on an `agent` node calling it, in the in-process worker; - `max_iterations` and `max_tool_calls` are compared with counters the runtime never - increments, so a limit above 0 never trips (a limit of 0 always does); `time_budget_secs` is never read; the error circuit - breaker never trips; nothing routes an autonomy escalation to a human or a - supervisor; and no tool compiles an `agents.yaml` into a card. See "What each - constraint does" in docs/guides/agent-model-guide.md. + - The card's limits are enforced too: see the next entry. + +- **An agent's autonomy limits are enforced.** `max_iterations` and `max_tool_calls` + were compared with counters nothing incremented, `time_budget_secs` was never read, + the circuit breaker never tripped, and `token_budget` and `cost_budget_usd` were + checked only for a `bounded_autonomous` card on an `agent` node in the in-process + worker. Every limit is now checked at the shared guard, at every autonomy level, on + the HTTP claim route and the in-process worker, for every node acting for the agent + and for an `agent` node calling it. + - The counts come from the run's event log, not from in-memory or snapshot counters: + each admitted step is recorded as a new `agent_step_admitted` event (its agents and + tool calls), and tokens, cost and failures are read from the step's `node_completed` + / `node_failed`. They survive a restart and a continue-as-new rollover, and a node's + `state_patch` cannot reset them. The step is recorded in a compare-and-set that + re-checks the limits, so two claims cannot both take the last step. + - An iteration is a step: a node acting for the agent that was let fire, retries + included. A tool call is one per step whose node uses a tool, plus each call of an + ADK dispatch batch; model steps count none. The time budget runs from the agent's + first step in the run. The circuit breaker counts the agent's failed steps since its + last completed one (new card field `max_consecutive_errors`, default 3) and applies + even to a card with no constraints. + - On a limit the node is refused and fails for good, recording + `autonomy_limit_reached` or `circuit_breaker_tripped`. New card field + `on_limit_exceeded: escalate_to_human` holds the node for an ordinary approval + instead (same digest binding, the policy chain's deadline, the limit shown under + `autonomy_limit`). Each approval of such a request releases one step: a retry of a failed step or a duplicate item asks again, and an approval given for another gate before the limit was reached does not count. `escalate_to_supervisor` + is not implemented and a card that asks for it is refused. + - `POST /work-items/:id/complete` accepts `cost_usd` (a negative or non-finite one is + 400), and records a result for a node acting for an agent only for an item the claim + route admitted (409 otherwise). The Python worker forwards `input_tokens`, + `output_tokens` and `cost_usd` from a node's output, and `JamjetClient.complete_work_item` + takes them. No executor in this runtime computes a cost, so a cost budget counts + only what workers report. + - Behaviour changes to plan for: a card that never tripped can now refuse nodes, + including through the default circuit breaker; an item claimed by an earlier build + for a node acting for an agent cannot be completed over HTTP (its lease lapses and + the node is claimed again); the worker no longer runs its own autonomy check, and + `BudgetState`'s iteration, tool call and error counters are documented as unused. + A build that predates this release reads `agent_step_admitted` as an unknown event + and stops advancing the run, so do not roll back past it on a store it has written. + See "Autonomy limits" in docs/guides/agent-model-guide.md. - **The audit log can be signed with Ed25519, so a third party can verify it.** HMAC signing uses one shared key, so anyone who can verify the log can also forge it. With diff --git a/docs/architecture/agent-model.md b/docs/architecture/agent-model.md index 5cf68b75..5f37e4bb 100644 --- a/docs/architecture/agent-model.md +++ b/docs/architecture/agent-model.md @@ -103,7 +103,7 @@ agent: | `bounded_autonomous` | Self-directs within iteration, token, cost limits | Configurable | | `fully_autonomous` | Free operation with guardrails only | Requires explicit policy | -What the runtime enforces of the level and its constraints is listed in the [Agent Model Guide](../guides/agent-model-guide.md#what-each-constraint-does). In short: a card's tool and delegation lists apply to every node that acts for the agent, at every level; for a `bounded_autonomous` card, an `agent` node calling it is refused (with `autonomy_limit_reached` and `escalation_required` events) once the run reaches the card's token or cost budget, in the in-process worker only. Iteration and tool-call limits are not counted, and nothing routes an escalation to a supervisor agent or a human approval node. +What the runtime enforces of the level and its constraints is listed in the [Agent Model Guide](../guides/agent-model-guide.md#what-each-constraint-does). In short: a card's tool and delegation lists and its limits (iterations, tool calls, tokens, cost, time and consecutive errors) apply to every node that acts for the agent, at every level, on both worker transports, counted from the run's event log. On a limit the node fails, or, with `on_limit_exceeded: escalate_to_human`, is held for a human approval. Nothing routes an escalation to a supervisor agent. --- diff --git a/docs/guides/agent-model-guide.md b/docs/guides/agent-model-guide.md index 24c88700..725cd5a3 100644 --- a/docs/guides/agent-model-guide.md +++ b/docs/guides/agent-model-guide.md @@ -20,6 +20,9 @@ agents: max_tool_calls: 50 token_budget: 100000 cost_budget_usd: 5.00 + time_budget_secs: 600 + max_consecutive_errors: 3 + on_limit_exceeded: escalate_to_human allowed_tools: [search_web, read_document, analyze_data] blocked_tools: ["delete_*", "write_to_production_*"] require_approval_for: [publish, external_api_write] @@ -33,9 +36,8 @@ The runtime reads these settings from the Agent Card registered with `POST /agen |------------|----------|-------| | `allowed_tools`, `blocked_tools`, `require_approval_for` | Yes | Every node that acts for the agent (`agent_ref`), on the HTTP claim route, the in-process worker and the completion gate. See below. | | `allowed_delegations` | Yes | Same nodes: `agent`, `agent_tool` and `a2a_task` targets. See below. | -| `token_budget`, `cost_budget_usd` | Partly | Only for a `bounded_autonomous` card, only on an `agent` node that calls it (the card named by that node's own `agent_ref` field, which says whom it calls, not the workflow's attribution), and only in the in-process worker, not on the claim route. Compared with the run's token and cost totals as the in-process worker accumulates them (results reported through the HTTP worker protocol add nothing); at or over the limit, or when the registry cannot be read, the node does not run and `autonomy_limit_reached` and `escalation_required` events are recorded (see Budget Enforcement below). | -| `max_iterations`, `max_tool_calls` | No, in effect | Read in the same place as the budgets, but compared with counters the runtime never increments, so a limit above 0 never trips (a limit of 0 always does). | -| `time_budget_secs` | No | Stored, never read. | +| `max_iterations`, `max_tool_calls`, `token_budget`, `cost_budget_usd`, `time_budget_secs`, `max_consecutive_errors` | Yes | Every node that acts for the agent, and an `agent` node calling it, at every autonomy level, on the HTTP claim route and the in-process worker. Counted from the run's event log. See [Autonomy limits](#autonomy-limits). | +| `on_limit_exceeded` | Yes | `fail` (default) or `escalate_to_human`. See [Autonomy limits](#autonomy-limits). | ### Enforcing `allowed_tools` @@ -65,12 +67,12 @@ From the Python SDK, `Agent("researcher", ..., agent_ref="jamjet://acme/research | Level | What it means | What the runtime enforces for it | |-------|--------------|------| -| `deterministic` | Follows the workflow graph exactly, no self-direction | Nothing beyond the card's tool and delegation lists | -| `guided` | Follows the graph, but chooses how to approach each node (default) | Nothing beyond the card's tool and delegation lists | -| `bounded_autonomous` | Self-directs within iteration, token, and cost budgets | The token and cost budgets, as described in the table above | -| `fully_autonomous` | Operates freely within guardrails | Nothing beyond the card's tool and delegation lists | +| `deterministic` | Follows the workflow graph exactly, no self-direction | The card's lists and limits | +| `guided` | Follows the graph, but chooses how to approach each node (default) | The card's lists and limits | +| `bounded_autonomous` | Self-directs within iteration, token, and cost budgets | The card's lists and limits | +| `fully_autonomous` | Operates freely within guardrails | The card's lists and limits | -The tool and delegation lists apply at every level. A circuit breaker for repeated errors is defined for every level, but the runtime does not count consecutive errors, so it never trips. +The level does not change what is enforced: every level gets the tool and delegation lists and every limit the card declares, and the circuit breaker applies even to a card that declares no constraints. Guided mode does not hold every tool call for approval; name the tools in `require_approval_for` for that. --- @@ -129,10 +131,31 @@ When an `agent` node runs, JamJet invokes the referenced agent with the given in --- -## Budget Enforcement +## Autonomy limits -When an `agent` node calls a `bounded_autonomous` agent whose card's `token_budget` or `cost_budget_usd` the run has already reached (in-process worker only): -1. `autonomy_limit_reached` and `escalation_required` (target `human_approval`) events are recorded -2. The node does not run. The worker leaves the item to its lease, so it is retried, and refused again, until its attempts run out. Nothing routes to a human or a supervisor agent automatically. +A node that acts for an agent (the workflow's `agent_ref`, or the node's own) is one step of that agent. So is an `agent` node, for the agent it calls; that agent's tool and delegation lists do not apply to the node, only its limits. When the runtime lets such a node fire, it records an `agent_step_admitted` event naming the agents and the tool calls the step makes. Every limit is checked before the step fires, against what the agent used in the run so far, all of it read from the run's event log: + +| Limit | What is counted | Refused when | +|-------|-----------------|--------------| +| `max_iterations` | Steps admitted, retries included | The agent has taken that many steps | +| `max_tool_calls` | One per step whose node uses a tool (`tool`, `mcp_tool`, `finalizer`, an `a2a_task` with a skill, `agent`, `agent_tool`, function nodes, and nodes that pick a target at run time), plus each call of an ADK dispatch batch. Model and condition steps count none | This step's calls would take the total past the limit (a step that makes no call is not limited) | +| `token_budget` | `input_tokens + output_tokens` on the agent's completed steps | The total is at or over the budget | +| `cost_budget_usd` | `cost_usd` on the agent's completed steps | The total is at or over the budget | +| `time_budget_secs` | Wall-clock time since the agent's first step in the run | That long has passed | +| `max_consecutive_errors` (default 3) | The agent's failed steps since its last completed step | That many in a row have failed | + +- Counted over the whole run: a continue-as-new segment and an engine restart keep the counts. A node's output or `state_patch` cannot change them. +- A limit of 0 refuses the first step (for `time_budget_secs`, the second). +- Token and cost totals grow only when a step reports. The step that crosses a budget is not refused (its usage is unknown until it finishes), so a total can end above the budget by that step's usage. A time budget is checked when a step starts; a running step is not interrupted. +- Usage is what the step reports. The in-process model executor reports tokens. No executor in this runtime reports a cost, so cost is counted only when a worker reports `cost_usd` on `POST /work-items/:id/complete` (the Python worker forwards `input_tokens`, `output_tokens` and `cost_usd` from a node's output); with nothing reported, the cost total stays 0 and a budget above 0 never trips. A worker that under-reports, or reports a failed call as completed, is believed. +- Only steps admitted while the guard records them count. A step running when the runtime is upgraded to this release has no record and adds nothing. +- Two claims cannot both take the last step: the step is recorded in the same compare-and-set that checks the limits. +- Over HTTP, a result for a node acting for an agent is recorded only for a work item the claim route admitted (409 otherwise), and only once per admission. +- A step is recorded when the guard lets the node fire. In the in-process worker a few things can still stop it after that (a spent workflow-level budget, an idempotency key another worker holds); such a step counts though it did not run. This errs toward refusing. +- Only an `agent` node charges the agent it calls. An `agent_tool`, `a2a_task` or coordinator that reaches a registered agent does not charge that agent's limits; attribute the node with `acts_for` to count it. An in-process worker built without an agent registry does not apply a called agent's limits (nodes acting for an agent are refused there). + +On a limit, with `on_limit_exceeded: fail` (the default), the node is refused and fails for good, which fails the run. The runtime records `autonomy_limit_reached` (`limit_type` is `max_iterations`, `max_tool_calls`, `token_budget`, `cost_budget` or `time_budget`) or `circuit_breaker_tripped`, and the node's `node_failed` error names the agent and the limit (it contains `agent autonomy`). + +With `on_limit_exceeded: escalate_to_human`, the node is held for an ordinary approval instead: the request is bound to the action like any other (the same digest), carries the policy chain's deadline, and shows the limit under `autonomy_limit` in its context. `autonomy_limit_reached` (or `circuit_breaker_tripped`) and `escalation_required` (target `human_approval`) are recorded. An approval releases that one step, once: a retry of a step that failed, a duplicate item, or the next step past the limit asks again. Only a request raised for the limit (one whose context shows `autonomy_limit`) can release it, so an approval given for a policy or card gate before the limit was reached does not. A failing card's limit reached while the approved step was being decided still refuses it. Without a deadline in the policy chain the request does not expire. A rejection or an expired request fails the node. If a node acts for several agents and one card says `fail`, the node is refused. RFC-005's `escalate_to_supervisor` is not implemented, and a card that sets it is refused at registration. A workflow-level `token_budget` or `cost_budget_usd` in the IR is a separate check, made by the in-process worker for the nodes it runs; the HTTP claim route does not make it. diff --git a/docs/guides/yaml-reference.md b/docs/guides/yaml-reference.md index 19077606..6e4447b5 100644 --- a/docs/guides/yaml-reference.md +++ b/docs/guides/yaml-reference.md @@ -194,7 +194,7 @@ agents: token_env: PARTNER_API_KEY ``` -`autonomy` and `constraints` mirror the Agent Card's fields. The runtime reads them from the card registered with `POST /agents`, not from this file, and enforces only some of them: see [What each constraint does](agent-model-guide.md#what-each-constraint-does). +`autonomy` and `constraints` mirror the Agent Card's fields. The runtime reads them from the card registered with `POST /agents`, not from this file: see [What each constraint does](agent-model-guide.md#what-each-constraint-does). --- diff --git a/docs/rfcs/RFC-005-agent-model.md b/docs/rfcs/RFC-005-agent-model.md index 0c3fa0cf..d9720376 100644 --- a/docs/rfcs/RFC-005-agent-model.md +++ b/docs/rfcs/RFC-005-agent-model.md @@ -82,7 +82,7 @@ interface AutonomyConstraints { } ``` -Implementation status (2026-10): the card stores `time_budget` as `time_budget_secs`. Which constraints the runtime enforces, and where, is listed in the [Agent Model Guide](../guides/agent-model-guide.md#what-each-constraint-does). +Implementation status (2026-10): the card stores `time_budget` as `time_budget_secs`, and adds `max_consecutive_errors` (the circuit breaker's threshold, default 3) and `on_limit_exceeded` (`fail` or `escalate_to_human`). Which constraints the runtime enforces, and where, is listed in the [Agent Model Guide](../guides/agent-model-guide.md#what-each-constraint-does). ### 3. Agent Lifecycle @@ -157,6 +157,8 @@ When any constraint is exceeded: - `escalate_to_supervisor` — delegates to a configured supervisor agent - `fail` — marks the node/workflow as failed +Implementation status (2026-10): the counts are folded from the run's event log per agent, not kept per execution, and the handler is the card's `on_limit_exceeded`. A node over a limit records `autonomy_limit_reached` or `circuit_breaker_tripped`; `fail` fails it, `escalate_to_human` holds it for an approval. `escalate_to_supervisor` is not implemented. See the [Agent Model Guide](../guides/agent-model-guide.md#autonomy-limits). + ### 6. Hot Reload When an agent's definition is updated (new version): diff --git a/runtime/agents/src/card.rs b/runtime/agents/src/card.rs index d2dd3b84..9c2ad18d 100644 --- a/runtime/agents/src/card.rs +++ b/runtime/agents/src/card.rs @@ -105,6 +105,34 @@ pub struct AutonomyConstraints { /// Operations that require explicit approval before execution. pub require_approval_for: Vec, pub time_budget_secs: Option, + /// Consecutive failed steps of this agent before its circuit breaker + /// trips. Absent means [`DEFAULT_MAX_CONSECUTIVE_ERRORS`]. + pub max_consecutive_errors: Option, + /// What happens when a limit is reached or the circuit breaker trips. + /// Absent means [`OnLimitExceeded::Fail`]. + pub on_limit_exceeded: Option, +} + +/// The circuit breaker's threshold when a card sets none. +pub const DEFAULT_MAX_CONSECUTIVE_ERRORS: u32 = 3; + +/// What the runtime does with a node that would take an agent past one of +/// its limits (or past a tripped circuit breaker). +/// +/// RFC-005 also names `escalate_to_supervisor`. Nothing implements it, so it +/// is not accepted: a card that asks for it fails to register rather than +/// being silently treated as `fail`. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum OnLimitExceeded { + /// Refuse the node and fail it for good. + #[default] + Fail, + /// Hold the node for a human approval of that one step, through the + /// ordinary approval path (bound to the action, with the policy chain's + /// deadline). Each approval releases one step, once: a retry, a duplicate + /// item or the next step past the limit needs its own. + EscalateToHuman, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/runtime/api/src/routes.rs b/runtime/api/src/routes.rs index 6560fdda..5deaa7a2 100644 --- a/runtime/api/src/routes.rs +++ b/runtime/api/src/routes.rs @@ -1709,6 +1709,7 @@ async fn gate_claimed_item( remote_agents: Some(&remote_agents), registry: Some(agents), }, + Some(wi.id), ) .await { @@ -1882,6 +1883,7 @@ async fn completion_gate( remote_agents: Some(&remote_agents), registry: Some(agents), }, + Some(wi.id), ) .await { @@ -2017,6 +2019,10 @@ struct CompleteWorkItemRequest { /// Finish reason (e.g. "stop", "length", "tool_calls"). #[serde(default)] finish_reason: Option, + /// Estimated USD cost of this node, as the worker reports it. Counted + /// toward an agent card's `cost_budget_usd`. + #[serde(default)] + cost_usd: Option, } /// `POST /work-items/:id/complete` — mark a work item as completed and emit NodeCompleted event. @@ -2081,6 +2087,14 @@ async fn complete_work_item( // write. The budget enforcer reads them back, and the approval digest // leaves them out as engine-written, so a reported value would be a budget // reset and a change to later nodes' input that no approval covers. + // A reported cost is summed into agent cost budgets, so it must be a + // spend: a negative one would refund the budget. + if body.cost_usd.is_some_and(|c| !(c.is_finite() && c >= 0.0)) { + return Err(ApiError::BadRequest( + "cost_usd must be a non-negative number".to_string(), + )); + } + let state_patch = jamjet_state::approvals::action_state(&body.state_patch); if state_patch != body.state_patch { warn!(work_item_id = %id, "complete: dropped engine bookkeeping keys from the reported state_patch"); @@ -2102,7 +2116,7 @@ async fn complete_work_item( input_tokens: body.input_tokens, output_tokens: body.output_tokens, finish_reason: body.finish_reason.clone(), - cost_usd: None, + cost_usd: body.cost_usd, provenance: None, idempotency_key: derived_key.clone(), }, diff --git a/runtime/api/tests/agent_autonomy_limits.rs b/runtime/api/tests/agent_autonomy_limits.rs new file mode 100644 index 00000000..5a436d09 --- /dev/null +++ b/runtime/api/tests/agent_autonomy_limits.rs @@ -0,0 +1,1405 @@ +//! An agent card's autonomy limits are enforced on every node acting for +//! the agent, at the shared guard, from the run's durable event log. +//! +//! Before this, `max_iterations` and `max_tool_calls` were compared with +//! counters nothing incremented, `time_budget_secs` was never read, the +//! circuit breaker never tripped, and the token and cost budgets were checked +//! only for a `bounded_autonomous` card on an `agent` node in the in-process +//! worker. Every test here drives both transports (the HTTP claim route and +//! the in-process worker) on both backends (memory and SQLite). + +use async_trait::async_trait; +use jamjet_agents::{InMemoryAgentRegistry, SqliteAgentRegistry}; +use jamjet_api::state::AppState; +use jamjet_core::workflow::ExecutionId; +use jamjet_state::backend::{StateBackend, WorkItem}; +use jamjet_state::event::EventKind; +use jamjet_worker::executor::{ExecutionResult, ExecutorError, NodeExecutor}; +use serde_json::{json, Value}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; + +mod support; + +use support::hooked::{AfterWrite, HookedBackend, Write}; + +const AGENT: &str = "jamjet://acme/researcher"; +const QUEUES: [&str; 4] = ["python_tool", "tool", "general", "model"]; + +fn card(constraints: Value) -> Value { + json!({ + "id": AGENT, + "uri": AGENT, + "name": "researcher", + "description": "research agent", + "version": "1.0.0", + "capabilities": { "skills": [], "protocols": [], "tools_provided": [], "tools_consumed": [] }, + // The limits apply at every level; this one used to switch them off. + "autonomy": "fully_autonomous", + "constraints": constraints, + "auth": { "type": "none" } + }) +} + +fn mcp_tool() -> Value { + json!({ "type": "mcp_tool", "server": "srv", "tool": "search_web", "input_mapping": {}, "output_schema": "" }) +} + +fn model_node() -> Value { + json!({ "type": "model", "model_ref": "fake", "prompt_ref": "", "output_schema": "" }) +} + +fn adk_dispatch() -> Value { + json!({ + "type": "python_fn", + "module": "jamjet.agents.tool_runtime", + "function": "dispatch_tool_calls", + "output_schema": "", + "agent_tool_dispatch": true + }) +} + +fn calls(n: usize) -> Value { + let tool_calls: Vec = (0..n) + .map(|i| json!({ "id": format!("c{i}"), "name": "search_web", "arguments": { "i": i } })) + .collect(); + json!({ "tool_calls": tool_calls, "tools": {} }) +} + +/// A one-node workflow. `workflow_agent` / `node_agent` attribute it. +fn ir( + workflow_id: &str, + kind: Value, + workflow_agent: Option<&str>, + node_agent: Option<&str>, + policy: Option, +) -> Value { + let mut node = json!({ + "id": "n1", + "kind": kind, + "retry_policy": null, + "node_timeout_secs": null, + "description": null, + "labels": {} + }); + if let Some(a) = node_agent { + node["agent_ref"] = json!(a); + } + let mut ir = json!({ + "workflow_id": workflow_id, + "version": "0.1.0", + "name": null, + "description": null, + "state_schema": "", + "start_node": "n1", + "nodes": { "n1": node }, + "edges": [ { "from": "n1", "to": "end", "condition": null } ], + "retry_policies": {}, + "timeouts": {}, + "models": {}, + "tools": {}, + "mcp_servers": {}, + "remote_agents": {}, + "labels": {} + }); + if let Some(a) = workflow_agent { + ir["agent_ref"] = json!(a); + } + if let Some(p) = policy { + ir["policy"] = p; + } + ir +} + +#[derive(Clone, Copy, Debug, PartialEq)] +enum Transport { + Claim, + Worker, +} + +const TRANSPORTS: [Transport; 2] = [Transport::Claim, Transport::Worker]; + +/// How the step's work ends. +#[derive(Clone, Copy, Debug)] +enum Act { + /// Completes, reporting these tokens and (over HTTP only) this cost. + Done { tokens: u64, cost: Option }, + /// Fails. + Fail, +} + +const DONE: Act = Act::Done { + tokens: 0, + cost: None, +}; + +/// Runs the node as told, counting runs. +#[derive(Clone)] +struct Scripted { + act: Act, + runs: Arc, +} + +#[async_trait] +impl NodeExecutor for Scripted { + async fn execute(&self, _item: &WorkItem) -> Result { + self.runs.fetch_add(1, Ordering::SeqCst); + match self.act { + Act::Done { tokens, .. } => Ok(ExecutionResult { + output: json!({}), + state_patch: json!({}), + duration_ms: 0, + gen_ai_system: None, + gen_ai_model: None, + input_tokens: Some(tokens), + output_tokens: Some(0), + finish_reason: None, + }), + Act::Fail => Err(ExecutorError::Fatal("tool down".into())), + } + } +} + +/// One started run and what is needed to give its node more work. +struct Run { + state: AppState, + backend: Arc, + eid: ExecutionId, + ir: Value, + wf: String, + input: Value, + /// The start item has been taken; later steps enqueue their own. + started: bool, +} + +/// A state over `backend`, with an agent registry of `tb`'s kind. +async fn state_for(tb: &support::TestBackend, backend: Arc) -> AppState { + let mut state = support::make_state(backend); + if let support::TestBackend::Sqlite { url, .. } = tb { + state.agents = Arc::new(SqliteAgentRegistry::connect(url).await.unwrap()); + } else { + state.agents = Arc::new(InMemoryAgentRegistry::new()); + } + state +} + +async fn register(state: &AppState, cards: &[Value]) { + for c in cards { + let (st, body) = support::post(state, "/agents", c.clone()).await; + assert!(st.is_success(), "register agent: {st} {body}"); + } +} + +async fn start_on( + tb: &support::TestBackend, + backend: Arc, + cards: &[Value], + ir: Value, + input: Value, +) -> Run { + let state = state_for(tb, backend.clone()).await; + register(&state, cards).await; + let wf = ir["workflow_id"].as_str().unwrap().to_string(); + let eid = support::register_and_start(&state, ir.clone(), input.clone()).await; + let eid = + ExecutionId(uuid::Uuid::parse_str(eid.strip_prefix("exec_").unwrap_or(&eid)).unwrap()); + Run { + state, + backend, + eid, + ir, + wf, + input, + started: false, + } +} + +async fn start(tb: &support::TestBackend, cards: &[Value], ir: Value, input: Value) -> Run { + start_on(tb, tb.dyn_backend(), cards, ir, input).await +} + +impl Run { + fn eid_path(&self) -> String { + format!("exec_{}", self.eid.0) + } + + /// Give the node one more item, as the scheduler would on a loop. + async fn enqueue(&self) { + let ir_t: jamjet_ir::WorkflowIr = serde_json::from_value(self.ir.clone()).unwrap(); + let payload = jamjet_scheduler::build_work_item_payload( + &ir_t, + "n1", + ir_t.node("n1").unwrap(), + self.input.as_object().unwrap(), + &self.wf, + "0.1.0", + ); + let queue = ir_t + .node("n1") + .unwrap() + .kind + .queue_type() + .as_str() + .to_string(); + let (st, body) = support::post( + &self.state, + "/work-items", + json!({ "execution_id": self.eid_path(), "node_id": "n1", "queue_type": queue, "payload": payload }), + ) + .await; + assert!(st.is_success(), "enqueue: {st} {body}"); + } + + /// One step of the node over `t`: true if its work was handed out + /// (claim) or executed (worker). + async fn step(&mut self, t: Transport, act: Act) -> bool { + if self.started { + self.enqueue().await; + } + self.started = true; + match t { + Transport::Claim => { + let (_, claimed) = support::post( + &self.state, + "/work-items/claim", + json!({ "worker_id": "w", "queue_types": QUEUES }), + ) + .await; + if claimed["claimed"] != json!(true) { + return false; + } + let wi = &claimed["work_item"]; + let id = wi["id"].as_str().unwrap(); + let fence = wi["lease_fence"].clone(); + match act { + Act::Done { tokens, cost } => { + let mut body = json!({ + "output": {}, "state_patch": {}, "lease_fence": fence, + "input_tokens": tokens, "output_tokens": 0 + }); + if let Some(c) = cost { + body["cost_usd"] = json!(c); + } + let (st, body) = + support::post(&self.state, &format!("/work-items/{id}/complete"), body) + .await; + assert!(st.is_success(), "complete: {st} {body}"); + assert_eq!(body["completed"], json!(true), "{body}"); + } + Act::Fail => { + let (st, body) = support::post( + &self.state, + &format!("/work-items/{id}/fail"), + json!({ "error": "tool down", "lease_fence": fence }), + ) + .await; + assert!(st.is_success(), "fail: {st} {body}"); + // The engine's retry of a failed item: the test gives + // the node its next item itself, so drop the retry. + self.drop_pending().await; + } + } + true + } + Transport::Worker => { + let exec = Scripted { + act, + runs: Arc::new(AtomicUsize::new(0)), + }; + let worker = jamjet_worker::Worker::new( + "w".into(), + self.backend.clone(), + QUEUES.iter().map(|q| q.to_string()).collect(), + ) + .with_agents(self.state.agents.clone()) + .register_executor("mcp_tool", Arc::new(exec.clone())) + .register_executor("python_fn", Arc::new(exec.clone())) + .register_executor("model", Arc::new(exec.clone())) + .register_executor("agent", Arc::new(exec.clone())); + let _ = worker.poll_and_execute().await; + exec.runs.load(Ordering::SeqCst) > 0 + } + } + } + + /// Settle every pending item of the run (the retry a failed item left). + async fn drop_pending(&self) { + loop { + let Some(wi) = self + .backend + .claim_work_item_with_lease("sweeper", &QUEUES, Some(30)) + .await + .unwrap() + else { + return; + }; + self.backend + .complete_work_item_fenced(wi.id, wi.lease_fence) + .await + .unwrap(); + } + } + + async fn events(&self) -> Vec { + self.backend + .get_events(&self.eid) + .await + .unwrap() + .into_iter() + .map(|e| e.kind) + .collect() + } + + /// `(limit_type, limit_value, actual_value)` of every `AutonomyLimitReached`. + async fn limits(&self) -> Vec<(String, Value, Value)> { + self.events() + .await + .into_iter() + .filter_map(|k| match k { + EventKind::AutonomyLimitReached { + limit_type, + limit_value, + actual_value, + agent_ref, + .. + } => { + assert_eq!(agent_ref, AGENT); + Some((limit_type, limit_value, actual_value)) + } + _ => None, + }) + .collect() + } + + async fn admitted(&self) -> usize { + self.events() + .await + .iter() + .filter(|k| matches!(k, EventKind::AgentStepAdmitted { .. })) + .count() + } + + /// The last `NodeFailed`'s error, if it failed for good. + async fn failed_for_good(&self) -> Option { + self.events().await.into_iter().rev().find_map(|k| match k { + EventKind::NodeFailed { + error, + retryable: false, + .. + } => Some(error), + _ => None, + }) + } +} + +/// Every backend and transport, each with a fresh pair of backends: the +/// SQLite one holds the registered cards, and URIs are unique. +async fn each() -> Vec<(String, support::TestBackend, Transport, String)> { + let mut out = Vec::new(); + for t in TRANSPORTS { + for tb in support::backends().await { + let ctx = format!("{} over {t:?}", tb.name()); + let wf = format!("wf-{}-{t:?}", tb.name()).to_lowercase(); + out.push((ctx, tb, t, wf)); + } + } + out +} + +/// The node was refused by `limit`, recorded, and failed for good. +async fn assert_limit(run: &Run, ran: bool, limit: &str, value: Value, actual: Value, ctx: &str) { + assert!(!ran, "{ctx}: the step past the limit ran"); + let limits = run.limits().await; + assert!( + limits.contains(&(limit.to_string(), value.clone(), actual.clone())), + "{ctx}: no {limit} = {value} (actual {actual}) refusal: {limits:?}" + ); + let error = run.failed_for_good().await; + assert!( + error.as_deref().is_some_and(|e| e.contains(limit)), + "{ctx}: the node did not fail for good on {limit}: {error:?}" + ); +} + +// ── max_iterations ────────────────────────────────────────────────────────── + +#[tokio::test] +async fn iterations_trip_at_the_boundary() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_iterations": 3 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + for i in 1..=3 { + assert!(run.step(t, DONE).await, "{ctx}: step {i} of 3 did not run"); + } + assert!(run.limits().await.is_empty(), "{ctx}"); + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "max_iterations", json!(3), json!(3), &ctx).await; + assert_eq!(run.admitted().await, 3, "{ctx}"); + } +} + +/// A node's own `acts_for` (`agent_ref` on the node) attributes it too. +#[tokio::test] +async fn a_node_acting_for_the_agent_is_counted() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_iterations": 1 }))], + ir(&wf, mcp_tool(), None, Some(AGENT), None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "max_iterations", json!(1), json!(1), &ctx).await; + } +} + +fn agent_node(target: &str) -> Value { + json!({ "type": "agent", "agent_ref": target, "input_mapping": {}, "output_schema": "" }) +} + +/// An `agent` node is a step of the agent it calls, attributed or not. +#[tokio::test] +async fn an_agent_node_charges_the_agent_it_calls() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_iterations": 1 }))], + ir(&wf, agent_node(AGENT), None, None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "max_iterations", json!(1), json!(1), &ctx).await; + } +} + +/// A callee with no registered card has no limits to apply. +#[tokio::test] +async fn an_agent_node_calling_an_unregistered_agent_runs() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[], + ir(&wf, agent_node("jamjet://acme/nobody"), None, None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(run.step(t, DONE).await, "{ctx}"); + assert_eq!(run.admitted().await, 0, "{ctx}"); + } +} + +// ── max_tool_calls ────────────────────────────────────────────────────────── + +#[tokio::test] +async fn tool_calls_trip_at_the_boundary() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_tool_calls": 2 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(run.step(t, DONE).await, "{ctx}"); + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "max_tool_calls", json!(2), json!(3), &ctx).await; + } +} + +/// An ADK dispatch batch counts each of its calls. +#[tokio::test] +async fn a_dispatch_batch_counts_each_call() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_tool_calls": 4 }))], + ir(&wf, adk_dispatch(), Some(AGENT), None, None), + calls(2), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}: 2 of 4"); + assert!(run.step(t, DONE).await, "{ctx}: 4 of 4"); + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "max_tool_calls", json!(4), json!(6), &ctx).await; + } +} + +/// A model step calls no tool, so a spent tool call budget does not stop it. +#[tokio::test] +async fn a_model_step_is_not_a_tool_call() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_tool_calls": 0 }))], + ir(&wf, model_node(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(run.step(t, DONE).await, "{ctx}"); + assert_eq!(run.admitted().await, 2, "{ctx}"); + } +} + +// ── token_budget and cost_budget_usd ──────────────────────────────────────── + +#[tokio::test] +async fn the_token_budget_trips_once_reached() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "token_budget": 100 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + let spend = |tokens| Act::Done { tokens, cost: None }; + assert!(run.step(t, spend(60)).await, "{ctx}"); + assert!(run.step(t, spend(39)).await, "{ctx}: 60 of 100"); + // 99 of 100: one more step may run. + assert!(run.step(t, spend(1)).await, "{ctx}: 99 of 100"); + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "token_budget", json!(100), json!(100), &ctx).await; + } +} + +/// Cost is recorded only when a worker reports it on an HTTP completion; it +/// is enforced on both transports. +#[tokio::test] +async fn the_cost_budget_trips_once_reached() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "cost_budget_usd": 0.5 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + let spend = |cost| Act::Done { + tokens: 0, + cost: Some(cost), + }; + // Binary-exact amounts, so the boundary is exact. + assert!(run.step(Transport::Claim, spend(0.25)).await, "{ctx}"); + assert!(run.step(Transport::Claim, spend(0.125)).await, "{ctx}"); + assert!(run.step(t, DONE).await, "{ctx}: 0.375 of 0.5"); + assert!( + run.step(Transport::Claim, spend(0.125)).await, + "{ctx}: 0.375 of 0.5" + ); + let ran = run.step(t, DONE).await; + assert!(!ran, "{ctx}: ran past the cost budget"); + let limits = run.limits().await; + assert_eq!(limits.len(), 1, "{ctx}: {limits:?}"); + assert_eq!(limits[0].0, "cost_budget", "{ctx}"); + assert_eq!(limits[0].1, json!(0.5), "{ctx}"); + assert_eq!(limits[0].2, json!(0.5), "{ctx}"); + } +} + +/// A reported cost is a sum the budget reads: a negative one would refund +/// it. +#[tokio::test] +async fn a_negative_cost_is_refused() { + for tb in support::backends().await { + let mut run = start( + &tb, + &[card(json!({ "cost_budget_usd": 0.5 }))], + ir("wf-neg", mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + run.started = true; + let (_, claimed) = support::post( + &run.state, + "/work-items/claim", + json!({ "worker_id": "w", "queue_types": QUEUES }), + ) + .await; + let wi = &claimed["work_item"]; + let (st, body) = support::post( + &run.state, + &format!("/work-items/{}/complete", wi["id"].as_str().unwrap()), + json!({ "output": {}, "state_patch": {}, "lease_fence": wi["lease_fence"], "cost_usd": -1.0 }), + ) + .await; + assert_eq!( + st, + axum::http::StatusCode::BAD_REQUEST, + "{}: {body}", + tb.name() + ); + assert!( + !run.events() + .await + .iter() + .any(|k| matches!(k, EventKind::NodeCompleted { .. })), + "{}", + tb.name() + ); + } +} + +// ── time_budget_secs ──────────────────────────────────────────────────────── + +#[tokio::test] +async fn the_time_budget_runs_from_the_first_step() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "time_budget_secs": 1 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + // The clock runs from the FIRST step, not the latest: 600 ms after + // the first step is inside the budget, 600 ms after that is not, + // though it is only 600 ms after the step before it. + tokio::time::sleep(std::time::Duration::from_millis(600)).await; + assert!(run.step(t, DONE).await, "{ctx}: inside the budget"); + tokio::time::sleep(std::time::Duration::from_millis(600)).await; + let ran = run.step(t, DONE).await; + assert!(!ran, "{ctx}: ran past the time budget"); + let limits = run.limits().await; + assert_eq!(limits.len(), 1, "{ctx}: {limits:?}"); + assert_eq!(limits[0].0, "time_budget", "{ctx}"); + assert_eq!(limits[0].1, json!(1), "{ctx}"); + assert!(run.failed_for_good().await.is_some(), "{ctx}"); + } +} + +// ── circuit breaker ───────────────────────────────────────────────────────── + +fn tripped(events: &[EventKind]) -> Vec<(u32, u32)> { + events + .iter() + .filter_map(|k| match k { + EventKind::CircuitBreakerTripped { + consecutive_errors, + threshold, + agent_ref, + .. + } => { + assert_eq!(agent_ref, AGENT); + Some((*consecutive_errors, *threshold)) + } + _ => None, + }) + .collect() +} + +#[tokio::test] +async fn consecutive_failures_trip_the_breaker() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_consecutive_errors": 2 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, Act::Fail).await, "{ctx}"); + assert!(run.step(t, Act::Fail).await, "{ctx}: 1 of 2 errors"); + let ran = run.step(t, DONE).await; + assert!(!ran, "{ctx}: ran with the breaker tripped"); + assert_eq!(tripped(&run.events().await), vec![(2, 2)], "{ctx}"); + assert!( + run.failed_for_good() + .await + .is_some_and(|e| e.contains("circuit breaker")), + "{ctx}" + ); + } +} + +/// A completed step resets the count: errors must be consecutive. +#[tokio::test] +async fn a_completed_step_resets_the_breaker() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_consecutive_errors": 2 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, Act::Fail).await, "{ctx}"); + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(run.step(t, Act::Fail).await, "{ctx}"); + assert!(run.step(t, DONE).await, "{ctx}: one error since the reset"); + assert!(tripped(&run.events().await).is_empty(), "{ctx}"); + } +} + +/// With no threshold on the card, the breaker trips at three. +#[tokio::test] +async fn the_breaker_defaults_to_three() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({}))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + for i in 1..=3 { + assert!(run.step(t, Act::Fail).await, "{ctx}: failure {i}"); + } + assert!(!run.step(t, DONE).await, "{ctx}"); + assert_eq!(tripped(&run.events().await), vec![(3, 3)], "{ctx}"); + } +} + +// ── durability ────────────────────────────────────────────────────────────── + +/// The counters are folded from the event log, so a new engine over the +/// same store (a restart) still refuses. +#[tokio::test] +async fn the_counters_survive_a_restart() { + for (ctx, tb, t, wf) in each().await { + let c = card(json!({ "max_iterations": 2 })); + let mut run = start( + &tb, + std::slice::from_ref(&c), + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(run.step(t, DONE).await, "{ctx}"); + + // A new engine: a new state and worker over the same store. SQLite + // is reopened from its file; memory keeps its backend (it is the + // store) and gets a new registry, so the card is registered again. + let backend: Arc = match &tb { + support::TestBackend::Sqlite { url, .. } => { + Arc::new(jamjet_state::SqliteBackend::open(url).await.unwrap()) + } + support::TestBackend::Memory(b) => b.clone(), + }; + let state = state_for(&tb, backend.clone()).await; + if matches!(tb, support::TestBackend::Memory(_)) { + register(&state, std::slice::from_ref(&c)).await; + } + run.state = state; + run.backend = backend; + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "max_iterations", json!(2), json!(2), &ctx).await; + } +} + +/// A continue-as-new segment is the same run: it does not start the +/// agent's counters again. +#[tokio::test] +async fn a_new_segment_keeps_the_counters() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_iterations": 1 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + + let carried = jamjet_state::materializer::materialize(run.backend.as_ref(), &run.eid) + .await + .unwrap(); + let next = jamjet_state::segment::start_next_segment( + run.backend.as_ref(), + &run.eid, + &carried, + &wf, + "0.1.0", + 1, + "n1", + "tool", + "default", + ) + .await + .unwrap(); + // The parent is closed, as a rollover leaves it. + run.backend + .update_execution_status(&run.eid, jamjet_core::workflow::WorkflowStatus::Completed) + .await + .unwrap(); + run.eid = next; + let ran = run.step(t, DONE).await; + assert_limit(&run, ran, "max_iterations", json!(1), json!(1), &ctx).await; + } +} + +// ── escalation ────────────────────────────────────────────────────────────── + +/// `on_limit_exceeded: escalate_to_human` holds the step for an ordinary +/// approval instead of failing it, and an approval releases exactly that +/// step. +#[tokio::test] +async fn escalate_routes_the_step_to_an_approval() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card( + json!({ "max_iterations": 1, "on_limit_exceeded": "escalate_to_human" }), + )], + ir( + &wf, + mcp_tool(), + Some(AGENT), + None, + Some(json!({ "approval_timeout_secs": 600 })), + ), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!( + !run.step(t, DONE).await, + "{ctx}: ran past the limit unapproved" + ); + + let events = run.events().await; + let requests: Vec<&Value> = events + .iter() + .filter_map(|k| match k { + EventKind::ToolApprovalRequired { context, .. } => Some(context), + _ => None, + }) + .collect(); + assert_eq!(requests.len(), 1, "{ctx}: {events:?}"); + let request = requests[0].clone(); + assert_eq!( + request["autonomy_limit"]["limit_type"], + json!("max_iterations"), + "{ctx}: {request}" + ); + assert_eq!( + request["autonomy_limit"]["agent_ref"], + json!(AGENT), + "{ctx}" + ); + // The same binding and deadline as any other request. + assert!(request["action_digest"].is_string(), "{ctx}: {request}"); + assert!(request["expires_at"].is_string(), "{ctx}: {request}"); + assert!( + events.iter().any(|k| matches!( + k, + EventKind::EscalationRequired { escalation_target, .. } if escalation_target == "human_approval" + )), + "{ctx}" + ); + // A duplicate item while the request is pending is held without a + // second record of the stop. + assert!(!run.step(t, DONE).await, "{ctx}"); + assert_eq!(run.limits().await.len(), 1, "{ctx}"); + assert_eq!( + run.events() + .await + .iter() + .filter(|k| matches!(k, EventKind::EscalationRequired { .. })) + .count(), + 1, + "{ctx}" + ); + assert!( + run.failed_for_good().await.is_none(), + "{ctx}: failed instead of escalating" + ); + + let (st, body) = support::post( + &run.state, + &format!("/executions/{}/approve", run.eid_path()), + json!({ "decision": "approved", "node_id": "n1", "action_digest": request["action_digest"] }), + ) + .await; + assert!(st.is_success(), "{ctx}: approve: {st} {body}"); + assert!( + run.step(t, DONE).await, + "{ctx}: the approved step did not run" + ); + + // The next step past the limit needs its own approval. + assert!(!run.step(t, DONE).await, "{ctx}"); + let requests = run + .events() + .await + .into_iter() + .filter(|k| matches!(k, EventKind::ToolApprovalRequired { .. })) + .count(); + assert_eq!(requests, 2, "{ctx}"); + } +} + +/// RFC-005's `escalate_to_supervisor` has no implementation: a card asking +/// for it is refused rather than treated as `fail`. +#[tokio::test] +async fn a_card_asking_for_a_supervisor_is_refused() { + for tb in support::backends().await { + let state = state_for(&tb, tb.dyn_backend()).await; + let (st, body) = support::post( + &state, + "/agents", + card(json!({ "max_iterations": 1, "on_limit_exceeded": "escalate_to_supervisor" })), + ) + .await; + assert!(!st.is_success(), "{}: {st} {body}", tb.name()); + } +} + +/// When the node acts for two agents and one card escalates while the other +/// fails, the node is refused: a refusal is the stricter answer. +#[tokio::test] +async fn a_failing_card_wins_over_an_escalating_one() { + const OTHER: &str = "jamjet://acme/reviewer"; + for (ctx, tb, t, wf) in each().await { + let mut other = card(json!({ "max_iterations": 1 })); + other["id"] = json!(OTHER); + other["uri"] = json!(OTHER); + let mut run = start( + &tb, + &[ + card(json!({ "max_iterations": 1, "on_limit_exceeded": "escalate_to_human" })), + other, + ], + ir(&wf, mcp_tool(), Some(AGENT), Some(OTHER), None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(!run.step(t, DONE).await, "{ctx}"); + let events = run.events().await; + assert!( + !events + .iter() + .any(|k| matches!(k, EventKind::ToolApprovalRequired { .. })), + "{ctx}: escalated although one card fails" + ); + assert!(run.failed_for_good().await.is_some(), "{ctx}"); + } +} + +/// A log holding an event this build cannot read may hold a step: the +/// counts cannot be known, so the step is withheld. +#[tokio::test] +async fn an_unreadable_log_withholds_the_step() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_iterations": 5 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + run.backend + .append_event(jamjet_state::Event::new( + run.eid.clone(), + 0, + EventKind::Unknown, + )) + .await + .unwrap(); + assert!(!run.step(t, DONE).await, "{ctx}: ran on an unreadable log"); + assert_eq!(run.admitted().await, 0, "{ctx}"); + } +} + +// ── unaffected runs ───────────────────────────────────────────────────────── + +/// A run whose nodes act for no agent is not counted or limited. +#[tokio::test] +async fn a_run_without_a_card_is_unaffected() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ "max_iterations": 0, "max_tool_calls": 0 }))], + ir(&wf, mcp_tool(), None, None, None), + json!({}), + ) + .await; + for i in 1..=4 { + assert!(run.step(t, Act::Fail).await, "{ctx}: step {i}"); + } + assert!(run.step(t, DONE).await, "{ctx}"); + assert_eq!(run.admitted().await, 0, "{ctx}"); + assert!(run.limits().await.is_empty(), "{ctx}"); + } +} + +// ── the completion gate ───────────────────────────────────────────────────── + +/// A result over HTTP is recorded only for an item the guard admitted: an +/// item completed without a claim would be a step no limit saw. +#[tokio::test] +async fn a_result_for_an_unadmitted_item_is_refused() { + for tb in support::backends().await { + let run = start( + &tb, + &[card(json!({ "max_iterations": 5 }))], + ir("wf-unadmitted", mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + // The start item, never claimed. + let items = run + .backend + .claim_work_item_with_lease("peek", &QUEUES, Some(30)) + .await + .unwrap() + .expect("start item"); + // Hand it back unclaimed in the engine's eyes: complete it over + // HTTP without the claim route. + let (st, body) = support::post( + &run.state, + &format!("/work-items/{}/complete", items.id), + json!({ "output": {}, "state_patch": {}, "lease_fence": items.lease_fence }), + ) + .await; + assert_eq!( + st, + axum::http::StatusCode::CONFLICT, + "{}: {body}", + tb.name() + ); + assert!( + !run.events() + .await + .iter() + .any(|k| matches!(k, EventKind::NodeCompleted { .. })), + "{}", + tb.name() + ); + } +} + +// ── concurrency ───────────────────────────────────────────────────────────── + +/// Another claim admits a step between this claim's check and its record. +struct RaceAdmission { + inner: Arc, + armed: AtomicBool, + eid: Mutex>, + /// The agents the racing step charges. + agents: Vec, + /// The approval the racing step used, if any. + approval_sequence: Mutex>, +} + +#[async_trait] +impl AfterWrite for RaceAdmission { + async fn after(&self, write: Write) { + if let Write::BeforeAppendEventIf(eid) = write { + if Some(&eid) == self.eid.lock().unwrap().as_ref() + && self.armed.swap(false, Ordering::SeqCst) + { + let approval_sequence = *self.approval_sequence.lock().unwrap(); + self.inner + .append_event(jamjet_state::Event::new( + eid, + 0, + EventKind::AgentStepAdmitted { + node_id: "other".into(), + work_item_id: None, + agents: self.agents.clone(), + tool_calls: 1, + approval_sequence, + }, + )) + .await + .unwrap(); + } + } + } +} + +/// The last step cannot be taken twice: the step is recorded only if the +/// limits still allow it in the same compare-and-set. +#[tokio::test] +async fn two_claims_cannot_both_take_the_last_step() { + for (ctx, tb, t, wf) in each().await { + let inner = tb.dyn_backend(); + let hook = Arc::new(RaceAdmission { + inner: inner.clone(), + armed: AtomicBool::new(false), + eid: Mutex::new(None), + agents: vec![AGENT.into()], + approval_sequence: Mutex::new(None), + }); + let hooked: Arc = + Arc::new(HookedBackend::new(inner.clone(), hook.clone())); + let mut run = start_on( + &tb, + hooked, + &[card(json!({ "max_iterations": 1 }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + *hook.eid.lock().unwrap() = Some(run.eid.clone()); + hook.armed.store(true, Ordering::SeqCst); + let ran = run.step(t, DONE).await; + assert!(!ran, "{ctx}: both claims took the last step"); + assert_eq!(run.admitted().await, 1, "{ctx}: only the racing step"); + } +} + +// ── escalation approvals are single-use ───────────────────────────────────── + +async fn approve_latest(run: &Run, ctx: &str) { + let digest = run + .events() + .await + .into_iter() + .rev() + .find_map(|k| match k { + EventKind::ToolApprovalRequired { context, .. } => { + Some(context["action_digest"].clone()) + } + _ => None, + }) + .expect("a request"); + let (st, body) = support::post( + &run.state, + &format!("/executions/{}/approve", run.eid_path()), + json!({ "decision": "approved", "node_id": "n1", "action_digest": digest }), + ) + .await; + assert!(st.is_success(), "{ctx}: approve: {st} {body}"); +} + +async fn requests(run: &Run) -> Vec { + run.events() + .await + .into_iter() + .filter_map(|k| match k { + EventKind::ToolApprovalRequired { context, .. } => Some(context), + _ => None, + }) + .collect() +} + +/// An approval releases ONE step past the limit. A step that fails does not +/// advance the node's occurrence, so its retry carries the same digest; it +/// must still ask again, or one approval would release every retry (and +/// silence the circuit breaker). +#[tokio::test] +async fn one_escalation_approval_releases_one_step_even_when_it_fails() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card(json!({ + "max_iterations": 1, + "max_consecutive_errors": 10, + "on_limit_exceeded": "escalate_to_human" + }))], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(!run.step(t, DONE).await, "{ctx}"); + approve_latest(&run, &ctx).await; + assert!(run.step(t, Act::Fail).await, "{ctx}: the approved step"); + assert!( + !run.step(t, Act::Fail).await, + "{ctx}: a retry ran on the same approval" + ); + assert_eq!(requests(&run).await.len(), 2, "{ctx}: no fresh request"); + assert_eq!(run.admitted().await, 2, "{ctx}"); + } +} + +/// An approval given before the limit was reached (a policy gate's, which +/// showed the reviewer no limit) does not release a step past the limit: +/// the step asks again, with the limit in the request. +#[tokio::test] +async fn an_approval_granted_before_the_limit_does_not_cover_it() { + for (ctx, tb, t, wf) in each().await { + let mut run = start( + &tb, + &[card( + json!({ "max_iterations": 1, "on_limit_exceeded": "escalate_to_human" }), + )], + ir( + &wf, + mcp_tool(), + Some(AGENT), + None, + Some(json!({ "require_approval_for": ["search_web"] })), + ), + json!({}), + ) + .await; + assert!(!run.step(t, DONE).await, "{ctx}: policy gate"); + assert!( + requests(&run).await[0].get("autonomy_limit").is_none(), + "{ctx}" + ); + // A parallel branch of the same agent takes its last step. + run.backend + .append_event(jamjet_state::Event::new( + run.eid.clone(), + 0, + EventKind::AgentStepAdmitted { + node_id: "other".into(), + work_item_id: None, + agents: vec![AGENT.into()], + tool_calls: 1, + approval_sequence: None, + }, + )) + .await + .unwrap(); + approve_latest(&run, &ctx).await; + assert!( + !run.step(t, DONE).await, + "{ctx}: ran past the limit unasked" + ); + let reqs = requests(&run).await; + assert_eq!(reqs.len(), 2, "{ctx}"); + assert_eq!( + reqs[1]["autonomy_limit"]["limit_type"], + json!("max_iterations"), + "{ctx}" + ); + } +} + +/// An approved escalation for one agent does not excuse another agent's +/// `fail` limit reached while the step was being decided. +#[tokio::test] +async fn an_approved_escalation_still_rechecks_failing_cards() { + const OTHER: &str = "jamjet://acme/reviewer"; + for (ctx, tb, t, wf) in each().await { + let inner = tb.dyn_backend(); + let hook = Arc::new(RaceAdmission { + inner: inner.clone(), + armed: AtomicBool::new(false), + eid: Mutex::new(None), + agents: vec![OTHER.into()], + approval_sequence: Mutex::new(None), + }); + let hooked: Arc = + Arc::new(HookedBackend::new(inner.clone(), hook.clone())); + let mut other = card(json!({ "max_iterations": 2 })); + other["id"] = json!(OTHER); + other["uri"] = json!(OTHER); + let mut run = start_on( + &tb, + hooked, + &[ + card(json!({ "max_iterations": 1, "on_limit_exceeded": "escalate_to_human" })), + other, + ], + ir(&wf, mcp_tool(), Some(AGENT), Some(OTHER), None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(!run.step(t, DONE).await, "{ctx}: escalated"); + approve_latest(&run, &ctx).await; + *hook.eid.lock().unwrap() = Some(run.eid.clone()); + hook.armed.store(true, Ordering::SeqCst); + assert!( + !run.step(t, DONE).await, + "{ctx}: ran although the other card's limit was reached" + ); + } +} + +/// One admitted item backs one result. SQLite's legacy unfenced completion +/// settles an already-settled item again; the second result must not be +/// recorded, or it would end a newer step of the node and reset the breaker. +#[tokio::test] +async fn one_admission_backs_one_result() { + for tb in support::backends().await { + let run = start( + &tb, + &[card(json!({ "max_iterations": 5 }))], + ir("wf-one", mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + let (_, claimed) = support::post( + &run.state, + "/work-items/claim", + json!({ "worker_id": "w", "queue_types": QUEUES }), + ) + .await; + let id = claimed["work_item"]["id"].as_str().unwrap().to_string(); + for _ in 0..3 { + let _ = support::post( + &run.state, + &format!("/work-items/{id}/complete"), + json!({ "output": {}, "state_patch": {} }), + ) + .await; + } + let completed = run + .events() + .await + .iter() + .filter(|k| matches!(k, EventKind::NodeCompleted { .. })) + .count(); + assert_eq!(completed, 1, "{}", tb.name()); + } +} + +/// Two items holding the same approved escalation: only one may use it, also +/// when the other uses it between this one's check and its record. +#[tokio::test] +async fn an_escalation_approval_cannot_be_used_twice_in_a_race() { + for (ctx, tb, t, wf) in each().await { + let inner = tb.dyn_backend(); + let hook = Arc::new(RaceAdmission { + inner: inner.clone(), + armed: AtomicBool::new(false), + eid: Mutex::new(None), + agents: vec![AGENT.into()], + approval_sequence: Mutex::new(None), + }); + let hooked: Arc = + Arc::new(HookedBackend::new(inner.clone(), hook.clone())); + let mut run = start_on( + &tb, + hooked, + &[card( + json!({ "max_iterations": 1, "on_limit_exceeded": "escalate_to_human" }), + )], + ir(&wf, mcp_tool(), Some(AGENT), None, None), + json!({}), + ) + .await; + assert!(run.step(t, DONE).await, "{ctx}"); + assert!(!run.step(t, DONE).await, "{ctx}"); + approve_latest(&run, &ctx).await; + let seq = run + .backend + .get_events(&run.eid) + .await + .unwrap() + .into_iter() + .rev() + .find(|e| matches!(e.kind, EventKind::ToolApprovalRequired { .. })) + .unwrap() + .sequence; + *hook.approval_sequence.lock().unwrap() = Some(seq); + *hook.eid.lock().unwrap() = Some(run.eid.clone()); + hook.armed.store(true, Ordering::SeqCst); + assert!( + !run.step(t, DONE).await, + "{ctx}: the approval was used twice" + ); + } +} diff --git a/runtime/policy/src/autonomy.rs b/runtime/policy/src/autonomy.rs index 2c8fa941..277a0311 100644 --- a/runtime/policy/src/autonomy.rs +++ b/runtime/policy/src/autonomy.rs @@ -1,293 +1,366 @@ -//! Autonomy enforcement for JamJet agents (Phase 4 — Tier 1). +//! An agent card's autonomy limits, decided from what the agent has used. //! -//! Enforces the three autonomy levels: -//! - `Guided`: every tool call pauses for human approval -//! - `BoundedAutonomous`: runs within declared limits (max_iterations, cost, tokens, tool calls) -//! - `FullyAutonomous`: only the circuit breaker applies +//! The limits are the card's `max_iterations`, `max_tool_calls`, +//! `token_budget`, `cost_budget_usd`, `time_budget_secs` and the circuit +//! breaker (`max_consecutive_errors`, default +//! [`DEFAULT_MAX_CONSECUTIVE_ERRORS`]). They apply at every autonomy level. //! -//! The circuit breaker trips when an agent emits N consecutive errors (default: 3), -//! regardless of autonomy level. +//! This module only decides. What the agent has used ([`AgentUsage`]) is +//! folded from the run's durable event log by the shared node guard +//! (`jamjet_worker::node_guard`), which also records the refusal and settles +//! the node. Keeping the decision a pure function of the usage means the +//! guard can run the same check inside the compare-and-set that records the +//! step, so two concurrent claims cannot both take the last step. +//! +//! Every check is made before a step fires, against what earlier steps used: +//! +//! - `max_iterations`: a step is refused once the agent has taken that many. +//! - `max_tool_calls`: a step is refused if its own calls would take the +//! agent past the limit (a step that makes no tool call is not limited). +//! - `token_budget`, `cost_budget_usd`: a step is refused once the recorded +//! usage is at or over the budget. The step that crosses it is not known +//! to cross it until it reports, so the total can end above the budget by +//! at most that one step's usage. +//! - `time_budget_secs`: a step is refused once that long has passed since +//! the agent's first step in the run. A step already running is not +//! interrupted. +//! - The circuit breaker: a step is refused once that many of the agent's +//! steps in a row have failed. A completed step resets the count. +//! +//! A limit of zero refuses the first step. -use jamjet_agents::card::{AutonomyConstraints, AutonomyLevel}; -use tracing::warn; +use jamjet_agents::card::{AutonomyConstraints, DEFAULT_MAX_CONSECUTIVE_ERRORS}; -/// Current runtime context for an agent node. -#[derive(Debug, Clone)] -pub struct AutonomyContext { - /// Identifies the agent for logging / event emission. - pub agent_ref: String, - pub autonomy_level: AutonomyLevel, - pub constraints: Option, - /// Accumulated iterations across the full execution. - pub current_iterations: u32, - /// Accumulated tool calls across the full execution. - pub current_tool_calls: u32, - /// Accumulated USD cost across the full execution. - pub current_cost_usd: f64, - /// Accumulated tokens across the full execution. - pub current_tokens: u64, - /// Consecutive error count (for circuit breaker). +/// What one agent has used so far in a run, across its segments. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct AgentUsage { + /// Steps admitted: nodes acting for the agent that were allowed to fire. + pub iterations: u64, + /// Tool calls those steps made. + pub tool_calls: u64, + /// Input plus output tokens recorded on the agent's completed steps. + pub tokens: u64, + /// USD cost recorded on the agent's completed steps. + pub cost_usd: f64, + /// The agent's steps that failed since its last completed step. pub consecutive_errors: u32, - /// Max consecutive errors before the circuit breaker trips (default: 3). - pub circuit_breaker_threshold: u32, + /// Milliseconds since the agent's first step in the run; `None` before + /// its first step. + pub elapsed_ms: Option, } -impl Default for AutonomyContext { - fn default() -> Self { - Self { - agent_ref: String::new(), - autonomy_level: AutonomyLevel::Guided, - constraints: None, - current_iterations: 0, - current_tool_calls: 0, - current_cost_usd: 0.0, - current_tokens: 0, - consecutive_errors: 0, - circuit_breaker_threshold: 3, - } - } +/// The step about to fire. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct Step { + /// Tool calls it makes. + pub tool_calls: u32, } -/// Decision returned by `AutonomyEnforcer::check`. -#[derive(Debug)] +/// The decision for one agent and one step. +#[derive(Debug, Clone, PartialEq)] pub enum AutonomyDecision { - /// Proceed — no autonomy constraint blocks this node. + /// No limit stops the step. Proceed, - /// Guided mode: every tool call pauses for human approval. - RequireToolApproval { tool_name: String }, - /// A declared limit was hit — escalate. - EscalateLimit { - limit_type: String, + /// The step would take the agent past a limit. + LimitReached { + /// `max_iterations`, `max_tool_calls`, `token_budget`, `cost_budget` + /// or `time_budget`. + limit_type: &'static str, limit_value: serde_json::Value, + /// What the agent had used (for `max_tool_calls`, what it would use + /// with this step's calls). actual_value: serde_json::Value, - escalation_target: EscalationTarget, }, - /// The circuit breaker tripped. - TripCircuitBreaker { + /// Too many of the agent's steps in a row have failed. + CircuitBreakerTripped { consecutive_errors: u32, threshold: u32, - escalation_target: EscalationTarget, }, } -/// Where to escalate when a limit or circuit breaker fires. -#[derive(Debug, Clone)] -pub enum EscalationTarget { - SupervisorAgent { agent_id: String }, - HumanApproval, -} +/// Decide whether `step` may fire for an agent whose card sets +/// `constraints` (`None`: a card with no constraints, where only the +/// default circuit breaker applies) and which has used `usage`. +pub fn check( + constraints: Option<&AutonomyConstraints>, + usage: &AgentUsage, + step: Step, +) -> AutonomyDecision { + let threshold = constraints + .and_then(|c| c.max_consecutive_errors) + .unwrap_or(DEFAULT_MAX_CONSECUTIVE_ERRORS); + if usage.consecutive_errors >= threshold { + return AutonomyDecision::CircuitBreakerTripped { + consecutive_errors: usage.consecutive_errors, + threshold, + }; + } + let Some(c) = constraints else { + return AutonomyDecision::Proceed; + }; + let reached = |limit_type, limit_value, actual_value| AutonomyDecision::LimitReached { + limit_type, + limit_value, + actual_value, + }; -impl EscalationTarget { - pub fn as_str(&self) -> String { - match self { - Self::SupervisorAgent { agent_id } => format!("supervisor_agent:{agent_id}"), - Self::HumanApproval => "human_approval".to_string(), + if let Some(max) = c.max_iterations { + if usage.iterations >= u64::from(max) { + return reached("max_iterations", max.into(), usage.iterations.into()); } } -} - -/// Checks autonomy constraints before a node executes. -pub struct AutonomyEnforcer; - -impl AutonomyEnforcer { - /// Evaluate autonomy constraints for the given context. - /// - /// Call this before executing an agent or tool node to determine whether - /// the node should proceed, pause for approval, or escalate. - pub fn check(&self, ctx: &AutonomyContext, tool_name: Option<&str>) -> AutonomyDecision { - // Circuit breaker fires for all autonomy levels. - if ctx.consecutive_errors >= ctx.circuit_breaker_threshold { - warn!( - agent_ref = %ctx.agent_ref, - consecutive_errors = ctx.consecutive_errors, - threshold = ctx.circuit_breaker_threshold, - "Circuit breaker tripped" + if let Some(max) = c.max_tool_calls { + let would_be = usage.tool_calls + u64::from(step.tool_calls); + if step.tool_calls > 0 && would_be > u64::from(max) { + return reached("max_tool_calls", max.into(), would_be.into()); + } + } + if let Some(max) = c.token_budget { + if usage.tokens >= max { + return reached("token_budget", max.into(), usage.tokens.into()); + } + } + if let Some(max) = c.cost_budget_usd { + // Stated positively: only a cost known to be under the budget may + // proceed, so a NaN total (or budget) refuses. + let under = usage.cost_usd < max; + if !under { + return reached( + "cost_budget", + serde_json::json!(max), + serde_json::json!(usage.cost_usd), ); - return AutonomyDecision::TripCircuitBreaker { - consecutive_errors: ctx.consecutive_errors, - threshold: ctx.circuit_breaker_threshold, - escalation_target: EscalationTarget::HumanApproval, - }; } + } + if let (Some(secs), Some(elapsed_ms)) = (c.time_budget_secs, usage.elapsed_ms) { + if elapsed_ms >= secs.saturating_mul(1_000) { + return reached("time_budget", secs.into(), elapsed_ms.into()); + } + } + AutonomyDecision::Proceed +} - match &ctx.autonomy_level { - AutonomyLevel::Guided => { - // Every tool call requires approval in guided mode. - if let Some(name) = tool_name { - return AutonomyDecision::RequireToolApproval { - tool_name: name.to_string(), - }; - } - AutonomyDecision::Proceed - } - - AutonomyLevel::BoundedAutonomous => { - let constraints = match &ctx.constraints { - Some(c) => c, - None => return AutonomyDecision::Proceed, - }; +#[cfg(test)] +mod tests { + use super::*; - // Check max_iterations. - if let Some(max) = constraints.max_iterations { - if ctx.current_iterations >= max { - return AutonomyDecision::EscalateLimit { - limit_type: "max_iterations".to_string(), - limit_value: serde_json::json!(max), - actual_value: serde_json::json!(ctx.current_iterations), - escalation_target: EscalationTarget::HumanApproval, - }; - } - } + fn limits(v: serde_json::Value) -> AutonomyConstraints { + serde_json::from_value(v).unwrap() + } - // Check max_tool_calls. - if let Some(max) = constraints.max_tool_calls { - if ctx.current_tool_calls >= max { - return AutonomyDecision::EscalateLimit { - limit_type: "max_tool_calls".to_string(), - limit_value: serde_json::json!(max), - actual_value: serde_json::json!(ctx.current_tool_calls), - escalation_target: EscalationTarget::HumanApproval, - }; - } - } + fn used(f: impl FnOnce(&mut AgentUsage)) -> AgentUsage { + let mut u = AgentUsage::default(); + f(&mut u); + u + } - // Check cost budget. - if let Some(max_cost) = constraints.cost_budget_usd { - if ctx.current_cost_usd >= max_cost { - return AutonomyDecision::EscalateLimit { - limit_type: "cost_budget".to_string(), - limit_value: serde_json::json!(max_cost), - actual_value: serde_json::json!(ctx.current_cost_usd), - escalation_target: EscalationTarget::HumanApproval, - }; - } - } + fn limit_of(d: &AutonomyDecision) -> Option<&'static str> { + match d { + AutonomyDecision::LimitReached { limit_type, .. } => Some(limit_type), + _ => None, + } + } - // Check token budget. - if let Some(max_tokens) = constraints.token_budget { - if ctx.current_tokens >= max_tokens { - return AutonomyDecision::EscalateLimit { - limit_type: "token_budget".to_string(), - limit_value: serde_json::json!(max_tokens), - actual_value: serde_json::json!(ctx.current_tokens), - escalation_target: EscalationTarget::HumanApproval, - }; - } - } + const ONE_CALL: Step = Step { tool_calls: 1 }; + const NO_CALL: Step = Step { tool_calls: 0 }; - AutonomyDecision::Proceed + #[test] + fn iterations_trip_at_the_limit_and_not_before() { + let c = limits(serde_json::json!({ "max_iterations": 3 })); + let below = used(|u| u.iterations = 2); + let at = used(|u| u.iterations = 3); + assert_eq!(check(Some(&c), &below, NO_CALL), AutonomyDecision::Proceed); + let d = check(Some(&c), &at, NO_CALL); + assert_eq!(limit_of(&d), Some("max_iterations"), "{d:?}"); + assert_eq!( + d, + AutonomyDecision::LimitReached { + limit_type: "max_iterations", + limit_value: serde_json::json!(3), + actual_value: serde_json::json!(3), } + ); + } - // Fully autonomous: only circuit breaker (checked above). - AutonomyLevel::FullyAutonomous | AutonomyLevel::Deterministic => { - AutonomyDecision::Proceed + #[test] + fn tool_calls_count_the_calls_this_step_would_make() { + let c = limits(serde_json::json!({ "max_tool_calls": 5 })); + // 3 used + 2 now = 5: exactly at the limit, allowed. + let three = used(|u| u.tool_calls = 3); + assert_eq!( + check(Some(&c), &three, Step { tool_calls: 2 }), + AutonomyDecision::Proceed + ); + // 3 used + 3 now = 6: over. + let d = check(Some(&c), &three, Step { tool_calls: 3 }); + assert_eq!( + d, + AutonomyDecision::LimitReached { + limit_type: "max_tool_calls", + limit_value: serde_json::json!(5), + actual_value: serde_json::json!(6), } - } + ); + // A step that calls no tool is not limited by the tool call count. + let five = used(|u| u.tool_calls = 5); + assert_eq!(check(Some(&c), &five, NO_CALL), AutonomyDecision::Proceed); + assert_eq!( + limit_of(&check(Some(&c), &five, ONE_CALL)), + Some("max_tool_calls") + ); + // Past the limit (an approved escalation took it there), a step that + // calls no tool still is not limited by the tool call count. + let six = used(|u| u.tool_calls = 6); + assert_eq!(check(Some(&c), &six, NO_CALL), AutonomyDecision::Proceed); } -} -#[cfg(test)] -mod tests { - use super::*; - use jamjet_agents::card::AutonomyConstraints; + #[test] + fn token_budget_trips_once_reached() { + let c = limits(serde_json::json!({ "token_budget": 100 })); + assert_eq!( + check(Some(&c), &used(|u| u.tokens = 99), ONE_CALL), + AutonomyDecision::Proceed + ); + assert_eq!( + limit_of(&check(Some(&c), &used(|u| u.tokens = 100), ONE_CALL)), + Some("token_budget") + ); + } - fn bounded_ctx(iterations: u32, cost: f64, tokens: u64) -> AutonomyContext { - AutonomyContext { - agent_ref: "test-agent".into(), - autonomy_level: AutonomyLevel::BoundedAutonomous, - constraints: Some(AutonomyConstraints { - max_iterations: Some(5), - max_tool_calls: Some(10), - token_budget: Some(1000), - cost_budget_usd: Some(1.0), - allowed_tools: vec![], - blocked_tools: vec![], - allowed_delegations: vec![], - require_approval_for: vec![], - time_budget_secs: None, - }), - current_iterations: iterations, - current_tool_calls: 0, - current_cost_usd: cost, - current_tokens: tokens, - consecutive_errors: 0, - circuit_breaker_threshold: 3, - } + #[test] + fn cost_budget_trips_once_reached() { + let c = limits(serde_json::json!({ "cost_budget_usd": 0.5 })); + assert_eq!( + check(Some(&c), &used(|u| u.cost_usd = 0.49), ONE_CALL), + AutonomyDecision::Proceed + ); + let d = check(Some(&c), &used(|u| u.cost_usd = 0.5), ONE_CALL); + assert_eq!( + d, + AutonomyDecision::LimitReached { + limit_type: "cost_budget", + limit_value: serde_json::json!(0.5), + actual_value: serde_json::json!(0.5), + } + ); } + /// A cost that is not a number (no JSON source can produce one, but the + /// fold adds floats) must not walk through a `<` guard. #[test] - fn guided_requires_approval_for_tool() { - let enforcer = AutonomyEnforcer; - let ctx = AutonomyContext { - agent_ref: "a".into(), - autonomy_level: AutonomyLevel::Guided, - ..Default::default() - }; - assert!(matches!( - enforcer.check(&ctx, Some("search")), - AutonomyDecision::RequireToolApproval { .. } - )); + fn a_nan_cost_is_refused() { + let c = limits(serde_json::json!({ "cost_budget_usd": 0.5 })); + assert_eq!( + limit_of(&check(Some(&c), &used(|u| u.cost_usd = f64::NAN), ONE_CALL)), + Some("cost_budget") + ); } #[test] - fn bounded_under_limits_proceeds() { - let enforcer = AutonomyEnforcer; - let ctx = bounded_ctx(2, 0.5, 500); - assert!(matches!( - enforcer.check(&ctx, None), + fn time_budget_trips_once_elapsed() { + let c = limits(serde_json::json!({ "time_budget_secs": 2 })); + // Before the first step there is no clock to run out. + assert_eq!( + check(Some(&c), &AgentUsage::default(), ONE_CALL), + AutonomyDecision::Proceed + ); + assert_eq!( + check(Some(&c), &used(|u| u.elapsed_ms = Some(1_999)), ONE_CALL), AutonomyDecision::Proceed - )); + ); + assert_eq!( + limit_of(&check( + Some(&c), + &used(|u| u.elapsed_ms = Some(2_000)), + ONE_CALL + )), + Some("time_budget") + ); } #[test] - fn bounded_iteration_limit_escalates() { - let enforcer = AutonomyEnforcer; - let ctx = bounded_ctx(5, 0.0, 0); - assert!(matches!( - enforcer.check(&ctx, None), - AutonomyDecision::EscalateLimit { limit_type, .. } if limit_type == "max_iterations" - )); + fn zero_limits_refuse_the_first_step() { + for (field, limit) in [ + ("max_iterations", "max_iterations"), + ("max_tool_calls", "max_tool_calls"), + ("token_budget", "token_budget"), + ("cost_budget_usd", "cost_budget"), + ] { + let c = limits(serde_json::json!({ field: 0 })); + assert_eq!( + limit_of(&check(Some(&c), &AgentUsage::default(), ONE_CALL)), + Some(limit), + "{field}" + ); + } + // A zero time budget refuses once the clock has started: the first + // step starts it at zero elapsed. + let c = limits(serde_json::json!({ "time_budget_secs": 0 })); + assert_eq!( + limit_of(&check( + Some(&c), + &used(|u| u.elapsed_ms = Some(0)), + ONE_CALL + )), + Some("time_budget") + ); } #[test] - fn bounded_cost_limit_escalates() { - let enforcer = AutonomyEnforcer; - let ctx = bounded_ctx(0, 1.5, 0); - assert!(matches!( - enforcer.check(&ctx, None), - AutonomyDecision::EscalateLimit { limit_type, .. } if limit_type == "cost_budget" - )); + fn the_circuit_breaker_defaults_to_three_and_applies_without_constraints() { + assert_eq!( + check(None, &used(|u| u.consecutive_errors = 2), ONE_CALL), + AutonomyDecision::Proceed + ); + assert_eq!( + check(None, &used(|u| u.consecutive_errors = 3), ONE_CALL), + AutonomyDecision::CircuitBreakerTripped { + consecutive_errors: 3, + threshold: 3 + } + ); + let c = limits(serde_json::json!({ "max_consecutive_errors": 1 })); + assert_eq!( + check(Some(&c), &AgentUsage::default(), ONE_CALL), + AutonomyDecision::Proceed + ); + assert_eq!( + check(Some(&c), &used(|u| u.consecutive_errors = 1), ONE_CALL), + AutonomyDecision::CircuitBreakerTripped { + consecutive_errors: 1, + threshold: 1 + } + ); } #[test] - fn circuit_breaker_fires_at_threshold() { - let enforcer = AutonomyEnforcer; - let ctx = AutonomyContext { - agent_ref: "a".into(), - autonomy_level: AutonomyLevel::FullyAutonomous, - consecutive_errors: 3, - circuit_breaker_threshold: 3, - ..Default::default() - }; - assert!(matches!( - enforcer.check(&ctx, None), - AutonomyDecision::TripCircuitBreaker { .. } - )); + fn no_limits_proceed_whatever_was_used() { + let u = used(|u| { + u.iterations = 1_000; + u.tool_calls = 1_000; + u.tokens = 1_000_000; + u.cost_usd = 1_000.0; + u.elapsed_ms = Some(1_000_000); + }); + assert_eq!( + check(Some(&AutonomyConstraints::default()), &u, ONE_CALL), + AutonomyDecision::Proceed + ); + assert_eq!(check(None, &u, ONE_CALL), AutonomyDecision::Proceed); } #[test] - fn fully_autonomous_below_cb_threshold_proceeds() { - let enforcer = AutonomyEnforcer; - let ctx = AutonomyContext { - agent_ref: "a".into(), - autonomy_level: AutonomyLevel::FullyAutonomous, - consecutive_errors: 2, - circuit_breaker_threshold: 3, - ..Default::default() - }; - assert!(matches!( - enforcer.check(&ctx, None), - AutonomyDecision::Proceed - )); + fn escalate_to_supervisor_is_not_accepted() { + let r: Result = serde_json::from_value( + serde_json::json!({ "on_limit_exceeded": "escalate_to_supervisor" }), + ); + assert!(r.is_err()); + let ok: AutonomyConstraints = + serde_json::from_value(serde_json::json!({ "on_limit_exceeded": "escalate_to_human" })) + .unwrap(); + assert_eq!( + ok.on_limit_exceeded, + Some(jamjet_agents::card::OnLimitExceeded::EscalateToHuman) + ); } } diff --git a/runtime/policy/tests/trading_agent_autonomy.rs b/runtime/policy/tests/trading_agent_autonomy.rs index 4182d763..bfcf7b47 100644 --- a/runtime/policy/tests/trading_agent_autonomy.rs +++ b/runtime/policy/tests/trading_agent_autonomy.rs @@ -1,326 +1,122 @@ -//! # Real-World Example: Algorithmic Trading Agent — Autonomy Enforcement +//! # Example: an algorithmic trading research agent's autonomy limits //! -//! ## Scenario -//! A quantitative hedge fund deploys AI agents for pre-trade research: -//! - **Research Agent**: scans filings, news, social sentiment -//! - **Risk Agent**: validates position sizes against risk limits -//! - **Execution Agent**: submits orders to the broker API +//! A fund runs a research agent that may take 10 steps, spend 100k tokens and +//! $2 per analysis, and must stop after 3 consecutive errors (a broker or +//! filings API that is down should not be hammered). These are the card's +//! `max_iterations`, `token_budget`, `cost_budget_usd` and +//! `max_consecutive_errors`; the runtime folds what the agent used from the +//! run's event log and asks [`check`] before every step. //! -//! The fund's CRO (Chief Risk Officer) mandates: -//! -//! 1. **Research agent runs in `bounded_autonomous` mode** — it may iterate up -//! to 10 times (tool calls: search APIs, parse filings) but must escalate -//! if it exceeds $2 in LLM costs per analysis. The CRO saw a runaway agent -//! spend $47 in one night iterating over SEC filings. -//! -//! 2. **Execution agent runs in `guided` mode** — every single trade order -//! requires a human trader's explicit approval. No exceptions. A rogue agent -//! placing live orders could cost millions before anyone notices. -//! -//! 3. **Circuit breaker** — if any agent hits 3 consecutive errors (e.g., the -//! broker API is down), stop immediately and alert the operations team. -//! In March 2025, a competitor lost $3M when their agent retried failed -//! orders 200+ times, each time with slightly different parameters that -//! the broker interpreted as new orders. -//! -//! ## Why This Matters -//! Without autonomy enforcement: -//! - A research agent could consume $1,000+/day in LLM costs chasing rabbit holes -//! - An execution agent could place unauthorized trades -//! - Error loops could generate thousands of duplicate orders - -use jamjet_agents::card::{AutonomyConstraints, AutonomyLevel}; -use jamjet_policy::autonomy::{AutonomyContext, AutonomyDecision, AutonomyEnforcer}; - -fn research_agent_context( - iterations: u32, - cost: f64, - tokens: u64, - consecutive_errors: u32, -) -> AutonomyContext { - AutonomyContext { - agent_ref: "quant-research-v2".to_string(), - autonomy_level: AutonomyLevel::BoundedAutonomous, - constraints: Some(AutonomyConstraints { - max_iterations: Some(10), - max_tool_calls: Some(50), - token_budget: Some(100_000), // 100k tokens per analysis - cost_budget_usd: Some(2.0), // $2 max per analysis run - allowed_tools: vec![], - blocked_tools: vec![], - allowed_delegations: vec![], - require_approval_for: vec![], - time_budget_secs: Some(300), // 5 minute timeout - }), - current_iterations: iterations, - current_tool_calls: 0, - current_cost_usd: cost, - current_tokens: tokens, - consecutive_errors, - circuit_breaker_threshold: 3, - } +//! The autonomy level does not switch any of this off: a +//! `fully_autonomous` card's declared limits apply too. Guided mode does not +//! gate every call; use `require_approval_for` for that. + +use jamjet_agents::card::AutonomyConstraints; +use jamjet_policy::autonomy::{check, AgentUsage, AutonomyDecision, Step}; + +fn research_card() -> AutonomyConstraints { + serde_json::from_value(serde_json::json!({ + "max_iterations": 10, + "max_tool_calls": 50, + "token_budget": 100_000, + "cost_budget_usd": 2.0, + "time_budget_secs": 300, + "max_consecutive_errors": 3 + })) + .unwrap() } -fn execution_agent_context(consecutive_errors: u32) -> AutonomyContext { - AutonomyContext { - agent_ref: "trade-executor-v1".to_string(), - autonomy_level: AutonomyLevel::Guided, - constraints: None, // Guided mode doesn't need constraints — everything is gated - current_iterations: 0, - current_tool_calls: 0, - current_cost_usd: 0.0, - current_tokens: 0, +fn usage(iterations: u64, cost_usd: f64, tokens: u64, consecutive_errors: u32) -> AgentUsage { + AgentUsage { + iterations, + tool_calls: iterations, + tokens, + cost_usd, consecutive_errors, - circuit_breaker_threshold: 3, + elapsed_ms: Some(60_000), } } -// ── Test 1: Research agent within limits proceeds normally ────────────── - -#[test] -fn research_agent_within_limits_proceeds() { - // Story: It's 8am. The research agent starts analyzing NVDA's latest - // 10-K filing. It has used 3 iterations, $0.80, and 40k tokens. - // All within limits — agent should proceed. - - let enforcer = AutonomyEnforcer; - let ctx = research_agent_context(3, 0.80, 40_000, 0); - - let decision = enforcer.check(&ctx, None); - assert!(matches!(decision, AutonomyDecision::Proceed)); - println!("PROCEED: Research agent at iteration 3, $0.80, 40k tokens — within all limits"); -} - -// ── Test 2: Research agent hits iteration limit → escalate ───────────── - -#[test] -fn research_agent_iteration_limit_escalates() { - // Story: The research agent has been analyzing a complex restructuring. - // It's now on iteration 10 — the maximum. The LLM keeps wanting to "dig - // deeper into subsidiary filings" but the CRO says: no. Escalate. - - let enforcer = AutonomyEnforcer; - let ctx = research_agent_context(10, 1.50, 80_000, 0); +const SEARCH: Step = Step { tool_calls: 1 }; - let decision = enforcer.check(&ctx, None); - - match &decision { - AutonomyDecision::EscalateLimit { - limit_type, - limit_value, - actual_value, - .. - } => { - assert_eq!(limit_type, "max_iterations"); - println!("ESCALATED: Research agent hit iteration limit"); - println!(" Limit: {limit_value}, Actual: {actual_value}"); - println!(" -> Analyst notified. They can review partial findings and decide"); - println!(" whether to increase the limit or accept current analysis."); - } - other => panic!("Expected EscalateLimit for iterations, got: {other:?}"), +fn limit(d: &AutonomyDecision) -> &'static str { + match d { + AutonomyDecision::LimitReached { limit_type, .. } => limit_type, + other => panic!("expected a limit, got {other:?}"), } } -// ── Test 3: Research agent hits cost budget → escalate ────────────────── - #[test] -fn research_agent_cost_budget_escalates() { - // Story: The agent is only on iteration 5, but it's been using long-context - // prompts (full 10-K text = 150k tokens each). Cost has hit $2.10 — over - // the $2.00 budget. Without this check, last week's runaway spent $47. - - let enforcer = AutonomyEnforcer; - let ctx = research_agent_context(5, 2.10, 95_000, 0); - - let decision = enforcer.check(&ctx, None); - - match &decision { - AutonomyDecision::EscalateLimit { - limit_type, - limit_value, - actual_value, - .. - } => { - assert_eq!(limit_type, "cost_budget"); - println!("ESCALATED: Research agent exceeded cost budget"); - println!(" Budget: ${limit_value}, Spent: ${actual_value}"); - println!(" -> Operations team notified before costs spiral further."); - } - other => panic!("Expected EscalateLimit for cost, got: {other:?}"), - } +fn within_limits_proceeds() { + let d = check(Some(&research_card()), &usage(3, 0.80, 40_000, 0), SEARCH); + assert_eq!(d, AutonomyDecision::Proceed); } -// ── Test 4: Research agent hits token budget → escalate ───────────────── - #[test] -fn research_agent_token_budget_escalates() { - // Story: The agent has consumed 100k+ tokens processing multiple filings. - // Token budget enforcement catches this even though cost might still be - // under limit (e.g., using a cheap model). - - let enforcer = AutonomyEnforcer; - let ctx = research_agent_context(4, 0.50, 100_000, 0); - - let decision = enforcer.check(&ctx, None); - - match &decision { - AutonomyDecision::EscalateLimit { limit_type, .. } => { - assert_eq!(limit_type, "token_budget"); - println!("ESCALATED: Research agent exceeded token budget (100k tokens)"); - } - other => panic!("Expected EscalateLimit for tokens, got: {other:?}"), - } +fn the_eleventh_step_is_refused() { + let card = research_card(); + assert_eq!( + check(Some(&card), &usage(9, 1.50, 80_000, 0), SEARCH), + AutonomyDecision::Proceed + ); + assert_eq!( + limit(&check(Some(&card), &usage(10, 1.50, 80_000, 0), SEARCH)), + "max_iterations" + ); } -// ── Test 5: Execution agent — EVERY trade requires approval ──────────── - #[test] -fn execution_agent_guided_mode_requires_approval() { - // Story: The execution agent wants to submit a market order: - // buy 1000 shares of NVDA @ market - // In guided mode, this call is PAUSED until a human trader approves. - - let enforcer = AutonomyEnforcer; - let ctx = execution_agent_context(0); - - let decision = enforcer.check(&ctx, Some("submit_market_order")); - - match &decision { - AutonomyDecision::RequireToolApproval { tool_name } => { - assert_eq!(tool_name, "submit_market_order"); - println!("APPROVAL REQUIRED: {tool_name}"); - println!(" -> Trade desk notified. Trader reviews order parameters."); - println!(" -> No order submitted until explicit approval via /approve."); - } - other => panic!("Expected RequireToolApproval, got: {other:?}"), - } +fn a_spent_cost_budget_is_refused() { + let d = check(Some(&research_card()), &usage(5, 2.10, 95_000, 0), SEARCH); + assert_eq!(limit(&d), "cost_budget"); } -// ── Test 6: Execution agent — even cancel orders need approval ───────── - #[test] -fn execution_agent_cancel_also_requires_approval() { - // Story: The agent decides to cancel a pending limit order. Even cancellations - // require approval — the CRO saw an agent cancel a hedging position that was - // protecting against downside risk. - - let enforcer = AutonomyEnforcer; - let ctx = execution_agent_context(0); - - let decision = enforcer.check(&ctx, Some("cancel_order")); - assert!( - matches!(decision, AutonomyDecision::RequireToolApproval { .. }), - "Cancel orders must also require approval in guided mode" - ); - println!("APPROVAL REQUIRED: cancel_order (even cancellations are gated)"); +fn a_spent_token_budget_is_refused_even_when_cost_is_low() { + let d = check(Some(&research_card()), &usage(4, 0.50, 100_000, 0), SEARCH); + assert_eq!(limit(&d), "token_budget"); } -// ── Test 7: Circuit breaker — broker API goes down ───────────────────── - #[test] -fn circuit_breaker_trips_after_consecutive_errors() { - // Story: At 2:47pm, the broker's API starts returning 503 errors. - // The execution agent fails 3 times in a row trying to check positions. - // The circuit breaker trips — agent STOPS and alerts operations. - // - // Without this: the agent at a competitor kept retrying for 2 hours, - // creating 200+ duplicate position queries that the broker rate-limited, - // which cascaded into a 24-hour trading suspension. - - let enforcer = AutonomyEnforcer; - - // After 1 error: proceed - let ctx1 = execution_agent_context(1); - assert!(matches!( - enforcer.check(&ctx1, None), - AutonomyDecision::Proceed - )); - println!("Error 1/3: Proceed (might be transient)"); - - // After 2 errors: proceed - let ctx2 = execution_agent_context(2); - assert!(matches!( - enforcer.check(&ctx2, None), - AutonomyDecision::Proceed - )); - println!("Error 2/3: Proceed (still within threshold)"); - - // After 3 errors: TRIP - let ctx3 = execution_agent_context(3); - let decision = enforcer.check(&ctx3, None); - - match &decision { - AutonomyDecision::TripCircuitBreaker { - consecutive_errors, - threshold, - .. - } => { - assert_eq!(*consecutive_errors, 3); - assert_eq!(*threshold, 3); - println!("CIRCUIT BREAKER TRIPPED: {consecutive_errors} consecutive errors (threshold: {threshold})"); - println!(" -> Agent STOPPED. Operations team alerted."); - println!(" -> No further API calls until human resets via /approve."); - println!(" -> Potential duplicate order storm PREVENTED."); - } - other => panic!("Expected TripCircuitBreaker, got: {other:?}"), +fn the_breaker_trips_on_the_third_consecutive_error() { + let card = research_card(); + for errors in [1, 2] { + assert_eq!( + check(Some(&card), &usage(2, 0.5, 20_000, errors), SEARCH), + AutonomyDecision::Proceed, + "{errors} errors" + ); } + assert_eq!( + check(Some(&card), &usage(2, 0.5, 20_000, 3), SEARCH), + AutonomyDecision::CircuitBreakerTripped { + consecutive_errors: 3, + threshold: 3 + } + ); } -// ── Test 8: Circuit breaker overrides even bounded_autonomous ────────── - #[test] -fn circuit_breaker_overrides_bounded_autonomous_limits() { - // Story: The research agent is within all its budgets (iteration 2, $0.50), - // but the SEC EDGAR API is down and it's hit 3 consecutive errors. - // Circuit breaker fires BEFORE checking budget limits — safety first. - - let enforcer = AutonomyEnforcer; - let ctx = research_agent_context(2, 0.50, 20_000, 3); // 3 errors but within budgets - - let decision = enforcer.check(&ctx, None); +fn the_breaker_is_checked_before_the_budgets() { + // Within every budget, but the filings API has failed three times. + let d = check(Some(&research_card()), &usage(2, 0.50, 20_000, 3), SEARCH); assert!( - matches!(decision, AutonomyDecision::TripCircuitBreaker { .. }), - "Circuit breaker should fire regardless of budget limits" + matches!(d, AutonomyDecision::CircuitBreakerTripped { .. }), + "{d:?}" ); - println!("CIRCUIT BREAKER: fires even though agent is within budget limits"); - println!(" -> Error recovery takes priority over continued analysis."); } -// ── Test 9: Fully autonomous agent — only circuit breaker applies ────── - #[test] -fn fully_autonomous_agent_only_has_circuit_breaker() { - // Story: The market data ingestion agent runs in fully_autonomous mode — - // it streams price data 24/7 with no human interaction needed. The ONLY - // safety net is the circuit breaker. - - let enforcer = AutonomyEnforcer; - - let ctx = AutonomyContext { - agent_ref: "market-data-ingest-v1".to_string(), - autonomy_level: AutonomyLevel::FullyAutonomous, - constraints: Some(AutonomyConstraints { - max_iterations: Some(1000), // high limit - max_tool_calls: Some(10000), - token_budget: Some(1_000_000), - cost_budget_usd: Some(100.0), - allowed_tools: vec![], - blocked_tools: vec![], - allowed_delegations: vec![], - require_approval_for: vec![], - time_budget_secs: None, - }), - current_iterations: 500, // Way past what bounded would allow - current_tool_calls: 5000, - current_cost_usd: 50.0, // Way past typical budgets - current_tokens: 500_000, - consecutive_errors: 0, - circuit_breaker_threshold: 3, - }; - - // Should proceed — fully autonomous ignores iteration/cost/token limits - let decision = enforcer.check(&ctx, None); - assert!(matches!(decision, AutonomyDecision::Proceed)); - println!("PROCEED: Fully autonomous agent at 500 iterations, $50 — no limits enforced"); - println!(" -> Only the circuit breaker (3 consecutive errors) would stop this agent."); +fn high_limits_still_bind() { + // A market data agent with high limits: still refused once past them. + let card: AutonomyConstraints = + serde_json::from_value(serde_json::json!({ "max_iterations": 1000 })).unwrap(); + assert_eq!( + check(Some(&card), &usage(999, 50.0, 500_000, 0), SEARCH), + AutonomyDecision::Proceed + ); + assert_eq!( + limit(&check(Some(&card), &usage(1000, 50.0, 500_000, 0), SEARCH)), + "max_iterations" + ); } diff --git a/runtime/state/src/budget.rs b/runtime/state/src/budget.rs index 8f5dcea0..04c3a7d8 100644 --- a/runtime/state/src/budget.rs +++ b/runtime/state/src/budget.rs @@ -16,13 +16,17 @@ pub struct BudgetState { pub total_input_tokens: u64, pub total_output_tokens: u64, pub total_cost_usd: f64, - /// Number of agent reasoning iterations across all agent nodes. + /// Not maintained by the runtime and not read by any enforcement. An + /// agent card's limits (iterations, tool calls, consecutive errors) are + /// folded from the event log instead (`jamjet_worker::agent_limits`), so + /// nothing a node writes into its state can reset them. Kept so stored + /// snapshots keep their shape. pub iteration_count: u32, - /// Number of tool calls across all nodes. + /// Not maintained or read; see `iteration_count`. pub tool_call_count: u32, - /// Consecutive error count for circuit-breaker logic. + /// Not maintained or read; see `iteration_count`. pub consecutive_error_count: u32, - /// True once the circuit breaker has fired for this execution. + /// Not maintained or read; see `iteration_count`. pub circuit_breaker_tripped: bool, } diff --git a/runtime/state/src/event.rs b/runtime/state/src/event.rs index a1bcff88..c754e60e 100644 --- a/runtime/state/src/event.rs +++ b/runtime/state/src/event.rs @@ -246,7 +246,8 @@ pub enum EventKind { AutonomyLimitReached { node_id: NodeId, agent_ref: String, - /// "max_iterations" | "cost_budget" | "token_budget" | "max_tool_calls" + /// "max_iterations" | "max_tool_calls" | "token_budget" | "cost_budget" + /// | "time_budget" limit_type: String, limit_value: serde_json::Value, actual_value: serde_json::Value, @@ -257,6 +258,32 @@ pub enum EventKind { consecutive_errors: u32, threshold: u32, }, + /// A node acting for one or more agents (`agent_ref`) was admitted to + /// fire: one step of each of `agents`, using `tool_calls` tool calls. + /// Appended by the shared guard (`node_guard`) on both transports, only + /// when every agent's limits allow the step (or an approval covers it), + /// in the same compare-and-set that checked them. The autonomy counters + /// (iterations, tool calls, time, and through the node's later + /// `NodeCompleted` / `NodeFailed`, tokens, cost and consecutive errors) + /// are folded from these, so they survive a restart and a rollover. + AgentStepAdmitted { + node_id: NodeId, + /// The work item that was admitted. A result reported over HTTP is + /// recorded only for an admitted item. + #[serde(default, skip_serializing_if = "Option::is_none")] + work_item_id: Option, + /// Every agent the node acts for (the workflow's and the node's + /// `agent_ref`), each charged one step. + agents: Vec, + /// Tool calls this step makes: one for a node that uses a tool, the + /// number of calls in an agent tool dispatch batch. + tool_calls: u32, + /// For a step past a limit released by a human (`escalate_to_human`): + /// the sequence of the approval request it used. Each request + /// releases one step. + #[serde(default, skip_serializing_if = "Option::is_none")] + approval_sequence: Option, + }, EscalationRequired { node_id: NodeId, agent_ref: String, @@ -499,6 +526,7 @@ impl EventKind { | Self::CostBudgetExceeded { node_id, .. } | Self::AutonomyLimitReached { node_id, .. } | Self::CircuitBreakerTripped { node_id, .. } + | Self::AgentStepAdmitted { node_id, .. } | Self::EscalationRequired { node_id, .. } | Self::PolicyViolation { node_id, .. } | Self::ToolApprovalRequired { node_id, .. } diff --git a/runtime/workers/src/agent_limits.rs b/runtime/workers/src/agent_limits.rs new file mode 100644 index 00000000..f2b0f2de --- /dev/null +++ b/runtime/workers/src/agent_limits.rs @@ -0,0 +1,578 @@ +//! An agent card's autonomy limits, enforced from the run's event log. +//! +//! A node that acts for an agent (the workflow's `agent_ref` or the node's +//! own) is one STEP of that agent. When the shared guard +//! ([`crate::node_guard::guard_item`]) lets such a node fire, it records +//! [`EventKind::AgentStepAdmitted`] naming the agents and the tool calls the +//! step makes. Everything the limits read is folded from the log: +//! +//! - iterations: admitted steps; +//! - tool calls: the admitted steps' `tool_calls` (one for a node that uses +//! a tool, each call of an agent tool dispatch batch, none for a model or +//! condition node); +//! - tokens and cost: `input_tokens + output_tokens` and `cost_usd` on the +//! `NodeCompleted` that ends an admitted step; +//! - consecutive errors: admitted steps ending in `NodeFailed` since the +//! agent's last completed step; +//! - time: since the agent's first admitted step. +//! +//! The fold covers every segment of the run (`segment_chain`), so a restart +//! or a continue-as-new rollover does not start the counters again, and +//! nothing a node returns (its `state_patch`) can reach them. The step is +//! recorded with `append_event_if`, whose condition re-runs the same check on +//! the log as it stands at the append, so two claims cannot both take the +//! last step. +//! +//! What a step reports about itself (tokens, cost, success or failure) is the +//! executor's or, over HTTP, the worker's word: a worker that under-reports +//! or reports a failed call as completed is believed, as it is everywhere +//! else in the worker protocol. + +use chrono::{DateTime, Utc}; +use jamjet_agents::card::{AutonomyConstraints, OnLimitExceeded}; +use jamjet_core::workflow::ExecutionId; +use jamjet_policy::autonomy::{self, AgentUsage, AutonomyDecision, Step}; +use jamjet_state::backend::{ConditionalAppend, StateBackend}; +use jamjet_state::event::EventKind; +use jamjet_state::Event; +use serde_json::Value; +use std::collections::HashMap; +use tracing::warn; + +use crate::dispatch_guard::DispatchGuardOutcome as GuardOutcome; + +/// Every rule recorded for a refusal by an agent's autonomy limits starts +/// with this. +pub const AUTONOMY_RULE: &str = "agent autonomy"; + +/// The agents a node acts for, each with its card's constraints (`None`: a +/// card with none, where only the default circuit breaker applies). +pub(crate) type Agents = Vec<(String, Option)>; + +/// What one agent has used, as folded from the log. +#[derive(Debug, Clone, Default)] +struct Spent { + iterations: u64, + tool_calls: u64, + tokens: u64, + cost_usd: f64, + consecutive_errors: u32, + first_step_at: Option>, +} + +/// The fold of a run's log into each agent's usage. +#[derive(Debug, Clone, Default)] +pub(crate) struct Ledger { + agents: HashMap, + /// Admitted steps that have not ended yet, by node: the agents each was + /// admitted for, oldest first. A `NodeCompleted` or `NodeFailed` of the + /// node ends the oldest. + open: HashMap>>, +} + +/// The log holds an event this build cannot read: it may be a step, so the +/// counters cannot be known. +#[derive(Debug)] +pub(crate) struct Unreadable; + +impl Ledger { + /// Fold one segment's events. Steps are matched to their ends within a + /// segment, so the open steps are dropped when the segment changes. + fn apply_segment(&mut self, events: &[Event]) -> Result<(), Unreadable> { + self.open.clear(); + for e in events { + self.apply(e)?; + } + Ok(()) + } + + fn apply(&mut self, e: &Event) -> Result<(), Unreadable> { + match &e.kind { + EventKind::AgentStepAdmitted { + node_id, + agents, + tool_calls, + .. + } => { + for a in agents { + let s = self.agents.entry(a.clone()).or_default(); + s.iterations = s.iterations.saturating_add(1); + s.tool_calls = s.tool_calls.saturating_add(u64::from(*tool_calls)); + if s.first_step_at.is_none() { + s.first_step_at = Some(e.created_at); + } + } + self.open + .entry(node_id.clone()) + .or_default() + .push(agents.clone()); + } + EventKind::NodeCompleted { + node_id, + input_tokens, + output_tokens, + cost_usd, + .. + } => { + let Some(agents) = self.end(node_id) else { + return Ok(()); + }; + let tokens = input_tokens + .unwrap_or(0) + .saturating_add(output_tokens.unwrap_or(0)); + // A cost that is not a finite number cannot be summed into + // something the budget can trust: count it as spending all + // of it. A negative cost is refused where it is reported and + // would refund the budget here, so it counts as nothing. + let cost = match cost_usd { + Some(c) if !c.is_finite() => f64::INFINITY, + Some(c) if *c > 0.0 => *c, + _ => 0.0, + }; + for a in agents { + let s = self.agents.entry(a).or_default(); + s.tokens = s.tokens.saturating_add(tokens); + s.cost_usd += cost; + s.consecutive_errors = 0; + } + } + EventKind::NodeFailed { node_id, .. } => { + if let Some(agents) = self.end(node_id) { + for a in agents { + let s = self.agents.entry(a).or_default(); + s.consecutive_errors = s.consecutive_errors.saturating_add(1); + } + } + } + EventKind::Unknown => return Err(Unreadable), + _ => {} + } + Ok(()) + } + + /// End the node's oldest open step, returning its agents. + fn end(&mut self, node_id: &str) -> Option> { + let open = self.open.get_mut(node_id)?; + if open.is_empty() { + return None; + } + Some(open.remove(0)) + } + + fn usage(&self, agent: &str, now: DateTime) -> AgentUsage { + let s = self.agents.get(agent).cloned().unwrap_or_default(); + AgentUsage { + iterations: s.iterations, + tool_calls: s.tool_calls, + tokens: s.tokens, + cost_usd: s.cost_usd, + consecutive_errors: s.consecutive_errors, + elapsed_ms: s + .first_step_at + .map(|t| u64::try_from((now - t).num_milliseconds()).unwrap_or(0)), + } + } +} + +/// The first agent whose limits stop `step`, with what its card says to do +/// about it. A `fail` card wins over an `escalate_to_human` one: a refusal +/// is the stricter answer. +fn first_stop( + ledger: &Ledger, + agents: &Agents, + step: Step, + now: DateTime, +) -> Option<(String, AutonomyDecision, OnLimitExceeded)> { + let mut escalation = None; + for (agent, constraints) in agents { + let decision = autonomy::check(constraints.as_ref(), &ledger.usage(agent, now), step); + if decision == AutonomyDecision::Proceed { + continue; + } + let on_limit = constraints + .as_ref() + .and_then(|c| c.on_limit_exceeded) + .unwrap_or_default(); + match on_limit { + OnLimitExceeded::Fail => return Some((agent.clone(), decision, on_limit)), + OnLimitExceeded::EscalateToHuman => { + if escalation.is_none() { + escalation = Some((agent.clone(), decision, on_limit)); + } + } + } + } + escalation +} + +/// Whether any agent whose card says `fail` stops `step`. An approved +/// escalation releases a step past an escalating card's limit, never past a +/// failing card's. +fn fail_stop(ledger: &Ledger, agents: &Agents, step: Step, now: DateTime) -> bool { + agents.iter().any(|(agent, constraints)| { + let on_limit = constraints + .as_ref() + .and_then(|c| c.on_limit_exceeded) + .unwrap_or_default(); + on_limit == OnLimitExceeded::Fail + && autonomy::check(constraints.as_ref(), &ledger.usage(agent, now), step) + != AutonomyDecision::Proceed + }) +} + +/// The answered request of `node_id` that can release one step past an +/// agent's limit, by sequence: approved, raised for that limit (its context +/// carries `autonomy_limit`, so the reviewer saw it), and not used by an +/// earlier step. `None` when there is none. +pub(crate) fn unused_escalation(events: &[Event], node_id: &str) -> Option { + use jamjet_state::approvals::{answered_request, node_approval_status, NodeApprovalStatus}; + if !matches!( + node_approval_status(events, node_id), + NodeApprovalStatus::Approved { .. } + ) { + return None; + } + let request = answered_request(events, node_id)?; + request.context.get(ESCALATION_KEY)?; + let used = events.iter().any(|e| { + matches!(&e.kind, EventKind::AgentStepAdmitted { approval_sequence: Some(s), .. } if *s == request.sequence) + }); + (!used).then_some(request.sequence) +} + +/// The key under which an escalation request's context shows the limit. +pub(crate) const ESCALATION_KEY: &str = "autonomy_limit"; + +/// An over-limit step whose card escalates: it may fire only under an +/// approval of that step. +#[derive(Debug, Clone)] +pub(crate) struct Escalation { + pub agent: String, + pub decision: AutonomyDecision, +} + +impl Escalation { + /// What the reviewer is asked about, in the request's context. + pub(crate) fn context(&self) -> Value { + let mut v = serde_json::json!({ "agent_ref": self.agent }); + match &self.decision { + AutonomyDecision::LimitReached { + limit_type, + limit_value, + actual_value, + } => { + v["limit_type"] = Value::from(*limit_type); + v["limit_value"] = limit_value.clone(); + v["actual_value"] = actual_value.clone(); + } + AutonomyDecision::CircuitBreakerTripped { + consecutive_errors, + threshold, + } => { + v["limit_type"] = Value::from("circuit_breaker"); + v["limit_value"] = Value::from(*threshold); + v["actual_value"] = Value::from(*consecutive_errors); + } + AutonomyDecision::Proceed => {} + } + v + } +} + +/// A step the limits allow (or escalate), ready to be recorded once the +/// rest of the guard allows it too. +pub(crate) struct Plan { + agents: Agents, + step: Step, + /// The fold of the run's earlier segments. They are closed, so only the + /// current segment's log is re-read at the append. + earlier: Ledger, + now: DateTime, + pub escalation: Option, +} + +/// Decide whether a step of `agents` making `tool_calls` tool calls may +/// fire. `Ok(plan)`: it may, possibly only under an approval +/// (`plan.escalation`); record it with [`admit`] if the rest of the guard +/// allows the node. `Err`: refused (a limit with `on_limit_exceeded: fail`, +/// recorded here) or undecidable (`Unavailable`). +pub(crate) async fn plan( + backend: &dyn StateBackend, + execution_id: &ExecutionId, + node_id: &str, + agents: Agents, + tool_calls: u32, +) -> Result { + let unavailable = |what: String| GuardOutcome::Unavailable { + reason: format!( + "agent autonomy counters unavailable; refusing to run without checking them: {what}" + ), + }; + let chain = jamjet_state::segment_chain(backend, execution_id) + .await + .map_err(|e| unavailable(e.to_string()))?; + let mut earlier = Ledger::default(); + for id in chain.iter().filter(|id| *id != execution_id) { + let events = backend + .get_events(id) + .await + .map_err(|e| unavailable(e.to_string()))?; + earlier + .apply_segment(&events) + .map_err(|_| unavailable(format!("segment {id} holds an unreadable event")))?; + } + let events = backend + .get_events(execution_id) + .await + .map_err(|e| unavailable(e.to_string()))?; + let mut ledger = earlier.clone(); + ledger + .apply_segment(&events) + .map_err(|_| unavailable("the run's log holds an unreadable event".to_string()))?; + + let now = Utc::now(); + let step = Step { tool_calls }; + let escalation = match first_stop(&ledger, &agents, step, now) { + None => None, + Some((agent, decision, OnLimitExceeded::Fail)) => { + let reason = describe(&agent, &decision); + warn!(execution_id = %execution_id, node_id, %reason, "agent autonomy limit refused node"); + record_stop(backend, execution_id, node_id, &agent, &decision, None).await; + return Err(GuardOutcome::Blocked { reason }); + } + Some((agent, decision, OnLimitExceeded::EscalateToHuman)) => { + Some(Escalation { agent, decision }) + } + }; + Ok(Plan { + agents, + step, + earlier, + now, + escalation, + }) +} + +/// The refusal text, recorded on the node's `NodeFailed`. +fn describe(agent: &str, decision: &AutonomyDecision) -> String { + match decision { + AutonomyDecision::LimitReached { + limit_type, + limit_value, + actual_value, + } => format!( + "{AUTONOMY_RULE}: agent '{agent}' reached {limit_type} (limit {limit_value}, \ + at {actual_value})" + ), + AutonomyDecision::CircuitBreakerTripped { + consecutive_errors, + threshold, + } => format!( + "{AUTONOMY_RULE}: agent '{agent}' circuit breaker tripped ({consecutive_errors} \ + consecutive errors, threshold {threshold})" + ), + AutonomyDecision::Proceed => format!("{AUTONOMY_RULE}: agent '{agent}'"), + } +} + +/// Record the stop (`AutonomyLimitReached` or `CircuitBreakerTripped`) and, +/// for an escalation, `EscalationRequired`. Best effort: the refusal or the +/// hold stands whether or not these land. +pub(crate) async fn record_stop( + backend: &dyn StateBackend, + execution_id: &ExecutionId, + node_id: &str, + agent: &str, + decision: &AutonomyDecision, + escalated_to: Option<&str>, +) { + let stop = match decision { + AutonomyDecision::LimitReached { + limit_type, + limit_value, + actual_value, + } => EventKind::AutonomyLimitReached { + node_id: node_id.to_string(), + agent_ref: agent.to_string(), + limit_type: limit_type.to_string(), + limit_value: limit_value.clone(), + actual_value: actual_value.clone(), + }, + AutonomyDecision::CircuitBreakerTripped { + consecutive_errors, + threshold, + } => EventKind::CircuitBreakerTripped { + node_id: node_id.to_string(), + agent_ref: agent.to_string(), + consecutive_errors: *consecutive_errors, + threshold: *threshold, + }, + AutonomyDecision::Proceed => return, + }; + let reason = match decision { + AutonomyDecision::CircuitBreakerTripped { .. } => "circuit_breaker".to_string(), + AutonomyDecision::LimitReached { limit_type, .. } => format!("autonomy_limit:{limit_type}"), + AutonomyDecision::Proceed => return, + }; + let mut kinds = vec![stop]; + if let Some(target) = escalated_to { + kinds.push(EventKind::EscalationRequired { + node_id: node_id.to_string(), + agent_ref: agent.to_string(), + reason, + escalation_target: target.to_string(), + }); + } + for kind in kinds { + if let Err(e) = backend + .append_event(Event::new(execution_id.clone(), 0, kind)) + .await + { + warn!(execution_id = %execution_id, node_id, "could not record an autonomy stop: {e}"); + } + } +} + +/// Record the step, if the limits still allow it on the log as it stands +/// (`approved`: an approval covers this over-limit step, so only the record +/// is needed). `Allow` once recorded; anything else withholds the node +/// (`Unavailable`: the next attempt decides again). +pub(crate) async fn admit( + backend: &dyn StateBackend, + execution_id: &ExecutionId, + node_id: &str, + work_item_id: Option, + plan: Plan, +) -> GuardOutcome { + let Plan { + agents, + step, + earlier, + now, + escalation, + } = plan; + // An escalated step fires only on the approval that released it, read + // now, and re-checked in the append below. + let approval_sequence = match &escalation { + None => None, + Some(_) => match backend.get_events(execution_id).await { + Ok(events) => match unused_escalation(&events, node_id) { + Some(seq) => Some(seq), + None => { + return GuardOutcome::Unavailable { + reason: "the approval releasing this step past the agent's limit was \ + used or withdrawn; deciding again on the next attempt" + .to_string(), + } + } + }, + Err(e) => { + return GuardOutcome::Unavailable { + reason: format!("agent steps unreadable: {e}"), + } + } + }, + }; + let names = agents.iter().map(|(a, _)| a.clone()).collect(); + let node = node_id.to_string(); + let event = Event::new( + execution_id.clone(), + 0, + EventKind::AgentStepAdmitted { + node_id: node_id.to_string(), + work_item_id, + agents: names, + tool_calls: step.tool_calls, + approval_sequence, + }, + ); + let still_allowed = move |log: &[Event]| -> bool { + let mut ledger = earlier.clone(); + if ledger.apply_segment(log).is_err() { + return false; + } + match approval_sequence { + // Every limit still allows it. + None => first_stop(&ledger, &agents, step, now).is_none(), + // Past an escalating card's limit on that one approval, still + // unused, and no failing card's limit reached meanwhile. + Some(seq) => { + unused_escalation(log, &node) == Some(seq) + && !fail_stop(&ledger, &agents, step, now) + } + } + }; + match backend.append_event_if(event, &still_allowed).await { + Ok(ConditionalAppend::Appended(_)) => GuardOutcome::Allow, + Ok(ConditionalAppend::Refused) => GuardOutcome::Unavailable { + reason: "another step of the agent was recorded while this one was decided; \ + deciding again on the next attempt" + .to_string(), + }, + Ok(ConditionalAppend::ExecutionTerminal(status)) => GuardOutcome::Unavailable { + reason: format!("execution is {status:?}"), + }, + Err(e) => GuardOutcome::Unavailable { + reason: format!("agent step could not be recorded; refusing to run it unrecorded: {e}"), + }, + } +} + +/// `GuardMode::Complete`: a result for a node acting for an agent is +/// recorded only for a work item the guard admitted. Without this a result +/// reported for an item that never went through the claim gate would be a +/// step no limit saw. +pub(crate) async fn completion_admitted( + backend: &dyn StateBackend, + execution_id: &ExecutionId, + work_item_id: Option, +) -> Result<(), GuardOutcome> { + let Some(item) = work_item_id else { + return Err(GuardOutcome::Blocked { + reason: format!( + "{AUTONOMY_RULE}: a result for a node acting for an agent needs the work item \ + it was admitted as" + ), + }); + }; + let events = backend + .get_events(execution_id) + .await + .map_err(|e| GuardOutcome::Unavailable { + reason: format!("agent steps unreadable; refusing to record the result: {e}"), + })?; + // The item's step must still be open: each admission of the node from + // the item's latest one on is ended by one result, so a step whose + // result is already recorded backs no second one. + let admitted = match events.iter().rposition(|e| { + matches!(&e.kind, EventKind::AgentStepAdmitted { work_item_id: Some(id), .. } if *id == item) + }) { + None => false, + Some(at) => { + let node = events[at].kind.node_id().unwrap_or_default(); + let (mut opened, mut ended) = (0usize, 0usize); + for e in &events[at..] { + match &e.kind { + EventKind::AgentStepAdmitted { node_id, .. } if node_id == node => opened += 1, + EventKind::NodeCompleted { node_id, .. } | EventKind::NodeFailed { node_id, .. } + if node_id == node => + { + ended += 1 + } + _ => {} + } + } + ended < opened + } + }; + if admitted { + Ok(()) + } else { + Err(GuardOutcome::Blocked { + reason: format!( + "{AUTONOMY_RULE}: work item {item} acts for an agent and was never admitted by \ + the claim gate, so its result cannot be recorded" + ), + }) + } +} diff --git a/runtime/workers/src/dispatch_guard.rs b/runtime/workers/src/dispatch_guard.rs index 2f1b1b6c..7528d0b1 100644 --- a/runtime/workers/src/dispatch_guard.rs +++ b/runtime/workers/src/dispatch_guard.rs @@ -1395,6 +1395,7 @@ mod tests { &json!({ "input": input_for(&["send_wire"]) }), Some(&sealer), Default::default(), + None, ) .await; assert!( diff --git a/runtime/workers/src/lib.rs b/runtime/workers/src/lib.rs index c49a4e80..78bb8764 100644 --- a/runtime/workers/src/lib.rs +++ b/runtime/workers/src/lib.rs @@ -7,6 +7,7 @@ //! - Emit heartbeats to renew leases //! - Report results back via the state backend +pub mod agent_limits; pub mod dispatch_guard; pub mod executor; pub mod executors; diff --git a/runtime/workers/src/node_guard.rs b/runtime/workers/src/node_guard.rs index 2c89c2ef..469789e7 100644 --- a/runtime/workers/src/node_guard.rs +++ b/runtime/workers/src/node_guard.rs @@ -149,6 +149,9 @@ struct CardGate { /// Names of an agent tool dispatch's calls that are gated: raised through /// the dispatch guard's per-call request, alongside any chain gates. calls: Vec, + /// Every agent the node acts for, with its card's constraints. Empty + /// when it acts for none. + agents: crate::agent_limits::Agents, } /// Decide whether `node_def` may fire (or, in `Complete` mode, whether its @@ -186,10 +189,13 @@ pub async fn guard_item( payload: &Value, sealer: Option<&dyn ApprovalSealer>, agents: AgentScope<'_>, + work_item_id: Option, ) -> GuardOutcome { + use crate::agent_limits; + // First, so a node the agent may not use never raises an approval request // and is refused even when no policy applies. - let card = match agent_card( + let mut card = match agent_card( backend, mode, execution_id, @@ -203,6 +209,123 @@ pub async fn guard_item( Ok(gate) => gate, Err(refused) => return refused, }; + + // Then the agents' autonomy limits (see `agent_limits`). A firing check + // decides from the run's log and, if the node may fire, records the step + // below; a completion only has to belong to an admitted step. + let mut agents_of_node = std::mem::take(&mut card.agents); + match callee(agents.registry, &node_def.kind).await { + Ok(Some(c)) if !agents_of_node.iter().any(|(a, _)| *a == c.0) => agents_of_node.push(c), + Ok(_) => {} + Err(unavailable) => return unavailable, + } + let plan = if agents_of_node.is_empty() { + None + } else { + match mode { + GuardMode::Fire => match agent_limits::plan( + backend, + execution_id, + node_id, + agents_of_node, + match step_tool_calls(node_def, payload) { + Ok(calls) => calls, + Err(rule) => return GuardOutcome::Blocked { reason: rule }, + }, + ) + .await + { + Ok(plan) => Some(plan), + Err(refused) => return refused, + }, + GuardMode::Complete => { + if let Err(refused) = + agent_limits::completion_admitted(backend, execution_id, work_item_id).await + { + return refused; + } + None + } + } + }; + let escalation = plan.as_ref().and_then(|p| p.escalation.clone()); + // Where the log stood before deciding: an escalation's stop is recorded + // only when this decision raised a new request for it. + let before = match &escalation { + Some(_) => match backend.latest_sequence(execution_id).await { + Ok(seq) => Some(seq), + Err(e) => { + return GuardOutcome::Unavailable { + reason: format!("run log unreadable: {e}"), + } + } + }, + None => None, + }; + + let outcome = guard_policy( + backend, + mode, + execution_id, + node_id, + tenant_id, + workflow_policy, + node_def, + payload, + sealer, + card, + escalation.as_ref(), + ) + .await; + + if let (Some(e), Some(before)) = (&escalation, before) { + if matches!(outcome, GuardOutcome::Held { .. }) { + let raised = backend + .get_events_since(execution_id, before) + .await + .map(|evs| { + evs.iter().any(|ev| { + matches!(&ev.kind, EventKind::ToolApprovalRequired { node_id: n, .. } if n == node_id) + }) + }) + .unwrap_or(true); + if raised { + agent_limits::record_stop( + backend, + execution_id, + node_id, + &e.agent, + &e.decision, + Some("human_approval"), + ) + .await; + } + } + } + match (outcome, plan) { + (GuardOutcome::Allow, Some(plan)) => { + agent_limits::admit(backend, execution_id, node_id, work_item_id, plan).await + } + (outcome, _) => outcome, + } +} + +/// [`guard_item`] after the agent card: the policy chain, the card's gates, +/// an autonomy escalation, and the approvals that answer them. +#[allow(clippy::too_many_arguments)] +async fn guard_policy( + backend: &dyn StateBackend, + mode: GuardMode, + execution_id: &ExecutionId, + node_id: &str, + tenant_id: &str, + workflow_policy: Option<&PolicySetIr>, + node_def: &NodeDef, + payload: &Value, + sealer: Option<&dyn ApprovalSealer>, + card: CardGate, + escalation: Option<&crate::agent_limits::Escalation>, +) -> GuardOutcome { let card_calls = card.calls.clone(); let card_gate = card.node; let is_dispatch = dispatch_guard::is_agent_tool_dispatch(&node_def.kind); @@ -279,13 +402,16 @@ pub async fn guard_item( } else { PolicyEvaluator.evaluate(&ctx, &sets) }; - // A card gate turns an allow into an approval; a block stays a block, and - // a chain gate is already one. - let decision = match (decision, &card_gate) { - (PolicyDecision::Allow, Some(_)) => PolicyDecision::RequireApproval { - approver: "human".to_string(), - }, - (d, _) => d, + // A card gate or an autonomy escalation turns an allow into an approval; + // a block stays a block, and a chain gate is already one (whose approval + // also covers the escalation: both approve this one action). + let decision = match decision { + PolicyDecision::Allow if card_gate.is_some() || escalation.is_some() => { + PolicyDecision::RequireApproval { + approver: "human".to_string(), + } + } + d => d, }; match decision { PolicyDecision::Allow if is_dispatch => per_call(false).await, @@ -339,10 +465,17 @@ pub async fn guard_item( payload, tool_name, approver, - serde_json::json!({ "node_id": node_id }), + { + let mut c = serde_json::json!({ "node_id": node_id }); + if let Some(e) = escalation { + c[crate::agent_limits::ESCALATION_KEY] = e.context(); + } + c + }, deadline, tenant_id, sealer, + escalation.is_some(), ) .await } @@ -359,6 +492,61 @@ pub async fn guard_item( } } +/// The agent an `agent` node calls, with its card's constraints: that +/// agent takes a step when the node runs, so its autonomy limits apply (its +/// tool and delegation lists do not: they restrict what the agent itself +/// uses, and the node is not acting for it). `Ok(None)`: not an `agent` node, +/// no registry, or the callee has no registered card. An unreadable registry +/// refuses. +async fn callee( + registry: Option<&dyn AgentRegistry>, + kind: &NodeKind, +) -> Result)>, GuardOutcome> { + let (NodeKind::Agent { agent_ref, .. }, Some(registry)) = (kind, registry) else { + return Ok(None); + }; + match registry.get_by_uri(agent_ref).await { + Ok(Some(agent)) => Ok(Some((agent_ref.clone(), agent.card.constraints))), + Ok(None) => Ok(None), + Err(e) => Err(GuardOutcome::Unavailable { + reason: format!( + "agent registry unavailable; refusing to run without checking the autonomy \ + limits of agent '{agent_ref}': {e}" + ), + }), + } +} + +/// The tool calls a node makes if it fires, for `max_tool_calls`: one for a +/// node that uses a tool (or picks its target at run time), plus each call +/// of an agent tool dispatch batch. The ADK dispatcher itself is not a tool +/// use, only its calls are. An unreadable batch is an error (its audit rule), +/// never zero calls: the limit must not depend on a later guard refusing it. +fn step_tool_calls(node_def: &NodeDef, payload: &Value) -> Result { + let is_dispatch = dispatch_guard::is_agent_tool_dispatch(&node_def.kind); + let calls_only = match &node_def.kind { + NodeKind::PythonFn { + module, function, .. + } => is_dispatch && jamjet_ir::is_adk_dispatch_coordinates(module, function), + _ => false, + }; + let ctx = EvaluationContext::from_node_kind(&node_def.id, &node_def.kind); + let own_call = !calls_only + && (ctx.tool_name.is_some() + || jamjet_policy::allowed_tools::RUN_TIME_TARGET_TAGS + .contains(&ctx.node_kind_tag.as_str())); + let batch = if is_dispatch { + jamjet_policy::dispatch::read_pending_tool_calls(&dispatch_guard::payload_input(payload)) + .map(|c| c.len()) + .map_err(dispatch_guard::unreadable_rule)? + } else { + 0 + }; + Ok(u32::try_from(batch) + .unwrap_or(u32::MAX) + .saturating_add(u32::from(own_call))) +} + /// The agent card half of the decision. `Ok(None)`: every agent the node /// acts for allows it outright (or it acts for none). `Ok(Some(name))`: some /// card's `require_approval_for` names it (as `name`), so the node must be @@ -461,6 +649,8 @@ async fn agent_card( }); } }; + gate.agents + .push((agent_ref.to_string(), card.constraints.clone())); let Some(c) = card.constraints.as_ref() else { continue; }; @@ -596,6 +786,7 @@ pub async fn require_approval( deadline: Option, tenant_id: &str, sealer: Option<&dyn ApprovalSealer>, + escalation: bool, ) -> GuardOutcome { let gated = vec![tool_name.clone()]; let events = match backend.get_events(execution_id).await { @@ -639,7 +830,12 @@ pub async fn require_approval( }; let current = basis.digest(node_id); - if approval_covers(&status, ¤t) { + // A step past an agent's limit needs an approval of a request that + // showed the limit, and each such approval releases one step: a failed + // step's retry has the same digest, and must ask again. + let covers = approval_covers(&status, ¤t) + && (!escalation || crate::agent_limits::unused_escalation(&events, node_id).is_some()); + if covers { info!(execution_id = %execution_id, node_id, "Approval satisfied — proceeding"); return GuardOutcome::Allow; } @@ -842,3 +1038,44 @@ pub(crate) fn basis_preview(payload: &Value, step: u64, node_def: &NodeDef) -> V }, }) } + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + fn dispatch_node() -> NodeDef { + serde_json::from_value(json!({ "id": "n1", "kind": { + "type": "python_fn", + "module": "jamjet.agents.tool_runtime", + "function": "dispatch_tool_calls", + "output_schema": "", + "agent_tool_dispatch": true + }})) + .expect("dispatch node must deserialize") + } + + #[test] + fn a_readable_batch_counts_each_call() { + let payload = json!({"input": {"tool_calls": [ + {"id": "c1", "name": "lookup", "arguments": {}}, + {"id": "c2", "name": "lookup", "arguments": {}} + ]}}); + assert_eq!(step_tool_calls(&dispatch_node(), &payload), Ok(2)); + } + + #[test] + fn an_unreadable_batch_is_refused_not_counted_as_zero() { + for input in [ + json!(null), + json!({"tool_calls": "send_wire"}), + json!({"tool_calls": [null]}), + ] { + let payload = json!({ "input": input }); + assert!( + step_tool_calls(&dispatch_node(), &payload).is_err(), + "unreadable batch {payload} must not count as a number of calls" + ); + } + } +} diff --git a/runtime/workers/src/worker.rs b/runtime/workers/src/worker.rs index 815b1044..be9f2834 100644 --- a/runtime/workers/src/worker.rs +++ b/runtime/workers/src/worker.rs @@ -9,10 +9,8 @@ use chrono::Utc; const MAX_RETRY_AFTER_SECS: u64 = 3_600; use jamjet_agents::AgentRegistry; -use jamjet_core::node::NodeKind; use jamjet_core::workflow::ExecutionId; use jamjet_ir::workflow::{NodeDef, WorkflowIr}; -use jamjet_policy::autonomy::{AutonomyContext, AutonomyDecision, AutonomyEnforcer}; use jamjet_policy::engine::node_kind_tag; use jamjet_state::backend::{StateBackend, StateBackendError, WorkItem, WorkItemId}; use jamjet_state::budget::BudgetState; @@ -267,26 +265,6 @@ impl Worker { return r; } - // Autonomy check. - if let Some(r) = self - .check_autonomy( - &execution_id, - item_id, - lease_fence, - attempt, - &node_id, - &tenant_id, - node_def, - &item.payload, - kind, - &budget, - ) - .await - { - heartbeat.abort(); - return r; - } - // Compute deterministic idempotency key: content_hash({run, segment, step, node, input_hash}). // step = count of NodeCompleted events for this node_id in the event log. // Retries do not complete, so the count is stable across retries and advances @@ -728,73 +706,6 @@ impl Worker { Ok(()) } - // ── Approval hold helper ────────────────────────────────────────────────── - - /// Park a node pending human approval (or let it run if already approved). - /// - /// Consults the event log. The gate is satisfied, and the caller proceeds - /// to execute (`None`), only when the latest request was `Approved` AND the - /// approval's `approved_digest` equals the action digest computed NOW - /// (`approvals::derive_action_digest`: run, occurrence step, node, - /// definition, dispatched payload). An approval therefore covers exactly the action - /// that was reviewed: if the state changed after the decision, or this is a - /// later visit of the node (a loop or any re-enqueue), or the approval is a - /// legacy one with no digest, it does not cover this action, and the node - /// asks again with a fresh request, settled exactly like a first request. - /// - /// Otherwise ensure exactly one outstanding `ToolApprovalRequired` exists - /// and settle the work item cleanly so the lease never expires into the - /// retry/dead-letter path. The node stays parked because it remains in the - /// scheduler fold's `scheduled` set (so it is never re-dispatched); the - /// fold's `held` set only gates which decision events are acted on. - /// The helper requires the FULL event log because a partial fold could - /// misreport NotRequested for an already-approved node. - /// - /// Every read that the decision depends on fails CLOSED: an error means the - /// node does not run. - /// The decision itself is `node_guard::require_approval`, shared with the - /// HTTP claim route; this method only settles the item. - #[allow(clippy::too_many_arguments)] - async fn hold_for_approval( - &self, - execution_id: &ExecutionId, - node_id: &str, - node_def: &NodeDef, - payload: &serde_json::Value, - item_id: WorkItemId, - lease_fence: i64, - attempt: u32, - tool_name: String, - approver: String, - context: serde_json::Value, - tenant_id: &str, - ) -> Option>> { - // An autonomy hold is not raised by a policy, so it has no deadline. - let outcome = node_guard::require_approval( - self.backend.as_ref(), - execution_id, - node_id, - node_def, - payload, - tool_name, - approver, - context, - None, - tenant_id, - self.sealer.as_deref(), - ) - .await; - self.settle_outcome( - execution_id, - node_id, - item_id, - lease_fence, - attempt, - outcome, - ) - .await - } - /// Turn a guard decision into the worker's lifecycle answer: `None` to run, /// `Some(Err)` to fail the attempt, `Some(Ok)` once the item is settled /// (held, or failed for good on a policy block). @@ -1011,6 +922,7 @@ impl Worker { remote_agents: Some(&ir.remote_agents), registry: self.agents.as_deref(), }, + Some(item_id), ) .await; self.settle_outcome( @@ -1024,187 +936,6 @@ impl Worker { .await } - // ── Autonomy check ──────────────────────────────────────────────────────── - - #[allow(clippy::too_many_arguments)] - async fn check_autonomy( - &self, - execution_id: &ExecutionId, - item_id: WorkItemId, - lease_fence: i64, - attempt: u32, - node_id: &str, - tenant_id: &str, - node_def: &NodeDef, - payload: &serde_json::Value, - kind: &NodeKind, - budget: &BudgetState, - ) -> Option>> { - let kind_tag = node_kind_tag(kind); - if kind_tag != "agent" { - return None; - } - let agents = self.agents.as_ref()?; - - let agent_ref = serde_json::to_value(kind) - .ok() - .and_then(|v| { - v.get("agent_ref") - .and_then(|a| a.as_str()) - .map(|s| s.to_string()) - }) - .unwrap_or_else(|| node_id.to_string()); - - // An unregistered callee has no card to enforce. An unreadable - // registry refuses: the item is left to its lease and retried, never - // run on a guess (the event appends below are best effort for the - // same reason, and never decide whether the node runs). - let agent = match agents.get_by_uri(&agent_ref).await { - Ok(Some(a)) => a, - Ok(None) => return None, - Err(e) => { - warn!(execution_id = %execution_id, node_id, %agent_ref, "agent registry unreadable; refusing the agent node"); - return Some(Err(format!("agent registry unavailable: {e}").into())); - } - }; - let card = agent.card; - - let ctx = AutonomyContext { - agent_ref: agent_ref.clone(), - autonomy_level: card.autonomy.clone(), - constraints: card.constraints.clone(), - current_iterations: budget.iteration_count, - current_tool_calls: budget.tool_call_count, - current_cost_usd: budget.total_cost_usd, - current_tokens: budget.total_tokens(), - consecutive_errors: budget.consecutive_error_count, - circuit_breaker_threshold: 3, - }; - - let decision = AutonomyEnforcer.check(&ctx, None); - - match decision { - AutonomyDecision::Proceed => None, - - // Currently unreachable at runtime: `AutonomyEnforcer.check` above - // is called with `tool_name: None` (pre-existing), so it never - // returns RequireToolApproval. The wiring is in place for when - // pending tool names are threaded through. - AutonomyDecision::RequireToolApproval { tool_name } => { - self.hold_for_approval( - execution_id, - node_id, - node_def, - payload, - item_id, - lease_fence, - attempt, - tool_name, - "human".to_string(), - serde_json::json!({ "agent_ref": agent_ref }), - tenant_id, - ) - .await - } - - AutonomyDecision::EscalateLimit { - limit_type, - limit_value, - actual_value, - escalation_target, - } => { - warn!(execution_id = %execution_id, node_id, %limit_type, "Autonomy limit reached"); - let seq = self - .backend - .latest_sequence(execution_id) - .await - .unwrap_or(0) - + 1; - let _ = self - .backend - .append_event(jamjet_state::Event::new( - execution_id.clone(), - seq, - EventKind::AutonomyLimitReached { - node_id: node_id.to_string(), - agent_ref: agent_ref.clone(), - limit_type: limit_type.clone(), - limit_value, - actual_value, - }, - )) - .await; - let seq = self - .backend - .latest_sequence(execution_id) - .await - .unwrap_or(0) - + 1; - let _ = self - .backend - .append_event(jamjet_state::Event::new( - execution_id.clone(), - seq, - EventKind::EscalationRequired { - node_id: node_id.to_string(), - agent_ref, - reason: format!("autonomy_limit:{limit_type}"), - escalation_target: escalation_target.as_str(), - }, - )) - .await; - Some(Err(format!("autonomy limit: {limit_type}").into())) - } - - AutonomyDecision::TripCircuitBreaker { - consecutive_errors, - threshold, - escalation_target, - } => { - warn!(execution_id = %execution_id, node_id, consecutive_errors, "Circuit breaker tripped"); - let seq = self - .backend - .latest_sequence(execution_id) - .await - .unwrap_or(0) - + 1; - let _ = self - .backend - .append_event(jamjet_state::Event::new( - execution_id.clone(), - seq, - EventKind::CircuitBreakerTripped { - node_id: node_id.to_string(), - agent_ref: agent_ref.clone(), - consecutive_errors, - threshold, - }, - )) - .await; - let seq = self - .backend - .latest_sequence(execution_id) - .await - .unwrap_or(0) - + 1; - let _ = self - .backend - .append_event(jamjet_state::Event::new( - execution_id.clone(), - seq, - EventKind::EscalationRequired { - node_id: node_id.to_string(), - agent_ref, - reason: "circuit_breaker".to_string(), - escalation_target: escalation_target.as_str(), - }, - )) - .await; - Some(Err("circuit breaker tripped".into())) - } - } - } - // ── Budget enforcement ──────────────────────────────────────────────────── async fn load_budget_state(&self, execution_id: &ExecutionId) -> BudgetState { diff --git a/sdk/python/jamjet/cli/main.py b/sdk/python/jamjet/cli/main.py index e24c6667..45c9e5c4 100644 --- a/sdk/python/jamjet/cli/main.py +++ b/sdk/python/jamjet/cli/main.py @@ -30,6 +30,7 @@ import enum import importlib.metadata import json +import math import re import sys import time @@ -1757,6 +1758,24 @@ async def _workers() -> None: asyncio.run(_workers()) +def _reported_usage(output: dict[str, Any]) -> dict[str, Any]: + """Token usage and cost a node reports in its output, for the engine to + charge to the agent the node acts for. Only values that are a spend are + forwarded: non-negative integers below 2**64 for tokens (the engine reads + a u64 and refuses the whole result otherwise), a non-negative finite number + for cost (a negative cost would refund an agent's budget, and the engine + refuses one). Booleans are not numbers here.""" + usage: dict[str, Any] = {} + for key in ("input_tokens", "output_tokens"): + v = output.get(key) + if isinstance(v, int) and not isinstance(v, bool) and 0 <= v < 2**64: + usage[key] = v + cost = output.get("cost_usd") + if isinstance(cost, (int, float)) and not isinstance(cost, bool) and math.isfinite(cost) and cost >= 0: + usage["cost_usd"] = float(cost) + return usage + + async def _worker_loop( client: Any, worker_id: str, @@ -1926,9 +1945,11 @@ async def _run_handler() -> Any: # Forward optional GenAI telemetry if the tool surfaced it. gen_ai_model_field: str | None = None finish_reason_field: str | None = None + usage: dict[str, Any] = {} if isinstance(output, dict): gen_ai_model_field = output.get("gen_ai_model") or None finish_reason_field = output.get("finish_reason") or None + usage = _reported_usage(output) try: await client.complete_work_item( @@ -1945,6 +1966,7 @@ async def _run_handler() -> Any: # 0/absent, keeping the unfenced fallback backward-compatible. lease_fence=lease["fence"], idempotency_key=idempotency_key, + **usage, ) except Exception as complete_exc: # A 409 means our echoed fence no longer matches: the lease was diff --git a/sdk/python/jamjet/client.py b/sdk/python/jamjet/client.py index 6fa1bc76..6eb15176 100644 --- a/sdk/python/jamjet/client.py +++ b/sdk/python/jamjet/client.py @@ -310,9 +310,17 @@ async def complete_work_item( finish_reason: str | None = None, lease_fence: int | None = None, idempotency_key: str | None = None, + input_tokens: int | None = None, + output_tokens: int | None = None, + cost_usd: float | None = None, ) -> None: """Mark a work item as completed and emit a NodeCompleted event. + ``input_tokens``, ``output_tokens`` and ``cost_usd`` are recorded on + the event and count toward the token and cost budgets of the agent + the node acts for. ``cost_usd`` must be a non-negative number (the + runtime answers 400 otherwise). + When ``lease_fence`` is provided (and non-zero), it is echoed back so the runtime can prove this worker still holds the lease: a stale fence (the item was reclaimed/replayed) is rejected with 409 and no NodeCompleted is @@ -331,6 +339,12 @@ async def complete_work_item( body["gen_ai_model"] = gen_ai_model if finish_reason is not None: body["finish_reason"] = finish_reason + if input_tokens is not None: + body["input_tokens"] = input_tokens + if output_tokens is not None: + body["output_tokens"] = output_tokens + if cost_usd is not None: + body["cost_usd"] = cost_usd if lease_fence: body["lease_fence"] = lease_fence if idempotency_key is not None: diff --git a/sdk/python/tests/test_worker.py b/sdk/python/tests/test_worker.py index 193fe2c6..d6124489 100644 --- a/sdk/python/tests/test_worker.py +++ b/sdk/python/tests/test_worker.py @@ -52,6 +52,20 @@ async def _with_genai(input: dict) -> dict: } +async def _with_usage(input: dict) -> dict: + """Return output that reports token usage and cost, plus values the + worker must not forward (negative, non-numeric, boolean).""" + return {"result": "answer", "input_tokens": 12, "output_tokens": 3, "cost_usd": 0.25} + + +async def _with_bad_usage(input: dict) -> dict: + return {"result": "answer", "input_tokens": True, "output_tokens": "3", "cost_usd": -1.0} + + +async def _with_bool_cost(input: dict) -> dict: + return {"result": "answer", "cost_usd": True, "input_tokens": 2**64} + + @tool(name="echo") async def echo_tool(text: str) -> str: """A @tool resolved by dispatch_tool_calls through the tools map, which @@ -86,9 +100,15 @@ async def complete_work_item( finish_reason: str | None = None, lease_fence: int | None = None, idempotency_key: str | None = None, + input_tokens: int | None = None, + output_tokens: int | None = None, + cost_usd: float | None = None, ) -> None: self.complete_calls.append( { + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "cost_usd": cost_usd, "item_id": item_id, "execution_id": execution_id, "node_id": node_id, @@ -274,6 +294,40 @@ async def test_worker_forwards_gen_ai_fields() -> None: assert call["output"]["result"] == "answer" +def _usage_item(function: str) -> dict: + return { + "id": "wi-usage", + "execution_id": "exec_usage", + "node_id": "n1", + "queue_type": "python_tool", + "payload": {"module": "tests.test_worker", "function": function, "input": {}}, + } + + +async def test_worker_forwards_usage_for_agent_budgets() -> None: + """Token usage and cost a tool reports reach the engine, which charges + them to an agent card's token and cost budgets.""" + stub = _StubClient(claimed_item=_usage_item("_with_usage")) + await _worker_loop(stub, "test-worker", ["python_tool"], once=True) + call = stub.complete_calls[0] + assert call["input_tokens"] == 12 + assert call["output_tokens"] == 3 + assert call["cost_usd"] == 0.25 + + +async def test_worker_drops_usage_that_is_not_a_spend() -> None: + stub = _StubClient(claimed_item=_usage_item("_with_bad_usage")) + await _worker_loop(stub, "test-worker", ["python_tool"], once=True) + call = stub.complete_calls[0] + assert call["input_tokens"] is None + assert call["output_tokens"] is None + assert call["cost_usd"] is None + stub = _StubClient(claimed_item=_usage_item("_with_bool_cost")) + await _worker_loop(stub, "test-worker", ["python_tool"], once=True) + assert stub.complete_calls[0]["cost_usd"] is None + assert stub.complete_calls[0]["input_tokens"] is None + + _DISPATCH_ITEM: dict = { "id": "wi-005", "execution_id": "exec_dispatch", @@ -623,3 +677,29 @@ async def _no_loop(*args: Any, lease_secs: int | None = None, **kwargs: Any) -> result = CliRunner().invoke(cli.app, ["worker", "--once", "--lease-secs", "600"]) assert result.exit_code == 0, result.output assert started == [600] + + +async def test_client_sends_usage_on_complete() -> None: + """`JamjetClient.complete_work_item` puts usage on the body the engine + records on `node_completed`, and leaves out what it was not given.""" + import json as _json + + from jamjet.client import JamjetClient + + bodies: list[dict] = [] + + def handler(request: httpx.Request) -> httpx.Response: + bodies.append(_json.loads(request.content)) + return httpx.Response(200, json={"completed": True}) + + client = JamjetClient("http://runtime") + client._client = httpx.AsyncClient(base_url="http://runtime", transport=httpx.MockTransport(handler)) + await client.complete_work_item( + "wi", None, None, {}, {}, input_tokens=12, output_tokens=3, cost_usd=0.25, lease_fence=7 + ) + await client.complete_work_item("wi", None, None, {}, {}, lease_fence=7) + assert bodies[0]["input_tokens"] == 12 + assert bodies[0]["output_tokens"] == 3 + assert bodies[0]["cost_usd"] == 0.25 + for key in ("input_tokens", "output_tokens", "cost_usd"): + assert key not in bodies[1]