From 829e18630c27642ed6c1e7175797922f8859657e Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Sun, 19 Jul 2026 08:30:23 +1200 Subject: [PATCH 1/4] feat: add llm reflector --- CHANGELOG.md | 24 ++ Cargo.toml | 2 + README.md | 1 + src/engine/bare.rs | 3 + src/engine/bare/dispatch.rs | 11 +- src/reflection.rs | 514 +++++++++++++++++++++-- src/reflection/llm.rs | 783 ++++++++++++++++++++++++++++++++++++ 7 files changed, 1308 insertions(+), 30 deletions(-) create mode 100644 src/reflection/llm.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 5e576dc..3a71663 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -65,9 +65,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. compiling a tool registry's schemas into a sampler grammar. - `grammar` feature flag (depends on `providers`): opt-in grammar / sampler support for the `Grammar` mode of `ToolConstraint`. +- `LlmReflector` (`reflection::llm` module): a `Reflector` that asks the + model to classify failed tool calls and suggest corrections via + `request_structured::`. First in-tree consumer of + `StructuredOutput`. Opt-in via `BareLoop::set_reflector`; the default + stays `NoopReflector`. Each analysed failure triggers one model + round-trip (see its rustdoc for the latency/cost note). +- `impl StructuredOutput for FailureAnalysis` (`reflection` module) with a + hand-written JSON Schema covering the 5 fields and the nested + `CorrectionType` snake_case enum. +- `schema_validation` feature flag (pulls `jsonschema` as an optional + dependency): when enabled, `LlmReflector` validates the model's + `Correction::modified_input` against the failing tool's `input_schema` + and returns `ReflectionError::Internal` on a mismatch. When disabled, + validation is skipped. ### Changed +- **Breaking:** `Reflector::analyze` gains a new `tool_schema: + Option<&ToolSchema>` parameter between `tool_input` and `context`. The + engine's call site now resolves the failing tool's schema from the + registry (passing `None` when the tool isn't found). Every `Reflector` + impl must add the new parameter; `NoopReflector` and the trait-doc + example have been updated. + Migration: add `_tool_schema: Option<&loopctl::tool::ToolSchema>` to + your `analyze` signature. Ignore it if your reflector does not validate + suggested corrections; otherwise use it to validate `modified_input` + before returning the analysis. - `OpenAiClient`, `AnthropicClient`, and `GeminiClient` now honor `RequestOptions::tool_constraint`. Under `Strict`, each tool's schema is tightened (recursive `additionalProperties: false` and full `required`); diff --git a/Cargo.toml b/Cargo.toml index 357f82f..b134fc2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,7 @@ parking_lot = "0.12" reqwest = { version = "0.13", default-features = false, features = ["json", "stream", "rustls"], optional = true } async-stream = { version = "0.3", optional = true } httpdate = { version = "1", optional = true } +jsonschema = { version = "0.30", optional = true } [dev-dependencies] tokio = { version = "1.52.3", features = ["rt-multi-thread", "macros", "sync", "time", "signal"] } @@ -53,6 +54,7 @@ grok = ["providers", "openai"] gemini = ["providers"] zai = ["providers", "anthropic"] grammar = ["providers"] +schema_validation = ["dep:jsonschema"] [[example]] name = "hello-cli" diff --git a/README.md b/README.md index 3a9485e..05a3b08 100644 --- a/README.md +++ b/README.md @@ -158,6 +158,7 @@ let agent = BareLoop::new( | `gemini` | No | `providers` | Google Gemini API client (`provider::gemini`) | | `zai` | No | `providers`, `anthropic` | Z.AI API client (Anthropic-compatible) | | `grammar` | No | `providers` | Tool-call grammar providers for grammar-aware samplers (vLLM `guided_json`); enables the `Grammar` mode of `ToolConstraint` | +| `schema_validation` | No | — | JSON Schema validation of `Correction::modified_input` in `LlmReflector` (pulls `jsonschema`); when off, validation is skipped | ## Architecture diff --git a/src/engine/bare.rs b/src/engine/bare.rs index 1aeac72..15ebe1c 100644 --- a/src/engine/bare.rs +++ b/src/engine/bare.rs @@ -2602,6 +2602,7 @@ mod tests { error: &str, tool_name: &str, _tool_input: &serde_json::Value, + _tool_schema: Option<&crate::tool::ToolSchema>, _context: &crate::reflection::ReflectionContext, ) -> Pin< Box< @@ -2913,6 +2914,7 @@ mod tests { error: &str, tool_name: &str, _tool_input: &serde_json::Value, + _tool_schema: Option<&crate::tool::ToolSchema>, _context: &crate::reflection::ReflectionContext, ) -> Pin< Box< @@ -3625,6 +3627,7 @@ mod tests { error: &str, tool_name: &str, _tool_input: &serde_json::Value, + _tool_schema: Option<&crate::tool::ToolSchema>, _context: &crate::reflection::ReflectionContext, ) -> Pin< Box< diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 7333694..f7ccfea 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -1049,12 +1049,18 @@ impl BareLoop { max_attempts: Self::MAX_RECOVERY_ATTEMPTS, }; + let tool_schema = self.tools.get(&tc.tool).map(crate::tool::Tool::schema); let Ok(analysis) = self .reflector - .analyze(&error_msg, &tc.tool, &tc.input, &context) + .analyze( + &error_msg, + &tc.tool, + &tc.input, + tool_schema.as_ref(), + &context, + ) .await else { - // Reflector failed — conservatively fail. return (RecoveryAction::Fail(error_msg), None); }; @@ -1555,6 +1561,7 @@ mod tests { error: &str, tool_name: &str, _tool_input: &Value, + _tool_schema: Option<&crate::tool::ToolSchema>, _context: &crate::reflection::ReflectionContext, ) -> Pin< Box< diff --git a/src/reflection.rs b/src/reflection.rs index fbddf4f..ce20362 100644 --- a/src/reflection.rs +++ b/src/reflection.rs @@ -37,6 +37,9 @@ pub mod backoff; pub use backoff::ExponentialBackoffRecovery; +pub mod llm; +pub use llm::LlmReflector; + use serde::{Deserialize, Serialize}; use std::fmt; use std::future::Future; @@ -54,6 +57,20 @@ use std::time::Duration; /// issue (e.g., a transient network blip) is more retryable than a /// `Critical` one (e.g., invalid API key). /// +/// # Ordering and comparison +/// +/// Derives [`Ord`], so variants can be compared directly: `Low < Medium < +/// High < Critical`. Strategies commonly write thresholds like +/// `analysis.severity >= FailureSeverity::High` to gate retries. +/// +/// # Serialization +/// +/// Serializes to and from the lowercase snake-case form of the variant +/// name (`"low"`, `"medium"`, `"high"`, `"critical"`) via +/// `#[serde(rename_all = "snake_case")]`. The same four strings are the +/// `enum` values in the [`FailureAnalysis`] JSON Schema, so model output +/// round-trips through deserialization without renaming. +/// /// # Example /// /// ```rust @@ -61,18 +78,54 @@ use std::time::Duration; /// /// assert!(FailureSeverity::Low < FailureSeverity::Critical); /// ``` +/// +/// [`FailureAnalysis`]: crate::reflection::FailureAnalysis #[derive( Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, serde::Serialize, serde::Deserialize, )] #[serde(rename_all = "snake_case")] pub enum FailureSeverity { /// Minor issue — a retry will likely fix it. + /// + /// Typical causes: a transient network blip, a momentary rate-limit, + /// or a tool that succeeded on a second attempt without any input + /// changes. Strategies usually retry immediately on `Low` severity. Low, + /// Moderate issue — may need a correction before retrying. + /// + /// The failure looks recoverable, but a bare retry is less likely to + /// succeed than at `Low`: the input may have a small mistake (a typo + /// in a path, a slightly wrong argument), or the tool's + /// preconditions may need a step first. Strategies commonly consult + /// [`FailureAnalysis::correction`] before retrying at this severity. + /// + /// [`FailureAnalysis::correction`]: crate::reflection::FailureAnalysis::correction Medium, + /// Serious issue — retrying without changes is unlikely to help. + /// + /// The call is fundamentally off: the wrong tool was chosen, the + /// argument type is incorrect, or the task itself is misframed. A + /// retry that doesn't first apply a [`Correction`] is probably + /// wasted. Strategies may still retry when a correction is supplied + /// and the attempt budget allows, but should treat `High` as a + /// signal to slow down rather than retry blindly. + /// + /// [`Correction`]: crate::reflection::Correction High, + /// Unrecoverable — the agent should stop or escalate. + /// + /// No correction can rescue this turn. Typical causes: an invalid + /// API key, a permissions failure the agent can't self-resolve, or a + /// bug in the tool itself. Strategies usually map `Critical` to + /// [`RecoveryAction::Fail`] (or + /// [`RecoveryAction::AskUser`] when interactive recovery is an + /// option) rather than retry. + /// + /// [`RecoveryAction::Fail`]: crate::reflection::RecoveryAction::Fail + /// [`RecoveryAction::AskUser`]: crate::reflection::RecoveryAction::AskUser Critical, } @@ -93,9 +146,24 @@ impl fmt::Display for FailureSeverity { /// Context provided to [`Reflector::analyze()`] describing the retry state. /// -/// Built by the framework before invoking the reflector. Contains -/// information about what the agent was doing and how many attempts -/// have been made so far. +/// Built by the framework before invoking the reflector. Carries what the +/// agent was trying to do and where it is in the retry budget, so a +/// reflector can factor both into its analysis — e.g., be more conservative +/// with suggested corrections on the last permitted attempt. +/// +/// # Lifecycle +/// +/// The engine constructs a fresh `ReflectionContext` for each failure +/// (see `recover_tool_error` in `engine/bare/dispatch.rs`) and passes it by +/// reference to [`Reflector::analyze`]. It is not stored across calls; +/// reflectors that want to track cross-failure history must keep their own +/// state. +/// +/// # Attempt indexing +/// +/// `attempt` is 0-indexed: the first try of a tool is `attempt = 0`. A +/// reflector rendering the value for a model prompt should add 1 (the +/// framework's built-in `LlmReflector` does this — "Attempt: 1 of N"). /// /// # Example /// @@ -108,14 +176,37 @@ impl fmt::Display for FailureSeverity { /// max_attempts: 5, /// }; /// assert_eq!(context.attempt, 2); +/// assert_eq!(context.max_attempts, 5); /// ``` #[derive(Debug, Clone, Default)] pub struct ReflectionContext { - /// What the agent was trying to accomplish. + /// What the agent was trying to accomplish when the failure occurred. + /// + /// Free-form text — typically the original user message, a summary of + /// the current step, or an empty string when the engine has no task + /// description to share. A reflector may include this in its prompt so + /// the model can reason about whether the failure is relevant to the + /// stated goal. pub task: String, - /// Current attempt number (0-indexed). + + /// Current attempt number for this tool call, 0-indexed. + /// + /// `0` is the first attempt; the framework increments this on each + /// retry within the recovery loop. Compare against + /// [`max_attempts`](Self::max_attempts) to know how much budget + /// remains. Render as `attempt + 1` when showing the value to a user + /// or model. pub attempt: u32, - /// Maximum attempts allowed before giving up. + + /// Maximum attempts allowed before the framework gives up on this + /// tool call. + /// + /// Set by the engine to its configured recovery ceiling + /// (`BareLoop::MAX_RECOVERY_ATTEMPTS`). When `attempt >= max_attempts`, + /// a [`RecoveryStrategy`] should typically return + /// [`RecoveryAction::Fail`] rather than schedule another retry. + /// + /// [`RecoveryAction::Fail`]: crate::reflection::RecoveryAction::Fail pub max_attempts: u32, } @@ -169,23 +260,84 @@ pub enum CorrectionType { /// A correction produced by the reflection system. /// /// When a tool call fails and reflection is enabled (via configuration), -/// the agent analyzes the error and produces a `Correction` that describes how to fix -/// the problem. The framework applies the correction and retries. +/// the agent analyzes the error and produces a `Correction` that describes +/// how to fix the problem. The framework applies the correction and +/// retries — `recover_tool_error` in `engine/bare/dispatch.rs` clones the +/// `Correction` out of the [`FailureAnalysis`] and feeds it back into the +/// retry loop. +/// +/// # Which fields apply when +/// +/// The fields are deliberately permissive (four `Option`s + one enum) so a +/// single shape covers every [`CorrectionType`] strategy. The convention is +/// that only the fields named by the variant's doc are meaningful for a +/// given `correction_type`; consumers should consult `correction_type` +/// first and read the relevant fields accordingly rather than treating +/// every `Some` as load-bearing. There is no runtime enforcement of this +/// pairing — a reflector that fills `modified_input` while declaring +/// `correction_type: Escalate` will not be rejected. /// /// # Serialization /// -/// Implements `Serialize` and `Deserialize` for persistence and observability. +/// Implements `Serialize` and `Deserialize` for persistence (e.g., writing +/// analyses to a session log) and so an [`LlmReflector`] can request it +/// back as a nested object inside a [`FailureAnalysis`] via the +/// [`StructuredOutput`] trait. +/// +/// [`LlmReflector`]: crate::reflection::LlmReflector +/// [`StructuredOutput`]: crate::structured::StructuredOutput +/// [`FailureAnalysis`]: crate::reflection::FailureAnalysis #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Correction { - /// See [`CorrectionType`] for available strategies. + /// Which fix strategy this correction represents. + /// + /// Drives which of the remaining fields the framework will actually + /// consume on retry. See [`CorrectionType`] for the five strategies + /// and the field each one pairs with. pub correction_type: CorrectionType, - /// Explains *what* went wrong and *how* the correction addresses it. + + /// Human-readable explanation of *what* went wrong and *how* this + /// correction addresses it. + /// + /// Always populated — even an [`CorrectionType::Escalate`] correction + /// carries a description so the framework can surface it to the user + /// or a higher-level handler. The string is free-form; some + /// reflectors include a short root-cause summary here in addition to + /// [`FailureAnalysis::root_cause`]. + /// + /// [`FailureAnalysis::root_cause`]: crate::reflection::FailureAnalysis::root_cause pub description: String, - /// Corrected JSON input when [`CorrectionType::InputFix`]. `None` otherwise. + + /// Corrected JSON input to pass on retry, when the fix is to change + /// the arguments rather than the tool. + /// + /// Set only for [`CorrectionType::InputFix`] corrections, where the + /// retry should substitute this value for the original + /// [`MessagePart::ToolCall::input`]. The shape must conform to the + /// tool's `input_schema`; an `LlmReflector` with the + /// `schema_validation` feature enabled will reject a non-conforming + /// suggestion before it reaches the retry. + /// + /// [`MessagePart::ToolCall::input`]: crate::message::MessagePart::ToolCall pub modified_input: Option, - /// Alternative tool name when [`CorrectionType::ToolChange`]. `None` otherwise. + + /// Name of a different tool to call instead, when the fix is to swap + /// tools rather than rewrite arguments. + /// + /// Set only for [`CorrectionType::ToolChange`] corrections. Must + /// match a tool name the registry knows; the retry loop will fail + /// normally if it does not. Leave `None` for all other correction + /// types. pub alternative_tool: Option, - /// Extra context or instructions to help avoid the same failure. + + /// Free-form instructions to help avoid the same failure on a future + /// turn. + /// + /// Used by [`CorrectionType::ApproachChange`] (high-level + /// re-strategizing) and optionally by other variants as a sidecar + /// note. Not consumed mechanically by the retry loop — it is + /// advisory, typically surfaced to a user or appended to context for + /// the next model turn. pub guidance: Option, } @@ -223,8 +375,26 @@ pub enum CorrectionResult { /// Result of analysing a failure via [`Reflector::analyze()`]. /// /// Describes what went wrong, how severe it is, whether it's worth -/// retrying, and optionally provides a [`Correction`] the agent can -/// apply before retrying. +/// retrying, and optionally provides a [`Correction`] the agent can apply +/// before retrying. Produced by a [`Reflector`] and consumed by a +/// [`RecoveryStrategy`] to decide the next action. +/// +/// # How the engine consumes it +/// +/// `recover_tool_error` in `engine/bare/dispatch.rs` calls +/// [`Reflector::analyze`] to get a `FailureAnalysis`, hands it to +/// [`RecoveryStrategy::decide`] for the action, and clones `correction` +/// out separately so the retry loop can apply it. If the reflector +/// itself errors, the engine conservatively fails the turn rather than +/// guessing — see [`ReflectionError`]. +/// +/// # Structured output +/// +/// Implements [`StructuredOutput`] so an [`LlmReflector`] can request it +/// back from a model via [`request_structured`] with a guaranteed-schema +/// response. The hand-written schema enumerates the five fields below, +/// pins [`FailureSeverity`] as a four-value string enum, and embeds the +/// [`Correction`] shape under `correction`. /// /// # Example /// @@ -240,20 +410,132 @@ pub enum CorrectionResult { /// }; /// assert!(analysis.is_recoverable); /// ``` +/// +/// [`StructuredOutput`]: crate::structured::StructuredOutput +/// [`LlmReflector`]: crate::reflection::LlmReflector +/// [`request_structured`]: crate::structured::request_structured +/// [`RecoveryStrategy::decide`]: crate::reflection::RecoveryStrategy::decide #[derive(Debug, Clone, serde::Serialize, serde::Deserialize)] pub struct FailureAnalysis { /// Whether the failure can be recovered from. + /// + /// The headline signal. A [`RecoveryStrategy`] typically maps `false` + /// to [`RecoveryAction::Fail`] (or [`RecoveryAction::Skip`] when the + /// step is optional) and `true` to [`RecoveryAction::Retry`] when the + /// attempt budget allows. This is independent of [`severity`](Self::severity): + /// a `Low`-severity failure may still be non-recoverable if the + /// reflector can't suggest a fix, and a `Critical`-severity failure + /// may technically be recoverable if a correction is supplied. + /// + /// [`RecoveryAction::Fail`]: crate::reflection::RecoveryAction::Fail + /// [`RecoveryAction::Skip`]: crate::reflection::RecoveryAction::Skip + /// [`RecoveryAction::Retry`]: crate::reflection::RecoveryAction::Retry pub is_recoverable: bool, + /// Description of what went wrong. + /// + /// Free-form text identifying the root cause — e.g., `"file not + /// found"`, `"401 Unauthorized"`, `"tool input did not match schema"`. + /// Often echoes the tool's error message but may be rephrased or + /// sharpened by the reflector. Surfaced to users in failure output + /// and included by an [`LlmReflector`] in its analysis of subsequent + /// failures. + /// + /// [`LlmReflector`]: crate::reflection::LlmReflector pub root_cause: String, + /// How severe the failure is. + /// + /// See [`FailureSeverity`] for the four levels and their typical + /// recovery implications. Strategies commonly use this to gate + /// retries — e.g., refusing to retry `Critical` even when + /// [`is_recoverable`](Self::is_recoverable) is `true`. pub severity: FailureSeverity, + /// Suggested correction for the agent to apply before retrying. + /// + /// `None` when the reflector has no concrete fix to suggest (the + /// failure is either non-recoverable, or recoverable by a bare retry + /// with no input changes). When `Some`, the engine clones the + /// [`Correction`] out and feeds it into the retry loop, which + /// substitutes `modified_input` / `alternative_tool` as the + /// correction directs. Reflectors that populate this field should + /// keep [`Correction::correction_type`] consistent with the fields + /// they fill. pub correction: Option, - /// Additional context (e.g., environment state at time of failure). + + /// Additional context the reflector wants the framework or a future + /// turn to see. + /// + /// Free-form; common uses are environment state at the time of + /// failure (cwd, available tools), the original task description, or + /// a short note about what the reflector considered. The framework + /// does not parse this field — it is advisory, typically logged + /// alongside the analysis or surfaced to a user when the turn fails. pub context: String, } +impl crate::structured::StructuredOutput for FailureAnalysis { + fn name() -> &'static str { + "failure_analysis" + } + + fn schema() -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { + "is_recoverable": {"type": "boolean"}, + "root_cause": {"type": "string"}, + "severity": { + "type": "string", + "enum": ["low", "medium", "high", "critical"] + }, + "correction": { + "anyOf": [ + { + "type": "object", + "properties": { + "correction_type": { + "type": "string", + "enum": [ + "input_fix", + "tool_change", + "prerequisite_fix", + "approach_change", + "escalate" + ] + }, + "description": {"type": "string"}, + "modified_input": {"type": ["object", "null"]}, + "alternative_tool": {"type": ["string", "null"]}, + "guidance": {"type": ["string", "null"]} + }, + "required": [ + "correction_type", + "description", + "modified_input", + "alternative_tool", + "guidance" + ], + "additionalProperties": false + }, + {"type": "null"} + ] + }, + "context": {"type": "string"} + }, + "required": [ + "is_recoverable", + "root_cause", + "severity", + "correction", + "context" + ], + "additionalProperties": false + }) + } +} + // =================================================== // ReflectionError // =================================================== @@ -286,8 +568,26 @@ pub enum ReflectionError { /// What the framework should do after a failure. /// -/// Produced by [`RecoveryStrategy::decide()`]. Each variant maps to a -/// different action in the agent loop. +/// Produced by [`RecoveryStrategy::decide()`] after the +/// [`Reflector::analyze()`] step. Each variant maps to a different action +/// in the agent loop — the strategy decides which one; the engine +/// executes it. +/// +/// # Variants in rough order of severity +/// +/// [`Retry`] is the most permissive (try again, optionally with a +/// correction); [`Skip`] continues past the failed step; [`AskUser`] +/// yields control for human input; [`Fail`] terminates the operation and +/// propagates the error. A typical strategy progresses through these as +/// the attempt budget drains: early attempts → `Retry`, late attempts → +/// `Fail` or `AskUser`. +/// +/// # Equality and ordering +/// +/// Derives [`Eq`] so two actions compare equal when their payloads +/// match (same `delay`, same string). Useful in tests that assert a +/// strategy picked a specific action; not meaningful for runtime +/// prioritization — there is no `Ord` impl, by design. /// /// # Example /// @@ -304,35 +604,76 @@ pub enum ReflectionError { /// let fail = RecoveryAction::Fail("unrecoverable".to_string()); /// assert!(fail.is_fail()); /// ``` +/// +/// [`Retry`]: Self::Retry +/// [`Skip`]: Self::Skip +/// [`AskUser`]: Self::AskUser +/// [`Fail`]: Self::Fail #[derive(Debug, Clone, PartialEq, Eq)] pub enum RecoveryAction { /// Retry the failed operation. /// - /// Wait for `delay` before retrying. If `correction` is `Some`, - /// apply it before the retry. + /// Wait for `delay` before retrying. If the + /// [`FailureAnalysis::correction`] that produced this action was + /// `Some`, the engine applies it (substituting + /// [`Correction::modified_input`] or swapping to + /// [`Correction::alternative_tool`]) before re-dispatching the tool + /// call. The strategy is responsible for choosing a sensible `delay` + /// — typically backoff that grows with the attempt number. + /// + /// [`FailureAnalysis::correction`]: crate::reflection::FailureAnalysis::correction + /// [`Correction::modified_input`]: crate::reflection::Correction::modified_input + /// [`Correction::alternative_tool`]: crate::reflection::Correction::alternative_tool Retry { - /// Duration to wait before retrying. + /// Duration to wait before retrying, chosen by the strategy. + /// + /// Strategies commonly use exponential backoff here. A `delay` + /// of zero is permitted (immediate retry) but should be reserved + /// for cases where the failure is known to be transient and the + /// retry is cheap. delay: Duration, }, - /// Skip the failed operation and continue. + /// Skip the failed operation and continue with the rest of the turn. /// - /// The framework should log the reason and move to the next turn. + /// The framework logs the reason and moves on rather than + /// retrying. Use when the step is optional or the failure is + /// non-fatal — e.g., a metric-emitting tool that the agent can + /// safely proceed without. The carried string is the human-readable + /// reason to log. Skip(String), - /// Ask the user for input. + /// Ask the user for input before continuing. /// - /// The framework should return control to the caller with a prompt. + /// Yields control to the caller with a prompt string. The framework + /// surfaces this in whatever interaction model it runs under + /// (headless: prints the prompt and waits on stdin; TUI: renders a + /// prompt and waits for input). Use when the strategy cannot decide + /// autonomously — e.g., a permission-style failure that a human + /// should adjudicate, or a `ToolChange` correction that requires a + /// choice between plausible alternatives. AskUser(String), /// Fail the operation and propagate the error. /// - /// No further retries — the framework should report this failure. + /// No further retries — the framework should report this failure and + /// stop the recovery loop. The carried string is the error message + /// to surface. Use when the failure is unrecoverable, when the + /// attempt budget is exhausted, or when a [`Reflector`] returned + /// [`ReflectionError::Internal`] (the engine conservatively maps + /// reflector failure to `Fail`). + /// + /// [`ReflectionError::Internal`]: crate::reflection::ReflectionError::Internal Fail(String), } impl RecoveryAction { /// Returns the retry delay, if this is a [`Retry`](Self::Retry) action. + /// + /// Lets callers branch on the wait without a full `match`. Returns + /// `None` for the other three variants, so a strategy can write + /// `action.delay().unwrap_or(Duration::ZERO)` to default a non-retry + /// action to immediate handling. #[must_use] pub fn delay(&self) -> Option { match self { @@ -342,24 +683,39 @@ impl RecoveryAction { } /// Returns `true` if this is a [`Retry`](Self::Retry) action. + /// + /// Convenience predicate; equivalent to + /// `matches!(action, RecoveryAction::Retry { .. })`. Useful in + /// engine code that gates on retry vs. non-retry without caring + /// about the delay. #[must_use] pub fn is_retry(&self) -> bool { matches!(self, Self::Retry { .. }) } /// Returns `true` if this is a [`Fail`](Self::Fail) action. + /// + /// Convenience predicate. Engine code commonly checks this to decide + /// whether to terminate the recovery loop and propagate the error. #[must_use] pub fn is_fail(&self) -> bool { matches!(self, Self::Fail(_)) } /// Returns `true` if this is a [`Skip`](Self::Skip) action. + /// + /// Convenience predicate. Use to distinguish "move on silently" from + /// the harder-failure variants when logging. #[must_use] pub fn is_skip(&self) -> bool { matches!(self, Self::Skip(_)) } /// Returns `true` if this is an [`AskUser`](Self::AskUser) action. + /// + /// Convenience predicate. Engine code checks this to know it must + /// yield control to the caller (headless: stdin; TUI: prompt) rather + /// than continue autonomously. #[must_use] pub fn is_ask_user(&self) -> bool { matches!(self, Self::AskUser(_)) @@ -409,6 +765,7 @@ impl fmt::Display for RecoveryAction { /// error: &str, /// tool_name: &str, /// _tool_input: &serde_json::Value, +/// _tool_schema: Option<&loopctl::tool::ToolSchema>, /// _context: &ReflectionContext, /// ) -> Pin> + Send + '_>> { /// let error = error.to_string(); @@ -443,6 +800,11 @@ pub trait Reflector: Send + Sync { /// - `error` — The error message from the failed call. /// - `tool_name` — Which tool was called. /// - `tool_input` — The JSON input that was passed. + /// - `tool_schema` — The schema of the tool that failed, when the + /// engine can resolve it. `None` if the tool isn't in the registry + /// or the schema is otherwise unavailable. Reflectors that want to + /// validate a suggested `modified_input` should skip validation + /// when this is `None`. /// - `context` — Retry state and task description. /// /// # Errors @@ -454,6 +816,7 @@ pub trait Reflector: Send + Sync { error: &str, tool_name: &str, tool_input: &serde_json::Value, + tool_schema: Option<&crate::tool::ToolSchema>, context: &ReflectionContext, ) -> Pin> + Send + '_>>; } @@ -553,6 +916,7 @@ impl Reflector for NoopReflector { error: &str, _tool_name: &str, _tool_input: &serde_json::Value, + _tool_schema: Option<&crate::tool::ToolSchema>, _context: &ReflectionContext, ) -> Pin> + Send + '_>> { let root_cause = error.to_string(); @@ -723,7 +1087,7 @@ mod tests { max_attempts: 3, }; let analysis = reflector - .analyze("some error", "tool", &serde_json::json!({}), &ctx) + .analyze("some error", "tool", &serde_json::json!({}), None, &ctx) .await .unwrap(); assert!(!analysis.is_recoverable); @@ -738,4 +1102,98 @@ mod tests { let debug = format!("{reflector:?}"); assert!(debug.contains("NoopReflector")); } + + // ---- StructuredOutput impl tests (1-4) ---- + + #[test] + fn failure_analysis_structured_round_trip() { + use crate::structured::StructuredOutput; + let v = serde_json::json!({ + "is_recoverable": true, + "root_cause": "timeout", + "severity": "low", + "correction": { + "correction_type": "input_fix", + "description": "fix the path", + "modified_input": {"path": "/x"}, + "alternative_tool": null, + "guidance": null + }, + "context": "open call" + }); + let analysis = FailureAnalysis::from_value(v).expect("should deserialize"); + assert!(analysis.is_recoverable); + assert_eq!(analysis.root_cause, "timeout"); + assert_eq!(analysis.severity, FailureSeverity::Low); + let correction = analysis.correction.expect("correction"); + assert_eq!(correction.description, "fix the path"); + assert_eq!( + correction.modified_input, + Some(serde_json::json!({"path": "/x"})) + ); + } + + #[test] + fn failure_analysis_schema_is_valid_json() { + let schema = ::schema(); + let obj = schema.as_object().expect("schema must be a JSON object"); + // Five top-level properties. + assert_eq!(obj["type"], "object"); + let required = obj["required"] + .as_array() + .expect("required must be an array"); + assert_eq!(required.len(), 5); + } + + #[test] + fn failure_analysis_schema_correction_type_enum() { + let schema = ::schema(); + let enum_values = schema + .pointer("/properties/correction/anyOf/0/properties/correction_type/enum") + .expect("correction_type enum must be present") + .as_array() + .expect("enum must be an array"); + let values: Vec<&str> = enum_values.iter().map(|v| v.as_str().unwrap()).collect(); + // Pin the 5 snake_case variants matching CorrectionType's serde rename. + assert_eq!( + values, + vec![ + "input_fix", + "tool_change", + "prerequisite_fix", + "approach_change", + "escalate" + ] + ); + } + + #[test] + fn failure_analysis_name_is_stable() { + use crate::structured::StructuredOutput; + let name = FailureAnalysis::name(); + assert_eq!(name, "failure_analysis"); + assert!( + name.chars() + .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'), + "name must match ^[a-zA-Z0-9_-]+$: {name}" + ); + } + + #[tokio::test] + async fn noop_reflector_accepts_new_signature() { + // Pins the breaking trait change: NoopReflector implements the + // 5-arg analyze and its semantics are unchanged. + let reflector = NoopReflector; + let ctx = ReflectionContext { + task: "t".to_string(), + attempt: 0, + max_attempts: 1, + }; + let analysis = reflector + .analyze("err", "tool", &serde_json::json!({}), None, &ctx) + .await + .unwrap(); + assert!(!analysis.is_recoverable); + assert_eq!(analysis.severity, FailureSeverity::Medium); + } } diff --git a/src/reflection/llm.rs b/src/reflection/llm.rs new file mode 100644 index 0000000..9d93435 --- /dev/null +++ b/src/reflection/llm.rs @@ -0,0 +1,783 @@ +//! LLM-powered failure reflection. +//! +//! [`LlmReflector`] asks the model to classify a failed tool call and +//! suggest a correction, returning a typed [`FailureAnalysis`] via the +//! [`StructuredOutput`](crate::structured::StructuredOutput) trait. This is +//! the first in-tree consumer of structured output, and it makes tool-error +//! recovery semantically intelligent instead of heuristic. +//! +//! # Construction +//! +//! ```rust,ignore +//! use loopctl::reflection::LlmReflector; +//! use std::sync::Arc; +//! +//! let reflector = LlmReflector::new(client); +//! agent.set_reflector(Arc::new(reflector)); +//! ``` +//! +//! # Latency and cost +//! +//! Each analysed tool failure triggers one model round-trip. The reflector +//! is opt-in — the framework's default reflector is +//! [`NoopReflector`](super::NoopReflector), which performs no I/O. Only +//! callers that explicitly install an `LlmReflector` pay the per-failure +//! round-trip cost. + +use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; + +use crate::api::ApiClient; +use crate::message::Message; +use crate::reflection::{FailureAnalysis, ReflectionContext, ReflectionError, Reflector}; +use crate::structured::request_structured; +use crate::tool::ToolSchema; + +/// The default system prompt used when the caller does not supply one via +/// [`LlmReflector::with_system_prompt`]. +const DEFAULT_PROMPT: &str = "\ +You are a tool-failure analyst for an LLM agent loop. Given a tool name, \ +the JSON input that was passed, the tool's input schema (if provided), and \ +the error that resulted, classify the failure and suggest a correction.\n\ +\n\ +Respond with a single JSON object matching this exact shape:\n\ +{\n\ + \"is_recoverable\": ,\n\ + \"root_cause\": ,\n\ + \"severity\": \"low\" | \"medium\" | \"high\" | \"critical\",\n\ + \"correction\": {\n\ + \"correction_type\": \"input_fix\" | \"tool_change\" | \"prerequisite_fix\" | \"approach_change\" | \"escalate\",\n\ + \"description\": ,\n\ + \"modified_input\": ,\n\ + \"alternative_tool\": ,\n\ + \"guidance\": \n\ + } | null,\n\ + \"context\": \n\ +}\n\ +\n\ +Only set \"modified_input\" when \"correction_type\" is \"input_fix\", and only \ +when you can produce an input that conforms to the tool's schema. Prefer \ +\"is_recoverable\": false over inventing a correction."; + +/// A [`Reflector`] that asks an LLM to classify tool failures and suggest +/// corrections. +/// +/// Holds a shared API client handle and an optional system-prompt override. +/// On each call to [`analyze`](Reflector::analyze) it builds a single user +/// message describing the failure, calls +/// [`request_structured::`](crate::structured::request_structured), +/// and returns the typed analysis. +/// +/// When the `schema_validation` feature is enabled and the engine supplies +/// the failing tool's schema, the reflector validates the model's +/// `Correction::modified_input` against that schema and returns +/// [`ReflectionError::Internal`] on a mismatch. +pub struct LlmReflector { + /// The shared API client used to make the structured-output call on each + /// failure analysis. + /// + /// Held as an `Arc` so the reflector can outlive the + /// borrow that [`Reflector::analyze`] is called under: the returned + /// future borrows `&self`, but `request_structured` takes `&dyn + /// ApiClient` — the `Arc` lets the future reach the client without + /// capturing `&self`'s borrow chain. The trait's `'static + Send + + /// Sync` bound rules out borrowing the client by lifetime. + client: Arc, + + /// The system prompt sent with every analysis request, instructing the + /// model to return JSON conforming to the [`FailureAnalysis`] schema. + /// + /// Defaults to [`DEFAULT_PROMPT`] when the reflector is constructed via + /// [`new`](Self::new); replace it with + /// [`with_system_prompt`](Self::with_system_prompt). Cloned per + /// `analyze` call so the future owns its copy. + /// + /// [`FailureAnalysis`]: crate::reflection::FailureAnalysis + system_prompt: String, +} + +impl LlmReflector { + /// Construct an `LlmReflector` backed by the given shared client. + /// + /// Uses the module's default system prompt. Override with + /// [`with_system_prompt`](Self::with_system_prompt). + #[must_use] + pub fn new(client: Arc) -> Self { + Self { + client, + system_prompt: DEFAULT_PROMPT.to_string(), + } + } + + /// Replace the default system prompt with a caller-supplied one. + /// + /// Builder-style; consumes and returns `self`. The prompt should still + /// instruct the model to return JSON conforming to the + /// `FailureAnalysis` schema, since the response is parsed via + /// [`StructuredOutput`](crate::structured::StructuredOutput). + #[must_use] + pub fn with_system_prompt(mut self, prompt: impl Into) -> Self { + self.system_prompt = prompt.into(); + self + } +} + +impl std::fmt::Debug for LlmReflector { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("LlmReflector") + .field("client", &"") + .field( + "system_prompt", + &format!("{} chars", self.system_prompt.len()), + ) + .finish() + } +} + +impl Reflector for LlmReflector { + fn analyze( + &self, + error: &str, + tool_name: &str, + tool_input: &serde_json::Value, + tool_schema: Option<&ToolSchema>, + context: &ReflectionContext, + ) -> Pin> + Send + '_>> { + let user_message = build_user_message(error, tool_name, tool_input, context); + let system = self.system_prompt.clone(); + let client = std::sync::Arc::clone(&self.client); + let schema_value = tool_schema.map(|s| s.input_schema.clone()); + + Box::pin(async move { + let analysis = request_structured::( + &*client, + vec![Message::user(user_message)], + Some(system), + ) + .await + .map_err(|e| ReflectionError::Internal(format!("{e}")))?; + + validate_modified_input(&analysis, schema_value.as_ref())?; + Ok(analysis) + }) + } +} + +/// Build the single user message describing the failure. +/// +/// Carries the error message, the tool name, the serialized tool input, +/// and the task description from the reflection context so the model has +/// everything it needs to produce a typed `FailureAnalysis`. +fn build_user_message( + error: &str, + tool_name: &str, + tool_input: &serde_json::Value, + context: &ReflectionContext, +) -> String { + format!( + "Tool: {tool_name}\n\ + Input: {tool_input}\n\ + Error: {error}\n\ + Task: {task}\n\ + Attempt: {attempt} of {max}", + tool_name = tool_name, + tool_input = tool_input, + error = error, + task = context.task, + attempt = context.attempt.saturating_add(1), + max = context.max_attempts, + ) +} + +/// Validate the model's suggested `modified_input` against the failing +/// tool's input schema. +/// +/// Returns `Ok(())` when there is nothing to validate (no correction, no +/// `modified_input`, or no schema supplied) — these are all legitimate +/// "skip validation" cases. When a schema is supplied and validation is +/// enabled, returns `Ok(())` on a match or +/// [`ReflectionError::Internal`](crate::reflection::ReflectionError::Internal) +/// on a mismatch. +/// +/// # Feature gating +/// +/// The schema check itself only runs when the `schema_validation` feature +/// is enabled. Without it, this function is always `Ok(())` once the +/// early-return skips have passed — the analysis is returned unchanged. +/// Callers who want validation must enable the feature. +/// +/// # Errors +/// +/// See the body — returns `ReflectionError::Internal` only under +/// `schema_validation` + supplied schema + non-conforming `modified_input`. +fn validate_modified_input( + analysis: &FailureAnalysis, + tool_schema: Option<&serde_json::Value>, +) -> Result<(), ReflectionError> { + #[cfg(not(feature = "schema_validation"))] + { + let _ = (modified_input, schema); + } + + let Some(correction) = &analysis.correction else { + return Ok(()); + }; + let Some(modified_input) = &correction.modified_input else { + return Ok(()); + }; + let Some(schema) = tool_schema else { + return Ok(()); + }; + + #[cfg(feature = "schema_validation")] + { + if !jsonschema::is_valid(schema, modified_input) { + return Err(ReflectionError::Internal( + "corrected input does not match the tool's schema".to_string(), + )); + } + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::api::error::ApiError; + use crate::message::MessagePart; + use crate::reflection::{Correction, CorrectionType, FailureSeverity}; + use crate::structured::RequestOptions; + use crate::tool::ToolSchema; + use std::future::Future; + use std::pin::Pin; + use std::sync::Mutex; + + // ---- Mock clients ---- + + /// A capture of what the reflector sent to the client, plus the canned + /// response to return. + #[derive(Clone)] + struct Captured { + system: Option, + user: Option, + } + + /// Mock returning a canned `FailureAnalysis`-shaped JSON value. + struct CannedMock { + response: serde_json::Value, + captured: Arc>>, + } + + impl ApiClient for CannedMock { + fn model(&self) -> String { + "test".to_string() + } + fn stream_messages( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn futures::Stream> + + Send + + 'static, + >, + > { + Box::pin(futures::stream::empty()) + } + fn create_message( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin> + Send + '_>> + { + Box::pin(async { Ok(serde_json::json!({})) }) + } + fn create_message_with_options( + &self, + messages: Vec, + system: Option, + _tools: Option>, + _options: RequestOptions, + ) -> Pin> + Send + '_>> + { + let user = messages.first().and_then(|m| { + m.parts.first().and_then(|p| match p { + MessagePart::Text { text } => Some(text.clone()), + _ => None, + }) + }); + *self.captured.lock().unwrap() = Some(Captured { system, user }); + let response = self.response.clone(); + Box::pin(async move { Ok(response) }) + } + fn extract_structured(&self, raw: &serde_json::Value) -> serde_json::Value { + raw.clone() + } + } + + /// Mock returning a structured-output API error. + struct ErrorMock; + impl ApiClient for ErrorMock { + fn model(&self) -> String { + "test".to_string() + } + fn stream_messages( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin< + Box< + dyn futures::Stream> + + Send + + 'static, + >, + > { + Box::pin(futures::stream::empty()) + } + fn create_message( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + ) -> Pin> + Send + '_>> + { + Box::pin(async { Ok(serde_json::json!({})) }) + } + fn create_message_with_options( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + _options: RequestOptions, + ) -> Pin> + Send + '_>> + { + Box::pin(async { Err(ApiError::http("upstream 500".to_string())) }) + } + } + + /// Mock returning prose (a valid JSON string, but not an object matching + /// the FailureAnalysis schema) to exercise the deserialize-error path. + struct ProseMock(ErrorMock); + impl ApiClient for ProseMock { + fn model(&self) -> String { + self.0.model() + } + fn stream_messages( + &self, + m: Vec, + s: Option, + t: Option>, + ) -> Pin< + Box< + dyn futures::Stream> + + Send + + 'static, + >, + > { + self.0.stream_messages(m, s, t) + } + fn create_message( + &self, + m: Vec, + s: Option, + t: Option>, + ) -> Pin> + Send + '_>> + { + self.0.create_message(m, s, t) + } + fn create_message_with_options( + &self, + _messages: Vec, + _system: Option, + _tools: Option>, + _options: RequestOptions, + ) -> Pin> + Send + '_>> + { + Box::pin(async { Ok(serde_json::json!("I cannot produce that.")) }) + } + fn extract_structured(&self, raw: &serde_json::Value) -> serde_json::Value { + raw.clone() + } + } + + fn fixture_analysis() -> serde_json::Value { + serde_json::json!({ + "is_recoverable": true, + "root_cause": "file not found", + "severity": "medium", + "correction": { + "correction_type": "input_fix", + "description": "fix the path", + "modified_input": {"path": "/correct/path"}, + "alternative_tool": null, + "guidance": null + }, + "context": "open() call" + }) + } + + fn ctx() -> ReflectionContext { + ReflectionContext { + task: "fix the bug".to_string(), + attempt: 0, + max_attempts: 3, + } + } + + #[tokio::test] + async fn llm_reflector_returns_typed_analysis() { + let captured = Arc::new(Mutex::new(None)); + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: captured.clone(), + }); + let reflector = LlmReflector::new(client); + let analysis = reflector + .analyze( + "open: file not found", + "read", + &serde_json::json!({"path": "/wrong"}), + None, + &ctx(), + ) + .await + .expect("should succeed"); + assert!(analysis.is_recoverable); + assert_eq!(analysis.root_cause, "file not found"); + assert_eq!(analysis.severity, FailureSeverity::Medium); + let correction = analysis.correction.expect("correction present"); + assert_eq!(correction.correction_type, CorrectionType::InputFix); + assert_eq!(correction.description, "fix the path"); + assert_eq!( + correction.modified_input, + Some(serde_json::json!({"path": "/correct/path"})) + ); + } + + #[tokio::test] + async fn llm_reflector_api_error_maps_to_internal() { + let client: Arc = Arc::new(ErrorMock); + let reflector = LlmReflector::new(client); + let err = reflector + .analyze("e", "t", &serde_json::json!({}), None, &ctx()) + .await + .expect_err("should fail"); + assert!( + matches!(err, ReflectionError::Internal(ref msg) if msg.contains("upstream 500")), + "got: {err:?}" + ); + } + + #[tokio::test] + async fn llm_reflector_prose_maps_to_internal() { + let client: Arc = Arc::new(ProseMock(ErrorMock)); + let reflector = LlmReflector::new(client); + let err = reflector + .analyze("e", "t", &serde_json::json!({}), None, &ctx()) + .await + .expect_err("should fail"); + // Prose → deserialize error → Internal. + assert!(matches!(err, ReflectionError::Internal(_)), "got: {err:?}"); + } + + #[tokio::test] + async fn llm_reflector_uses_default_prompt() { + let captured = Arc::new(Mutex::new(None)); + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: captured.clone(), + }); + let reflector = LlmReflector::new(client); + let result = reflector + .analyze("e", "t", &serde_json::json!({}), None, &ctx()) + .await; + assert!(result.is_ok(), "analyze should succeed: {:?}", result.err()); + let cap = captured.lock().unwrap().clone().expect("captured"); + assert_eq!(cap.system.as_deref(), Some(DEFAULT_PROMPT)); + } + + #[tokio::test] + async fn llm_reflector_with_system_prompt_overrides() { + let captured = Arc::new(Mutex::new(None)); + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: captured.clone(), + }); + let reflector = LlmReflector::new(client).with_system_prompt("custom analyst prompt"); + let result = reflector + .analyze("e", "t", &serde_json::json!({}), None, &ctx()) + .await; + assert!(result.is_ok(), "analyze should succeed: {:?}", result.err()); + let cap = captured.lock().unwrap().clone().expect("captured"); + assert_eq!(cap.system.as_deref(), Some("custom analyst prompt")); + } + + #[tokio::test] + async fn llm_reflector_user_message_carries_all_fields() { + let captured = Arc::new(Mutex::new(None)); + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: captured.clone(), + }); + let reflector = LlmReflector::new(client); + let result = reflector + .analyze( + "the error text", + "the_tool", + &serde_json::json!({"k": "v"}), + None, + &ctx(), + ) + .await; + assert!(result.is_ok(), "analyze should succeed: {:?}", result.err()); + let cap = captured.lock().unwrap().clone().expect("captured"); + let user = cap.user.expect("user message"); + assert!(user.contains("the error text"), "user: {user}"); + assert!(user.contains("the_tool"), "user: {user}"); + assert!(user.contains("\"k\""), "user: {user}"); + assert!(user.contains("fix the bug"), "user: {user}"); + } + + #[cfg(feature = "schema_validation")] + #[tokio::test] + async fn llm_reflector_validates_modified_input_pass() { + let schema = serde_json::json!({ + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"], + "additionalProperties": false + }); + let tool_schema = ToolSchema { + tool: "read".into(), + description: "Read a file".into(), + input_schema: schema, + }; + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: Arc::new(Mutex::new(None)), + }); + let reflector = LlmReflector::new(client); + // modified_input is {"path": "/correct/path"}, which matches. + let result = reflector + .analyze( + "e", + "read", + &serde_json::json!({}), + Some(&tool_schema), + &ctx(), + ) + .await; + assert!( + result.is_ok(), + "valid modified_input should pass: {:?}", + result.err() + ); + } + + #[cfg(feature = "schema_validation")] + #[tokio::test] + async fn llm_reflector_validates_modified_input_fail() { + // Tool schema requires {path: string}, but the fixture's + // modified_input is {"path": "/correct/path"} — that matches. + // Construct a schema it *doesn't* match (requires a number). + let schema = serde_json::json!({ + "type": "object", + "properties": {"n": {"type": "number"}}, + "required": ["n"], + "additionalProperties": false + }); + let tool_schema = ToolSchema { + tool: "calc".into(), + description: "Calculate".into(), + input_schema: schema, + }; + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: Arc::new(Mutex::new(None)), + }); + let reflector = LlmReflector::new(client); + let err = reflector + .analyze( + "e", + "calc", + &serde_json::json!({}), + Some(&tool_schema), + &ctx(), + ) + .await + .expect_err("invalid modified_input should fail"); + assert!( + matches!(err, ReflectionError::Internal(ref m) if m.contains("does not match")), + "got: {err:?}" + ); + } + + #[cfg(feature = "schema_validation")] + #[tokio::test] + async fn llm_reflector_skips_validation_when_no_schema() { + // tool_schema = None; even if modified_input is weird, we return Ok. + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: Arc::new(Mutex::new(None)), + }); + let reflector = LlmReflector::new(client); + let result = reflector + .analyze("e", "t", &serde_json::json!({}), None, &ctx()) + .await; + assert!(result.is_ok()); + } + + #[test] + fn validate_modified_input_noop_without_correction() { + let analysis = FailureAnalysis { + is_recoverable: false, + root_cause: "x".to_string(), + severity: FailureSeverity::Low, + correction: None, + context: String::new(), + }; + // With a schema but no correction → Ok. + assert!(validate_modified_input(&analysis, Some(&serde_json::json!({}))).is_ok()); + } + + #[test] + fn validate_modified_input_noop_without_modified_input() { + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "x".to_string(), + severity: FailureSeverity::Low, + correction: Some(Correction { + correction_type: CorrectionType::ApproachChange, + description: "no modified input".to_string(), + modified_input: None, + alternative_tool: None, + guidance: None, + }), + context: String::new(), + }; + assert!(validate_modified_input(&analysis, Some(&serde_json::json!({}))).is_ok()); + } + + #[test] + fn validate_modified_input_skips_when_no_schema() { + // Even with a modified_input present, no schema → Ok (the engine + // passes None when the tool isn't in the registry). + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "x".to_string(), + severity: FailureSeverity::Low, + correction: Some(Correction { + correction_type: CorrectionType::InputFix, + description: "fix".to_string(), + modified_input: Some(serde_json::json!({"anything": true})), + alternative_tool: None, + guidance: None, + }), + context: String::new(), + }; + assert!(validate_modified_input(&analysis, None).is_ok()); + } + + #[test] + fn validate_modified_input_passes_without_feature() { + // Without `schema_validation`, even a deliberately-mismatched + // modified_input must return Ok (validation is gated off). + let analysis = FailureAnalysis { + is_recoverable: true, + root_cause: "x".to_string(), + severity: FailureSeverity::Low, + correction: Some(Correction { + correction_type: CorrectionType::InputFix, + description: "fix".to_string(), + // Mismatched shape — schema requires a number, input is a string. + modified_input: Some(serde_json::json!({"wrong": "shape"})), + alternative_tool: None, + guidance: None, + }), + context: String::new(), + }; + let schema = serde_json::json!({ + "type": "object", + "properties": {"n": {"type": "number"}}, + "required": ["n"], + "additionalProperties": false + }); + let result = validate_modified_input(&analysis, Some(&schema)); + // Under `schema_validation` this fails; without it, it's Ok. Pin + // the no-feature behavior here; the under-feature behavior is + // covered by llm_reflector_validates_modified_input_fail. + #[cfg(feature = "schema_validation")] + assert!( + matches!(result, Err(ReflectionError::Internal(_))), + "with schema_validation the mismatch should fail: {result:?}" + ); + #[cfg(not(feature = "schema_validation"))] + assert!( + result.is_ok(), + "without schema_validation validation is skipped: {result:?}" + ); + } + + // ---- build_user_message direct unit tests ---- + + #[test] + fn build_user_message_contains_all_fields() { + let msg = build_user_message( + "the error", + "the_tool", + &serde_json::json!({"k": "v"}), + &ReflectionContext { + task: "the task".to_string(), + attempt: 1, + max_attempts: 4, + }, + ); + assert!(msg.contains("the_tool"), "tool name missing: {msg}"); + assert!(msg.contains("the error"), "error missing: {msg}"); + assert!(msg.contains("\"k\""), "input missing: {msg}"); + assert!(msg.contains("the task"), "task missing: {msg}"); + } + + #[test] + fn build_user_message_attempt_is_one_indexed() { + // attempt is 0-indexed in ReflectionContext; the message should + // render it 1-indexed ("Attempt: 1 of N" for attempt=0). + let msg = build_user_message( + "e", + "t", + &serde_json::json!({}), + &ReflectionContext { + task: "x".to_string(), + attempt: 0, + max_attempts: 3, + }, + ); + assert!( + msg.contains("Attempt: 1 of 3"), + "expected 1-indexed attempt in: {msg}" + ); + } + + #[test] + fn build_user_message_saturates_attempt_overflow() { + // u32::MAX + 1 must saturate rather than panic. + let msg = build_user_message( + "e", + "t", + &serde_json::json!({}), + &ReflectionContext { + task: "x".to_string(), + attempt: u32::MAX, + max_attempts: u32::MAX, + }, + ); + // Should contain u32::MAX (saturated, not overflowed). + assert!(msg.contains(&u32::MAX.to_string())); + } +} From 1658023073d505427ec1b2a376666d0b0daf3b5f Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Sun, 19 Jul 2026 14:03:27 +1200 Subject: [PATCH 2/4] chore: use resolved tool name --- CHANGELOG.md | 2 +- src/engine/bare/dispatch.rs | 12 ++-- src/reflection.rs | 10 ++- src/reflection/llm.rs | 127 +++++++++++++++++++++++++++++++++--- 4 files changed, 129 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a71663..e673cf5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -69,7 +69,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/2.0.0. model to classify failed tool calls and suggest corrections via `request_structured::`. First in-tree consumer of `StructuredOutput`. Opt-in via `BareLoop::set_reflector`; the default - stays `NoopReflector`. Each analysed failure triggers one model + stays `NoopReflector`. Each analyzed failure triggers one model round-trip (see its rustdoc for the latency/cost note). - `impl StructuredOutput for FailureAnalysis` (`reflection` module) with a hand-written JSON Schema covering the 5 fields and the nested diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index f7ccfea..4b3bb73 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -1049,7 +1049,13 @@ impl BareLoop { max_attempts: Self::MAX_RECOVERY_ATTEMPTS, }; - let tool_schema = self.tools.get(&tc.tool).map(crate::tool::Tool::schema); + // Look up the schema under the name a routing middleware may have + // redirected the call to; `tc.tool` (passed below as `tool_name`) + // stays as the originally requested name. + let tool_schema = self + .tools + .get(&result.resolved_tool_name) + .map(crate::tool::Tool::schema); let Ok(analysis) = self .reflector .analyze( @@ -1236,8 +1242,6 @@ mod tests { } } - // ----- ToolDependencyGraph unit tests ----- - fn make_call(id: &str, tool: &str, input: Value) -> ToolCall { ToolCall { id: id.into(), @@ -1409,8 +1413,6 @@ mod tests { assert_eq!(plan.waves[0], vec![0, 1, 2, 3]); } - // ----- dispatch_tools_parallel integration tests ----- - fn make_parallel_loop(tools: ToolRegistry) -> BareLoop { let mut config = LoopConfig::default(); config.parallel_tool_dispatch.mode = crate::config::ParallelMode::Parallel; diff --git a/src/reflection.rs b/src/reflection.rs index ce20362..567f066 100644 --- a/src/reflection.rs +++ b/src/reflection.rs @@ -372,7 +372,7 @@ pub enum CorrectionResult { // FailureAnalysis // =================================================== -/// Result of analysing a failure via [`Reflector::analyze()`]. +/// Result of analyzing a failure via [`Reflector::analyze()`]. /// /// Describes what went wrong, how severe it is, whether it's worth /// retrying, and optionally provides a [`Correction`] the agent can apply @@ -547,7 +547,7 @@ impl crate::structured::StructuredOutput for FailureAnalysis { /// during analysis. #[derive(Debug, thiserror::Error)] pub enum ReflectionError { - /// The reflector opted out of analysing this failure. + /// The reflector opted out of analyzing this failure. /// /// The framework should fall back to its default error handling. #[error("reflection skipped: {0}")] @@ -555,7 +555,7 @@ pub enum ReflectionError { /// The reflector itself encountered an error. /// - /// Distinct from the tool failure being analysed — it means + /// Distinct from the tool failure being analyzed — it means /// the reflector's own logic broke (e.g., an LLM call for /// summarisation failed). #[error("reflection internal error: {0}")] @@ -825,7 +825,7 @@ pub trait Reflector: Send + Sync { // RecoveryStrategy trait // =================================================== -/// Decides what to do after a failure has been analysed. +/// Decides what to do after a failure has been analyzed. /// /// Takes a [`FailureAnalysis`] and the current retry state, returns a /// [`RecoveryAction`]. Implement this trait to provide custom recovery @@ -1103,8 +1103,6 @@ mod tests { assert!(debug.contains("NoopReflector")); } - // ---- StructuredOutput impl tests (1-4) ---- - #[test] fn failure_analysis_structured_round_trip() { use crate::structured::StructuredOutput; diff --git a/src/reflection/llm.rs b/src/reflection/llm.rs index 9d93435..0217559 100644 --- a/src/reflection/llm.rs +++ b/src/reflection/llm.rs @@ -18,7 +18,7 @@ //! //! # Latency and cost //! -//! Each analysed tool failure triggers one model round-trip. The reflector +//! Each analyzed tool failure triggers one model round-trip. The reflector //! is opt-in — the framework's default reflector is //! [`NoopReflector`](super::NoopReflector), which performs no I/O. Only //! callers that explicitly install an `LlmReflector` pay the per-failure @@ -144,10 +144,11 @@ impl Reflector for LlmReflector { tool_schema: Option<&ToolSchema>, context: &ReflectionContext, ) -> Pin> + Send + '_>> { - let user_message = build_user_message(error, tool_name, tool_input, context); + let schema_value = tool_schema.map(|s| s.input_schema.clone()); + let user_message = + build_user_message(error, tool_name, tool_input, schema_value.as_ref(), context); let system = self.system_prompt.clone(); let client = std::sync::Arc::clone(&self.client); - let schema_value = tool_schema.map(|s| s.input_schema.clone()); Box::pin(async move { let analysis = request_structured::( @@ -167,22 +168,33 @@ impl Reflector for LlmReflector { /// Build the single user message describing the failure. /// /// Carries the error message, the tool name, the serialized tool input, -/// and the task description from the reflection context so the model has -/// everything it needs to produce a typed `FailureAnalysis`. +/// the tool's input schema (when available), and the task description +/// from the reflection context so the model has everything it needs to +/// produce a typed `FailureAnalysis`. The schema block is omitted +/// entirely when `tool_schema` is `None` (the engine passes `None` when +/// the tool isn't in the registry) so the model isn't shown a misleading +/// empty placeholder. fn build_user_message( error: &str, tool_name: &str, tool_input: &serde_json::Value, + tool_schema: Option<&serde_json::Value>, context: &ReflectionContext, ) -> String { + let schema_line = match tool_schema { + Some(schema) => format!("Schema: {schema}\n"), + None => String::new(), + }; format!( "Tool: {tool_name}\n\ Input: {tool_input}\n\ + {schema_line}\ Error: {error}\n\ Task: {task}\n\ Attempt: {attempt} of {max}", tool_name = tool_name, tool_input = tool_input, + schema_line = schema_line, error = error, task = context.task, attempt = context.attempt.saturating_add(1), @@ -215,9 +227,14 @@ fn validate_modified_input( analysis: &FailureAnalysis, tool_schema: Option<&serde_json::Value>, ) -> Result<(), ReflectionError> { + // Without the schema_validation feature, validation never runs — bail + // out early so we don't bind `modified_input` / `schema` only to drop + // them on the floor. The signature is unchanged; the early return + // keeps the function a no-op under the default feature set. #[cfg(not(feature = "schema_validation"))] { - let _ = (modified_input, schema); + let _ = (analysis, tool_schema); + return Ok(()); } let Some(correction) = &analysis.correction else { @@ -227,6 +244,8 @@ fn validate_modified_input( return Ok(()); }; let Some(schema) = tool_schema else { + // No schema available — the engine couldn't resolve the tool. + // Skip validation rather than reject a possibly-correct fix. return Ok(()); }; @@ -254,8 +273,6 @@ mod tests { use std::pin::Pin; use std::sync::Mutex; - // ---- Mock clients ---- - /// A capture of what the reflector sent to the client, plus the canned /// response to return. #[derive(Clone)] @@ -544,6 +561,48 @@ mod tests { assert!(user.contains("fix the bug"), "user: {user}"); } + #[tokio::test] + async fn llm_reflector_user_message_includes_schema() { + // When the engine supplies a tool schema, the reflector should + // forward it in the user message so the model can produce a + // schema-conforming modified_input. + let captured = Arc::new(Mutex::new(None)); + let client = Arc::new(CannedMock { + response: fixture_analysis(), + captured: captured.clone(), + }); + let reflector = LlmReflector::new(client); + let tool_schema = ToolSchema { + tool: "read".into(), + description: "Read a file".into(), + input_schema: serde_json::json!({ + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"] + }), + }; + let result = reflector + .analyze( + "e", + "read", + &serde_json::json!({"path": "/wrong"}), + Some(&tool_schema), + &ctx(), + ) + .await; + assert!(result.is_ok(), "analyze should succeed: {:?}", result.err()); + let cap = captured.lock().unwrap().clone().expect("captured"); + let user = cap.user.expect("user message"); + assert!( + user.contains("Schema:"), + "expected the schema to appear in the user message: {user}" + ); + assert!( + user.contains("\"required\""), + "expected schema content in the user message: {user}" + ); + } + #[cfg(feature = "schema_validation")] #[tokio::test] async fn llm_reflector_validates_modified_input_pass() { @@ -724,14 +783,13 @@ mod tests { ); } - // ---- build_user_message direct unit tests ---- - #[test] fn build_user_message_contains_all_fields() { let msg = build_user_message( "the error", "the_tool", &serde_json::json!({"k": "v"}), + None, &ReflectionContext { task: "the task".to_string(), attempt: 1, @@ -752,6 +810,7 @@ mod tests { "e", "t", &serde_json::json!({}), + None, &ReflectionContext { task: "x".to_string(), attempt: 0, @@ -771,6 +830,7 @@ mod tests { "e", "t", &serde_json::json!({}), + None, &ReflectionContext { task: "x".to_string(), attempt: u32::MAX, @@ -780,4 +840,51 @@ mod tests { // Should contain u32::MAX (saturated, not overflowed). assert!(msg.contains(&u32::MAX.to_string())); } + + #[test] + fn build_user_message_includes_schema_when_present() { + let schema = serde_json::json!({ + "type": "object", + "properties": {"path": {"type": "string"}}, + "required": ["path"] + }); + let msg = build_user_message( + "e", + "read", + &serde_json::json!({"path": "/wrong"}), + Some(&schema), + &ReflectionContext { + task: "x".to_string(), + attempt: 0, + max_attempts: 3, + }, + ); + assert!( + msg.contains("Schema:"), + "expected a Schema: line when schema is supplied: {msg}" + ); + assert!( + msg.contains("\"required\""), + "schema content missing from message: {msg}" + ); + } + + #[test] + fn build_user_message_omits_schema_line_when_none() { + let msg = build_user_message( + "e", + "t", + &serde_json::json!({}), + None, + &ReflectionContext { + task: "x".to_string(), + attempt: 0, + max_attempts: 3, + }, + ); + assert!( + !msg.contains("Schema:"), + "no Schema: line expected when schema is None: {msg}" + ); + } } From 7cd6ebac1395502d24f71a31c65389f7922862ef Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Sun, 19 Jul 2026 14:04:16 +1200 Subject: [PATCH 3/4] chore: improve docs --- src/stream/heartbeat.rs | 123 ++++++++-- src/tool/shield.rs | 515 +++++++++++++++++++++++++++++++++------- 2 files changed, 533 insertions(+), 105 deletions(-) diff --git a/src/stream/heartbeat.rs b/src/stream/heartbeat.rs index 5efe242..c6b24cb 100644 --- a/src/stream/heartbeat.rs +++ b/src/stream/heartbeat.rs @@ -13,18 +13,19 @@ //! //! # Architecture //! -//! ```text -//! ┌────────────────────────────────────┐ -//! │ HeartbeatStream │ -//! │ │ -//! │ poll_next(): │ -//! │ 1. Check heartbeat interval │ -//! │ └─ fire callback if elapsed │ -//! │ 2. Check hard timeout │ -//! │ └─ return ApiError if hit │ -//! │ 3. Delegate to inner stream │ -//! └────────────────────────────────────┘ -//! ``` +//! On each `poll_next`, `HeartbeatStream` runs three checks in order before +//! delegating to the inner stream: +//! +//! 1. **Heartbeat interval** — if the configured interval has elapsed +//! since the last beat, fire the heartbeat callback with the elapsed +//! time and current timeout status. +//! 2. **Hard timeout** — if the total elapsed time has exceeded the +//! configured maximum, return an [`ApiError`] without consulting the +//! inner stream. +//! 3. **Delegate** — otherwise, forward to the inner stream's `poll_next` +//! and pass through its result. +//! +//! [`ApiError`]: crate::api::error::ApiError //! //! # Quick Start //! @@ -59,7 +60,10 @@ use std::time::{Duration, Instant}; /// Data emitted on each heartbeat callback. /// /// Passed to the callback registered in [`HeartbeatConfig`] at each -/// heartbeat interval. +/// heartbeat interval. Carries a snapshot of how long the stream has +/// been running and whether it has crossed its configured hard-timeout +/// deadline, so a UI or metrics collector can render progress without +/// owning a clock itself. /// /// # Example /// @@ -75,9 +79,25 @@ use std::time::{Duration, Instant}; /// ``` #[derive(Debug, Clone)] pub struct HeartbeatData { - /// Time elapsed since stream start. + /// Time elapsed since the wrapped stream was constructed. + /// + /// Monotonic — measured from the [`Instant`] captured in + /// [`HeartbeatStream::new`], not wall-clock time. Useful for + /// progress UIs ("streaming for 45s"), metrics, and detecting + /// stalls without each consumer holding its own start instant. + /// + /// [`Instant`]: std::time::Instant pub elapsed: Duration, - /// Whether the stream has exceeded its configured timeout. + + /// Whether the stream has exceeded its configured hard timeout. + /// + /// Set to `elapsed > config.timeout` at the moment the heartbeat + /// fires. Note this is *advisory*: the heartbeat may observe + /// `is_timeout == true` a beat after the deadline actually crossed, + /// because callbacks only fire on `poll_next`. The next + /// [`HeartbeatStream::poll_next`] after the deadline will return + /// the hard-timeout error regardless of when the callback last + /// fired. pub is_timeout: bool, } @@ -116,10 +136,30 @@ pub type HeartbeatCallback = Box; /// ``` pub struct HeartbeatConfig { /// Interval between heartbeat callbacks. + /// + /// Checked on every `poll_next`: when `last_heartbeat.elapsed()` + /// reaches this value, the callback fires and `last_heartbeat` is + /// reset. There is no background timer — heartbeats only fire while + /// the stream is actively being polled. heartbeat_interval: Duration, + /// Maximum total stream duration before triggering a hard timeout. + /// + /// Once `start.elapsed()` exceeds this value, the next `poll_next` + /// returns an [`ApiError`] instead of delegating to the inner + /// stream. Choose a value comfortably above the expected p99 + /// response time so transient slowness doesn't trip it. + /// + /// [`ApiError`]: crate::api::error::ApiError timeout: Duration, - /// Callback fired at each heartbeat interval. + + /// Callback invoked at each heartbeat interval. + /// + /// A `Box` so it can be shared + /// across the runtime and mutate captured state (typically an + /// `Arc>`-protected metrics struct or channel sender). + /// Called synchronously from `poll_next` — keep the body cheap to + /// avoid stalling the stream's task. on_heartbeat: HeartbeatCallback, } @@ -169,12 +209,19 @@ impl HeartbeatConfig { } /// Returns the configured heartbeat interval. + /// + /// Exposed so callers (e.g. a metrics reporter that wants to align + /// its own cadence with the heartbeat) can read the value the + /// stream was constructed with. #[must_use] pub fn heartbeat_interval(&self) -> Duration { self.heartbeat_interval } /// Returns the configured hard timeout. + /// + /// Exposed so callers can render the deadline ("stream will time out + /// after 600s") or compute remaining budget from the elapsed time. #[must_use] pub fn timeout(&self) -> Duration { self.timeout @@ -219,26 +266,62 @@ impl HeartbeatConfig { /// ``` pub struct HeartbeatStream { /// The inner stream being wrapped. + /// + /// All `poll_next` calls that survive the heartbeat check and the + /// hard-timeout check are delegated to this stream. The wrapper + /// does not buffer or transform items — it passes them through + /// verbatim, including errors from the inner stream. inner: S, + /// Heartbeat and timeout configuration. + /// + /// Holds the interval, the timeout, and the callback. Owned by + /// the stream (not shared) so the wrapper can call the callback + /// without synchronization. config: HeartbeatConfig, + /// Time of the last heartbeat callback. + /// + /// Compared against `Instant::now()` on every `poll_next` to + /// decide whether the interval has elapsed. Reset to `Instant::now()` + /// (not `start + n*interval`) after each fire so drift from + /// irregular polling doesn't accumulate. last_heartbeat: Instant, + /// Time the stream was created. + /// + /// The reference for all `elapsed` calculations — heartbeat + /// `elapsed`, and the hard-timeout comparison. Captured once in + /// [`new`](Self::new) and never mutated for the life of the stream. start: Instant, - /// A Sleep that fires at the hard timeout deadline. + + /// A `Sleep` future that fires at the hard-timeout deadline. /// /// Ensures the runtime wakes this task when the timeout expires, - /// even if the inner stream is Pending and nobody re-polls. + /// even if the inner stream is `Pending` and nobody re-polls. + /// Polled proactively in `poll_next` so a ready `Sleep` short- + /// circuits to the timeout error without delegating to the inner + /// stream first. timeout_sleep: std::pin::Pin>, } impl HeartbeatStream { /// Create a new heartbeat stream wrapping the given inner stream. /// - /// The heartbeat timer starts immediately upon construction. + /// The heartbeat timer starts immediately upon construction: `start` + /// and `last_heartbeat` are captured at this moment, and the + /// hard-timeout `Sleep` is armed against `start + config.timeout`. /// The first heartbeat callback fires after `heartbeat_interval` - /// elapses (checked on each `poll_next`). + /// elapses (checked on each `poll_next` — there is no background + /// timer). + /// + /// # Timeout overflow + /// + /// `start + config.timeout` is computed via `checked_add`. For any + /// realistic `Duration` this always succeeds; the fallback (a + /// deadline 30 years in the future) only triggers for extreme values + /// near `Duration::MAX`, where the deadline is effectively "never" + /// either way. /// /// # Example /// diff --git a/src/tool/shield.rs b/src/tool/shield.rs index fe1a410..077a84d 100644 --- a/src/tool/shield.rs +++ b/src/tool/shield.rs @@ -22,21 +22,16 @@ //! //! # Architecture //! -//! ```text -//! ┌────────────────────────────────────────────┐ -//! │ Agent Loop (middleware) │ -//! │ │ -//! │ tool_call ──▶ ToolSafetyShield::evaluate()│ -//! │ │ │ -//! │ ┌────────┴────────┐ │ -//! │ │ SafetyDecision │ │ -//! │ │ Allow / Warn / │ │ -//! │ │ Block │ │ -//! │ └────────┬────────┘ │ -//! │ │ │ -//! │ tool_result ──▶ record_invocation() │ -//! └────────────────────────────────────────────┘ -//! ``` +//! The shield sits in the agent-loop middleware path and is consulted twice +//! per tool call: +//! +//! 1. **Before dispatch** — +//! [`ToolSafetyShield::evaluate()`] inspects the `tool_call` and returns +//! a [`SafetyDecision`] of `Allow`, `Warn`, or `Block`, which the loop +//! honors before running the tool. +//! 2. **After dispatch** — `record_invocation()` is fed the `tool_result` +//! so the shield can update its multi-turn state (call history, +//! combination tracking) for future `evaluate()` calls. //! //! # Provided Implementations //! @@ -79,18 +74,61 @@ use serde_json::Value; /// /// Each dimension of the shield (single-turn, multi-turn, combination) /// scores into a float in `[0.0, 1.0]`. The aggregate score is then -/// mapped to a [`RiskLevel`] using configurable thresholds. +/// mapped to a [`RiskLevel`] using configurable thresholds +/// (`warn_threshold`, `block_threshold`) on [`UnixShield`]. +/// +/// # Ordering +/// +/// Variants are written in increasing severity. There is intentionally no +/// `Ord` derive — risk levels are labels produced by threshold +/// comparisons, not a totally-ordered set; `RiskLevel::Medium` is not +/// "less than" `RiskLevel::High` in any numeric sense callers should rely +/// on. Compare the underlying aggregate score if you need ordering. +/// +/// [`UnixShield`]: crate::tool::shield::UnixShield #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum RiskLevel { - /// No risk detected. Aggregate score < 0.2. + /// No risk detected. + /// + /// Aggregate score below `0.2`. The default decision mapping for + /// `Safe` is [`SafetyAction::Allow`] with no warning. Safe, - /// Low risk. Aggregate score in [0.2, `warn_threshold`). + + /// Low risk — below the warn threshold but not negligible. + /// + /// Aggregate score in `[0.2, warn_threshold)`. Like [`Safe`], this + /// maps to [`SafetyAction::Allow`]; the level exists so callers + /// inspecting a [`SafetyDecision`] can distinguish "nothing matched" + /// from "minor patterns matched, but below the warn bar." + /// + /// [`Safe`]: Self::Safe + /// [`SafetyAction::Allow`]: crate::tool::shield::SafetyAction::Allow Low, - /// Moderate risk. Aggregate score in [`warn_threshold`, `block_threshold`). + + /// Moderate risk — at or above the warn threshold. + /// + /// Aggregate score in `[warn_threshold, block_threshold)`. Maps to + /// [`SafetyAction::Warn`]: the call proceeds, but the middleware + /// surfaces a reason and category to observers / logs. + /// + /// [`SafetyAction::Warn`]: crate::tool::shield::SafetyAction::Warn Medium, - /// High risk. Aggregate score in [`block_threshold`, 0.9). + + /// High risk — at or above the block threshold. + /// + /// Aggregate score in `[block_threshold, 0.9)`. Maps to + /// [`SafetyAction::Block`]: the call does not proceed. + /// + /// [`SafetyAction::Block`]: crate::tool::shield::SafetyAction::Block High, - /// Critical risk. Aggregate score ≥ 0.9. + + /// Critical risk — the most dangerous category. + /// + /// Aggregate score at or above `0.9`. Reserved for patterns the + /// shield treats as maximally dangerous (e.g. `rm -rf /`, `curl … | + /// sh`). Maps to [`SafetyAction::Block`]. + /// + /// [`SafetyAction::Block`]: crate::tool::shield::SafetyAction::Block Critical, } @@ -110,14 +148,49 @@ impl std::fmt::Display for RiskLevel { // SafetyAction // =================================================== -/// The action the shield recommends. +/// The action the shield recommends for a tool call. +/// +/// Produced as the [`SafetyDecision::action`] field. The middleware is +/// expected to honor it: proceed on [`Allow`], proceed-and-log on +/// [`Warn`], refuse to dispatch on [`Block`]. +/// +/// There is intentionally no `Ord` derive — these are categorical +/// recommendations, not an ordered severity scale. Use [`RiskLevel`] for +/// severity comparisons. +/// +/// [`Allow`]: Self::Allow +/// [`Warn`]: Self::Warn +/// [`Block`]: Self::Block +/// [`SafetyDecision::action`]: crate::tool::shield::SafetyDecision::action #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum SafetyAction { /// Allow the tool call to proceed. + /// + /// The shield found no concerning patterns, or every pattern it found + /// scored below the warn threshold. No metadata is required on the + /// accompanying [`SafetyDecision`]. + /// + /// [`SafetyDecision`]: crate::tool::shield::SafetyDecision Allow, - /// Allow but emit a warning. + + /// Allow the tool call to proceed, but emit a warning. + /// + /// The shield detected moderate risk (score in the warn band). The + /// call is permitted to run — warnings are advisory, not blocking — + /// but the [`SafetyDecision`] carries a human-readable `reason` and + /// a machine-readable `category` for observers, logs, and TUIs. + /// + /// [`SafetyDecision`]: crate::tool::shield::SafetyDecision Warn, + /// Block the tool call entirely. + /// + /// The shield detected high or critical risk (score at or above the + /// block threshold). The middleware must refuse to dispatch the call; + /// the [`SafetyDecision`] carries a `reason` and `category` + /// describing what was matched, for surfacing to the user. + /// + /// [`SafetyDecision`]: crate::tool::shield::SafetyDecision Block, } @@ -127,20 +200,48 @@ pub enum SafetyAction { /// The shield's decision for a single tool invocation. /// -/// Produced by [`ToolSafetyShield::evaluate`]. Carries the action, an -/// optional human-readable reason, and a machine-readable category tag. +/// Produced by [`ToolSafetyShield::evaluate`]. Carries the action the +/// middleware should take, an optional human-readable reason explaining +/// the decision, and an optional machine-readable category tag for +/// programmatic handling / filtering in logs. +/// +/// # Construction +/// +/// Use the [`allow`](Self::allow), [`warn`](Self::warn), and +/// [`block`](Self::block) constructors rather than building the struct +/// directly — they keep the reason/category fields consistent with the +/// chosen action (e.g. `allow` always has `None` for both). #[derive(Debug, Clone)] pub struct SafetyDecision { - /// The recommended action. + /// The recommended action for this tool call. + /// + /// Always set. See [`SafetyAction`] for the three variants and the + /// middleware's responsibility for each. pub action: SafetyAction, - /// Human-readable explanation (set for `Warn` and `Block`). + + /// Human-readable explanation of the decision. + /// + /// Set for [`SafetyAction::Warn`] and [`SafetyAction::Block`] + /// decisions (the constructors populate it from the caller-supplied + /// string); `None` for [`SafetyAction::Allow`]. Suitable for + /// surfacing to a user or writing to a log line. pub reason: Option, - /// Machine-readable category (e.g. `"safety_evaluation"`, `"pattern_match"`). + + /// Machine-readable category tag for the decision. + /// + /// Set for `Warn` and `Block`; `None` for `Allow`. Free-form but + /// conventionally a stable identifier like `"safety_evaluation"` or + /// `"pattern_match"` so downstream consumers can route on it without + /// parsing the `reason` text. pub category: Option, } impl SafetyDecision { - /// Create an `Allow` decision with no attached metadata. + /// Create an [`SafetyAction::Allow`] decision with no attached + /// metadata. + /// + /// Both `reason` and `category` are `None` — an allowed call has + /// nothing to warn or block about. #[must_use] pub fn allow() -> Self { Self { @@ -150,7 +251,12 @@ impl SafetyDecision { } } - /// Create a `Warn` decision with reason and category. + /// Create a [`SafetyAction::Warn`] decision with a reason and + /// category. + /// + /// The call proceeds, but the middleware should surface `reason` to + /// observers / logs. `category` should be a stable machine-readable + /// identifier consumers can filter on. #[must_use] pub fn warn(reason: String, category: &str) -> Self { Self { @@ -160,7 +266,12 @@ impl SafetyDecision { } } - /// Create a `Block` decision with reason and category. + /// Create a [`SafetyAction::Block`] decision with a reason and + /// category. + /// + /// The middleware must refuse to dispatch the call. `reason` should + /// explain what was matched so the user can understand why; `category` + /// should be a stable identifier for programmatic routing. #[must_use] pub fn block(reason: String, category: &str) -> Self { Self { @@ -170,19 +281,31 @@ impl SafetyDecision { } } - /// Whether the decision is `Allow`. + /// Returns `true` if the decision's action is + /// [`SafetyAction::Allow`]. + /// + /// Convenience predicate for middleware that wants to short-circuit + /// "this call is fine, dispatch it" without a full `match`. #[must_use] pub fn is_allowed(&self) -> bool { self.action == SafetyAction::Allow } - /// Whether the decision is `Block`. + /// Returns `true` if the decision's action is + /// [`SafetyAction::Block`]. + /// + /// Convenience predicate for middleware that wants to short-circuit + /// "refuse this call" without a full `match`. #[must_use] pub fn is_blocked(&self) -> bool { self.action == SafetyAction::Block } - /// Whether the decision is `Warn`. + /// Returns `true` if the decision's action is + /// [`SafetyAction::Warn`]. + /// + /// Convenience predicate for middleware that wants to branch on + /// "proceed but log" without a full `match`. #[must_use] pub fn is_warn(&self) -> bool { self.action == SafetyAction::Warn @@ -195,17 +318,53 @@ impl SafetyDecision { /// Context provided to the shield for evaluation. /// -/// Constructed by the middleware before each tool call. The shield uses -/// this to evaluate single-turn risk and multi-turn patterns. +/// Constructed by the middleware before each tool call. Carries what the +/// shield needs to evaluate both single-turn risk (the current call in +/// isolation) and multi-turn / combination risk (the current call in the +/// context of the recent call history). +/// +/// # Lifecycle +/// +/// A fresh `ShieldContext` is built per tool invocation and passed by +/// reference to [`ToolSafetyShield::evaluate`]. It is not stored across +/// calls — shields that need persistent history keep their own +/// (typically `Mutex>`-protected) state, updated in +/// [`ToolSafetyShield::record_invocation`]. #[derive(Debug, Clone)] pub struct ShieldContext { /// Name of the tool being invoked. + /// + /// Matches the registry key the engine used to dispatch the call. + /// Single-turn risk patterns are looked up under this name, so a + /// tool's risk configuration is keyed off the same identifier the + /// agent uses to call it. pub tool_name: String, - /// The JSON input to the tool. + + /// The JSON input passed to the tool. + /// + /// Shield patterns operate on the stringified form of this value + /// (substring matching), so any JSON-shaped input is admissible — + /// objects, arrays, primitives. The same value is later handed to + /// [`ToolSafetyShield::record_invocation`] so the shield can store + /// it for combination-rule matching. pub input: Value, - /// Current turn number in the agent session (0-indexed). + + /// Current turn number in the agent session, 0-indexed. + /// + /// Lets the shield factor recency into its scoring (e.g. weight + /// recent calls more heavily) and lets logs correlate shield + /// decisions with the turn that produced them. pub turn: usize, - /// Recent tool invocations in this session (tool name, turn). + + /// Recent tool invocations in this session, as `(tool_name, turn)` + /// pairs. + /// + /// A read-only snapshot provided by the middleware for shields that + /// prefer not to maintain their own history. Shields that *do* + /// maintain their own history (like [`UnixShield`]) can ignore this + /// field and consult their internal state instead. + /// + /// [`UnixShield`]: crate::tool::shield::UnixShield pub recent_calls: Vec<(String, usize)>, } @@ -213,18 +372,45 @@ pub struct ShieldContext { // RiskPattern // =================================================== -/// A named risk pattern matched against tool input. +/// A named risk pattern matched against a tool's input. /// /// Patterns are stored per tool name. When a tool call is evaluated, /// its stringified JSON input is checked against every pattern -/// registered for that tool name. +/// registered for that tool name; any pattern whose `pattern` substring +/// appears contributes its `score` to the single-turn risk dimension. +/// +/// # Matching +/// +/// Matching is plain substring search on `input.to_string()` — there is +/// no regex, word-boundary, or JSON-aware path matching. Patterns must +/// be chosen so that a substring match implies the intended risk (e.g. +/// `"rm -rf"` is distinctive enough; a bare `"rm"` would over-match). #[derive(Clone)] pub struct RiskPattern { /// Human-readable name for diagnostics and log messages. + /// + /// Appears in shield output and logs when this pattern is the + /// matched one. Should be unique within a tool's pattern list and + /// descriptive enough to identify the risk at a glance (e.g. + /// `"recursive_delete"`, `"write_ssh"`, `"curl_pipe_sh"`). pub name: &'static str, - /// Score contribution if matched, in `[0.0, 1.0]`. + + /// Score contributed to the single-turn dimension if this pattern + /// matches, in `[0.0, 1.0]`. + /// + /// When multiple patterns match a single input, the highest score + /// wins ([`UnixShield::assess_single_tool_risk`] takes the `max`). + /// Convention: `0.9` for critical patterns (`rm -rf`, `curl | sh`), + /// `0.5`–`0.8` for serious patterns, `< 0.5` for minor ones. + /// + /// [`UnixShield::assess_single_tool_risk`]: crate::tool::shield::UnixShield::assess_single_tool_risk pub score: f32, + /// Substring to search for in the stringified input. + /// + /// Matched verbatim (no regex, no case folding). Picked so that a + /// substring hit reliably indicates the intended risk — see the + /// type-level docs on matching. pub pattern: &'static str, } @@ -234,17 +420,49 @@ pub struct RiskPattern { /// A rule that scores a dangerous *sequence* of tool calls. /// -/// A combination rule triggers when **all** of its [`triggers`](CombinationRule::triggers) -/// appear in recent history or the current call. Each trigger is a -/// `(tool_name, optional_substring)` pair. +/// A combination rule triggers when **all** of its [`triggers`](Self::triggers) +/// appear in recent history or the current call, in the order specified. +/// Each trigger is a `(tool_name, optional_substring)` pair: the +/// `tool_name` must match exactly, and when the substring is present the +/// trigger only fires if that substring is in the call's stringified +/// input. +/// +/// # Example +/// +/// A rule with triggers `&[("Bash", Some("curl")), ("Bash", Some("| sh"))]` +/// fires only when a `curl` call is followed by a `| sh` call in the +/// session — catching the classic "download and execute" pattern that +/// neither call would flag in isolation. #[derive(Clone)] pub struct CombinationRule { /// Human-readable description for diagnostics and log messages. + /// + /// Surfaces in shield output when the rule fires. Should name the + /// dangerous sequence (e.g. `"download then execute"`, + /// `"write then chmod +x"`) so a user reading the warning + /// understands what the shield saw. pub description: &'static str, - /// Score contribution if matched. + + /// Score contributed to the combination dimension if this rule + /// triggers, in `[0.0, 1.0]`. + /// + /// When multiple rules fire on the same call, the highest score + /// wins ([`UnixShield::assess_combination`] takes the `max`). + /// Convention: `0.8+` for sequences that are almost always + /// adversarial; `0.5`–`0.7` for suspicious-but-defensible ones. + /// + /// [`UnixShield::assess_combination`]: crate::tool::shield::UnixShield::assess_combination pub score: f32, - /// Pairs of `(tool_name, optional_substring)` that must all appear - /// for the rule to trigger. + + /// Pairs of `(tool_name, optional_substring)` that must all appear, + /// in order, for the rule to trigger. + /// + /// The match is chronological against the candidate sequence + /// (recorded history followed by the current call). Each trigger + /// advances only when both the tool name matches and, if a + /// substring is supplied, that substring appears in the call's + /// stringified input. A `None` substring means "any call to this + /// tool". pub triggers: &'static [(&'static str, Option<&'static str>)], } @@ -276,20 +494,40 @@ pub struct CombinationRule { pub trait ToolSafetyShield: Send + Sync { /// Evaluate whether the tool call described by `ctx` should be /// allowed, warned about, or blocked. + /// + /// Called by the middleware before dispatch. The shield inspects the + /// call (and, if it maintains its own history, prior calls) and + /// returns a [`SafetyDecision`] naming the action to take, with a + /// reason/category populated for `Warn` and `Block`. + /// + /// Implementations must be deterministic for a given `(self, ctx)` + /// pair so that replaying the same call produces the same decision; + /// non-determinism breaks the VCR / cassette test path. fn evaluate(&self, ctx: &ShieldContext) -> SafetyDecision; - /// Called after the tool executes. Allows the shield to update - /// internal state (e.g. call history for multi-turn analysis). + /// Called after the tool executes so the shield can update internal + /// state for multi-turn / combination analysis. + /// + /// The middleware calls this with the same `input` [`Value`] passed + /// to [`evaluate`](Self::evaluate), plus a `success` flag indicating + /// whether the tool returned an error. Shields that maintain call + /// history (e.g. [`UnixShield`]) append to it here; shields with no + /// multi-turn state (e.g. [`NullShield`]) no-op. /// - /// The `input` is the same [`Value`] passed to [`evaluate`](ToolSafetyShield::evaluate), - /// so the shield can store it for later combination-rule matching. + /// [`UnixShield`]: crate::tool::shield::UnixShield + /// [`NullShield`]: crate::tool::shield::NullShield fn record_invocation(&self, tool_name: &str, input: &Value, success: bool); /// Return the set of tool names this shield wants to inspect. /// - /// Optimization: if the shield returns an empty set, the middleware - /// can skip calling [`evaluate()`](ToolSafetyShield::evaluate) for - /// tools that the shield has no rules for. + /// Optimization: the middleware consults this before calling + /// [`evaluate`](Self::evaluate) and skips evaluation for tools not + /// in the set. Shields should return the keys of their per-tool + /// pattern database. An empty set means "no tools watched" and + /// causes the middleware to skip evaluation entirely (used by + /// [`NullShield`] to be truly zero-cost). + /// + /// [`NullShield`]: crate::tool::shield::NullShield fn watched_tools(&self) -> HashSet; } @@ -328,15 +566,42 @@ pub trait ToolSafetyShield: Send + Sync { /// [`assess_multi_turn`]: UnixShield::assess_multi_turn /// [`assess_combination`]: UnixShield::assess_combination pub struct UnixShield { - /// Score at which [`SafetyAction::Warn`] is returned. + /// Aggregate score at or above which [`SafetyAction::Warn`] is + /// returned. + /// + /// Configurable via [`with_thresholds`](Self::with_thresholds) or the + /// builder; defaults to `0.4`. Clamped to `[0.0, 1.0]` on + /// construction. warn_threshold: f32, - /// Score at which [`SafetyAction::Block`] is returned. + + /// Aggregate score at or above which [`SafetyAction::Block`] is + /// returned. + /// + /// Must be `>= warn_threshold` for sensible behavior; the shield does + /// not enforce this at construction time, so a misconfigured builder + /// can produce surprising results. Defaults to `0.7`. block_threshold: f32, - /// Per-turn invocation history: `(tool_name, input_string)`. + + /// Per-invocation call history: `(tool_name, stringified_input)`. + /// + /// Mutex-protected because [`ToolSafetyShield`] requires `Send + + /// Sync` and `evaluate`/`record_invocation` are `&self` methods. + /// Trimmed to the last 20 entries on each `record_invocation` to + /// bound memory growth in long sessions. turn_history: Mutex>, - /// Per-tool single-turn risk patterns. + + /// Per-tool single-turn risk patterns, keyed by tool name. + /// + /// Populated from [`unix_patterns`](Self::unix_patterns) by default + /// (`Bash`, `Write`, `Edit`); extendable via the builder. The keyset + /// also defines [`watched_tools`](ToolSafetyShield::watched_tools). patterns: HashMap<&'static str, Vec>, + /// Combination rules for dangerous sequences. + /// + /// Populated from [`unix_combination_rules`](Self::unix_combination_rules) + /// by default (write→execute, download→execute, chmod→write); + /// extendable via the builder. combination_rules: Vec, } @@ -356,9 +621,13 @@ impl UnixShield { } } - /// Set custom warn and block thresholds. + /// Override the default warn and block thresholds, builder-style. /// - /// Values are clamped to `[0.0, 1.0]`. + /// Both values are clamped to `[0.0, 1.0]` so an out-of-range + /// configuration cannot produce a shield that never warns or never + /// blocks. The shield does not enforce `block >= warn`; passing an + /// inverted pair will produce surprising decisions, so callers + /// should validate their own inputs. #[must_use] pub fn with_thresholds(mut self, warn: f32, block: f32) -> Self { self.warn_threshold = warn.clamp(0.0, 1.0); @@ -376,11 +645,18 @@ impl UnixShield { UnixShieldBuilder::new() } - /// Assess single-turn risk: match the tool input against known - /// dangerous patterns for that tool. + /// Assess single-turn risk by matching the tool input against the + /// patterns registered for that tool. + /// + /// Stringifies `input` and substring-matches it against every + /// [`RiskPattern`] keyed under `tool_name`. When multiple patterns + /// match, the highest `score` among them is returned (a single + /// dangerous pattern dominates several minor ones). Returns `0.0` + /// when the tool has no registered patterns or none of them match. /// - /// Returns the highest score among all matched patterns, or `0.0` - /// if the tool has no registered patterns or none matched. + /// This is the single-turn contribution to the aggregate score; + /// [`evaluate`](ToolSafetyShield::evaluate) weights it at `1.0` + /// (the highest-weighted dimension). pub fn assess_single_tool_risk(&self, tool_name: &str, input: &Value) -> f32 { let Some(patterns) = self.patterns.get(tool_name) else { return 0.0; @@ -395,11 +671,16 @@ impl UnixShield { max_score } - /// Assess multi-turn risk: look for repeated calls to the same - /// tool across recent turns. + /// Assess multi-turn risk by counting prior calls to the same tool + /// in the recorded history. /// - /// The score graduates with repetition: 0 calls → 0.0, 1 → 0.1, - /// 2 → 0.3, 3+ → 0.6. + /// The score graduates with repetition so that a single repeat is + /// mild but a long run of the same tool flags as suspicious: + /// Graduated: 0 calls = 0.0, 1 = 0.1, 2 = 0.3, 3+ = 0.6 + /// [`evaluate`](ToolSafetyShield::evaluate) weights this at `0.5`, + /// so even at the saturated `0.6` it cannot alone push the + /// aggregate past the warn threshold — but combined with a + /// single-turn hit it adds up. pub fn assess_multi_turn(&self, ctx: &ShieldContext) -> f32 { let history = self .turn_history @@ -409,7 +690,6 @@ impl UnixShield { .iter() .filter(|(name, _)| name == &ctx.tool_name) .count(); - // Graduated: 0 calls = 0.0, 1 = 0.1, 2 = 0.3, 3+ = 0.6 match same_tool_count { 0 => 0.0, 1 => 0.1, @@ -418,13 +698,23 @@ impl UnixShield { } } - /// Assess combination risk: detect dangerous tool sequences in - /// recent history. + /// Assess combination risk by detecting dangerous sequences of + /// tool calls in recorded history plus the current call. + /// + /// Each [`CombinationRule`] specifies an ordered sequence of + /// triggers (`(tool_name, optional_substring)` pairs). The match + /// walks the candidate sequence — recorded history followed by the + /// current call — and advances a per-rule trigger pointer only + /// when both the tool name matches and, if a substring is + /// supplied, that substring appears in the call's stringified + /// input. A rule fires when its pointer reaches the end of its + /// trigger list, contributing its `score`. /// - /// Each [`CombinationRule`] specifies an ordered sequence of triggers. - /// The match is chronological — triggers must appear in the order - /// specified, either in recorded history or the current call. - /// Returns the highest score among all matched rules. + /// When multiple rules fire, the highest score wins. Returns `0.0` + /// if no rule triggers. [`evaluate`](ToolSafetyShield::evaluate) + /// weights this dimension at `0.3` — the lowest-weighted, but the + /// only one that can detect adversarial *sequences* (download → + /// execute, write → chmod +x). pub fn assess_combination(&self, ctx: &ShieldContext) -> f32 { let history = self .turn_history @@ -463,6 +753,18 @@ impl UnixShield { max_risk } + /// Map an aggregate risk score to a [`RiskLevel`] using this + /// shield's configured thresholds. + /// + /// Bands (default thresholds): `< 0.2` → [`Safe`](RiskLevel::Safe), + /// `[0.2, 0.4)` → [`Low`](RiskLevel::Low), + /// `[warn_threshold, block_threshold)` → + /// [`Medium`](RiskLevel::Medium), + /// `[block_threshold, 0.9)` → [`High`](RiskLevel::High), + /// `>= 0.9` → [`Critical`](RiskLevel::Critical). + /// `0.9` is hardcoded for `Critical` so the most dangerous patterns + /// always land in their own bucket regardless of threshold + /// configuration. fn score_to_level(&self, score: f32) -> RiskLevel { if score >= 0.9 { RiskLevel::Critical @@ -662,9 +964,30 @@ impl ToolSafetyShield for UnixShield { /// .build(); /// ``` pub struct UnixShieldBuilder { + /// Aggregate score at or above which the built shield will return + /// [`SafetyAction::Warn`]. Defaults to `0.4`; override via + /// [`warn_threshold`](Self::warn_threshold). warn_threshold: f32, + + /// Aggregate score at or above which the built shield will return + /// [`SafetyAction::Block`]. Defaults to `0.7`; override via + /// [`block_threshold`](Self::block_threshold). block_threshold: f32, + + /// Per-tool single-turn risk patterns, keyed by tool name. + /// + /// Populated from [`UnixShield::unix_patterns`] by + /// [`new`](Self::new); empty under [`blank`](Self::blank); extended + /// via [`pattern`](Self::pattern). patterns: HashMap<&'static str, Vec>, + + /// Combination rules for dangerous sequences. + /// + /// Populated from [`UnixShield::unix_combination_rules`] by + /// [`new`](Self::new); empty under [`blank`](Self::blank); extended + /// via [`combination_rule`](Self::combination_rule). + /// + /// [`UnixShield::unix_combination_rules`]: crate::tool::shield::UnixShield::unix_combination_rules combination_rules: Vec, } @@ -718,19 +1041,29 @@ impl UnixShieldBuilder { self } - /// Register single-turn risk patterns for a tool. + /// Register single-turn risk patterns for a tool, builder-style. /// - /// Appends to any existing patterns for that tool name. + /// Appends to any patterns already registered for `tool_name` + /// rather than replacing them — call repeatedly to assemble a + /// tool's risk profile across multiple additions. The first call + /// for a given `tool_name` also adds it to the built shield's + /// [`watched_tools`](ToolSafetyShield::watched_tools) set. #[must_use] pub fn pattern(mut self, tool_name: &'static str, patterns: Vec) -> Self { self.patterns.entry(tool_name).or_default().extend(patterns); self } - /// Register a combination rule. + /// Register a combination rule for dangerous tool sequences, + /// builder-style. + /// + /// The rule is evaluated on every call to + /// [`UnixShield::evaluate`] and fires when all of its triggers + /// appear in the session history in order. Call repeatedly to add + /// multiple rules; rules are independent (the highest score among + /// triggered rules wins). /// - /// The rule is evaluated on every call to [`UnixShield::evaluate`] - /// and fires when all of its triggers appear in the session history. + /// [`UnixShield::evaluate`]: crate::tool::shield::UnixShield::evaluate #[must_use] pub fn combination_rule(mut self, rule: CombinationRule) -> Self { self.combination_rules.push(rule); @@ -764,11 +1097,23 @@ impl Default for UnixShieldBuilder { // NullShield // =================================================== -/// A no-op shield that allows everything. +/// A no-op [`ToolSafetyShield`] that allows every call and watches no +/// tools. +/// +/// Used as the default shield type when the `tool_shield` feature is +/// disabled, so engine code that holds `Arc` +/// compiles and runs without forcing the feature on downstream +/// consumers. +/// +/// # Zero-cost /// -/// Used when the `tool_shield` feature is disabled. Because all methods -/// return constants with no field access, the compiler can inline and -/// eliminate all calls to this type, making it truly zero-cost. +/// All three trait methods return constants with no field access +/// (`evaluate` returns a pre-built [`SafetyDecision::allow`], +/// `record_invocation` no-ops, `watched_tools` returns an empty +/// [`HashSet`]). Because the middleware short-circuits evaluation +/// entirely when `watched_tools` is empty, an installed `NullShield` +/// costs nothing at runtime — and the compiler can inline and elide +/// the constant returns. pub struct NullShield; impl ToolSafetyShield for NullShield { From 22e90c77a55da07d4dd350f27989a66b22c3f943 Mon Sep 17 00:00:00 2001 From: Iurii Bobrykov Date: Sun, 19 Jul 2026 14:37:36 +1200 Subject: [PATCH 4/4] fix: requested tool name fallback for schema lookup --- src/engine/bare/dispatch.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/engine/bare/dispatch.rs b/src/engine/bare/dispatch.rs index 4b3bb73..6eb3a37 100644 --- a/src/engine/bare/dispatch.rs +++ b/src/engine/bare/dispatch.rs @@ -1049,12 +1049,18 @@ impl BareLoop { max_attempts: Self::MAX_RECOVERY_ATTEMPTS, }; - // Look up the schema under the name a routing middleware may have - // redirected the call to; `tc.tool` (passed below as `tool_name`) - // stays as the originally requested name. + // Resolve the schema under the name a routing middleware may have + // redirected the call to, falling back to the requested name when + // the resolved name is empty or unknown to the registry. + let resolved_tool = if result.resolved_tool_name.is_empty() { + &tc.tool + } else { + &result.resolved_tool_name + }; let tool_schema = self .tools - .get(&result.resolved_tool_name) + .get(resolved_tool) + .or_else(|| self.tools.get(&tc.tool)) .map(crate::tool::Tool::schema); let Ok(analysis) = self .reflector