From d1f6e98f91b96c743c725ab7374376e65a7c1457 Mon Sep 17 00:00:00 2001 From: stephen Date: Mon, 10 Aug 2026 16:52:42 -0700 Subject: [PATCH 01/21] feat(agent): expose per-turn request preparation on Agent Add Agent::prepare_turn, returning the new public PreparedTurn/TurnTools: the agent's baseline configuration resolved into one turn's completion request plus the turn's executable and allowed tool-name sets, the synthetic output-tool name, and tool dispatch pinned to the turn's registry snapshot. Hand-driving the sans-IO AgentRun machine is a documented supported use case (custom provider transport, durable suspend/resume), but until now the only way to do it was to re-derive the agent's configuration by hand - our own agent_run_stepping example registered its tool twice and copy-pasted its preamble into the loop. A prepared turn is a configuration read plus a dispatch target, not a second execution path: AgentRunner remains the only path that executes an agent with hooks, memory, retrieval policy, and telemetry, and a prepared turn reflects the baseline configuration with no hook patches applied. Design notes: - ToolRegistrySnapshot stays private. The public dispatch surface is the new TurnTools wrapper, whose execute() mirrors ToolServerHandle:: execute; retain_names and ToolDispatch remain internal. - The advertised sets and the dispatch snapshot are resolved together at one instant and carried as one value, so a hand-driver cannot mix a snapshot from one turn with name sets from another. - Impossible tool_choice/tool-set combinations fail at prepare time with no provider round-trip, including Required against an empty advertised set. - Dispatching the synthetic output tool is rejected with an error explaining that its call carries the final structured answer. Both hand-driven examples (agent_run_stepping, agent_with_durable_approval) are rewritten on the new API: each builds one Agent and drives AgentRun from it, with tool calls executing through the same snapshot the provider saw advertised. No behavior change on the runner or streaming paths; the shared build_prepared_completion_request serves both callers unchanged. --- MIGRATING.md | 9 + crates/rig-agent/CHANGELOG.md | 1 + crates/rig-agent/src/agent/completion.rs | 88 +++++ crates/rig-agent/src/agent/mod.rs | 2 + crates/rig-agent/src/agent/prepared_turn.rs | 359 ++++++++++++++++++ crates/rig-agent/src/agent/run/mod.rs | 13 +- examples/agent_run_stepping/src/main.rs | 60 ++- .../agent_with_durable_approval/src/main.rs | 57 +-- 8 files changed, 527 insertions(+), 62 deletions(-) create mode 100644 crates/rig-agent/src/agent/prepared_turn.rs diff --git a/MIGRATING.md b/MIGRATING.md index 5c51ba5540..933a801258 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1561,6 +1561,15 @@ response. For intentionally hook-free transport, start from for custom drivers. It holds no configured model, tools, memory, or hooks and is not an alternate execution path for configured agents. +Reading a configured agent's request state back out (the 0.40 public fields, +`agent.tool_server_handle`) has a supported replacement: +`Agent::prepare_turn(prompt, &history)` resolves the agent's configuration into +one turn's completion request plus its tool sets and dispatch snapshot +(`PreparedTurn` / `TurnTools`), so a hand-driven `AgentRun` no longer restates +the preamble and tools. It is a configuration read, not an execution path — +hooks, memory, retrieval policy, and telemetry still run only under +`Agent::runner`. See `examples/agent_run_stepping`. + An `Agent`'s default model is set at construction. Per-run overrides now go through `runner(...).using_model(...)`, `Agent::set_model`, or a `ModelSelection` hook (see the "runtime model swapping" section for the diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 78f3a0e60f..ee4901b74f 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `AgentHook::on_reasoning_delta` with the Rig stream correlator, optional provider reasoning id, current fragment, and per-part aggregate; reasoning hooks share the existing observation-interest and stop-before-yield semantics used by other streaming deltas +- *(agent)* add `Agent::prepare_turn`, returning the new public `PreparedTurn`/`TurnTools`: the agent's baseline configuration resolved into one turn's completion request plus the turn's executable and allowed tool-name sets, the synthetic output-tool name, and tool dispatch pinned to the turn's registry snapshot — so hand-driven `AgentRun` loops (custom provider transport, durable suspend/resume) reuse the configured `Agent` instead of restating its preamble and tools. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. A prepared turn is a configuration read: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` - *(agent)* add opaque, cloneable `ModelHandle` values with by-value `ProviderCapabilities` snapshots, plus default replacement, per-run default override (`using_model`), and hook-driven per-call selection via `AgentHook::on_model_select` - *(agent)* add run-local extractor default-model overrides used across retries diff --git a/crates/rig-agent/src/agent/completion.rs b/crates/rig-agent/src/agent/completion.rs index 2332371b31..b642bc8fd5 100644 --- a/crates/rig-agent/src/agent/completion.rs +++ b/crates/rig-agent/src/agent/completion.rs @@ -668,6 +668,94 @@ impl Agent { ) -> Result, ToolServerError> { self.tool_server_handle.get_tool_defs(prompt).await } + + /// Prepare one model turn from this agent's configuration, for hand-driven + /// [`AgentRun`](super::run::AgentRun) loops. + /// + /// Resolves the agent's preamble, static context, model parameters, + /// `tool_choice`, output schema/mode, and tool registry into a + /// [`PreparedTurn`](super::PreparedTurn): the fully configured completion request plus a + /// [`TurnTools`](super::TurnTools) value carrying the turn's executable and + /// allowed tool-name sets, the synthetic output-tool name (if any), and a + /// dispatch target pinned to this turn's registry snapshot. Use it to drive + /// the sans-IO [`AgentRun`](super::run::AgentRun) machine — custom provider + /// transport, suspend/resume — without restating the agent's configuration. + /// + /// # Baseline, not hook-patched + /// + /// A prepared turn reflects the agent's **baseline** configuration: no + /// `CompletionCall` hooks run (so no per-turn `RequestPatch` or + /// `active_tools` narrowing), no model-selection hooks run (the agent's + /// current default model is used), and no memory or telemetry is involved. + /// [`Agent::runner`](Self::runner) remains the only path that executes an + /// agent with hooks, memory, retrieval policy, and telemetry. + /// + /// # Per-turn snapshot + /// + /// The tool definitions in the request and the dispatch target in + /// [`TurnTools`](super::TurnTools) come from one registry snapshot taken by + /// this call. Registry changes made afterwards (`add_tool` / `remove_tool`, + /// MCP refreshes) affect the next prepared turn, never this one — see + /// [`TurnTools`](super::TurnTools) for the semantics. Retrieval-selected + /// (dynamic) tools are resolved here too, using the prompt's text (or the + /// latest history text) as the retrieval query, exactly as the runner does. + /// + /// # Errors + /// + /// Fails locally — with no provider round-trip — when the configuration + /// cannot produce a valid request: a `tool_choice` that is impossible + /// against the advertised tool set (`Required` with no advertised tool, or + /// `Specific` naming an unadvertised tool), or a real tool colliding with + /// the structured-output tool name. + /// + /// # Example + /// + /// ```rust,ignore + /// let mut run = AgentRun::new("What is 2 + 5?").max_turns(2); + /// // ... in the AgentRunStep::CallModel arm: + /// let prepared = agent.prepare_turn(prompt, &history).await?; + /// let (request, tools) = prepared.into_parts(); + /// let response = request.send().await?; + /// run.model_response(ModelTurn::new( + /// response.message_id.clone(), + /// response.choice.clone(), + /// response.usage, + /// tools.executable_tool_names().clone(), + /// tools.allowed_tool_names().clone(), + /// ))?; + /// // ... in the AgentRunStep::CallTools arm, dispatch via `tools.execute(...)`. + /// ``` + pub async fn prepare_turn( + &self, + prompt: impl Into + WasmCompatSend, + history: &[Message], + ) -> Result { + let prepared = build_prepared_completion_request( + &self.model, + prompt.into(), + history, + self.preamble.as_deref(), + &self.static_context, + self.temperature, + self.max_tokens, + self.additional_params.as_ref(), + self.record_telemetry_content, + self.tool_choice.as_ref(), + &self.tool_server_handle, + self.output_schema.as_ref(), + &self.output_mode, + // Baseline preparation matches a fresh runner's defaults: no + // run-committed output tool, the default output-tool description, + // and output-mode preamble augmentation enabled. + None, + None, + true, + // No hook stack exists on this path, so there is no request patch. + None, + ) + .await?; + Ok(super::PreparedTurn::from_prepared(prepared)) + } } // Here, we need to ensure that usage of `.prompt` on agent uses these redefinitions on the opaque diff --git a/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index b4a7828df5..cf40f7a5bb 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -103,6 +103,7 @@ mod builder; mod completion; pub mod hook; pub mod model; +mod prepared_turn; pub(crate) mod prompt_request; pub mod run; pub mod runner; @@ -123,6 +124,7 @@ pub use hook::{ ToolCallAction, ToolCallDelta, ToolResultAction, ToolResultEvent, }; pub use model::ModelHandle; +pub use prepared_turn::{PreparedTurn, TurnTools}; pub use prompt_request::streaming::{ MultiTurnStreamItem, StreamingError, StreamingPromptRequest, StreamingResult, stream_to_stdout, }; diff --git a/crates/rig-agent/src/agent/prepared_turn.rs b/crates/rig-agent/src/agent/prepared_turn.rs new file mode 100644 index 0000000000..79886bdad0 --- /dev/null +++ b/crates/rig-agent/src/agent/prepared_turn.rs @@ -0,0 +1,359 @@ +//! One prepared model turn of a configured [`Agent`](super::Agent), for +//! hand-driven [`AgentRun`](super::run::AgentRun) loops. +//! +//! [`Agent::prepare_turn`](super::Agent::prepare_turn) reads the agent's +//! configuration — preamble, static context, temperature, `max_tokens`, +//! `additional_params`, `tool_choice`, output schema/mode, and the tool +//! registry — and resolves it into a [`PreparedTurn`]: the fully configured +//! completion request for one turn plus a [`TurnTools`] value that keeps the +//! turn's advertised tool set and its dispatch target consistent. +//! +//! This exists so a hand-driven `AgentRun` loop (custom provider transport, +//! suspend/resume across processes) can reuse the `Agent` it configured instead +//! of restating every piece of that configuration by hand. It is a +//! configuration read plus a dispatch target — not a second execution path: +//! [`Agent::runner`](super::Agent::runner) remains the only way to execute an +//! agent with hooks, memory, retrieval policy, and telemetry. + +use std::collections::BTreeSet; +use std::sync::Arc; + +use super::completion::PreparedCompletionRequest; +use super::model::ModelHandle; +use crate::completion::CompletionRequestBuilder; +use crate::tool::server::ToolRegistrySnapshot; +use crate::tool::{ToolContext, ToolDispatch, ToolExecutionError, ToolResult}; + +/// One prepared model turn: the configured completion request plus the turn's +/// tool sets and dispatch target, resolved together at one instant. +/// +/// Produced by [`Agent::prepare_turn`](super::Agent::prepare_turn). The request +/// half and the tool half were computed from the same registry snapshot, so the +/// tools advertised in the request are exactly the tools [`TurnTools`] +/// dispatches — even if the agent's live registry changes afterwards. +/// +/// A prepared turn reflects the agent's **baseline** configuration. No hooks +/// run: there is no `CompletionCall` request patch, no model-selection hook, +/// and no per-turn `active_tools` narrowing. Anyone who needs those wants +/// [`Agent::runner`](super::Agent::runner). +pub struct PreparedTurn { + builder: CompletionRequestBuilder, + tools: TurnTools, +} + +impl std::fmt::Debug for PreparedTurn { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PreparedTurn") + .field("tools", &self.tools) + .finish_non_exhaustive() + } +} + +impl PreparedTurn { + pub(crate) fn from_prepared(prepared: PreparedCompletionRequest) -> Self { + let PreparedCompletionRequest { + builder, + tool_snapshot, + executable_tool_names, + allowed_tool_names, + output_tool_name, + } = prepared; + Self { + builder, + tools: TurnTools { + snapshot: tool_snapshot, + executable_tool_names, + allowed_tool_names, + output_tool_name, + }, + } + } + + /// The turn's tool sets and dispatch target. + /// + /// [`TurnTools`] is cheap to clone; clone it before consuming the request + /// with [`into_parts`](Self::into_parts) if you only need the tool half + /// later. + pub fn tools(&self) -> &TurnTools { + &self.tools + } + + /// Split the prepared turn into the configured completion request builder + /// and the turn's [`TurnTools`]. + /// + /// The builder already carries the agent's preamble (with any output-mode + /// augmentation), static context, model parameters, `tool_choice`, and this + /// turn's tool definitions. Call + /// [`send`](crate::completion::CompletionRequestBuilder::send) on it to + /// issue the request against the agent's configured model, or + /// [`build`](crate::completion::CompletionRequestBuilder::build) it to hand + /// the raw request to a custom transport. + pub fn into_parts(self) -> (CompletionRequestBuilder, TurnTools) { + (self.builder, self.tools) + } +} + +/// One turn's advertised tool sets and their dispatch target. +/// +/// All four pieces were resolved together when the turn was prepared: +/// +/// - [`executable_tool_names`](Self::executable_tool_names) — the real +/// registry tools advertised to the provider this turn. Feed this to +/// [`ModelTurn::new`](super::run::ModelTurn::new). +/// - [`allowed_tool_names`](Self::allowed_tool_names) — the tools the active +/// `tool_choice` lets the model call this turn (including the synthetic +/// output tool, which is allowed but not executable). Feed this to +/// [`ModelTurn::new`](super::run::ModelTurn::new) as well. +/// - [`output_tool_name`](Self::output_tool_name) — the synthetic +/// structured-output tool, when Tool output mode is active. +/// - [`execute`](Self::execute) — dispatch against the registry **snapshot** +/// taken when the turn was prepared. +/// +/// # Snapshot semantics +/// +/// The dispatch target is a per-turn snapshot, not the live registry. Tools +/// added to or removed from the agent after the turn was prepared (including +/// MCP refreshes) take effect on the *next* prepared turn; calls belonging to +/// this turn dispatch through the exact implementations whose definitions the +/// provider received. This is the same guarantee the agent runner gives its own +/// turns, and it is why this type exists instead of a handle to the live +/// registry: advertising one implementation and dispatching another is the skew +/// this closes. +#[derive(Clone)] +pub struct TurnTools { + snapshot: Arc, + executable_tool_names: BTreeSet, + allowed_tool_names: BTreeSet, + output_tool_name: Option, +} + +impl std::fmt::Debug for TurnTools { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TurnTools") + .field("executable_tool_names", &self.executable_tool_names) + .field("allowed_tool_names", &self.allowed_tool_names) + .field("output_tool_name", &self.output_tool_name) + .finish_non_exhaustive() + } +} + +impl TurnTools { + /// The executable registry tools advertised to the provider this turn. + pub fn executable_tool_names(&self) -> &BTreeSet { + &self.executable_tool_names + } + + /// The tools the active `tool_choice` lets the model call this turn, + /// including the synthetic output tool when Tool output mode is active. + pub fn allowed_tool_names(&self) -> &BTreeSet { + &self.allowed_tool_names + } + + /// The synthetic structured-output tool advertised this turn, when the + /// agent's output schema resolved to Tool output mode. + /// + /// This tool is **allowed but not executable**: a model call to it *is* the + /// final structured answer. A hand-driven loop must intercept it by name — + /// consume its arguments as the run's structured output instead of + /// dispatching it (dispatching it through [`execute`](Self::execute) is + /// rejected). Drivers that validate the arguments against the output schema + /// should feed failures back to the model as an error tool result and + /// re-prompt, mirroring + /// [`AgentRun::with_output_validation`](super::run::AgentRun::with_output_validation). + pub fn output_tool_name(&self) -> Option<&str> { + self.output_tool_name.as_deref() + } + + /// Execute a tool call through the exact implementation advertised for this + /// turn. + /// + /// Mirrors [`ToolServerHandle::execute`](crate::tool::server::ToolServerHandle::execute): + /// the result carries success, failure, or refusal as a [`ToolResult`], and + /// the tool's result metadata is published back to `context`. A name that + /// was not advertised this turn — including one registered on the agent + /// *after* the turn was prepared — resolves to a not-found failure rather + /// than reaching the live registry. Calling the synthetic output tool is + /// rejected: it has no implementation, and its call is the final structured + /// answer the driver must consume itself (see + /// [`output_tool_name`](Self::output_tool_name)). + pub async fn execute( + &self, + tool_name: &str, + args: &str, + context: &mut ToolContext, + ) -> ToolResult { + if self.output_tool_name.as_deref() == Some(tool_name) { + return ToolResult::failed(ToolExecutionError::not_found(format!( + "`{tool_name}` is this turn's synthetic structured-output tool: it is advertised \ + to the model but has no implementation. Its call carries the final structured \ + answer; the driver must consume the call's arguments instead of dispatching it." + ))); + } + context.clear_dispatch_result(); + let ToolDispatch { + result, + context: dispatch_context, + } = self.snapshot.dispatch(tool_name, args, context).await; + context.accept_dispatch_result(dispatch_context); + result + } +} + +#[cfg(test)] +mod tests { + use crate::agent::AgentBuilder; + use crate::agent::run::OutputMode; + use crate::completion::{CompletionError, Message}; + use crate::test_utils::{MockAddTool, MockCompletionModel, MockSubtractTool}; + use crate::tool::{ToolContext, ToolErrorKind}; + use rig_core::message::ToolChoice; + use serde_json::json; + + fn schema(value: serde_json::Value) -> schemars::Schema { + serde_json::from_value(value).expect("valid schema") + } + + /// Criterion: advertise/dispatch consistency. A prepared turn keeps + /// dispatching the implementation it advertised even after the agent's + /// live registry mutates, and never sees tools registered afterwards. + #[tokio::test] + async fn prepared_turn_dispatches_the_implementation_it_advertised() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .preamble("prepared preamble") + .tool(MockAddTool) + .build(); + + let prepared = agent + .prepare_turn("add 1 and 2", &[]) + .await + .expect("prepare succeeds"); + let (request, tools) = prepared.into_parts(); + + // The request carries the agent's configuration: the preamble leads + // the history and the registered tool is advertised. + let request = request.build(); + assert!(matches!( + request.chat_history.first(), + Some(Message::System { content }) if content == "prepared preamble" + )); + assert!(request.tools.iter().any(|tool| tool.name == "add")); + assert!(tools.executable_tool_names().contains("add")); + assert_eq!(tools.executable_tool_names(), tools.allowed_tool_names()); + + // Mutate the live registry AFTER preparation: retire the advertised + // tool and register a different one. + agent.tool_server_handle.remove_tool("add").await; + agent.tool_server_handle.add_tool(MockSubtractTool).await; + + // The prepared turn still dispatches the implementation it advertised... + let mut context = ToolContext::new(); + let result = tools.execute("add", r#"{"x": 1, "y": 2}"#, &mut context).await; + assert!( + result.is_success(), + "snapshot dispatch must reach the advertised implementation" + ); + + // ...and does not see tools registered after the snapshot was taken. + let mut context = ToolContext::new(); + let result = tools + .execute("subtract", r#"{"x": 1, "y": 2}"#, &mut context) + .await; + assert!(result.is_error_kind(ToolErrorKind::NotFound)); + + // The live registry really did change underneath the snapshot. + let mut context = ToolContext::new(); + let result = agent + .tool_server_handle + .execute("add", r#"{"x": 1, "y": 2}"#, &mut context) + .await; + assert!(result.is_error_kind(ToolErrorKind::NotFound)); + } + + /// Criterion: under Tool output mode the synthetic output tool is allowed + /// and advertised but never executable, and dispatching it is rejected. + #[tokio::test] + async fn output_tool_is_allowed_but_never_executable() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .tool(MockAddTool) + .output_schema_raw(schema(json!({ + "type": "object", + "properties": { "value": { "type": "integer" } }, + "required": ["value"] + }))) + .output_mode(OutputMode::Tool) + .build(); + + let prepared = agent + .prepare_turn("compute", &[]) + .await + .expect("prepare succeeds"); + let (request, tools) = prepared.into_parts(); + let output_tool = tools + .output_tool_name() + .expect("Tool output mode advertises an output tool") + .to_owned(); + + assert!(tools.allowed_tool_names().contains(&output_tool)); + assert!(!tools.executable_tool_names().contains(&output_tool)); + // The provider sees the synthetic tool alongside the real one. + let request = request.build(); + assert!(request.tools.iter().any(|tool| tool.name == output_tool)); + assert!(request.tools.iter().any(|tool| tool.name == "add")); + + let mut context = ToolContext::new(); + let result = tools.execute(&output_tool, "{}", &mut context).await; + assert!(result.is_error_kind(ToolErrorKind::NotFound)); + let message = result.error().expect("rejected dispatch").message().to_owned(); + assert!( + message.contains("structured-output"), + "rejection must explain the synthetic tool: {message}" + ); + } + + /// Criterion: an impossible `ToolChoice` fails at prepare time, locally, + /// with no provider round-trip. + #[tokio::test] + async fn impossible_tool_choice_fails_locally_at_prepare_time() { + let model = MockCompletionModel::text("unused"); + let agent = AgentBuilder::new(model.clone()) + .tool(MockAddTool) + .tool_choice(ToolChoice::Specific { + function_names: vec!["missing".to_string()], + }) + .build(); + + let err = agent + .prepare_turn("go", &[]) + .await + .expect_err("a tool_choice naming an unadvertised tool must fail at prepare time"); + assert!(matches!( + err, + CompletionError::RequestError(inner) if inner.to_string().contains("missing") + )); + assert_eq!( + model.request_count(), + 0, + "local validation must not cost a provider round-trip" + ); + } + + /// `Required` forces a tool call, so an empty advertised set can never + /// satisfy it — that must fail at prepare time, not degrade silently. + #[tokio::test] + async fn required_tool_choice_with_no_tools_fails_locally() { + let model = MockCompletionModel::text("unused"); + let agent = AgentBuilder::new(model.clone()) + .tool_choice(ToolChoice::Required) + .build(); + + let err = agent + .prepare_turn("go", &[]) + .await + .expect_err("Required with no advertised tool must fail at prepare time"); + assert!(matches!( + err, + CompletionError::RequestError(inner) if inner.to_string().contains("Required") + )); + assert_eq!(model.request_count(), 0); + } +} diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index 763529be6d..d6211b35d8 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -24,10 +24,17 @@ //! //! `AgentRun` deliberately contains no model, tool registry, memory backend, or //! hook stack. Hand-driving it is a low-level provider integration: the caller -//! owns all IO and any lifecycle policy. To execute a configured [`Agent`](crate::agent::Agent) -//! with its hooks, tools, retrieval, and memory, use +//! owns all IO and any lifecycle policy. To drive a *configured* +//! [`Agent`](crate::agent::Agent) by hand, obtain each turn's completion +//! request and tool dispatch from +//! [`Agent::prepare_turn`](crate::agent::Agent::prepare_turn) — it resolves the +//! agent's configuration (preamble, tools, `tool_choice`, output mode, model +//! parameters) into a [`PreparedTurn`](crate::agent::PreparedTurn) instead of +//! the driver restating it. To execute an `Agent` with its hooks, memory, +//! retrieval policy, and telemetry, use //! [`Agent::runner`](crate::agent::Agent::runner); constructing an `AgentRun` -//! directly is not an alternate way to execute an `Agent`. +//! directly is not an alternate way to execute an `Agent`, and a prepared turn +//! is a configuration read, not a second execution path. //! //! [`crate::completion::Prompt::prompt`] and //! [`Agent::runner`](crate::agent::Agent::runner) drive this machine internally; diff --git a/examples/agent_run_stepping/src/main.rs b/examples/agent_run_stepping/src/main.rs index 05aa3d2d80..95d70cc3b0 100644 --- a/examples/agent_run_stepping/src/main.rs +++ b/examples/agent_run_stepping/src/main.rs @@ -7,6 +7,11 @@ //! because the machine is fully serializable between steps — pause a run while //! tool calls are pending and resume it later (even in another process). //! +//! The per-turn request comes from [`Agent::prepare_turn`]: the loop reuses the +//! configured `Agent`'s preamble, tools, and model parameters instead of +//! restating them, and dispatches tool calls through the same registry snapshot +//! the provider saw advertised. +//! //! ## Part 2 — high-level [`rig::agent::AgentRunner`] with hooks //! //! For the common case you don't need that level of control: attach an @@ -17,18 +22,16 @@ //! //! Requires `OPENAI_API_KEY`. -use std::collections::BTreeSet; - use anyhow::Result; +use rig::agent::TurnTools; use rig::agent::run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}; use rig::agent::{ AgentHook, HookContext, InvalidToolCallAction, ToolCall as ToolCallEvent, ToolCallAction, }; -use rig::completion::CompletionModel; use rig::message::UserContent; use rig::prelude::*; use rig::providers::openai; -use rig::tool::{Tool, ToolSet}; +use rig::tool::Tool; use serde::Deserialize; use serde_json::json; @@ -92,14 +95,16 @@ impl AgentHook for ToolLoggerHook { async fn main() -> Result<()> { let openai = openai::Client::from_env()?; let model = openai.completion_model(openai::GPT_4O); - let agent = rig::agent::AgentBuilder::new(model.clone()) + let agent = rig::agent::AgentBuilder::new(model) .preamble("You are a calculator. Always use the provided tools to compute results.") .tool(Add) .build(); - let local_tools = ToolSet::builder().static_tool(Add).build(); - let tool_definitions = local_tools.get_tool_definitions(); let mut run = AgentRun::new("What is 2 + 5?").max_turns(2); + // The tool sets and dispatch target of the most recent prepared turn. Tool + // calls always execute through the snapshot whose definitions the provider + // saw — the same guarantee the runner gives its own turns. + let mut turn_tools: Option = None; loop { match run.next_step()? { @@ -110,34 +115,21 @@ async fn main() -> Result<()> { } => { println!("→ model call #{turn}"); // A hand-driven `AgentRun` is a sans-IO protocol primitive, not - // execution of the configured `Agent`. Its transport is an - // explicit raw model request and therefore has no agent hooks. - let response = model - .completion_request(prompt) - .messages(history) - .preamble( - "You are a calculator. Always use the provided tools to compute results." - .to_string(), - ) - .tools(tool_definitions.clone()) - .send() - .await?; - - // The tools advertised to the provider for this turn. With - // static tools these are the agent's registered tools; agents - // with dynamic (RAG) tools would resolve them per turn. - let tool_names: BTreeSet = tool_definitions - .iter() - .map(|def| def.name.clone()) - .collect(); + // execution of the configured `Agent`: the driver owns the IO + // and no agent hooks run. `prepare_turn` supplies the request — + // preamble, tools, model parameters — from the agent's + // configuration instead of restating it here. + let (request, tools) = agent.prepare_turn(prompt, &history).await?.into_parts(); + let response = request.send().await?; let mut outcome = run.model_response(ModelTurn::new( response.message_id.clone(), response.choice.clone(), response.usage, - tool_names.clone(), - tool_names, + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), ))?; + turn_tools = Some(tools); while let ModelTurnOutcome::NeedsResolution(context) = outcome { eprintln!("model called unknown tool `{}`", context.tool_name); // Preserve the agent loop's default fail-fast behavior; a @@ -149,12 +141,18 @@ async fn main() -> Result<()> { // The whole run is serializable while tool calls are pending: // persist it here to pause for approval and resume later — // even in a process that never saw this step. The resumed run - // re-emits the pending tool calls from its own state. + // re-emits the pending tool calls from its own state. (Tool + // implementations are live objects: a genuinely separate + // process would rebuild the same `Agent` and prepare its own + // turns from there.) let suspended = serde_json::to_string(&run)?; let mut run_resumed: AgentRun = serde_json::from_str(&suspended)?; let AgentRunStep::CallTools { calls } = run_resumed.next_step()? else { anyhow::bail!("resumed run must re-emit the pending tool calls"); }; + let Some(tools) = turn_tools.as_ref() else { + anyhow::bail!("CallTools always follows a prepared CallModel turn"); + }; let mut results = Vec::new(); for call in calls { @@ -168,7 +166,7 @@ async fn main() -> Result<()> { let args = call.tool_call.function.arguments.to_string(); println!("→ executing {name}({args})"); let mut context = rig::tool::ToolContext::new(); - let result = local_tools.execute(name, args, &mut context).await; + let result = tools.execute(name, &args, &mut context).await; results.push(UserContent::tool_result_for( call.tool_call.id.clone(), call.tool_call.provider.clone(), diff --git a/examples/agent_with_durable_approval/src/main.rs b/examples/agent_with_durable_approval/src/main.rs index 19ccfbfbc4..a16a717a8a 100644 --- a/examples/agent_with_durable_approval/src/main.rs +++ b/examples/agent_with_durable_approval/src/main.rs @@ -28,16 +28,14 @@ //! Requires `OPENAI_API_KEY`. Run with: `cargo run -p agent_with_durable_approval` use anyhow::Result; -use rig::agent::InvalidToolCallAction; use rig::agent::run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}; -use rig::completion::CompletionModel; +use rig::agent::{InvalidToolCallAction, TurnTools}; use rig::message::{ToolResultContent, UserContent}; use rig::prelude::*; use rig::providers::openai; -use rig::tool::{Tool, ToolSet}; +use rig::tool::Tool; use serde::Deserialize; use serde_json::json; -use std::collections::BTreeSet; // --------------------------------------------------------------------------- // One read-only tool and one side-effecting tool worth gating. @@ -145,17 +143,19 @@ async fn ask(prompt: &str) -> Option { #[tokio::main] async fn main() -> Result<()> { - // A serializable `AgentRun` is a sans-IO protocol primitive. This example - // intentionally supplies raw model transport and tool dispatch explicitly; - // configured `Agent` execution instead always goes through `AgentRunner`. + // A hand-driven `AgentRun` is a sans-IO protocol primitive: this loop owns + // the IO and no agent hooks run. The per-turn request and the tool dispatch + // target both come from the configured `Agent` via `prepare_turn`, so + // nothing about the agent's configuration is restated below. let model = openai::Client::from_env()?.completion_model(openai::GPT_4O); - let preamble = "You are a banking assistant. Use the tools to carry out the user's request. \ - Call one tool at a time."; - let tools = ToolSet::builder() - .static_tool(GetBalance) - .static_tool(TransferFunds) + let agent = rig::agent::AgentBuilder::new(model) + .preamble( + "You are a banking assistant. Use the tools to carry out the user's request. \ + Call one tool at a time.", + ) + .tool(GetBalance) + .tool(TransferFunds) .build(); - let tool_definitions = tools.get_tool_definitions(); let prompt = "Check the balance of account A-1, then transfer $500 to account B-2."; println!("User: {prompt}"); @@ -165,6 +165,9 @@ async fn main() -> Result<()> { let _ = std::fs::remove_file(&state_path); let mut run = AgentRun::new(prompt).max_turns(10); + // Tool dispatch target of the most recent prepared turn: approved calls + // execute through the snapshot whose definitions the provider saw. + let mut turn_tools: Option = None; loop { match run.next_step()? { @@ -174,24 +177,16 @@ async fn main() -> Result<()> { turn, } => { println!("\n→ model call #{turn}"); - let response = model - .completion_request(prompt) - .messages(history) - .preamble(preamble.to_string()) - .tools(tool_definitions.clone()) - .send() - .await?; - let tool_names: BTreeSet = tool_definitions - .iter() - .map(|def| def.name.clone()) - .collect(); + let (request, tools) = agent.prepare_turn(prompt, &history).await?.into_parts(); + let response = request.send().await?; let mut outcome = run.model_response(ModelTurn::new( response.message_id.clone(), response.choice.clone(), response.usage, - tool_names.clone(), - tool_names, + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), ))?; + turn_tools = Some(tools); while let ModelTurnOutcome::NeedsResolution(context) = outcome { eprintln!("model called unknown tool `{}`", context.tool_name); outcome = run.resolve_invalid_tool_call(InvalidToolCallAction::fail())?; @@ -207,11 +202,17 @@ async fn main() -> Result<()> { println!("\n💾 run suspended to {}", state_path.display()); // ----- imagine the process exits here and resumes later ----- + // (Tool implementations are live objects: a genuinely separate + // process would rebuild the same `Agent` and prepare its own + // turns from there.) let mut resumed: AgentRun = serde_json::from_slice(&std::fs::read(&state_path)?)?; let AgentRunStep::CallTools { calls } = resumed.next_step()? else { anyhow::bail!("resumed run must re-emit the pending tool calls"); }; + let Some(tools) = turn_tools.as_ref() else { + anyhow::bail!("CallTools always follows a prepared CallModel turn"); + }; let mut results = Vec::new(); let mut aborted = false; @@ -233,7 +234,7 @@ async fn main() -> Result<()> { { Some("a") | Some("approve") => { let execution = tools - .execute(&name, args, &mut rig::tool::ToolContext::new()) + .execute(&name, &args, &mut rig::tool::ToolContext::new()) .await; results.push(UserContent::tool_result_for( id, @@ -252,7 +253,7 @@ async fn main() -> Result<()> { let execution = tools .execute( &name, - value.to_string(), + &value.to_string(), &mut rig::tool::ToolContext::new(), ) .await; From eba03b2bd25d587d6f09deb304a6bb5b2b2e6c65 Mon Sep 17 00:00:00 2001 From: stephen Date: Mon, 10 Aug 2026 18:47:01 -0700 Subject: [PATCH 02/21] refactor!(agent): replace prepare_turn with the AgentDriver shell Review of the unreleased prepare_turn API (PR #2278 round 1) verified ten findings; five of them - the unarmed output-tool intercept, the lost Tool-mode pin, the examples bailing on cross-process resume, the transposable ModelTurn incantation, and run construction restating agent config - were symptoms of one missing abstraction: the pairing between run-state (AgentRun) and turn-state (TurnTools) lived in every caller. With the API unreleased, the architecture is fixed instead of the symptoms. Agent::drive(prompt) / Agent::drive_run(run) return an AgentDriver that owns the pairing and nothing else. It seeds the run from the agent's configuration (default_max_turns, tool_choice, output schema), yields DriveStep::SendRequest with the fully configured request for the caller to send (or hand to a custom transport), pairs DriveStep::ExecuteTools with the exact registry snapshot that advertised the calls, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry remain AgentRunner-only, and it is built strictly on AgentRun's own step methods and the runner's shared build_agent_run/build_prepared_completion_request - no duplicated turn-decision logic. Resume in a fresh process re-derives a fresh dispatch snapshot from the rebuilt agent; a pending call whose tool is missing from this process's registry is surfaced as an error before dispatch (opt out with allow_missing_resumed_tools). PreparedTurn and Agent::prepare_turn are removed; the public surface is AgentDriver + DriveStep + TurnTools, with TurnTools now cheaply cloneable (Arc-shared sets) and gaining execute_call for pending-call dispatch with preresolved-result handling. Standalone review fixes in the same pass: - ToolErrorKind::NotExecutable (rig-core, non_exhaustive enum): the output-tool dispatch rejection is machine-readable instead of an ambiguous NotFound. - The clear -> dispatch -> publish sequence is single-sourced (ToolDispatch::publish_to + ToolRegistrySnapshot::execute); the rejection path can no longer skip clear_dispatch_result (regression test included). - PreparedCompletionRequest composes TurnTools directly; the 1:1 mirror struct and its repack are gone. - The "no telemetry" overclaim is corrected: the prepared request honors the agent's record_telemetry_content, documented with the opt-out. Both examples now contain zero pairing logic - the CallTools arm is decision logic plus execute_call - and demonstrate true fresh-driver resume. Runner and streaming paths are behaviorally unchanged; all 518 pre-existing rig-agent tests pass as-is. --- MIGRATING.md | 14 +- crates/rig-agent/CHANGELOG.md | 3 + crates/rig-agent/src/agent/completion.rs | 121 +-- crates/rig-agent/src/agent/driver.rs | 771 ++++++++++++++++++ crates/rig-agent/src/agent/mod.rs | 6 +- crates/rig-agent/src/agent/prepared_turn.rs | 359 -------- .../src/agent/prompt_request/streaming.rs | 11 +- crates/rig-agent/src/agent/run/mod.rs | 18 +- crates/rig-agent/src/agent/runner.rs | 15 +- crates/rig-agent/src/agent/turn_tools.rs | 262 ++++++ crates/rig-agent/src/tool/mod.rs | 21 +- crates/rig-agent/src/tool/server.rs | 25 +- crates/rig-core/CHANGELOG.md | 1 + crates/rig-core/src/tool/result.rs | 23 +- examples/agent_run_stepping/src/main.rs | 121 +-- .../agent_with_durable_approval/src/main.rs | 72 +- 16 files changed, 1205 insertions(+), 638 deletions(-) create mode 100644 crates/rig-agent/src/agent/driver.rs delete mode 100644 crates/rig-agent/src/agent/prepared_turn.rs create mode 100644 crates/rig-agent/src/agent/turn_tools.rs diff --git a/MIGRATING.md b/MIGRATING.md index 933a801258..ce1a2e793b 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1563,12 +1563,14 @@ not an alternate execution path for configured agents. Reading a configured agent's request state back out (the 0.40 public fields, `agent.tool_server_handle`) has a supported replacement: -`Agent::prepare_turn(prompt, &history)` resolves the agent's configuration into -one turn's completion request plus its tool sets and dispatch snapshot -(`PreparedTurn` / `TurnTools`), so a hand-driven `AgentRun` no longer restates -the preamble and tools. It is a configuration read, not an execution path — -hooks, memory, retrieval policy, and telemetry still run only under -`Agent::runner`. See `examples/agent_run_stepping`. +`Agent::drive(prompt)` returns an `AgentDriver` that hand-drives the sans-IO +`AgentRun` with the agent's own configuration — each `DriveStep::SendRequest` +carries the fully configured completion request for the caller to send, and +each `DriveStep::ExecuteTools` carries the `TurnTools` dispatch snapshot of +the turn that advertised the calls. It is a configuration and pairing layer, +not an execution path — hooks, memory, retrieval policy, and telemetry still +run only under `Agent::runner`. See `examples/agent_run_stepping` and +`examples/agent_with_durable_approval`. An `Agent`'s default model is set at construction. Per-run overrides now go through `runner(...).using_model(...)`, `Agent::set_model`, or a diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index ee4901b74f..39e5b9f490 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -32,6 +32,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `AgentHook::on_reasoning_delta` with the Rig stream correlator, optional provider reasoning id, current fragment, and per-part aggregate; reasoning hooks share the existing observation-interest and stop-before-yield semantics used by other streaming deltas - *(agent)* add `Agent::prepare_turn`, returning the new public `PreparedTurn`/`TurnTools`: the agent's baseline configuration resolved into one turn's completion request plus the turn's executable and allowed tool-name sets, the synthetic output-tool name, and tool dispatch pinned to the turn's registry snapshot — so hand-driven `AgentRun` loops (custom provider transport, durable suspend/resume) reuse the configured `Agent` instead of restating its preamble and tools. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. A prepared turn is a configuration read: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` + +- *(agent)* add `Agent::drive` / `Agent::drive_run`, returning the new public `AgentDriver` (with `DriveStep` and `TurnTools`): hand-drive the sans-IO `AgentRun` machine with the agent's own configuration while owning every side effect. The driver seeds the run from the agent (`default_max_turns`, `tool_choice`, output schema), yields each turn's fully configured completion request for the caller to send (or hand to a custom transport), pairs every pending tool batch with the exact registry snapshot that advertised it, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. Resuming a serialized run in a fresh process re-derives a fresh dispatch snapshot from the rebuilt agent, and surfaces missing pending tools as an error (opt out with `allow_missing_resumed_tools`). The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` +- *(tool)* add `ToolErrorKind::NotExecutable` (`ToolExecutionError::not_executable`): the tool is advertised to the model but not executable by the dispatcher — produced when dispatching the synthetic structured-output tool, whose call carries the final structured answer - *(agent)* add opaque, cloneable `ModelHandle` values with by-value `ProviderCapabilities` snapshots, plus default replacement, per-run default override (`using_model`), and hook-driven per-call selection via `AgentHook::on_model_select` - *(agent)* add run-local extractor default-model overrides used across retries diff --git a/crates/rig-agent/src/agent/completion.rs b/crates/rig-agent/src/agent/completion.rs index b642bc8fd5..8d2467f1c3 100644 --- a/crates/rig-agent/src/agent/completion.rs +++ b/crates/rig-agent/src/agent/completion.rs @@ -6,33 +6,18 @@ use super::runner::AgentRunner; use crate::{ agent::prompt_request::streaming::StreamingPromptRequest, completion::{ - Chat, CompletionError, CompletionModel, CompletionRequestBuilder, Document, Message, - Prompt, PromptError, ToolDefinition, TypedPrompt, + Chat, CompletionError, CompletionModel, Document, Message, Prompt, PromptError, + ToolDefinition, TypedPrompt, }, json_utils, streaming::{StreamingChat, StreamingPrompt}, - tool::server::{ToolRegistrySnapshot, ToolServerError, ToolServerHandle}, + tool::server::{ToolServerError, ToolServerHandle}, }; use rig_core::{message::ToolChoice, wasm_compat::WasmCompatSend}; use std::{collections::BTreeSet, sync::Arc}; use super::UNKNOWN_AGENT_NAME; - -/// A prepared completion request plus the executable Rig tool names advertised -/// to the provider for this turn. -pub(crate) struct PreparedCompletionRequest { - /// Builder carrying the selected model handle: request preparation ran - /// against this handle's captured capabilities, and the same handle - /// executes the prepared request. - pub(crate) builder: CompletionRequestBuilder, - /// Exact implementations behind this turn's provider definitions. - pub(crate) tool_snapshot: Arc, - pub(crate) executable_tool_names: BTreeSet, - pub(crate) allowed_tool_names: BTreeSet, - /// When Tool output mode is active, the name of the synthetic output tool - /// advertised to the model (allowed but not executable). See #1928. - pub(crate) output_tool_name: Option, -} +use super::turn_tools::{PreparedCompletionRequest, TurnTools}; /// Base name of the synthetic output tool used by [`OutputMode::Tool`]. const DEFAULT_OUTPUT_TOOL_NAME: &str = "final_result"; @@ -518,10 +503,12 @@ pub(crate) async fn build_prepared_completion_request( Ok(PreparedCompletionRequest { builder: completion_request, - tool_snapshot: Arc::new(tool_snapshot), - executable_tool_names, - allowed_tool_names, - output_tool_name, + tools: TurnTools { + snapshot: Arc::new(tool_snapshot), + executable_tool_names: Arc::new(executable_tool_names), + allowed_tool_names: Arc::new(allowed_tool_names), + output_tool_name, + }, }) } @@ -668,94 +655,6 @@ impl Agent { ) -> Result, ToolServerError> { self.tool_server_handle.get_tool_defs(prompt).await } - - /// Prepare one model turn from this agent's configuration, for hand-driven - /// [`AgentRun`](super::run::AgentRun) loops. - /// - /// Resolves the agent's preamble, static context, model parameters, - /// `tool_choice`, output schema/mode, and tool registry into a - /// [`PreparedTurn`](super::PreparedTurn): the fully configured completion request plus a - /// [`TurnTools`](super::TurnTools) value carrying the turn's executable and - /// allowed tool-name sets, the synthetic output-tool name (if any), and a - /// dispatch target pinned to this turn's registry snapshot. Use it to drive - /// the sans-IO [`AgentRun`](super::run::AgentRun) machine — custom provider - /// transport, suspend/resume — without restating the agent's configuration. - /// - /// # Baseline, not hook-patched - /// - /// A prepared turn reflects the agent's **baseline** configuration: no - /// `CompletionCall` hooks run (so no per-turn `RequestPatch` or - /// `active_tools` narrowing), no model-selection hooks run (the agent's - /// current default model is used), and no memory or telemetry is involved. - /// [`Agent::runner`](Self::runner) remains the only path that executes an - /// agent with hooks, memory, retrieval policy, and telemetry. - /// - /// # Per-turn snapshot - /// - /// The tool definitions in the request and the dispatch target in - /// [`TurnTools`](super::TurnTools) come from one registry snapshot taken by - /// this call. Registry changes made afterwards (`add_tool` / `remove_tool`, - /// MCP refreshes) affect the next prepared turn, never this one — see - /// [`TurnTools`](super::TurnTools) for the semantics. Retrieval-selected - /// (dynamic) tools are resolved here too, using the prompt's text (or the - /// latest history text) as the retrieval query, exactly as the runner does. - /// - /// # Errors - /// - /// Fails locally — with no provider round-trip — when the configuration - /// cannot produce a valid request: a `tool_choice` that is impossible - /// against the advertised tool set (`Required` with no advertised tool, or - /// `Specific` naming an unadvertised tool), or a real tool colliding with - /// the structured-output tool name. - /// - /// # Example - /// - /// ```rust,ignore - /// let mut run = AgentRun::new("What is 2 + 5?").max_turns(2); - /// // ... in the AgentRunStep::CallModel arm: - /// let prepared = agent.prepare_turn(prompt, &history).await?; - /// let (request, tools) = prepared.into_parts(); - /// let response = request.send().await?; - /// run.model_response(ModelTurn::new( - /// response.message_id.clone(), - /// response.choice.clone(), - /// response.usage, - /// tools.executable_tool_names().clone(), - /// tools.allowed_tool_names().clone(), - /// ))?; - /// // ... in the AgentRunStep::CallTools arm, dispatch via `tools.execute(...)`. - /// ``` - pub async fn prepare_turn( - &self, - prompt: impl Into + WasmCompatSend, - history: &[Message], - ) -> Result { - let prepared = build_prepared_completion_request( - &self.model, - prompt.into(), - history, - self.preamble.as_deref(), - &self.static_context, - self.temperature, - self.max_tokens, - self.additional_params.as_ref(), - self.record_telemetry_content, - self.tool_choice.as_ref(), - &self.tool_server_handle, - self.output_schema.as_ref(), - &self.output_mode, - // Baseline preparation matches a fresh runner's defaults: no - // run-committed output tool, the default output-tool description, - // and output-mode preamble augmentation enabled. - None, - None, - true, - // No hook stack exists on this path, so there is no request patch. - None, - ) - .await?; - Ok(super::PreparedTurn::from_prepared(prepared)) - } } // Here, we need to ensure that usage of `.prompt` on agent uses these redefinitions on the opaque diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs new file mode 100644 index 0000000000..ca466d926a --- /dev/null +++ b/crates/rig-agent/src/agent/driver.rs @@ -0,0 +1,771 @@ +//! Hand-drive a configured [`Agent`] while owning every side effect yourself. +//! +//! [`Agent::drive`] pairs the sans-IO [`AgentRun`] state machine with the +//! agent's configuration: each [`DriveStep::SendRequest`] carries the fully +//! configured completion request (the caller sends it — or hands it to a +//! custom transport), and each [`DriveStep::ExecuteTools`] carries the +//! [`TurnTools`] of the turn that advertised those calls, paired by +//! construction rather than by caller discipline. The driver itself performs +//! **no** provider IO and **no** tool dispatch, and runs no hooks, memory, +//! retrieval policy, or telemetry — it is the run/turn pairing logic the +//! runner keeps internally, exposed for callers who own the IO. To *execute* +//! an agent with hooks, memory, and telemetry, use +//! [`Agent::runner`](super::Agent::runner); the driver is not a second +//! execution path. +//! +//! # Why a driver and not a bag of getters +//! +//! Hand-driving `AgentRun` requires state that must stay mutually consistent: +//! the request a turn sent, the tool sets that validate the model's calls, the +//! snapshot that dispatches them, and the run's committed structured-output +//! tool. Leaving that pairing to callers is how configuration drift and +//! advertise/dispatch skew happen; the driver owns it in one place while every +//! side effect stays with the caller. +//! +//! # Durability +//! +//! The serializable state is still [`AgentRun`] — serialize +//! [`AgentDriver::run`] while tool calls are pending, and resume in another +//! process with [`Agent::drive_run`]. Tool implementations are live objects +//! and cannot be serialized: the resuming process rebuilds the same `Agent` +//! and the driver takes a **fresh** registry snapshot for the pending calls. +//! If that snapshot no longer contains a pending call's tool, the driver +//! surfaces an error instead of silently feeding a not-found result to the +//! model (see [`AgentDriver::allow_missing_resumed_tools`]) — the model chose +//! that tool from a registry this process no longer has, and re-prompting +//! cannot fix deployment drift. If you suspend runs across deploys, version +//! your agent definitions alongside the serialized run. + +use std::collections::BTreeSet; +use std::sync::Arc; + +use rig_core::message::UserContent; + +use super::completion::{Agent, allowed_tool_names_for_choice, build_prepared_completion_request}; +use super::model::ModelHandle; +use super::run::{AgentRun, AgentRunStep, ModelTurnOutcome, PendingToolCall}; +use super::runner::build_agent_run; +use super::turn_tools::{PreparedCompletionRequest, TurnTools}; +use crate::agent::hook::InvalidToolCallAction; +use crate::agent::prompt_request::PromptResponse; +use crate::completion::{ + CompletionError, CompletionRequestBuilder, CompletionResponse, Message, PromptError, +}; + +impl Agent { + /// Hand-drive this agent: build a driver whose run is seeded from the + /// agent's configuration. + /// + /// Seeding mirrors [`Agent::runner`]: the model-call budget comes from + /// `default_max_turns` (implicit budget of one when unset), and the run + /// inherits the agent's `tool_choice` and output schema (with the default + /// output-retry budget). Override per run with [`AgentDriver::max_turns`] + /// and [`AgentDriver::history`], or construct a custom [`AgentRun`] and + /// use [`Agent::drive_run`]. + pub fn drive(&self, prompt: impl Into) -> AgentDriver { + self.drive_run(build_agent_run( + prompt.into(), + self.default_max_turns.unwrap_or(1), + 0, + self.output_schema.as_ref(), + None, + self.tool_choice.clone(), + )) + } + + /// Hand-drive an existing [`AgentRun`] with this agent's configuration — + /// the resume path for a run deserialized in a new process, or the entry + /// point for a custom-configured run (which is taken as-is, not re-seeded). + pub fn drive_run(&self, run: AgentRun) -> AgentDriver { + AgentDriver { + agent: self.clone(), + run, + turn_tools: None, + allow_missing_resumed_tools: false, + } + } +} + +/// What the caller must do next to advance an [`AgentDriver`]. +/// +/// Deliberately exhaustive, like [`AgentRunStep`]: a driver loop must handle +/// every step, so adding a variant is a breaking change by design. +pub enum DriveStep { + /// Send this completion request to the model — via + /// [`send`](CompletionRequestBuilder::send), or + /// [`build`](CompletionRequestBuilder::build) it for a custom transport — + /// then feed the response back through [`AgentDriver::model_response`]. + SendRequest { + /// The fully configured request: the agent's preamble (with any + /// output-mode augmentation), static context, model parameters, + /// `tool_choice`, and this turn's tool definitions. It reflects the + /// agent's **baseline** configuration — no `CompletionCall` hooks run + /// on this path, so there is no per-turn request patch, model + /// selection, or `active_tools` narrowing. The request still honors + /// the agent's `record_telemetry_content` for provider-level spans; + /// call `.record_content_telemetry(false)` on it to opt a hand-driven + /// turn out. + request: Box>, + /// The turn's advertised tool sets — informational here (the driver + /// assembles the model turn itself); the same value arrives on the + /// following [`ExecuteTools`](Self::ExecuteTools) step for dispatch. + tools: TurnTools, + /// One-based index of this model call within the run. + turn: usize, + }, + /// Execute these tool calls — typically via + /// [`TurnTools::execute_call`] — and feed the results back through + /// [`AgentDriver::tool_results`]. `tools` is the dispatch target of the + /// turn that advertised `calls`, paired by construction. + ExecuteTools { + /// The pending tool calls of the current assistant turn, in emission + /// order. + calls: Vec, + /// The advertising turn's tool sets and snapshot dispatch target. + tools: TurnTools, + }, + /// The run is complete. + Done(PromptResponse), +} + +impl std::fmt::Debug for DriveStep { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::SendRequest { tools, turn, .. } => f + .debug_struct("SendRequest") + .field("turn", turn) + .field("tools", tools) + .finish_non_exhaustive(), + Self::ExecuteTools { calls, tools } => f + .debug_struct("ExecuteTools") + .field("calls", calls) + .field("tools", tools) + .finish(), + Self::Done(response) => f.debug_tuple("Done").field(response).finish(), + } + } +} + +/// Hand-drives one [`AgentRun`] with one [`Agent`]'s configuration. Built by +/// [`Agent::drive`] / [`Agent::drive_run`]; see the [module docs](self) for +/// the driving protocol and the boundary with [`AgentRunner`](super::AgentRunner). +pub struct AgentDriver { + agent: Agent, + run: AgentRun, + /// Tool state of the most recently prepared turn. `None` until the first + /// `SendRequest` — or in a process that resumed a serialized run, where + /// [`Self::resume_tools`] derives a fresh snapshot on demand. + turn_tools: Option, + allow_missing_resumed_tools: bool, +} + +impl AgentDriver { + /// Set the input chat history preceding the prompt. + pub fn history(mut self, history: Vec) -> Self { + self.run = self.run.with_history(history); + self + } + + /// Override the seeded total model-call budget for this run. + pub fn max_turns(mut self, max_turns: usize) -> Self { + self.run = self.run.max_turns(max_turns); + self + } + + /// Opt out of the resumed-run drift check: dispatch pending calls whose + /// tools are missing from this process's registry anyway, feeding the + /// resulting not-found errors to the model instead of surfacing the drift + /// to the caller. See the [module docs](self) on durability. + pub fn allow_missing_resumed_tools(mut self) -> Self { + self.allow_missing_resumed_tools = true; + self + } + + /// The sans-IO run state. Serialize this to suspend the run — for example + /// while tool calls are pending approval — and resume it elsewhere with + /// [`Agent::drive_run`]. + pub fn run(&self) -> &AgentRun { + &self.run + } + + /// Consume the driver, returning the run state. + pub fn into_run(self) -> AgentRun { + self.run + } + + /// Advance to the next step the caller must perform. + /// + /// Preparing a model turn reads the agent's configuration and tool + /// registry (one snapshot per turn) and maintains the run's committed + /// structured-output tool exactly as the runner does: the committed name + /// is re-advertised on every later turn, so Tool output mode cannot flip + /// or re-pick a name mid-run. Fails locally — with no provider + /// round-trip — when the configuration cannot produce a valid request + /// (e.g. a `tool_choice` impossible against the advertised tool set). + pub async fn next_step(&mut self) -> Result { + match self.run.next_step()? { + AgentRunStep::CallModel { + prompt, + history, + turn, + } => { + // Pin Tool output mode once committed (#1928), mirroring the + // runner: read the run's committed name into preparation, and + // store the resolved name back (fill-once). + let committed = self.run.output_tool_name().map(str::to_owned); + let prepared = build_prepared_completion_request( + &self.agent.model, + prompt, + &history, + self.agent.preamble.as_deref(), + &self.agent.static_context, + self.agent.temperature, + self.agent.max_tokens, + self.agent.additional_params.as_ref(), + self.agent.record_telemetry_content, + self.agent.tool_choice.as_ref(), + &self.agent.tool_server_handle, + self.agent.output_schema.as_ref(), + &self.agent.output_mode, + committed.as_deref(), + None, + true, + None, + ) + .await + .map_err(PromptError::CompletionError)?; + self.run + .set_output_tool_name(prepared.tools.output_tool_name.clone()); + let PreparedCompletionRequest { builder, tools } = prepared; + self.turn_tools = Some(tools.clone()); + Ok(DriveStep::SendRequest { + request: Box::new(builder), + tools, + turn, + }) + } + AgentRunStep::CallTools { calls } => { + let tools = match &self.turn_tools { + Some(tools) => tools.clone(), + // A process resuming a serialized run wakes here with no + // prepared turn: derive a fresh dispatch snapshot. + None => { + let tools = self.resume_tools().await?; + if !self.allow_missing_resumed_tools { + let missing: Vec<&str> = calls + .iter() + .filter(|call| call.preresolved_result.is_none()) + .map(|call| call.tool_call.function.name.as_str()) + .filter(|name| { + !tools.executable_tool_names.contains(*name) + && tools.output_tool_name() != Some(*name) + }) + .collect(); + if !missing.is_empty() { + return Err(PromptError::CompletionError( + CompletionError::RequestError( + format!( + "resumed run has pending tool calls {missing:?} that \ + are not registered in this process; register the \ + tools on the agent before resuming, or call \ + `allow_missing_resumed_tools()` to dispatch anyway \ + and feed not-found results to the model" + ) + .into(), + ), + )); + } + } + self.turn_tools = Some(tools.clone()); + tools + } + }; + Ok(DriveStep::ExecuteTools { calls, tools }) + } + AgentRunStep::Done(response) => Ok(DriveStep::Done(response)), + } + } + + /// Feed one completion response back into the run. The turn's tool sets + /// are supplied by the driver — the caller never assembles them. + /// + /// As with [`AgentRun::model_response`], a + /// [`ModelTurnOutcome::NeedsResolution`] outcome must be answered via + /// [`Self::resolve_invalid_tool_call`] before advancing. + pub fn model_response( + &mut self, + response: &CompletionResponse, + ) -> Result { + let Some(tools) = &self.turn_tools else { + return Err(PromptError::CompletionError(CompletionError::RequestError( + "model_response must follow a SendRequest step from this driver".into(), + ))); + }; + self.run.model_response(tools.model_turn(response)) + } + + /// Resolve a pending invalid tool call, exactly as + /// [`AgentRun::resolve_invalid_tool_call`]. + pub fn resolve_invalid_tool_call( + &mut self, + action: InvalidToolCallAction, + ) -> Result { + self.run.resolve_invalid_tool_call(action) + } + + /// Feed the results for the pending tool calls back into the run, exactly + /// as [`AgentRun::tool_results`]. + pub fn tool_results(&mut self, results: Vec) -> Result<(), PromptError> { + self.run.tool_results(results) + } + + /// Derive a fresh dispatch target for a resumed run's pending tool calls. + /// + /// Necessarily a **new** snapshot: implementations are live objects, so a + /// fresh process dispatches against its own registry state. The retrieval + /// query is re-derived from the run's history (matching preparation), and + /// the run's committed output tool — which is serialized with the run — + /// stays non-executable. + async fn resume_tools(&self) -> Result { + let query = self + .run + .full_history() + .iter() + .rev() + .find_map(|message| message.rag_text()); + let snapshot = self + .agent + .tool_server_handle + .snapshot_tool_defs(query) + .await + .map_err(|_| { + PromptError::CompletionError(CompletionError::RequestError( + "Failed to get tool definitions".into(), + )) + })?; + let executable: BTreeSet = snapshot + .definitions() + .iter() + .map(|tool| tool.name.clone()) + .collect(); + let output_tool_name = self.run.output_tool_name().map(str::to_owned); + let mut allowed = allowed_tool_names_for_choice( + &executable, + self.agent.tool_choice.as_ref(), + output_tool_name.as_deref(), + None, + ) + .map_err(PromptError::CompletionError)?; + if let Some(name) = &output_tool_name { + allowed.insert(name.clone()); + } + Ok(TurnTools { + snapshot: Arc::new(snapshot), + executable_tool_names: Arc::new(executable), + allowed_tool_names: Arc::new(allowed), + output_tool_name, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::AgentBuilder; + use crate::agent::run::OutputMode; + use crate::completion::Message; + use crate::test_utils::{MockAddTool, MockCompletionModel, MockSubtractTool, MockTurn}; + use crate::tool::{ToolContext, ToolErrorKind}; + use rig_core::message::ToolChoice; + use serde_json::json; + + fn schema(value: serde_json::Value) -> schemars::Schema { + serde_json::from_value(value).expect("valid schema") + } + + fn value_schema() -> schemars::Schema { + schema(json!({ + "type": "object", + "properties": { "value": { "type": "integer" } }, + "required": ["value"] + })) + } + + /// Expect the next step to be `SendRequest`, panicking otherwise. + macro_rules! expect_send { + ($driver:expr) => { + match $driver.next_step().await.expect("next_step succeeds") { + DriveStep::SendRequest { + request, + tools, + turn, + } => (request, tools, turn), + other => panic!("expected SendRequest, got {other:?}"), + } + }; + } + + macro_rules! expect_execute_tools { + ($driver:expr) => { + match $driver.next_step().await.expect("next_step succeeds") { + DriveStep::ExecuteTools { calls, tools } => (calls, tools), + other => panic!("expected ExecuteTools, got {other:?}"), + } + }; + } + + /// Criterion: the full drive loop reuses the agent's configuration and the + /// prepared turn keeps dispatching the implementation it advertised even + /// after the live registry mutates. + #[tokio::test] + async fn drive_loop_dispatches_the_advertised_implementation() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("call_1", "add", json!({"x": 1, "y": 2})), + MockTurn::text("3"), + ]); + let agent = AgentBuilder::new(model) + .preamble("driver preamble") + .default_max_turns(2) + .tool(MockAddTool) + .build(); + + let mut driver = agent.drive("add 1 and 2"); + + let (request, tools, turn) = expect_send!(driver); + assert_eq!(turn, 1); + assert!(tools.executable_tool_names().contains("add")); + assert_eq!(tools.executable_tool_names(), tools.allowed_tool_names()); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + + // Mutate the live registry AFTER the turn advertised its tools. + agent.tool_server_handle.remove_tool("add").await; + agent.tool_server_handle.add_tool(MockSubtractTool).await; + + let (calls, tools) = expect_execute_tools!(driver); + assert_eq!(calls.len(), 1); + let mut context = ToolContext::new(); + + // The snapshot still dispatches the advertised implementation... + let probe = tools + .execute("add", r#"{"x": 1, "y": 2}"#, &mut context) + .await; + assert!( + probe.is_success(), + "snapshot dispatch must reach the advertised implementation" + ); + // ...and does not see tools registered after it was taken. + let probe = tools + .execute("subtract", r#"{"x": 1, "y": 2}"#, &mut context) + .await; + assert!(probe.is_error_kind(ToolErrorKind::NotFound)); + + let mut results = Vec::new(); + for call in &calls { + results.push(tools.execute_call(call, &mut context).await); + } + driver.tool_results(results).expect("results accepted"); + + let (request, _, turn) = expect_send!(driver); + assert_eq!(turn, 2); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + match driver.next_step().await.expect("next_step succeeds") { + DriveStep::Done(response) => assert_eq!(response.output, "3"), + other => panic!("expected Done, got {other:?}"), + } + } + + /// The prepared request carries the agent's configuration — the preamble + /// leads the history and the registered tools are advertised. + #[tokio::test] + async fn prepared_request_carries_agent_configuration() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .preamble("driver preamble") + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("go"); + let (request, _, _) = expect_send!(driver); + let request = request.build(); + assert!(matches!( + request.chat_history.first(), + Some(Message::System { content }) if content == "driver preamble" + )); + assert!(request.tools.iter().any(|tool| tool.name == "add")); + } + + /// Criterion: under Tool output mode the synthetic output tool is allowed + /// and advertised but never executable, and dispatching it is rejected + /// with the machine-readable `NotExecutable` kind. + #[tokio::test] + async fn output_tool_is_allowed_but_never_executable() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .tool(MockAddTool) + .output_schema_raw(value_schema()) + .output_mode(OutputMode::Tool) + .build(); + let mut driver = agent.drive("compute"); + let (request, tools, _) = expect_send!(driver); + let output_tool = tools + .output_tool_name() + .expect("Tool output mode advertises an output tool") + .to_owned(); + + assert!(tools.allowed_tool_names().contains(&output_tool)); + assert!(!tools.executable_tool_names().contains(&output_tool)); + let request = request.build(); + assert!(request.tools.iter().any(|tool| tool.name == output_tool)); + + let mut context = ToolContext::new(); + let result = tools.execute(&output_tool, "{}", &mut context).await; + assert!(result.is_error_kind(ToolErrorKind::NotExecutable)); + } + + /// Criterion (finding 1): a Tool-output-mode agent driven by the + /// documented pattern finalizes with its structured answer — the run's + /// intercept is armed by the driver, so the output-tool call never + /// surfaces as a pending tool. + #[tokio::test] + async fn tool_mode_run_finalizes_via_the_output_tool_intercept() { + let model = MockCompletionModel::new([MockTurn::tool_call( + "call_1", + "final_result", + json!({"value": 7}), + )]); + let agent = AgentBuilder::new(model) + .tool(MockAddTool) + .output_schema_raw(value_schema()) + .output_mode(OutputMode::Tool) + .build(); + let mut driver = agent.drive("compute"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + match driver.next_step().await.expect("next_step succeeds") { + DriveStep::Done(response) => { + assert!( + response.output.contains('7'), + "structured output should carry the answer: {}", + response.output + ); + } + other => panic!("the output-tool call must finalize the run, got {other:?}"), + } + } + + /// Criterion (finding 2): the run's committed output tool stays pinned + /// across turns even when the tool set changes in between. + #[tokio::test] + async fn committed_output_tool_pins_across_turns() { + let model = MockCompletionModel::new([ + MockTurn::text("not structured output"), + MockTurn::tool_call("call_1", "final_result", json!({"value": 7})), + ]); + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool(MockAddTool) + .output_schema_raw(value_schema()) + .output_mode(OutputMode::Tool) + .build(); + let mut driver = agent.drive("compute"); + let (request, tools, _) = expect_send!(driver); + let committed = tools + .output_tool_name() + .expect("Tool mode commits a name on turn 1") + .to_owned(); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn processed"); + + // Change the tool set between turns: retire every executable tool. + agent.tool_server_handle.remove_tool("add").await; + + // Tool-mode validation re-prompts; the committed name must survive. + let (_, tools, turn) = expect_send!(driver); + assert_eq!(turn, 2); + assert_eq!( + tools.output_tool_name(), + Some(committed.as_str()), + "the committed output tool must stay pinned when the tool set changes" + ); + assert!(tools.allowed_tool_names().contains(&committed)); + } + + /// Criterion (finding 3): a run resumed from serialized state alone — + /// fresh driver, fresh process semantics — re-emits its pending calls and + /// dispatches them through a fresh snapshot. + #[tokio::test] + async fn resume_from_serialized_state_dispatches_via_a_fresh_snapshot() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("call_1", "add", json!({"x": 2, "y": 5})), + MockTurn::text("7"), + ]); + let agent = AgentBuilder::new(model.clone()) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("what is 2 + 5?"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + drop(driver); + drop(agent); + + // "Fresh process": rebuild the same agent, deserialize the run. + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = agent.drive_run(run); + let (calls, tools) = expect_execute_tools!(driver); + let mut context = ToolContext::new(); + let mut results = Vec::new(); + for call in &calls { + results.push(tools.execute_call(call, &mut context).await); + } + driver.tool_results(results).expect("results accepted"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + match driver.next_step().await.expect("next_step succeeds") { + DriveStep::Done(response) => assert_eq!(response.output, "7"), + other => panic!("expected Done, got {other:?}"), + } + } + + /// Criterion (finding 3, drift): a resumed pending call whose tool is + /// missing from this process's registry is surfaced as an error before + /// dispatch — and the opt-out downgrades it to a not-found tool result. + #[tokio::test] + async fn resumed_drift_is_loud_before_dispatch() { + let model = MockCompletionModel::new([MockTurn::tool_call( + "call_1", + "add", + json!({"x": 2, "y": 5}), + )]); + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("what is 2 + 5?"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + + // Resume against an agent that no longer registers the pending tool. + let bare_agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(2) + .build(); + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = bare_agent.drive_run(run); + let err = driver + .next_step() + .await + .expect_err("a missing pending tool must surface as drift"); + let message = err.to_string(); + assert!(message.contains("add"), "error names the tool: {message}"); + assert!( + message.contains("allow_missing_resumed_tools"), + "error names the opt-out: {message}" + ); + + // The opt-out dispatches anyway and yields a not-found result. + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = bare_agent.drive_run(run).allow_missing_resumed_tools(); + let (calls, tools) = expect_execute_tools!(driver); + let mut context = ToolContext::new(); + let result = tools + .execute( + &calls[0].tool_call.function.name, + &calls[0].tool_call.function.arguments.to_string(), + &mut context, + ) + .await; + assert!(result.is_error_kind(ToolErrorKind::NotFound)); + } + + /// Criterion (finding 11): the driven run inherits the agent's + /// configuration instead of the driver restating it. + #[tokio::test] + async fn drive_seeds_the_run_from_agent_configuration() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(3) + .tool_choice(ToolChoice::Auto) + .output_schema_raw(value_schema()) + .tool(MockAddTool) + .build(); + let driver = agent.drive("go"); + let state = serde_json::to_value(driver.run()).expect("run serializes"); + assert_eq!(state["max_turns"], 3, "seeded from default_max_turns"); + assert!( + !state["tool_choice"].is_null(), + "seeded from the agent's tool_choice" + ); + assert!( + state["output_schema"].is_object(), + "output validation seeded from the agent's schema" + ); + + // No configured budget seeds the implicit budget of one, exactly as + // `AgentRunner::from_agent` does. + let unconfigured = AgentBuilder::new(MockCompletionModel::text("unused")).build(); + let state = serde_json::to_value(unconfigured.drive("go").run()).expect("run serializes"); + assert_eq!(state["max_turns"], 1); + + // Per-run overrides layer on top of the seeding, and the run state can + // be reclaimed by value for suspension. + let run = unconfigured + .drive("go") + .history(vec![Message::user("earlier context")]) + .max_turns(5) + .into_run(); + let state = serde_json::to_value(&run).expect("run serializes"); + assert_eq!(state["max_turns"], 5); + assert!(state["chat_history"].is_array()); + } + + /// Criterion: an impossible `ToolChoice` fails at prepare time, locally, + /// with no provider round-trip. + #[tokio::test] + async fn impossible_tool_choice_fails_locally_at_prepare_time() { + let model = MockCompletionModel::text("unused"); + let agent = AgentBuilder::new(model.clone()) + .tool_choice(ToolChoice::Specific { + function_names: vec!["missing".to_string()], + }) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("go"); + let err = driver + .next_step() + .await + .expect_err("a tool_choice naming an unadvertised tool must fail at prepare time"); + assert!(err.to_string().contains("missing")); + assert_eq!( + model.request_count(), + 0, + "local validation must not cost a provider round-trip" + ); + } + + /// `Required` forces a tool call, so an empty advertised set can never + /// satisfy it — that must fail at prepare time, not degrade silently. + #[tokio::test] + async fn required_tool_choice_with_no_tools_fails_locally() { + let model = MockCompletionModel::text("unused"); + let agent = AgentBuilder::new(model.clone()) + .tool_choice(ToolChoice::Required) + .build(); + let mut driver = agent.drive("go"); + let err = driver + .next_step() + .await + .expect_err("Required with no advertised tool must fail at prepare time"); + assert!(err.to_string().contains("Required")); + assert_eq!(model.request_count(), 0); + } +} diff --git a/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index cf40f7a5bb..22e7c21d7d 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -101,13 +101,14 @@ //! ``` mod builder; mod completion; +pub mod driver; pub mod hook; pub mod model; -mod prepared_turn; pub(crate) mod prompt_request; pub mod run; pub mod runner; mod tool; +mod turn_tools; /// Fallback display name used in telemetry spans and logs when an agent has no /// configured name. @@ -123,8 +124,9 @@ pub use hook::{ RetryRequest, RunId, Scratchpad, StepEventKind, StreamResponseFinish, TextDelta, ToolCall, ToolCallAction, ToolCallDelta, ToolResultAction, ToolResultEvent, }; +pub use driver::{AgentDriver, DriveStep}; pub use model::ModelHandle; -pub use prepared_turn::{PreparedTurn, TurnTools}; +pub use turn_tools::TurnTools; pub use prompt_request::streaming::{ MultiTurnStreamItem, StreamingError, StreamingPromptRequest, StreamingResult, stream_to_stdout, }; diff --git a/crates/rig-agent/src/agent/prepared_turn.rs b/crates/rig-agent/src/agent/prepared_turn.rs deleted file mode 100644 index 79886bdad0..0000000000 --- a/crates/rig-agent/src/agent/prepared_turn.rs +++ /dev/null @@ -1,359 +0,0 @@ -//! One prepared model turn of a configured [`Agent`](super::Agent), for -//! hand-driven [`AgentRun`](super::run::AgentRun) loops. -//! -//! [`Agent::prepare_turn`](super::Agent::prepare_turn) reads the agent's -//! configuration — preamble, static context, temperature, `max_tokens`, -//! `additional_params`, `tool_choice`, output schema/mode, and the tool -//! registry — and resolves it into a [`PreparedTurn`]: the fully configured -//! completion request for one turn plus a [`TurnTools`] value that keeps the -//! turn's advertised tool set and its dispatch target consistent. -//! -//! This exists so a hand-driven `AgentRun` loop (custom provider transport, -//! suspend/resume across processes) can reuse the `Agent` it configured instead -//! of restating every piece of that configuration by hand. It is a -//! configuration read plus a dispatch target — not a second execution path: -//! [`Agent::runner`](super::Agent::runner) remains the only way to execute an -//! agent with hooks, memory, retrieval policy, and telemetry. - -use std::collections::BTreeSet; -use std::sync::Arc; - -use super::completion::PreparedCompletionRequest; -use super::model::ModelHandle; -use crate::completion::CompletionRequestBuilder; -use crate::tool::server::ToolRegistrySnapshot; -use crate::tool::{ToolContext, ToolDispatch, ToolExecutionError, ToolResult}; - -/// One prepared model turn: the configured completion request plus the turn's -/// tool sets and dispatch target, resolved together at one instant. -/// -/// Produced by [`Agent::prepare_turn`](super::Agent::prepare_turn). The request -/// half and the tool half were computed from the same registry snapshot, so the -/// tools advertised in the request are exactly the tools [`TurnTools`] -/// dispatches — even if the agent's live registry changes afterwards. -/// -/// A prepared turn reflects the agent's **baseline** configuration. No hooks -/// run: there is no `CompletionCall` request patch, no model-selection hook, -/// and no per-turn `active_tools` narrowing. Anyone who needs those wants -/// [`Agent::runner`](super::Agent::runner). -pub struct PreparedTurn { - builder: CompletionRequestBuilder, - tools: TurnTools, -} - -impl std::fmt::Debug for PreparedTurn { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("PreparedTurn") - .field("tools", &self.tools) - .finish_non_exhaustive() - } -} - -impl PreparedTurn { - pub(crate) fn from_prepared(prepared: PreparedCompletionRequest) -> Self { - let PreparedCompletionRequest { - builder, - tool_snapshot, - executable_tool_names, - allowed_tool_names, - output_tool_name, - } = prepared; - Self { - builder, - tools: TurnTools { - snapshot: tool_snapshot, - executable_tool_names, - allowed_tool_names, - output_tool_name, - }, - } - } - - /// The turn's tool sets and dispatch target. - /// - /// [`TurnTools`] is cheap to clone; clone it before consuming the request - /// with [`into_parts`](Self::into_parts) if you only need the tool half - /// later. - pub fn tools(&self) -> &TurnTools { - &self.tools - } - - /// Split the prepared turn into the configured completion request builder - /// and the turn's [`TurnTools`]. - /// - /// The builder already carries the agent's preamble (with any output-mode - /// augmentation), static context, model parameters, `tool_choice`, and this - /// turn's tool definitions. Call - /// [`send`](crate::completion::CompletionRequestBuilder::send) on it to - /// issue the request against the agent's configured model, or - /// [`build`](crate::completion::CompletionRequestBuilder::build) it to hand - /// the raw request to a custom transport. - pub fn into_parts(self) -> (CompletionRequestBuilder, TurnTools) { - (self.builder, self.tools) - } -} - -/// One turn's advertised tool sets and their dispatch target. -/// -/// All four pieces were resolved together when the turn was prepared: -/// -/// - [`executable_tool_names`](Self::executable_tool_names) — the real -/// registry tools advertised to the provider this turn. Feed this to -/// [`ModelTurn::new`](super::run::ModelTurn::new). -/// - [`allowed_tool_names`](Self::allowed_tool_names) — the tools the active -/// `tool_choice` lets the model call this turn (including the synthetic -/// output tool, which is allowed but not executable). Feed this to -/// [`ModelTurn::new`](super::run::ModelTurn::new) as well. -/// - [`output_tool_name`](Self::output_tool_name) — the synthetic -/// structured-output tool, when Tool output mode is active. -/// - [`execute`](Self::execute) — dispatch against the registry **snapshot** -/// taken when the turn was prepared. -/// -/// # Snapshot semantics -/// -/// The dispatch target is a per-turn snapshot, not the live registry. Tools -/// added to or removed from the agent after the turn was prepared (including -/// MCP refreshes) take effect on the *next* prepared turn; calls belonging to -/// this turn dispatch through the exact implementations whose definitions the -/// provider received. This is the same guarantee the agent runner gives its own -/// turns, and it is why this type exists instead of a handle to the live -/// registry: advertising one implementation and dispatching another is the skew -/// this closes. -#[derive(Clone)] -pub struct TurnTools { - snapshot: Arc, - executable_tool_names: BTreeSet, - allowed_tool_names: BTreeSet, - output_tool_name: Option, -} - -impl std::fmt::Debug for TurnTools { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("TurnTools") - .field("executable_tool_names", &self.executable_tool_names) - .field("allowed_tool_names", &self.allowed_tool_names) - .field("output_tool_name", &self.output_tool_name) - .finish_non_exhaustive() - } -} - -impl TurnTools { - /// The executable registry tools advertised to the provider this turn. - pub fn executable_tool_names(&self) -> &BTreeSet { - &self.executable_tool_names - } - - /// The tools the active `tool_choice` lets the model call this turn, - /// including the synthetic output tool when Tool output mode is active. - pub fn allowed_tool_names(&self) -> &BTreeSet { - &self.allowed_tool_names - } - - /// The synthetic structured-output tool advertised this turn, when the - /// agent's output schema resolved to Tool output mode. - /// - /// This tool is **allowed but not executable**: a model call to it *is* the - /// final structured answer. A hand-driven loop must intercept it by name — - /// consume its arguments as the run's structured output instead of - /// dispatching it (dispatching it through [`execute`](Self::execute) is - /// rejected). Drivers that validate the arguments against the output schema - /// should feed failures back to the model as an error tool result and - /// re-prompt, mirroring - /// [`AgentRun::with_output_validation`](super::run::AgentRun::with_output_validation). - pub fn output_tool_name(&self) -> Option<&str> { - self.output_tool_name.as_deref() - } - - /// Execute a tool call through the exact implementation advertised for this - /// turn. - /// - /// Mirrors [`ToolServerHandle::execute`](crate::tool::server::ToolServerHandle::execute): - /// the result carries success, failure, or refusal as a [`ToolResult`], and - /// the tool's result metadata is published back to `context`. A name that - /// was not advertised this turn — including one registered on the agent - /// *after* the turn was prepared — resolves to a not-found failure rather - /// than reaching the live registry. Calling the synthetic output tool is - /// rejected: it has no implementation, and its call is the final structured - /// answer the driver must consume itself (see - /// [`output_tool_name`](Self::output_tool_name)). - pub async fn execute( - &self, - tool_name: &str, - args: &str, - context: &mut ToolContext, - ) -> ToolResult { - if self.output_tool_name.as_deref() == Some(tool_name) { - return ToolResult::failed(ToolExecutionError::not_found(format!( - "`{tool_name}` is this turn's synthetic structured-output tool: it is advertised \ - to the model but has no implementation. Its call carries the final structured \ - answer; the driver must consume the call's arguments instead of dispatching it." - ))); - } - context.clear_dispatch_result(); - let ToolDispatch { - result, - context: dispatch_context, - } = self.snapshot.dispatch(tool_name, args, context).await; - context.accept_dispatch_result(dispatch_context); - result - } -} - -#[cfg(test)] -mod tests { - use crate::agent::AgentBuilder; - use crate::agent::run::OutputMode; - use crate::completion::{CompletionError, Message}; - use crate::test_utils::{MockAddTool, MockCompletionModel, MockSubtractTool}; - use crate::tool::{ToolContext, ToolErrorKind}; - use rig_core::message::ToolChoice; - use serde_json::json; - - fn schema(value: serde_json::Value) -> schemars::Schema { - serde_json::from_value(value).expect("valid schema") - } - - /// Criterion: advertise/dispatch consistency. A prepared turn keeps - /// dispatching the implementation it advertised even after the agent's - /// live registry mutates, and never sees tools registered afterwards. - #[tokio::test] - async fn prepared_turn_dispatches_the_implementation_it_advertised() { - let agent = AgentBuilder::new(MockCompletionModel::text("unused")) - .preamble("prepared preamble") - .tool(MockAddTool) - .build(); - - let prepared = agent - .prepare_turn("add 1 and 2", &[]) - .await - .expect("prepare succeeds"); - let (request, tools) = prepared.into_parts(); - - // The request carries the agent's configuration: the preamble leads - // the history and the registered tool is advertised. - let request = request.build(); - assert!(matches!( - request.chat_history.first(), - Some(Message::System { content }) if content == "prepared preamble" - )); - assert!(request.tools.iter().any(|tool| tool.name == "add")); - assert!(tools.executable_tool_names().contains("add")); - assert_eq!(tools.executable_tool_names(), tools.allowed_tool_names()); - - // Mutate the live registry AFTER preparation: retire the advertised - // tool and register a different one. - agent.tool_server_handle.remove_tool("add").await; - agent.tool_server_handle.add_tool(MockSubtractTool).await; - - // The prepared turn still dispatches the implementation it advertised... - let mut context = ToolContext::new(); - let result = tools.execute("add", r#"{"x": 1, "y": 2}"#, &mut context).await; - assert!( - result.is_success(), - "snapshot dispatch must reach the advertised implementation" - ); - - // ...and does not see tools registered after the snapshot was taken. - let mut context = ToolContext::new(); - let result = tools - .execute("subtract", r#"{"x": 1, "y": 2}"#, &mut context) - .await; - assert!(result.is_error_kind(ToolErrorKind::NotFound)); - - // The live registry really did change underneath the snapshot. - let mut context = ToolContext::new(); - let result = agent - .tool_server_handle - .execute("add", r#"{"x": 1, "y": 2}"#, &mut context) - .await; - assert!(result.is_error_kind(ToolErrorKind::NotFound)); - } - - /// Criterion: under Tool output mode the synthetic output tool is allowed - /// and advertised but never executable, and dispatching it is rejected. - #[tokio::test] - async fn output_tool_is_allowed_but_never_executable() { - let agent = AgentBuilder::new(MockCompletionModel::text("unused")) - .tool(MockAddTool) - .output_schema_raw(schema(json!({ - "type": "object", - "properties": { "value": { "type": "integer" } }, - "required": ["value"] - }))) - .output_mode(OutputMode::Tool) - .build(); - - let prepared = agent - .prepare_turn("compute", &[]) - .await - .expect("prepare succeeds"); - let (request, tools) = prepared.into_parts(); - let output_tool = tools - .output_tool_name() - .expect("Tool output mode advertises an output tool") - .to_owned(); - - assert!(tools.allowed_tool_names().contains(&output_tool)); - assert!(!tools.executable_tool_names().contains(&output_tool)); - // The provider sees the synthetic tool alongside the real one. - let request = request.build(); - assert!(request.tools.iter().any(|tool| tool.name == output_tool)); - assert!(request.tools.iter().any(|tool| tool.name == "add")); - - let mut context = ToolContext::new(); - let result = tools.execute(&output_tool, "{}", &mut context).await; - assert!(result.is_error_kind(ToolErrorKind::NotFound)); - let message = result.error().expect("rejected dispatch").message().to_owned(); - assert!( - message.contains("structured-output"), - "rejection must explain the synthetic tool: {message}" - ); - } - - /// Criterion: an impossible `ToolChoice` fails at prepare time, locally, - /// with no provider round-trip. - #[tokio::test] - async fn impossible_tool_choice_fails_locally_at_prepare_time() { - let model = MockCompletionModel::text("unused"); - let agent = AgentBuilder::new(model.clone()) - .tool(MockAddTool) - .tool_choice(ToolChoice::Specific { - function_names: vec!["missing".to_string()], - }) - .build(); - - let err = agent - .prepare_turn("go", &[]) - .await - .expect_err("a tool_choice naming an unadvertised tool must fail at prepare time"); - assert!(matches!( - err, - CompletionError::RequestError(inner) if inner.to_string().contains("missing") - )); - assert_eq!( - model.request_count(), - 0, - "local validation must not cost a provider round-trip" - ); - } - - /// `Required` forces a tool call, so an empty advertised set can never - /// satisfy it — that must fail at prepare time, not degrade silently. - #[tokio::test] - async fn required_tool_choice_with_no_tools_fails_locally() { - let model = MockCompletionModel::text("unused"); - let agent = AgentBuilder::new(model.clone()) - .tool_choice(ToolChoice::Required) - .build(); - - let err = agent - .prepare_turn("go", &[]) - .await - .expect_err("Required with no advertised tool must fail at prepare time"); - assert!(matches!( - err, - CompletionError::RequestError(inner) if inner.to_string().contains("Required") - )); - assert_eq!(model.request_count(), 0); - } -} diff --git a/crates/rig-agent/src/agent/prompt_request/streaming.rs b/crates/rig-agent/src/agent/prompt_request/streaming.rs index dd6b579ac5..8c89f454e1 100644 --- a/crates/rig-agent/src/agent/prompt_request/streaming.rs +++ b/crates/rig-agent/src/agent/prompt_request/streaming.rs @@ -4,7 +4,8 @@ use rig_core::{ }; use crate::{ - agent::completion::{PreparedCompletionRequest, build_prepared_completion_request}, + agent::completion::build_prepared_completion_request, + agent::turn_tools::PreparedCompletionRequest, agent::hook::{ AgentHook, HookContext, HookStack, InvalidToolCallAction, ModelSelection, ModelSelectionAction, ModelTurnFinished, ReasoningDelta, StepEventKind, @@ -586,8 +587,8 @@ where break 'outer; } }; - run.set_output_tool_name(prepared.output_tool_name.clone()); - let turn_tool_snapshot = prepared.tool_snapshot.clone(); + run.set_output_tool_name(prepared.tools.output_tool_name.clone()); + let turn_tool_snapshot = prepared.tools.snapshot.clone(); if runner.record_telemetry_content { let input_messages = prepared.builder.messages_for_telemetry(); rig_core::telemetry::record_model_input(&chat_span, &input_messages, true); @@ -1082,8 +1083,8 @@ impl TurnSource for StreamingTurnSource { let mut last_usage = crate::completion::Usage::new(); let mut assembler = StreamedTurnAssembler::new( - prepared.executable_tool_names.clone(), - prepared.allowed_tool_names.clone(), + (*prepared.tools.executable_tool_names).clone(), + (*prepared.tools.allowed_tool_names).clone(), ); let mut completion_call_emitted = false; let mut turn_abandoned = false; diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index d6211b35d8..aca8c883f2 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -25,16 +25,16 @@ //! `AgentRun` deliberately contains no model, tool registry, memory backend, or //! hook stack. Hand-driving it is a low-level provider integration: the caller //! owns all IO and any lifecycle policy. To drive a *configured* -//! [`Agent`](crate::agent::Agent) by hand, obtain each turn's completion -//! request and tool dispatch from -//! [`Agent::prepare_turn`](crate::agent::Agent::prepare_turn) — it resolves the -//! agent's configuration (preamble, tools, `tool_choice`, output mode, model -//! parameters) into a [`PreparedTurn`](crate::agent::PreparedTurn) instead of -//! the driver restating it. To execute an `Agent` with its hooks, memory, -//! retrieval policy, and telemetry, use +//! [`Agent`](crate::agent::Agent) by hand, use +//! [`Agent::drive`](crate::agent::Agent::drive) — the returned +//! [`AgentDriver`](crate::agent::AgentDriver) seeds this machine from the +//! agent's configuration and pairs each turn's completion request with its +//! tool dispatch snapshot, while every side effect (the provider call, tool +//! execution) stays with the caller. To execute an `Agent` with its hooks, +//! memory, retrieval policy, and telemetry, use //! [`Agent::runner`](crate::agent::Agent::runner); constructing an `AgentRun` -//! directly is not an alternate way to execute an `Agent`, and a prepared turn -//! is a configuration read, not a second execution path. +//! directly is not an alternate way to execute an `Agent`, and the driver is a +//! configuration and pairing layer, not a second execution path. //! //! [`crate::completion::Prompt::prompt`] and //! [`Agent::runner`](crate::agent::Agent::runner) drive this machine internally; diff --git a/crates/rig-agent/src/agent/runner.rs b/crates/rig-agent/src/agent/runner.rs index 23a357a058..9109d293ac 100644 --- a/crates/rig-agent/src/agent/runner.rs +++ b/crates/rig-agent/src/agent/runner.rs @@ -33,7 +33,8 @@ use futures::StreamExt; use tracing::{Instrument, info_span, span::Id}; use super::{ - completion::{Agent, PreparedCompletionRequest}, + completion::Agent, + turn_tools::PreparedCompletionRequest, hook::{ AgentHook, CompletionCall, CompletionCallAction, CompletionResponse as CompletionResponseEvent, HookContext, HookStack, @@ -49,9 +50,7 @@ use super::{ }, tool_result_output, }, - run::{ - AgentRun, DEFAULT_OUTPUT_RETRIES, ModelTurn, ModelTurnOutcome, OutputMode, PendingToolCall, - }, + run::{AgentRun, DEFAULT_OUTPUT_RETRIES, ModelTurnOutcome, OutputMode, PendingToolCall}, }; use rig_core::{ memory::ConversationMemory, @@ -935,13 +934,7 @@ impl TurnSource for UnaryTurnSource { } }; - let mut outcome = match run.model_response(ModelTurn::new( - resp.message_id.clone(), - resp.choice.clone(), - resp.usage, - prepared.executable_tool_names, - prepared.allowed_tool_names, - )) { + let mut outcome = match run.model_response(prepared.tools.model_turn(&resp)) { Ok(outcome) => outcome, Err(err) => { yield Err(Box::new(err).into()); diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs new file mode 100644 index 0000000000..43d4788e78 --- /dev/null +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -0,0 +1,262 @@ +//! One turn's advertised tool sets and their dispatch target. +//! +//! [`TurnTools`] is produced when a model turn is prepared — by +//! [`AgentDriver`](super::AgentDriver) for hand-driven runs, and internally by +//! the runner — and carries everything the turn's tool handling needs, resolved +//! together at one instant: the executable and allowed tool-name sets, the +//! synthetic output-tool name, and dispatch pinned to the registry snapshot +//! whose definitions the provider received. + +use std::collections::BTreeSet; +use std::sync::Arc; + +use rig_core::message::UserContent; + +use super::model::ModelHandle; +use super::run::{ModelTurn, PendingToolCall}; +use crate::completion::{CompletionRequestBuilder, CompletionResponse}; +use crate::tool::server::ToolRegistrySnapshot; +use crate::tool::{ToolContext, ToolExecutionError, ToolResult}; + +/// A prepared completion request: the configured builder plus the turn's tool +/// state, computed together from one registry snapshot. +pub(crate) struct PreparedCompletionRequest { + /// Builder carrying the selected model handle: request preparation ran + /// against this handle's captured capabilities, and the same handle + /// executes the prepared request. + pub(crate) builder: CompletionRequestBuilder, + /// The turn's tool sets and dispatch target. + pub(crate) tools: TurnTools, +} + +/// One turn's advertised tool sets and their dispatch target. +/// +/// All four pieces were resolved together when the turn was prepared: +/// +/// - [`executable_tool_names`](Self::executable_tool_names) — the real +/// registry tools advertised to the provider this turn. +/// - [`allowed_tool_names`](Self::allowed_tool_names) — the tools the active +/// `tool_choice` lets the model call this turn (including the synthetic +/// output tool, which is allowed but not executable). +/// - [`output_tool_name`](Self::output_tool_name) — the synthetic +/// structured-output tool, when Tool output mode is active. +/// - [`execute`](Self::execute) / [`execute_call`](Self::execute_call) — +/// dispatch against the registry **snapshot** taken when the turn was +/// prepared. +/// +/// Cloning is cheap: the sets and the snapshot are shared behind `Arc`s. +/// +/// # Snapshot semantics +/// +/// The dispatch target is a per-turn snapshot, not the live registry. Tools +/// added to or removed from the agent after the turn was prepared (including +/// MCP refreshes) take effect on the *next* prepared turn; calls belonging to +/// this turn dispatch through the exact implementations whose definitions the +/// provider received. This is the same guarantee the agent runner gives its +/// own turns: advertising one implementation and dispatching another is the +/// skew this type exists to close. +#[derive(Clone)] +pub struct TurnTools { + pub(crate) snapshot: Arc, + pub(crate) executable_tool_names: Arc>, + pub(crate) allowed_tool_names: Arc>, + pub(crate) output_tool_name: Option, +} + +impl std::fmt::Debug for TurnTools { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("TurnTools") + .field("executable_tool_names", &self.executable_tool_names) + .field("allowed_tool_names", &self.allowed_tool_names) + .field("output_tool_name", &self.output_tool_name) + .finish_non_exhaustive() + } +} + +impl TurnTools { + /// The executable registry tools advertised to the provider this turn. + pub fn executable_tool_names(&self) -> &BTreeSet { + &self.executable_tool_names + } + + /// The tools the active `tool_choice` lets the model call this turn, + /// including the synthetic output tool when Tool output mode is active. + pub fn allowed_tool_names(&self) -> &BTreeSet { + &self.allowed_tool_names + } + + /// The synthetic structured-output tool advertised this turn, when the + /// agent's output schema resolved to Tool output mode. + /// + /// This tool is **allowed but not executable**: a model call to it *is* + /// the final structured answer. A driver never sees such a call as a + /// pending tool — the run machine intercepts it and finalizes — so this + /// is informational for [`AgentDriver`](super::AgentDriver) users. + /// Dispatching it through [`execute`](Self::execute) is rejected with + /// [`ToolErrorKind::NotExecutable`](crate::tool::ToolErrorKind::NotExecutable). + pub fn output_tool_name(&self) -> Option<&str> { + self.output_tool_name.as_deref() + } + + /// Execute a tool call through the exact implementation advertised for + /// this turn. + /// + /// Mirrors [`ToolServerHandle::execute`](crate::tool::server::ToolServerHandle::execute): + /// the result carries success, failure, or refusal as a [`ToolResult`], + /// and the tool's result metadata is published back to `context`. A name + /// that was not advertised this turn — including one registered on the + /// agent *after* the turn was prepared — resolves to a + /// [`NotFound`](crate::tool::ToolErrorKind::NotFound) failure rather than + /// reaching the live registry. The synthetic output tool is rejected with + /// [`NotExecutable`](crate::tool::ToolErrorKind::NotExecutable): it has no + /// implementation, and its call carries the final structured answer. + pub async fn execute( + &self, + tool_name: &str, + args: &str, + context: &mut ToolContext, + ) -> ToolResult { + context.clear_dispatch_result(); + if self.output_tool_name.as_deref() == Some(tool_name) { + return ToolResult::failed(ToolExecutionError::not_executable(format!( + "`{tool_name}` is this turn's synthetic structured-output tool: it is advertised \ + to the model but has no implementation. Its call carries the final structured \ + answer; the driver must consume the call's arguments instead of dispatching it." + ))); + } + self.snapshot.execute(tool_name, args, context).await + } + + /// Execute one pending call and produce its tool-result content. + /// + /// Honors [`PendingToolCall::preresolved_result`] — a call suppressed by + /// invalid tool-call recovery returns its pre-resolved content without + /// executing anything or touching `context`. Otherwise the call dispatches + /// through this turn's snapshot via [`execute`](Self::execute) and the + /// result is assembled into the [`UserContent`] tool-result value that + /// [`AgentDriver::tool_results`](super::AgentDriver::tool_results) (or + /// [`AgentRun::tool_results`](super::run::AgentRun::tool_results)) expects. + /// + /// Approval flows compose around this: execute approved calls here, and + /// build denied calls' result content directly (see + /// `examples/agent_with_durable_approval`). + pub async fn execute_call( + &self, + call: &PendingToolCall, + context: &mut ToolContext, + ) -> UserContent { + if let Some(result) = &call.preresolved_result { + return result.clone(); + } + let name = &call.tool_call.function.name; + let args = call.tool_call.function.arguments.to_string(); + let result = self.execute(name, &args, context).await; + UserContent::tool_result_for( + call.tool_call.id.clone(), + call.tool_call.provider.clone(), + name.clone(), + result.output().clone().into_content(), + ) + } + + /// Assemble the [`ModelTurn`] for a completion response received on this + /// prepared turn. The single construction site for driver-facing turns, + /// so the two name sets can never be transposed by a caller. + pub(crate) fn model_turn(&self, response: &CompletionResponse) -> ModelTurn { + ModelTurn::new( + response.message_id.clone(), + response.choice.clone(), + response.usage, + (*self.executable_tool_names).clone(), + (*self.allowed_tool_names).clone(), + ) + } +} + +#[cfg(test)] +mod tests { + use crate::agent::{AgentBuilder, DriveStep}; + use crate::agent::run::OutputMode; + use crate::test_utils::MockCompletionModel; + use crate::tool::{Tool, ToolContext, ToolErrorKind}; + use serde_json::json; + + /// A tool that publishes result metadata into the caller's context. + struct MetadataTool; + + #[derive(Clone, Debug, PartialEq)] + struct Marker(&'static str); + + #[derive(Debug, thiserror::Error)] + #[error("metadata tool failed")] + struct MetadataToolError; + + impl Tool for MetadataTool { + const NAME: &'static str = "metadata"; + type Error = MetadataToolError; + type Args = serde_json::Value; + type Output = i32; + + fn description(&self) -> String { + "publishes a result marker".to_string() + } + + fn parameters(&self) -> serde_json::Value { + json!({ "type": "object", "properties": {} }) + } + + async fn call( + &self, + context: &mut ToolContext, + _args: Self::Args, + ) -> Result { + context.insert_result(Marker("published")); + Ok(1) + } + } + + /// Regression (finding 4): the output-tool rejection must clear the + /// previous dispatch's result metadata, exactly like every other + /// `execute` path. + #[tokio::test] + async fn rejected_output_tool_dispatch_clears_stale_context_metadata() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .tool(MetadataTool) + .output_schema_raw( + serde_json::from_value(json!({ + "type": "object", + "properties": { "value": { "type": "integer" } }, + "required": ["value"] + })) + .expect("valid schema"), + ) + .output_mode(OutputMode::Tool) + .build(); + let mut driver = agent.drive("go"); + let tools = match driver.next_step().await.expect("prepare succeeds") { + DriveStep::SendRequest { tools, .. } => tools, + other => panic!("expected SendRequest, got {other:?}"), + }; + let output_tool = tools + .output_tool_name() + .expect("Tool mode advertises an output tool") + .to_owned(); + + let mut context = ToolContext::new(); + let result = tools.execute("metadata", "{}", &mut context).await; + assert!(result.is_success()); + assert_eq!( + context.result::(), + Some(&Marker("published")), + "the first dispatch publishes its metadata" + ); + + let result = tools.execute(&output_tool, "{}", &mut context).await; + assert!(result.is_error_kind(ToolErrorKind::NotExecutable)); + assert_eq!( + context.result::(), + None, + "the rejection must not leave the previous dispatch's metadata behind" + ); + } +} diff --git a/crates/rig-agent/src/tool/mod.rs b/crates/rig-agent/src/tool/mod.rs index 7231c341e7..2e8ebe5f8b 100644 --- a/crates/rig-agent/src/tool/mod.rs +++ b/crates/rig-agent/src/tool/mod.rs @@ -570,6 +570,18 @@ pub(crate) struct ToolDispatch { pub(crate) context: ToolContext, } +impl ToolDispatch { + /// Publish this dispatch's result metadata to the caller's context and + /// yield the result — the single completion step of the + /// clear → dispatch → publish sequence every `execute` surface follows. + /// Keeping publication here (rather than open-coded per caller) is what + /// stops the surfaces from drifting on result-metadata semantics. + pub(crate) fn publish_to(self, context: &mut ToolContext) -> ToolResult { + context.accept_dispatch_result(self.context); + self.result + } +} + /// Execute a resolved registry entry through the single dispatch boundary. /// /// Every surface enters here with its caller-owned context. The helper clones @@ -726,12 +738,9 @@ impl ToolSet { ) -> ToolResult { context.clear_dispatch_result(); let tool = self.get(name).cloned(); - let ToolDispatch { - result, - context: dispatch_context, - } = dispatch_tool(name, args.into(), tool, context).await; - context.accept_dispatch_result(dispatch_context); - result + dispatch_tool(name, args.into(), tool, context) + .await + .publish_to(context) } /// Documents describing all registered tools. diff --git a/crates/rig-agent/src/tool/server.rs b/crates/rig-agent/src/tool/server.rs index 1a29fe5b7f..d6faf61900 100644 --- a/crates/rig-agent/src/tool/server.rs +++ b/crates/rig-agent/src/tool/server.rs @@ -62,6 +62,22 @@ impl ToolRegistrySnapshot { let tool = self.tools.get(tool_name).cloned(); dispatch_tool(tool_name, args.to_string(), tool, context).await } + + /// Execute through the snapshot's pinned implementation, publishing result + /// metadata back to `context`. The snapshot counterpart of + /// [`ToolServerHandle::execute`], sharing the same + /// clear → dispatch → publish sequence. + pub(crate) async fn execute( + &self, + tool_name: &str, + args: &str, + context: &mut ToolContext, + ) -> ToolResult { + context.clear_dispatch_result(); + self.dispatch(tool_name, args, context) + .await + .publish_to(context) + } } /// Shared state behind a `ToolServerHandle`. @@ -427,12 +443,9 @@ impl ToolServerHandle { context: &mut ToolContext, ) -> ToolResult { context.clear_dispatch_result(); - let ToolDispatch { - result, - context: dispatch_context, - } = self.dispatch(tool_name, args, context).await; - context.accept_dispatch_result(dispatch_context); - result + self.dispatch(tool_name, args, context) + .await + .publish_to(context) } /// Run one isolated dispatch and retain its full context for agent hooks. diff --git a/crates/rig-core/CHANGELOG.md b/crates/rig-core/CHANGELOG.md index 7370fda877..1f4dd8b58f 100644 --- a/crates/rig-core/CHANGELOG.md +++ b/crates/rig-core/CHANGELOG.md @@ -56,6 +56,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- *(tool)* `ToolErrorKind::NotExecutable` and `ToolExecutionError::not_executable`: the tool is advertised to the model but not executable by the dispatcher — its call must be handled by the host. Produced by rig-agent when dispatching the synthetic structured-output tool; also the natural classification for host-executed or deferred tools. The enum is `#[non_exhaustive]`, so this is additive - *(streaming)* `wire::classify_typed_event` extends the decode-then-validate policy to typed-transport wires (bedrock, candle, gemini-grpc): modeled variants are `Known`, the SDK's non-exhaustive/unrecognized variants are `Unknown`, SDK decode errors are `Corrupt` — a typed transport earns no policy exemption - *(streaming)* `WireAdapter` gains an associated `Frame` type so typed-event wires implement the same contract over their SDK events; `classify` now takes the frame by value - *(streaming)* the conformance corpus accepts typed-event input (`WireInput::{Bytes, Event}`), so typed wires run the shared scenarios events-first with no mock transport; frame-level scenarios a typed wire cannot spell report visible skips diff --git a/crates/rig-core/src/tool/result.rs b/crates/rig-core/src/tool/result.rs index 6b6d26e424..dc3d4817ea 100644 --- a/crates/rig-core/src/tool/result.rs +++ b/crates/rig-core/src/tool/result.rs @@ -19,6 +19,12 @@ pub enum ToolErrorKind { Cancelled, /// The requested tool or resource was not found. NotFound, + /// The tool is advertised to the model but is not executable by the + /// dispatcher — its call must be handled by the host. Produced for the + /// synthetic structured-output tool, whose call carries the final + /// structured answer; also the natural classification for host-executed + /// or deferred tools. + NotExecutable, /// An authorization or permission check failed. Intentional tool refusals /// use this normalized kind with a separate refusal disposition. PermissionDenied, @@ -40,6 +46,7 @@ impl ToolErrorKind { Self::Timeout => "timeout", Self::Cancelled => "cancelled", Self::NotFound => "not_found", + Self::NotExecutable => "not_executable", Self::PermissionDenied => "permission_denied", Self::RateLimited => "rate_limited", Self::Provider => "provider", @@ -51,9 +58,11 @@ impl ToolErrorKind { const fn default_retryable(self) -> Option { match self { Self::Timeout | Self::RateLimited | Self::Network => Some(true), - Self::InvalidArgs | Self::Cancelled | Self::NotFound | Self::PermissionDenied => { - Some(false) - } + Self::InvalidArgs + | Self::Cancelled + | Self::NotFound + | Self::NotExecutable + | Self::PermissionDenied => Some(false), Self::Provider | Self::Other => None, } } @@ -64,6 +73,7 @@ impl ToolErrorKind { Self::Timeout => "tool execution timed out", Self::Cancelled => "tool execution was cancelled", Self::NotFound => "the requested tool or resource was not found", + Self::NotExecutable => "the tool is not executable by the dispatcher", Self::PermissionDenied => "the tool denied the request", Self::RateLimited => "the tool was rate limited; try again later", Self::Provider => "the tool provider failed", @@ -139,6 +149,13 @@ impl ToolExecutionError { Self::new(ToolErrorKind::NotFound, message) } + /// Advertised but not executable by the dispatcher — the host must handle + /// the call (e.g. the synthetic structured-output tool, whose call is the + /// final structured answer). + pub fn not_executable(message: impl Into) -> Self { + Self::new(ToolErrorKind::NotExecutable, message) + } + /// An authorization or permission failure. /// /// This is an ordinary execution error. Use [`Self::refused`] when the tool diff --git a/examples/agent_run_stepping/src/main.rs b/examples/agent_run_stepping/src/main.rs index 95d70cc3b0..46ae4a7015 100644 --- a/examples/agent_run_stepping/src/main.rs +++ b/examples/agent_run_stepping/src/main.rs @@ -1,16 +1,14 @@ //! Two complementary ways to drive the agent loop. //! -//! ## Part 1 — hand-driven [`AgentRun`] state machine +//! ## Part 1 — hand-driving a configured agent with [`Agent::drive`] //! -//! `agent.prompt(...)` runs this machine internally; stepping it yourself lets -//! you inspect every model call, execute tools with your own policy, and — -//! because the machine is fully serializable between steps — pause a run while -//! tool calls are pending and resume it later (even in another process). -//! -//! The per-turn request comes from [`Agent::prepare_turn`]: the loop reuses the -//! configured `Agent`'s preamble, tools, and model parameters instead of -//! restating them, and dispatches tool calls through the same registry snapshot -//! the provider saw advertised. +//! `agent.prompt(...)` runs the sans-IO [`AgentRun`] machine internally; +//! driving it yourself lets you inspect every model call, own the provider +//! transport, execute tools with your own policy, and — because the run state +//! is fully serializable between steps — pause a run while tool calls are +//! pending and resume it later, even in another process. The driver owns the +//! run/turn pairing (request preparation, tool snapshots, structured-output +//! bookkeeping); every side effect stays in this loop. //! //! ## Part 2 — high-level [`rig::agent::AgentRunner`] with hooks //! @@ -23,12 +21,11 @@ //! Requires `OPENAI_API_KEY`. use anyhow::Result; -use rig::agent::TurnTools; -use rig::agent::run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}; +use rig::agent::run::{AgentRun, ModelTurnOutcome}; use rig::agent::{ - AgentHook, HookContext, InvalidToolCallAction, ToolCall as ToolCallEvent, ToolCallAction, + AgentHook, DriveStep, HookContext, InvalidToolCallAction, ToolCall as ToolCallEvent, + ToolCallAction, }; -use rig::message::UserContent; use rig::prelude::*; use rig::providers::openai; use rig::tool::Tool; @@ -97,87 +94,59 @@ async fn main() -> Result<()> { let model = openai.completion_model(openai::GPT_4O); let agent = rig::agent::AgentBuilder::new(model) .preamble("You are a calculator. Always use the provided tools to compute results.") + .default_max_turns(2) .tool(Add) .build(); - let mut run = AgentRun::new("What is 2 + 5?").max_turns(2); - // The tool sets and dispatch target of the most recent prepared turn. Tool - // calls always execute through the snapshot whose definitions the provider - // saw — the same guarantee the runner gives its own turns. - let mut turn_tools: Option = None; + // The driver seeds the run from the agent's configuration (turn budget, + // tool choice, output schema) and owns the run/turn pairing. Every side + // effect — the provider call, tool execution — stays in this loop, and no + // agent hooks run. + let mut driver = agent.drive("What is 2 + 5?"); loop { - match run.next_step()? { - AgentRunStep::CallModel { - prompt, - history, - turn, - } => { + match driver.next_step().await? { + DriveStep::SendRequest { request, turn, .. } => { println!("→ model call #{turn}"); - // A hand-driven `AgentRun` is a sans-IO protocol primitive, not - // execution of the configured `Agent`: the driver owns the IO - // and no agent hooks run. `prepare_turn` supplies the request — - // preamble, tools, model parameters — from the agent's - // configuration instead of restating it here. - let (request, tools) = agent.prepare_turn(prompt, &history).await?.into_parts(); let response = request.send().await?; - - let mut outcome = run.model_response(ModelTurn::new( - response.message_id.clone(), - response.choice.clone(), - response.usage, - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ))?; - turn_tools = Some(tools); + let mut outcome = driver.model_response(&response)?; while let ModelTurnOutcome::NeedsResolution(context) = outcome { eprintln!("model called unknown tool `{}`", context.tool_name); // Preserve the agent loop's default fail-fast behavior; a // driver could instead retry, repair, or skip here. - outcome = run.resolve_invalid_tool_call(InvalidToolCallAction::fail())?; + outcome = driver.resolve_invalid_tool_call(InvalidToolCallAction::fail())?; } } - AgentRunStep::CallTools { .. } => { - // The whole run is serializable while tool calls are pending: + DriveStep::ExecuteTools { .. } => { + // The run state is serializable while tool calls are pending: // persist it here to pause for approval and resume later — - // even in a process that never saw this step. The resumed run - // re-emits the pending tool calls from its own state. (Tool - // implementations are live objects: a genuinely separate - // process would rebuild the same `Agent` and prepare its own - // turns from there.) - let suspended = serde_json::to_string(&run)?; - let mut run_resumed: AgentRun = serde_json::from_str(&suspended)?; - let AgentRunStep::CallTools { calls } = run_resumed.next_step()? else { - anyhow::bail!("resumed run must re-emit the pending tool calls"); - }; - let Some(tools) = turn_tools.as_ref() else { - anyhow::bail!("CallTools always follows a prepared CallModel turn"); + // even in a process that never saw this step. Resuming + // rebuilds the driver from the same agent; the resumed driver + // re-emits the pending calls and re-derives its dispatch + // snapshot (tool implementations are live objects, so a fresh + // process dispatches against its own registry). + let suspended = serde_json::to_string(driver.run())?; + let resumed: AgentRun = serde_json::from_str(&suspended)?; + driver = agent.drive_run(resumed); + let DriveStep::ExecuteTools { calls, tools } = driver.next_step().await? else { + anyhow::bail!("a resumed run re-emits its pending tool calls"); }; + let mut context = rig::tool::ToolContext::new(); let mut results = Vec::new(); - for call in calls { - // Tool calls suppressed by invalid tool-call recovery come - // with a pre-resolved result and must not be executed. - if let Some(result) = call.preresolved_result { - results.push(result); - continue; - } - let name = &call.tool_call.function.name; - let args = call.tool_call.function.arguments.to_string(); - println!("→ executing {name}({args})"); - let mut context = rig::tool::ToolContext::new(); - let result = tools.execute(name, &args, &mut context).await; - results.push(UserContent::tool_result_for( - call.tool_call.id.clone(), - call.tool_call.provider.clone(), - name.clone(), - result.output().clone().into_content(), - )); + for call in &calls { + println!( + "→ executing {}({})", + call.tool_call.function.name, call.tool_call.function.arguments + ); + // `execute_call` honors pre-resolved results (from invalid + // tool-call recovery) and dispatches through the exact + // snapshot the provider saw advertised for this turn. + results.push(tools.execute_call(call, &mut context).await); } - run_resumed.tool_results(results)?; - run = run_resumed; + driver.tool_results(results)?; } - AgentRunStep::Done(response) => { + DriveStep::Done(response) => { println!("✓ {}", response.output); println!( " {} model call(s), {} total tokens", diff --git a/examples/agent_with_durable_approval/src/main.rs b/examples/agent_with_durable_approval/src/main.rs index a16a717a8a..d1541b130d 100644 --- a/examples/agent_with_durable_approval/src/main.rs +++ b/examples/agent_with_durable_approval/src/main.rs @@ -28,8 +28,8 @@ //! Requires `OPENAI_API_KEY`. Run with: `cargo run -p agent_with_durable_approval` use anyhow::Result; -use rig::agent::run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}; -use rig::agent::{InvalidToolCallAction, TurnTools}; +use rig::agent::run::{AgentRun, ModelTurnOutcome}; +use rig::agent::{DriveStep, InvalidToolCallAction}; use rig::message::{ToolResultContent, UserContent}; use rig::prelude::*; use rig::providers::openai; @@ -143,16 +143,17 @@ async fn ask(prompt: &str) -> Option { #[tokio::main] async fn main() -> Result<()> { - // A hand-driven `AgentRun` is a sans-IO protocol primitive: this loop owns - // the IO and no agent hooks run. The per-turn request and the tool dispatch - // target both come from the configured `Agent` via `prepare_turn`, so - // nothing about the agent's configuration is restated below. + // A hand-driven run is a sans-IO protocol primitive: this loop owns the IO + // and no agent hooks run. The driver pairs each turn's request and tool + // dispatch with the configured `Agent`, so nothing about the agent's + // configuration is restated below. let model = openai::Client::from_env()?.completion_model(openai::GPT_4O); let agent = rig::agent::AgentBuilder::new(model) .preamble( "You are a banking assistant. Use the tools to carry out the user's request. \ Call one tool at a time.", ) + .default_max_turns(10) .tool(GetBalance) .tool(TransferFunds) .build(); @@ -164,54 +165,38 @@ async fn main() -> Result<()> { let state_path = std::env::temp_dir().join("rig_durable_approval.json"); let _ = std::fs::remove_file(&state_path); - let mut run = AgentRun::new(prompt).max_turns(10); - // Tool dispatch target of the most recent prepared turn: approved calls - // execute through the snapshot whose definitions the provider saw. - let mut turn_tools: Option = None; + let mut driver = agent.drive(prompt); loop { - match run.next_step()? { - AgentRunStep::CallModel { - prompt, - history, - turn, - } => { + match driver.next_step().await? { + DriveStep::SendRequest { request, turn, .. } => { println!("\n→ model call #{turn}"); - let (request, tools) = agent.prepare_turn(prompt, &history).await?.into_parts(); let response = request.send().await?; - let mut outcome = run.model_response(ModelTurn::new( - response.message_id.clone(), - response.choice.clone(), - response.usage, - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ))?; - turn_tools = Some(tools); + let mut outcome = driver.model_response(&response)?; while let ModelTurnOutcome::NeedsResolution(context) = outcome { eprintln!("model called unknown tool `{}`", context.tool_name); - outcome = run.resolve_invalid_tool_call(InvalidToolCallAction::fail())?; + outcome = driver.resolve_invalid_tool_call(InvalidToolCallAction::fail())?; } } - AgentRunStep::CallTools { .. } => { - // DURABLE PAUSE. Persist the whole run, then reconstruct it from - // the file before deciding. The write→reload boundary below could - // be a separate process / request / much later — the resumed run - // re-emits the pending tool calls purely from serialized state. - std::fs::write(&state_path, serde_json::to_vec_pretty(&run)?)?; + DriveStep::ExecuteTools { .. } => { + // DURABLE PAUSE. Persist the run state, then rebuild the driver + // from the file before deciding. The write→reload boundary below + // could be a separate process / request / much later — the + // resumed driver re-emits the pending tool calls purely from + // serialized state, and re-derives its dispatch snapshot from + // the rebuilt agent (tool implementations are live objects; if + // a pending tool is missing from this process's registry the + // driver surfaces the drift as an error instead of dispatching). + std::fs::write(&state_path, serde_json::to_vec_pretty(driver.run())?)?; println!("\n💾 run suspended to {}", state_path.display()); // ----- imagine the process exits here and resumes later ----- - // (Tool implementations are live objects: a genuinely separate - // process would rebuild the same `Agent` and prepare its own - // turns from there.) - let mut resumed: AgentRun = serde_json::from_slice(&std::fs::read(&state_path)?)?; - let AgentRunStep::CallTools { calls } = resumed.next_step()? else { - anyhow::bail!("resumed run must re-emit the pending tool calls"); - }; - let Some(tools) = turn_tools.as_ref() else { - anyhow::bail!("CallTools always follows a prepared CallModel turn"); + let resumed: AgentRun = serde_json::from_slice(&std::fs::read(&state_path)?)?; + driver = agent.drive_run(resumed); + let DriveStep::ExecuteTools { calls, tools } = driver.next_step().await? else { + anyhow::bail!("a resumed run re-emits its pending tool calls"); }; let mut results = Vec::new(); @@ -308,12 +293,11 @@ async fn main() -> Result<()> { return Ok(()); } - resumed.tool_results(results)?; + driver.tool_results(results)?; let _ = std::fs::remove_file(&state_path); - run = resumed; } - AgentRunStep::Done(response) => { + DriveStep::Done(response) => { println!("\n✓ {}", response.output); return Ok(()); } From ae1934bd398b28e4b0822cec8f52ab230b26c82a Mon Sep 17 00:00:00 2001 From: stephen Date: Mon, 10 Aug 2026 19:43:38 -0700 Subject: [PATCH 03/21] style: cargo fmt --- crates/rig-agent/src/agent/mod.rs | 4 ++-- crates/rig-agent/src/agent/prompt_request/streaming.rs | 2 +- crates/rig-agent/src/agent/runner.rs | 2 +- crates/rig-agent/src/agent/turn_tools.rs | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index 22e7c21d7d..bf8ae0334e 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -116,6 +116,7 @@ pub(crate) const UNKNOWN_AGENT_NAME: &str = "Unnamed Agent"; pub use builder::{AgentBuilder, NoToolConfig, WithBuilderTools, WithToolServerHandle}; pub use completion::Agent; +pub use driver::{AgentDriver, DriveStep}; pub use hook::CompletionCall as CompletionCallEvent; pub use hook::{ AgentHook, CompletionCallAction, CompletionResponse as CompletionResponseEvent, HookContext, @@ -124,9 +125,7 @@ pub use hook::{ RetryRequest, RunId, Scratchpad, StepEventKind, StreamResponseFinish, TextDelta, ToolCall, ToolCallAction, ToolCallDelta, ToolResultAction, ToolResultEvent, }; -pub use driver::{AgentDriver, DriveStep}; pub use model::ModelHandle; -pub use turn_tools::TurnTools; pub use prompt_request::streaming::{ MultiTurnStreamItem, StreamingError, StreamingPromptRequest, StreamingResult, stream_to_stdout, }; @@ -137,3 +136,4 @@ pub use prompt_request::{ pub use rig_core::message::Text; pub use run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome, OutputMode, PendingToolCall}; pub use runner::AgentRunner; +pub use turn_tools::TurnTools; diff --git a/crates/rig-agent/src/agent/prompt_request/streaming.rs b/crates/rig-agent/src/agent/prompt_request/streaming.rs index 8c89f454e1..490601d937 100644 --- a/crates/rig-agent/src/agent/prompt_request/streaming.rs +++ b/crates/rig-agent/src/agent/prompt_request/streaming.rs @@ -5,7 +5,6 @@ use rig_core::{ use crate::{ agent::completion::build_prepared_completion_request, - agent::turn_tools::PreparedCompletionRequest, agent::hook::{ AgentHook, HookContext, HookStack, InvalidToolCallAction, ModelSelection, ModelSelectionAction, ModelTurnFinished, ReasoningDelta, StepEventKind, @@ -21,6 +20,7 @@ use crate::{ append_run_messages, build_chat_span, new_execute_tool_span, observe_action, resolve_completion_call, resolve_model_turn_action, run_single_tool, }, + agent::turn_tools::PreparedCompletionRequest, streaming::{StreamedAssistantContent, StreamedUserContent, ToolCallDeltaContent}, tool::{ToolContext, server::ToolRegistrySnapshot}, }; diff --git a/crates/rig-agent/src/agent/runner.rs b/crates/rig-agent/src/agent/runner.rs index 9109d293ac..9ef2ac45a8 100644 --- a/crates/rig-agent/src/agent/runner.rs +++ b/crates/rig-agent/src/agent/runner.rs @@ -34,7 +34,6 @@ use tracing::{Instrument, info_span, span::Id}; use super::{ completion::Agent, - turn_tools::PreparedCompletionRequest, hook::{ AgentHook, CompletionCall, CompletionCallAction, CompletionResponse as CompletionResponseEvent, HookContext, HookStack, @@ -51,6 +50,7 @@ use super::{ tool_result_output, }, run::{AgentRun, DEFAULT_OUTPUT_RETRIES, ModelTurnOutcome, OutputMode, PendingToolCall}, + turn_tools::PreparedCompletionRequest, }; use rig_core::{ memory::ConversationMemory, diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs index 43d4788e78..e2774f2d28 100644 --- a/crates/rig-agent/src/agent/turn_tools.rs +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -175,8 +175,8 @@ impl TurnTools { #[cfg(test)] mod tests { - use crate::agent::{AgentBuilder, DriveStep}; use crate::agent::run::OutputMode; + use crate::agent::{AgentBuilder, DriveStep}; use crate::test_utils::MockCompletionModel; use crate::tool::{Tool, ToolContext, ToolErrorKind}; use serde_json::json; From ce401a6d428e6253baf9f6a342150c4a5fe39536 Mon Sep 17 00:00:00 2001 From: stephen Date: Mon, 10 Aug 2026 23:25:29 -0700 Subject: [PATCH 04/21] refactor!(agent): make every step of a driven run a resume point Review of the AgentDriver shell (PR #2278 round 2) found four defects, three of them one structural fact: the driver held turn_tools outside the serialized AgentRun. Everything downstream followed from that. A model call was committed - turn consumed, state moved to AwaitingModel - before the fallible request preparation ran, so a transiently unreachable tool server left the run wedged: next_step rejected it as a pending response, model_response rejected it for having no prepared turn, and no third door existed. There was no AwaitingModel counterpart to resume_tools, so a run serialized after SendRequest - the natural suspension point for a caller who owns the transport - could never be resumed. resume_tools validated a tool_choice against a request it was not building, masking the actionable drift message and making allow_missing_resumed_tools unreachable, contradicting its documented contract. Separately, a custom run's tool_choice reached only the run's internal decisions and never the wire, while its sibling output_tool_name was threaded correctly. TurnTools splits at the durability line. TurnToolNames (new, public, Serialize) is data and is recorded on the run at commit; TurnTools::from_parts pairs it back with a live snapshot, which cannot be serialized in any design. AgentDriver keeps only that snapshot, as a cache: it exists so a turn prepared in this process dispatches through the implementations the provider was shown, and is rebuilt on demand after a resume. model_response reads the names from the run, so a resumed run validates its reply against the set the request actually carried rather than against whatever the registry holds now. AgentRun::next_step's PreparingRequest arm decomposes into a pure peek_model_call, an infallible commit_model_call, and advance() -> NeedsModelCall | CallTools | Done. reprompt_for_output no longer recurses into next_step - it parks the run and lets advance loop - which is what closes the path by which the machine could commit a model call the driver had not prepared a request for. next_step is advance plus an immediate commit, so the runner, the streaming path, the conformance suite and the module-doc example are unchanged; the driver uses the two halves and a preparation failure now leaves the run byte-identical, ready to retry. Per-turn configuration gets one seam. The driver was passing None for the request_patch parameter it already had; AgentDriver::request_patch / set_request_patch now supply it, giving a hand-driven run the per-turn preamble, sampling parameters, tool_choice, active_tools narrowing, extra context and substituted history the runner gets from its CompletionCall hooks. The patch's tool_choice is seeded from the new public AgentRun::tool_choice() when the caller sets none, so a custom run's choice reaches the provider. build_prepared_completion_ request's 17 positional parameters collapse into TurnBaseline and TurnRequest, retiring the transposition hazard for the parameters TurnTools had only fixed two of. Serialized runs carry $schemaVersion (RUN_SCHEMA_VERSION = "1.0") and a build reads only the version it writes, with no serde(default) so an untagged payload fails on the missing field. This is breaking: runs suspended by an earlier build cannot be resumed, and the prose warning it replaces could not fail a load. In exchange a run is now suspendable and resumable at every step boundary, model call in flight included. Five driver regression tests (one per defect plus the patch seam) and three format-tag tests; the pre-monoid fixture test is re-stamped and split, since asserting that an untagged payload loads is what the version tag deliberately reverses. 495 lib tests pass, clippy and fmt clean across the workspace, both examples check. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- MIGRATING.md | 17 + crates/rig-agent/CHANGELOG.md | 6 + crates/rig-agent/src/agent/completion.rs | 109 ++- crates/rig-agent/src/agent/driver.rs | 517 +++++++++---- .../src/agent/prompt_request/streaming.rs | 40 +- crates/rig-agent/src/agent/run/mod.rs | 697 ++++++++++++------ crates/rig-agent/src/agent/turn_tools.rs | 82 ++- 7 files changed, 1072 insertions(+), 396 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index ce1a2e793b..1232177fa4 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1572,6 +1572,23 @@ not an execution path — hooks, memory, retrieval policy, and telemetry still run only under `Agent::runner`. See `examples/agent_run_stepping` and `examples/agent_with_durable_approval`. +A hand-driven turn gets per-turn configuration through +`AgentDriver::request_patch(RequestPatch)` (or `set_request_patch` between +steps) — the same `RequestPatch` the runner merges from its `CompletionCall` +hooks, covering the per-turn preamble, sampling parameters, `tool_choice`, +`active_tools` narrowing, extra context, and substituted history. A custom +run's own `AgentRun::with_tool_choice` now also reaches the provider; before, +it governed only the run's internal decisions while the request carried the +agent's baseline. + +**Serialized `AgentRun` state does not survive the upgrade.** Payloads now +carry a `$schemaVersion` tag and a build reads only the version it writes, so a +run suspended by an earlier release fails to deserialize with a named error +rather than being reinterpreted against fields whose meaning moved. Drain +in-flight suspended runs before upgrading, or discard them. In exchange, a run +can now be suspended and resumed at *every* step boundary — including with a +model call in flight — not only while tool calls are pending. + An `Agent`'s default model is set at construction. Per-run overrides now go through `runner(...).using_model(...)`, `Agent::set_model`, or a `ModelSelection` hook (see the "runtime model swapping" section for the diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 39e5b9f490..2bede78126 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* [**breaking**] `PromptResponse` JSON serialized before the `content` field existed no longer deserializes — the missing-`content` reconstruction (and the serde shadow repr that carried it) is deleted and `content` is a required field; the JSON wire shape is unchanged +- *(agent)* [**breaking**] the serialized `AgentRun` carries a `$schemaVersion` tag (`rig_agent::agent::run::RUN_SCHEMA_VERSION`, currently `"1.0"`) and a build reads only the version it writes. A payload from another version — including one written before versioning existed, which has no tag — fails deserialization with a named error instead of being silently reinterpreted. Runs suspended by an earlier build cannot be resumed by this one; drain them before upgrading + - *(agent)* [**breaking**] message content is a plain `Vec`, following rig-core's removal of the non-empty container — every content list the agent constructs, inspects, or hands to hooks is a `Vec`, and the `OneOrMany` re-export is gone from this crate - *(agent)* [**behavior**] a genuinely empty assistant turn no longer cancels the run and is no longer padded with a fabricated empty-text part — with the non-empty container gone the turn is honestly representable as an empty list, and `is_empty_assistant_turn` neutralizes it instead of the agent inventing content or failing a run that used to succeed @@ -34,6 +36,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `Agent::prepare_turn`, returning the new public `PreparedTurn`/`TurnTools`: the agent's baseline configuration resolved into one turn's completion request plus the turn's executable and allowed tool-name sets, the synthetic output-tool name, and tool dispatch pinned to the turn's registry snapshot — so hand-driven `AgentRun` loops (custom provider transport, durable suspend/resume) reuse the configured `Agent` instead of restating its preamble and tools. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. A prepared turn is a configuration read: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` - *(agent)* add `Agent::drive` / `Agent::drive_run`, returning the new public `AgentDriver` (with `DriveStep` and `TurnTools`): hand-drive the sans-IO `AgentRun` machine with the agent's own configuration while owning every side effect. The driver seeds the run from the agent (`default_max_turns`, `tool_choice`, output schema), yields each turn's fully configured completion request for the caller to send (or hand to a custom transport), pairs every pending tool batch with the exact registry snapshot that advertised it, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. Resuming a serialized run in a fresh process re-derives a fresh dispatch snapshot from the rebuilt agent, and surfaces missing pending tools as an error (opt out with `allow_missing_resumed_tools`). The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` +- *(agent)* `AgentDriver` is resumable at **every** step boundary, not just while tool calls are pending. The turn's advertised tool names are recorded on the run (new public `TurnToolNames`, the serializable half of `TurnTools`), so a run serialized after `DriveStep::SendRequest` — the natural suspension point for a queued or long-running provider call — can be resumed in another process and fed its reply, validated against the set the request actually carried rather than against whatever the resuming registry holds. The driver now keeps no state of its own beyond a registry-snapshot cache +- *(agent)* a failed model-turn preparation costs nothing: the run advances only once a request exists, so an unreachable tool server or an impossible `tool_choice` leaves the run byte-identical — same state, same turn budget — and the step can be retried in place. Previously the turn was consumed and the run was left unresumable +- *(agent)* add `AgentDriver::request_patch` / `set_request_patch`, giving a hand-driven run the per-turn configuration the runner gets from its `CompletionCall` hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history. Because the caller owns the loop, per-turn variation needs no callback +- *(agent)* add `AgentRun::tool_choice()`. A custom run's own `tool_choice` now reaches the provider when the run is hand-driven, instead of only governing the run's internal decisions while the request silently carried the agent's baseline - *(tool)* add `ToolErrorKind::NotExecutable` (`ToolExecutionError::not_executable`): the tool is advertised to the model but not executable by the dispatcher — produced when dispatching the synthetic structured-output tool, whose call carries the final structured answer - *(agent)* add opaque, cloneable `ModelHandle` values with by-value `ProviderCapabilities` snapshots, plus default replacement, per-run default override (`using_model`), and hook-driven per-call selection via `AgentHook::on_model_select` - *(agent)* add run-local extractor default-model overrides used across retries diff --git a/crates/rig-agent/src/agent/completion.rs b/crates/rig-agent/src/agent/completion.rs index 8d2467f1c3..a1c656cef6 100644 --- a/crates/rig-agent/src/agent/completion.rs +++ b/crates/rig-agent/src/agent/completion.rs @@ -201,28 +201,101 @@ pub(crate) fn allowed_tool_names_for_choice( Ok(allowed) } +/// The configured baseline a turn is prepared against: everything that is a +/// property of the *agent* rather than of the turn. +/// +/// Grouping these is not cosmetic. Passed positionally they were a long tail of +/// same-typed arguments — three `Option<&str>`-ish, two numeric options, two +/// bools — that a caller could transpose without the compiler noticing. Named +/// fields make each one say what it is at the call site. +pub(crate) struct TurnBaseline<'a> { + /// The model this turn is prepared for. Preparation reads its captured + /// capabilities, and the same handle executes the prepared request — so a + /// caller selecting a model per turn passes the selected one here. + pub(crate) model: &'a ModelHandle, + pub(crate) preamble: Option<&'a str>, + pub(crate) static_context: &'a [Document], + pub(crate) temperature: Option, + pub(crate) max_tokens: Option, + pub(crate) additional_params: Option<&'a serde_json::Value>, + pub(crate) record_telemetry_content: bool, + pub(crate) tool_choice: Option<&'a ToolChoice>, + pub(crate) tool_server_handle: &'a ToolServerHandle, + pub(crate) output_schema: Option<&'a schemars::Schema>, + pub(crate) output_mode: &'a OutputMode, + /// Description advertised for the synthetic structured-output tool. + pub(crate) output_tool_description: Option<&'a str>, + /// Whether Tool output mode augments the preamble with output guidance. + pub(crate) augment_output_preamble: bool, +} + +impl<'a> TurnBaseline<'a> { + /// The baseline as an agent configures it. `output_tool_description` and + /// `augment_output_preamble` are runner-level knobs, so they take their + /// defaults here. + pub(crate) fn from_agent(agent: &'a Agent) -> Self { + Self { + model: &agent.model, + preamble: agent.preamble.as_deref(), + static_context: &agent.static_context, + temperature: agent.temperature, + max_tokens: agent.max_tokens, + additional_params: agent.additional_params.as_ref(), + record_telemetry_content: agent.record_telemetry_content, + tool_choice: agent.tool_choice.as_ref(), + tool_server_handle: &agent.tool_server_handle, + output_schema: agent.output_schema.as_ref(), + output_mode: &agent.output_mode, + output_tool_description: None, + augment_output_preamble: true, + } + } +} + +/// What distinguishes one turn from the next: the message to send, the history +/// behind it, the run's committed output tool, and the per-turn patch layered +/// over the baseline. +pub(crate) struct TurnRequest<'a> { + pub(crate) prompt: Message, + pub(crate) chat_history: &'a [Message], + /// The run's already-committed Tool-mode output name, re-advertised so the + /// mode cannot flip or re-pick a name mid-run (#1928). + pub(crate) committed_output_tool: Option<&'a str>, + /// Per-turn overrides — from `CompletionCall` hooks in the runner, from + /// [`AgentDriver::request_patch`](crate::agent::AgentDriver::request_patch) + /// when hand-driven. The single seam through which a turn diverges from + /// the baseline. + pub(crate) patch: Option<&'a RequestPatch>, +} + /// Helper function to build a completion request from agent components while /// preserving the executable Rig tool names sent to the provider. -#[allow(clippy::too_many_arguments)] pub(crate) async fn build_prepared_completion_request( - model: &ModelHandle, - prompt: Message, - chat_history: &[Message], - preamble: Option<&str>, - static_context: &[Document], - temperature: Option, - max_tokens: Option, - additional_params: Option<&serde_json::Value>, - record_telemetry_content: bool, - tool_choice: Option<&ToolChoice>, - tool_server_handle: &ToolServerHandle, - output_schema: Option<&schemars::Schema>, - output_mode: &OutputMode, - committed_output_tool: Option<&str>, - output_tool_description: Option<&str>, - augment_output_preamble: bool, - request_patch: Option<&RequestPatch>, + baseline: TurnBaseline<'_>, + turn: TurnRequest<'_>, ) -> Result { + let TurnBaseline { + model, + preamble, + static_context, + temperature, + max_tokens, + additional_params, + record_telemetry_content, + tool_choice, + tool_server_handle, + output_schema, + output_mode, + output_tool_description, + augment_output_preamble, + } = baseline; + let TurnRequest { + prompt, + chat_history, + committed_output_tool, + patch: request_patch, + } = turn; + // Apply a per-turn request patch (the merged patch from every `CompletionCall` // hook): each set field replaces the agent's configured value for this turn, // unset fields inherit it, `additional_params` is shallow-merged, and diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index ca466d926a..28bf0cf47e 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -22,35 +22,55 @@ //! advertise/dispatch skew happen; the driver owns it in one place while every //! side effect stays with the caller. //! +//! It owns that pairing without *holding* it. Everything durable lives on the +//! [`AgentRun`] — including the turn's advertised tool names — and the driver +//! keeps only a cache of the live registry snapshot, which cannot be +//! serialized in any design. That is what makes the durability guarantees +//! below hold at every step rather than at one of them. +//! //! # Durability //! -//! The serializable state is still [`AgentRun`] — serialize -//! [`AgentDriver::run`] while tool calls are pending, and resume in another -//! process with [`Agent::drive_run`]. Tool implementations are live objects -//! and cannot be serialized: the resuming process rebuilds the same `Agent` -//! and the driver takes a **fresh** registry snapshot for the pending calls. -//! If that snapshot no longer contains a pending call's tool, the driver -//! surfaces an error instead of silently feeding a not-found result to the -//! model (see [`AgentDriver::allow_missing_resumed_tools`]) — the model chose -//! that tool from a registry this process no longer has, and re-prompting -//! cannot fix deployment drift. If you suspend runs across deploys, version -//! your agent definitions alongside the serialized run. - -use std::collections::BTreeSet; +//! The serializable state is *all* of the state: the driver holds nothing it +//! could lose. Serialize [`AgentDriver::run`] at any step boundary — while +//! tool calls are pending, or while a model call is in flight with a +//! long-running or queued provider — and resume in another process with +//! [`Agent::drive_run`]. Every step is a resume point, including +//! [`DriveStep::SendRequest`]: the turn's advertised tool names travel with +//! the run, so the resuming process validates the model's reply against the +//! set the request actually carried rather than against whatever its registry +//! holds now. +//! +//! Preparation failures are equally survivable. Nothing advances until a +//! request exists, so a turn that fails to prepare — an unreachable tool +//! server, an impossible `tool_choice` — costs no turn from the budget and +//! leaves the run byte-identical, ready to retry. +//! +//! Tool *implementations* are live objects and cannot be serialized: the +//! resuming process rebuilds the same `Agent` and the driver takes a fresh +//! registry snapshot to dispatch pending calls through. If that snapshot no +//! longer contains a pending call's tool, the driver surfaces an error instead +//! of silently feeding a not-found result to the model (see +//! [`AgentDriver::allow_missing_resumed_tools`]) — the model chose that tool +//! from a registry this process no longer has, and re-prompting cannot fix +//! deployment drift. If you suspend runs across deploys, version your agent +//! definitions alongside the serialized run; the run's own format is versioned +//! by [`RUN_SCHEMA_VERSION`](super::run::RUN_SCHEMA_VERSION). + use std::sync::Arc; use rig_core::message::UserContent; -use super::completion::{Agent, allowed_tool_names_for_choice, build_prepared_completion_request}; +use super::completion::{Agent, TurnBaseline, TurnRequest, build_prepared_completion_request}; use super::model::ModelHandle; -use super::run::{AgentRun, AgentRunStep, ModelTurnOutcome, PendingToolCall}; +use super::run::{Advance, AgentRun, ModelCallInputs, ModelTurnOutcome, PendingToolCall}; use super::runner::build_agent_run; use super::turn_tools::{PreparedCompletionRequest, TurnTools}; -use crate::agent::hook::InvalidToolCallAction; +use crate::agent::hook::{InvalidToolCallAction, RequestPatch}; use crate::agent::prompt_request::PromptResponse; use crate::completion::{ CompletionError, CompletionRequestBuilder, CompletionResponse, Message, PromptError, }; +use crate::tool::server::ToolRegistrySnapshot; impl Agent { /// Hand-drive this agent: build a driver whose run is seeded from the @@ -80,7 +100,8 @@ impl Agent { AgentDriver { agent: self.clone(), run, - turn_tools: None, + snapshot: None, + request_patch: RequestPatch::new(), allow_missing_resumed_tools: false, } } @@ -98,13 +119,15 @@ pub enum DriveStep { SendRequest { /// The fully configured request: the agent's preamble (with any /// output-mode augmentation), static context, model parameters, - /// `tool_choice`, and this turn's tool definitions. It reflects the - /// agent's **baseline** configuration — no `CompletionCall` hooks run - /// on this path, so there is no per-turn request patch, model - /// selection, or `active_tools` narrowing. The request still honors - /// the agent's `record_telemetry_content` for provider-level spans; - /// call `.record_content_telemetry(false)` on it to opt a hand-driven - /// turn out. + /// `tool_choice`, and this turn's tool definitions — with the + /// driver's [`RequestPatch`](AgentDriver::request_patch) applied over + /// that baseline, which is how a hand-driven turn gets the per-turn + /// preamble, `tool_choice`, or `active_tools` narrowing the runner + /// gets from its `CompletionCall` hooks. No hooks run on this path; + /// the patch is the seam. The request still honors the agent's + /// `record_telemetry_content` for provider-level spans; call + /// `.record_content_telemetry(false)` on it to opt a hand-driven turn + /// out. request: Box>, /// The turn's advertised tool sets — informational here (the driver /// assembles the model turn itself); the same value arrives on the @@ -152,10 +175,17 @@ impl std::fmt::Debug for DriveStep { pub struct AgentDriver { agent: Agent, run: AgentRun, - /// Tool state of the most recently prepared turn. `None` until the first - /// `SendRequest` — or in a process that resumed a serialized run, where - /// [`Self::resume_tools`] derives a fresh snapshot on demand. - turn_tools: Option, + /// Live dispatch target for the current turn — a **cache**, never state. + /// + /// Everything the driver must not lose lives on [`Self::run`]; a tool + /// registry snapshot cannot, because implementations are live objects. It + /// is held only so that a turn prepared in *this* process dispatches + /// through the exact implementations the provider was shown, and is + /// rebuilt on demand when a resumed run reaches its pending tool calls. + snapshot: Option>, + /// Per-turn request configuration applied to every turn this driver + /// prepares. See [`Self::request_patch`]. + request_patch: RequestPatch, allow_missing_resumed_tools: bool, } @@ -172,6 +202,29 @@ impl AgentDriver { self } + /// Set the per-turn request configuration for the turns this driver + /// prepares. + /// + /// The driver runs no hooks, so this is how a hand-driven run gets what + /// the runner gets from its `CompletionCall` hooks: a per-turn preamble, + /// sampling parameters, `tool_choice`, `active_tools` narrowing, extra + /// context, or a substituted history. Each set field replaces the agent's + /// configured value for the turn; unset fields inherit it. + /// + /// The patch applies to every turn this driver prepares. Because the + /// caller owns the loop, per-turn variation needs no callback — call + /// [`Self::set_request_patch`] between steps. + pub fn request_patch(mut self, patch: RequestPatch) -> Self { + self.request_patch = patch; + self + } + + /// Replace the per-turn request configuration in place, so a driving loop + /// can vary it from turn to turn. See [`Self::request_patch`]. + pub fn set_request_patch(&mut self, patch: RequestPatch) { + self.request_patch = patch; + } + /// Opt out of the resumed-run drift check: dispatch pending calls whose /// tools are missing from this process's registry anyway, feeding the /// resulting not-found errors to the model instead of surfacing the drift @@ -202,87 +255,55 @@ impl AgentDriver { /// or re-pick a name mid-run. Fails locally — with no provider /// round-trip — when the configuration cannot produce a valid request /// (e.g. a `tool_choice` impossible against the advertised tool set). + /// + /// **Such a failure costs nothing.** Preparation runs entirely before the + /// run advances: the turn is committed only once a request exists + /// ([`AgentRun::commit_model_call`]), so an error here leaves the run + /// exactly as it was — same state, same turn budget — and the step can be + /// retried once the cause is fixed (a tool server that was briefly + /// unreachable, say), here or in another process. pub async fn next_step(&mut self) -> Result { - match self.run.next_step()? { - AgentRunStep::CallModel { - prompt, - history, - turn, - } => { + match self.run.advance()? { + Advance::NeedsModelCall => { + // Peek, prepare, *then* commit. Reading the inputs consumes + // nothing, so everything fallible below happens while the run + // is still fully intact. + let ModelCallInputs { + prompt, history, .. + } = self.run.peek_model_call()?; // Pin Tool output mode once committed (#1928), mirroring the // runner: read the run's committed name into preparation, and // store the resolved name back (fill-once). let committed = self.run.output_tool_name().map(str::to_owned); + let patch = self.effective_request_patch(); let prepared = build_prepared_completion_request( - &self.agent.model, - prompt, - &history, - self.agent.preamble.as_deref(), - &self.agent.static_context, - self.agent.temperature, - self.agent.max_tokens, - self.agent.additional_params.as_ref(), - self.agent.record_telemetry_content, - self.agent.tool_choice.as_ref(), - &self.agent.tool_server_handle, - self.agent.output_schema.as_ref(), - &self.agent.output_mode, - committed.as_deref(), - None, - true, - None, + TurnBaseline::from_agent(&self.agent), + TurnRequest { + prompt, + chat_history: &history, + committed_output_tool: committed.as_deref(), + patch: Some(&patch), + }, ) .await .map_err(PromptError::CompletionError)?; - self.run - .set_output_tool_name(prepared.tools.output_tool_name.clone()); + let PreparedCompletionRequest { builder, tools } = prepared; - self.turn_tools = Some(tools.clone()); + self.snapshot = Some(tools.snapshot.clone()); + let turn = self + .run + .commit_model_call(Some(tools.names()), tools.output_tool_name.clone()); Ok(DriveStep::SendRequest { request: Box::new(builder), tools, turn, }) } - AgentRunStep::CallTools { calls } => { - let tools = match &self.turn_tools { - Some(tools) => tools.clone(), - // A process resuming a serialized run wakes here with no - // prepared turn: derive a fresh dispatch snapshot. - None => { - let tools = self.resume_tools().await?; - if !self.allow_missing_resumed_tools { - let missing: Vec<&str> = calls - .iter() - .filter(|call| call.preresolved_result.is_none()) - .map(|call| call.tool_call.function.name.as_str()) - .filter(|name| { - !tools.executable_tool_names.contains(*name) - && tools.output_tool_name() != Some(*name) - }) - .collect(); - if !missing.is_empty() { - return Err(PromptError::CompletionError( - CompletionError::RequestError( - format!( - "resumed run has pending tool calls {missing:?} that \ - are not registered in this process; register the \ - tools on the agent before resuming, or call \ - `allow_missing_resumed_tools()` to dispatch anyway \ - and feed not-found results to the model" - ) - .into(), - ), - )); - } - } - self.turn_tools = Some(tools.clone()); - tools - } - }; + Advance::CallTools(calls) => { + let tools = self.dispatch_tools_for_turn(&calls).await?; Ok(DriveStep::ExecuteTools { calls, tools }) } - AgentRunStep::Done(response) => Ok(DriveStep::Done(response)), + Advance::Done(response) => Ok(DriveStep::Done(response)), } } @@ -296,12 +317,17 @@ impl AgentDriver { &mut self, response: &CompletionResponse, ) -> Result { - let Some(tools) = &self.turn_tools else { + // The advertised names come from the run, not from this driver — which + // is what lets a run serialized between `SendRequest` and the model's + // reply be resumed in another process, and what guarantees the + // response is validated against the set the request actually carried + // rather than whatever the registry holds now. + let Some(names) = self.run.advertised_tools().cloned() else { return Err(PromptError::CompletionError(CompletionError::RequestError( "model_response must follow a SendRequest step from this driver".into(), ))); }; - self.run.model_response(tools.model_turn(response)) + self.run.model_response(names.model_turn(response)) } /// Resolve a pending invalid tool call, exactly as @@ -319,22 +345,81 @@ impl AgentDriver { self.run.tool_results(results) } - /// Derive a fresh dispatch target for a resumed run's pending tool calls. + /// The turn's tool sets paired with a dispatch target. + /// + /// The names always come from the run. The snapshot comes from this + /// process: the one taken when the turn was prepared if the turn was + /// prepared here, otherwise a fresh one — implementations are live + /// objects, so a resumed process can only dispatch against its own + /// registry. + async fn dispatch_tools_for_turn( + &mut self, + calls: &[PendingToolCall], + ) -> Result { + let Some(names) = self.run.advertised_tools().cloned() else { + return Err(PromptError::CompletionError(CompletionError::RequestError( + "the run has no advertised tool set for the pending calls; drive the model turn \ + through this driver so the turn's tools are recorded on the run" + .into(), + ))); + }; + let output_tool_name = self.run.output_tool_name().map(str::to_owned); + + if let Some(snapshot) = &self.snapshot { + return Ok(TurnTools::from_parts( + snapshot.clone(), + names, + output_tool_name, + )); + } + + // Resumed run: rebuild the live half, then report deployment drift + // rather than silently feeding not-found results to the model. The + // model chose these tools from a registry this process no longer has, + // and re-prompting cannot fix that. + let snapshot = Arc::new(self.fresh_snapshot().await?); + if !self.allow_missing_resumed_tools { + let missing: Vec<&str> = calls + .iter() + .filter(|call| call.preresolved_result.is_none()) + .map(|call| call.tool_call.function.name.as_str()) + .filter(|name| { + output_tool_name.as_deref() != Some(*name) + && !snapshot + .definitions() + .iter() + .any(|tool| tool.name.as_str() == *name) + }) + .collect(); + if !missing.is_empty() { + return Err(PromptError::CompletionError(CompletionError::RequestError( + format!( + "resumed run has pending tool calls {missing:?} that are not registered \ + in this process; register the tools on the agent before resuming, or \ + call `allow_missing_resumed_tools()` to dispatch anyway and feed \ + not-found results to the model" + ) + .into(), + ))); + } + } + + self.snapshot = Some(snapshot.clone()); + Ok(TurnTools::from_parts(snapshot, names, output_tool_name)) + } + + /// Take a registry snapshot for a run resumed in this process. /// - /// Necessarily a **new** snapshot: implementations are live objects, so a - /// fresh process dispatches against its own registry state. The retrieval - /// query is re-derived from the run's history (matching preparation), and - /// the run's committed output tool — which is serialized with the run — - /// stays non-executable. - async fn resume_tools(&self) -> Result { + /// The retrieval query is re-derived from the run's history, matching what + /// request preparation would have used. + async fn fresh_snapshot(&self) -> Result { let query = self .run .full_history() .iter() .rev() .find_map(|message| message.rag_text()); - let snapshot = self - .agent + self.agent .tool_server_handle .snapshot_tool_defs(query) .await @@ -342,29 +427,23 @@ impl AgentDriver { PromptError::CompletionError(CompletionError::RequestError( "Failed to get tool definitions".into(), )) - })?; - let executable: BTreeSet = snapshot - .definitions() - .iter() - .map(|tool| tool.name.clone()) - .collect(); - let output_tool_name = self.run.output_tool_name().map(str::to_owned); - let mut allowed = allowed_tool_names_for_choice( - &executable, - self.agent.tool_choice.as_ref(), - output_tool_name.as_deref(), - None, - ) - .map_err(PromptError::CompletionError)?; - if let Some(name) = &output_tool_name { - allowed.insert(name.clone()); + }) + } + + /// The patch actually applied to the next turn's request. + /// + /// The run's own `tool_choice` is the driver's baseline — a run built with + /// [`AgentRun::with_tool_choice`] and handed to + /// [`Agent::drive_run`](super::Agent::drive_run) is taken as-is, so its + /// choice must reach the provider and not merely the run's internal + /// decisions. An explicit [`Self::request_patch`] outranks it, exactly as + /// a per-turn patch outranks the agent's baseline everywhere else. + fn effective_request_patch(&self) -> RequestPatch { + let mut patch = self.request_patch.clone(); + if patch.tool_choice.is_none() { + patch.tool_choice = self.run.tool_choice().cloned(); } - Ok(TurnTools { - snapshot: Arc::new(snapshot), - executable_tool_names: Arc::new(executable), - allowed_tool_names: Arc::new(allowed), - output_tool_name, - }) + patch } } @@ -768,4 +847,194 @@ mod tests { assert!(err.to_string().contains("Required")); assert_eq!(model.request_count(), 0); } + + /// A preparation failure must cost nothing. Preparation runs before the + /// run advances, so a turn that never reached the provider consumes no + /// budget and leaves the run drivable — the caller fixes the cause (here, + /// a tool the registry was missing; in practice a briefly unreachable tool + /// server) and retries the same step. + #[tokio::test] + async fn failed_preparation_leaves_the_run_intact_and_retryable() { + let model = MockCompletionModel::new([MockTurn::text("done")]); + let agent = AgentBuilder::new(model.clone()) + .default_max_turns(1) + .tool_choice(ToolChoice::Specific { + function_names: vec!["add".to_string()], + }) + .build(); + let mut driver = agent.drive("go"); + + driver + .next_step() + .await + .expect_err("a tool_choice naming an unregistered tool must fail at prepare time"); + assert_eq!( + driver.run().turn(), + 0, + "a request that never left the process must not consume a turn" + ); + assert_eq!(model.request_count(), 0); + + // Fix the cause and drive the very same step again. + agent.tool_server_handle.add_tool(MockAddTool).await; + let (request, tools, turn) = expect_send!(driver); + assert_eq!(turn, 1, "the retry takes the turn the failure did not"); + assert!(tools.executable_tool_names().contains("add")); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + } + + /// The natural suspension point for a caller that owns the transport is + /// *after* the request is sent and before the reply lands — a queued or + /// long-running provider call. The turn's advertised tool names travel + /// with the run, so the reply can be fed to a driver in another process. + #[tokio::test] + async fn run_suspended_awaiting_the_model_resumes_in_another_process() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("call_1", "add", json!({"x": 2, "y": 5})), + MockTurn::text("7"), + ]); + let agent = AgentBuilder::new(model.clone()) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + + let mut driver = agent.drive("what is 2 + 5?"); + let (request, _, _) = expect_send!(driver); + // Suspend with the model call in flight. + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + let response = request.send().await.expect("scripted turn"); + drop(driver); + drop(agent); + + // "Fresh process": rebuild the agent, deserialize, feed the reply. + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = agent.drive_run(run); + driver + .model_response(&response) + .expect("a run resumed mid-model-call accepts its reply"); + + let (calls, tools) = expect_execute_tools!(driver); + let mut context = ToolContext::new(); + let mut results = Vec::new(); + for call in &calls { + results.push(tools.execute_call(call, &mut context).await); + } + driver.tool_results(results).expect("results accepted"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + match driver.next_step().await.expect("next_step succeeds") { + DriveStep::Done(response) => assert_eq!(response.output, "7"), + other => panic!("expected Done, got {other:?}"), + } + } + + /// A custom run is taken as-is, so its own `tool_choice` must reach the + /// provider — not merely the run's internal decisions. Forbidding tools on + /// the run and having the model call one anyway is not a recoverable + /// situation; the request has to carry the constraint. + #[tokio::test] + async fn a_custom_runs_tool_choice_reaches_the_request() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive_run(AgentRun::new("go").with_tool_choice(ToolChoice::None)); + let (request, tools, _) = expect_send!(driver); + assert_eq!(request.build().tool_choice, Some(ToolChoice::None)); + assert!( + tools.allowed_tool_names().is_empty(), + "ToolChoice::None allows nothing to be called" + ); + } + + /// An explicit patch outranks the run's own choice, and reaches every + /// other per-turn field the runner's `CompletionCall` hooks reach. + #[tokio::test] + async fn a_request_patch_overrides_the_baseline_for_the_turn() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .preamble("baseline preamble") + .tool(MockAddTool) + .tool(MockSubtractTool) + .build(); + let mut driver = agent + .drive_run(AgentRun::new("go").with_tool_choice(ToolChoice::None)) + .request_patch( + RequestPatch::new() + .preamble("patched preamble") + .tool_choice(ToolChoice::Required) + .active_tools(["add"]), + ); + + let (request, tools, _) = expect_send!(driver); + assert!( + tools.executable_tool_names().contains("add") + && !tools.executable_tool_names().contains("subtract"), + "active_tools narrows the advertised set: {:?}", + tools.executable_tool_names() + ); + let request = request.build(); + assert_eq!(request.tool_choice, Some(ToolChoice::Required)); + assert!(matches!( + request.chat_history.first(), + Some(Message::System { content }) if content == "patched preamble" + )); + } + + /// Resuming is about dispatch, not about validating a request that will + /// never be built. A `tool_choice` the resuming process could not satisfy + /// must not pre-empt the drift report, and must not defeat the opt-out + /// that exists precisely for this situation. + #[tokio::test] + async fn resumed_dispatch_does_not_validate_an_unbuilt_requests_tool_choice() { + let model = MockCompletionModel::new([MockTurn::tool_call( + "call_1", + "add", + json!({"x": 2, "y": 5}), + )]); + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool_choice(ToolChoice::Required) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("what is 2 + 5?"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + + // Resume against a process whose registry lost the tool. `Required` + // is unsatisfiable there, but no request is being built. + let bare_agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(2) + .tool_choice(ToolChoice::Required) + .build(); + + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let err = bare_agent + .drive_run(run) + .next_step() + .await + .expect_err("the missing pending tool must surface"); + let message = err.to_string(); + assert!( + message.contains("add") && message.contains("allow_missing_resumed_tools"), + "drift must be reported, not the unsatisfiable tool choice: {message}" + ); + + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = bare_agent.drive_run(run).allow_missing_resumed_tools(); + let (calls, tools) = expect_execute_tools!(driver); + let mut context = ToolContext::new(); + let result = tools.execute_call(&calls[0], &mut context).await; + assert!( + matches!(result, UserContent::ToolResult(_)), + "the opt-out must dispatch and produce a tool result" + ); + } } diff --git a/crates/rig-agent/src/agent/prompt_request/streaming.rs b/crates/rig-agent/src/agent/prompt_request/streaming.rs index 490601d937..fb86b66916 100644 --- a/crates/rig-agent/src/agent/prompt_request/streaming.rs +++ b/crates/rig-agent/src/agent/prompt_request/streaming.rs @@ -4,7 +4,7 @@ use rig_core::{ }; use crate::{ - agent::completion::build_prepared_completion_request, + agent::completion::{TurnBaseline, TurnRequest, build_prepared_completion_request}, agent::hook::{ AgentHook, HookContext, HookStack, InvalidToolCallAction, ModelSelection, ModelSelectionAction, ModelTurnFinished, ReasoningDelta, StepEventKind, @@ -560,23 +560,27 @@ where // consistent even if the per-turn tool set changes (#1928). let committed_output_tool = run.output_tool_name().map(str::to_owned); let mut prepared = match build_prepared_completion_request( - &selected_model, - prompt.clone(), - &history, - runner.preamble.as_deref(), - &runner.static_context, - runner.temperature, - runner.max_tokens, - runner.additional_params.as_ref(), - runner.record_telemetry_content, - runner.tool_choice.as_ref(), - &runner.tool_server_handle, - runner.output_schema.as_ref(), - &runner.output_mode, - committed_output_tool.as_deref(), - runner.output_tool_description.as_deref(), - runner.augment_output_preamble, - request_patch.as_ref(), + TurnBaseline { + model: &selected_model, + preamble: runner.preamble.as_deref(), + static_context: &runner.static_context, + temperature: runner.temperature, + max_tokens: runner.max_tokens, + additional_params: runner.additional_params.as_ref(), + record_telemetry_content: runner.record_telemetry_content, + tool_choice: runner.tool_choice.as_ref(), + tool_server_handle: &runner.tool_server_handle, + output_schema: runner.output_schema.as_ref(), + output_mode: &runner.output_mode, + output_tool_description: runner.output_tool_description.as_deref(), + augment_output_preamble: runner.augment_output_preamble, + }, + TurnRequest { + prompt: prompt.clone(), + chat_history: &history, + committed_output_tool: committed_output_tool.as_deref(), + patch: request_patch.as_ref(), + }, ) .await { diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index aca8c883f2..fb406ebf68 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -15,12 +15,23 @@ //! //! Because the machine never awaits anything, it is runtime-agnostic and the //! whole run state is `Serialize + Deserialize`: a driver can serialize a run -//! between steps (for example while tool calls are pending), persist it, and -//! resume it later in another process. Note that serialized run state embeds -//! the full conversation accumulated so far — persisting it inherits whatever -//! sensitivity the conversation content has — and the serialization format -//! carries no cross-version stability guarantee yet: resume with the same rig -//! version that suspended the run. +//! between *any* two steps — while tool calls are pending, or while a model +//! call is in flight — persist it, and resume it later in another process. +//! Note that serialized run state embeds the full conversation accumulated so +//! far, so persisting it inherits whatever sensitivity the conversation +//! content has. +//! +//! # Serialization format +//! +//! Every payload carries a `$schemaVersion` tag ([`RUN_SCHEMA_VERSION`]), and +//! a build reads only the version it writes. Forward *and* backward +//! compatibility are deliberately fail-fast: a mismatched or absent tag is a +//! deserialization error, never a silent reinterpretation of fields whose +//! meaning moved. Version your agent definitions alongside suspended runs. +//! +//! | version | change | +//! | --- | --- | +//! | `1.0` | Initial versioned format. Records the turn's advertised tool names so a run suspended mid-model-call can be resumed. | //! //! `AgentRun` deliberately contains no model, tool registry, memory backend, or //! hook stack. Hand-driving it is a low-level provider integration: the caller @@ -87,6 +98,7 @@ use crate::{ assistant_text_from_choice, build_full_history, build_history_for_request, invalid_tool_retry_user_message, is_empty_assistant_turn, tool_result_message, }, + agent::turn_tools::TurnToolNames, completion::{Message, PromptError, Usage}, json_utils, }; @@ -115,6 +127,42 @@ fn unknown_tool_call_error( } } +/// The serialization format [`AgentRun`] writes, and the only one it reads. +/// +/// Every bump must add a line to the table in the [module docs](self). Forward +/// compatibility is deliberately fail-fast: a build refuses to deserialize a +/// run written by any other version rather than silently reinterpreting fields +/// whose meaning moved. +pub const RUN_SCHEMA_VERSION: &str = "1.0"; + +/// Fail-fast schema tag on the serialized [`AgentRun`]. +/// +/// Serializes as [`RUN_SCHEMA_VERSION`] and refuses to deserialize anything +/// else. The field carries no `serde(default)`, so a payload written before +/// versioning existed fails with a missing-field error instead of being +/// silently accepted. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct SchemaVersion; + +impl Serialize for SchemaVersion { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(RUN_SCHEMA_VERSION) + } +} + +impl<'de> Deserialize<'de> for SchemaVersion { + fn deserialize>(deserializer: D) -> Result { + let found = String::deserialize(deserializer)?; + if found == RUN_SCHEMA_VERSION { + return Ok(Self); + } + Err(serde::de::Error::custom(format!( + "serialized agent run has schema version `{found}`, but this build reads and writes \ + `{RUN_SCHEMA_VERSION}`; resume the run with the rig version that suspended it" + ))) + } +} + /// Default number of times Tool output mode re-prompts the model for valid /// structured output before finalizing best-effort (see #1928). Mirrors /// pydantic-ai's default output-retry budget of 1. @@ -292,6 +340,9 @@ enum RunState { /// driving protocol. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AgentRun { + /// Fail-fast format tag; see [`RUN_SCHEMA_VERSION`]. + #[serde(rename = "$schemaVersion")] + schema_version: SchemaVersion, max_turns: usize, max_invalid_tool_call_retries: usize, tool_choice: Option, @@ -327,14 +378,49 @@ pub struct AgentRun { /// [`AgentRunStep::CallModel`] is emitted. #[serde(default)] streamed_completion_call_recorded: bool, + /// Tool names advertised on the most recently committed model call. + /// + /// Recorded by [`AgentRun::commit_model_call`] and kept for the whole turn + /// — through validation, resolution and tool execution — so a driver can + /// rebuild the turn's [`TurnTools`](crate::agent::TurnTools) in *any* + /// state, including in a process that only just deserialized the run. + /// `None` for runs hand-driven through [`AgentRun::next_step`], where the + /// advertised names arrive with the [`ModelTurn`] instead. + #[serde(default)] + advertised_tools: Option, state: RunState, } +/// The inputs for a model call the run is ready to make, read without +/// advancing anything. See [`AgentRun::peek_model_call`]. +#[derive(Debug, Clone)] +pub(crate) struct ModelCallInputs { + pub(crate) prompt: Message, + pub(crate) history: Vec, +} + +/// What [`AgentRun::advance`] resolved to. +/// +/// Splits "the run wants another model call" from every step that needs no +/// model call, so a caller with fallible request preparation can do that work +/// *before* anything advances (see [`AgentRun::commit_model_call`]). +pub(crate) enum Advance { + /// The run is parked in `PreparingRequest`: read the inputs with + /// [`AgentRun::peek_model_call`], build the request, and commit it with + /// [`AgentRun::commit_model_call`]. Nothing has advanced yet. + NeedsModelCall, + /// Execute these tool calls, then feed [`AgentRun::tool_results`]. + CallTools(Vec), + /// The run is complete. + Done(PromptResponse), +} + impl AgentRun { /// Create a run for one prompt with no input history, a one-model-call /// budget, and no invalid tool-call retries. pub fn new(prompt: impl Into) -> Self { Self { + schema_version: SchemaVersion, max_turns: 1, max_invalid_tool_call_retries: 0, tool_choice: None, @@ -351,6 +437,7 @@ impl AgentRun { invalid_tool_call_retries: 0, rollback_pending: false, streamed_completion_call_recorded: false, + advertised_tools: None, state: RunState::PreparingRequest, } } @@ -425,12 +512,12 @@ impl AgentRun { /// Roll the run back to re-prompt for valid output (#1928). The caller must /// have already appended the assistant turn and the corrective feedback - /// message to the history. Consumes one output-retry, then emits the retry - /// [`AgentRunStep::CallModel`]. - fn reprompt_for_output(&mut self) -> Result { + /// message to the history. Consumes one output-retry and parks the run in + /// `PreparingRequest`; the caller loops back through [`Self::advance`], + /// which reports the pending model call without committing it. + fn reprompt_for_output(&mut self) { self.output_retries += 1; self.state = RunState::PreparingRequest; - self.next_step() } /// Set the retry budget for [`InvalidToolCallAction::Retry`] @@ -449,6 +536,28 @@ impl AgentRun { self } + /// The tool choice active for this run, if one was set. + /// + /// A driver preparing this run's requests must honor it: it is the run's + /// own policy, not the agent's baseline, and it governs both the request + /// sent to the provider and the run's internal decisions (see + /// [`Self::with_tool_choice`]). + pub fn tool_choice(&self) -> Option<&ToolChoice> { + self.tool_choice.as_ref() + } + + /// The tool names advertised on the current turn's model call, when the + /// call was committed with them (see [`Self::commit_model_call`]). + /// + /// This is the durable half of the turn's + /// [`TurnTools`](crate::agent::TurnTools): pairing it with a registry + /// snapshot reconstitutes the whole thing, which is what lets a driver + /// resume a run suspended at *any* step rather than only while tool calls + /// are pending. + pub(crate) fn advertised_tools(&self) -> Option<&TurnToolNames> { + self.advertised_tools.as_ref() + } + /// Set the synthetic output-tool name for Tool output mode (see #1928). /// When a model turn calls this tool, the run finalizes with the call's /// arguments (serialized JSON) as the response. @@ -629,230 +738,325 @@ impl AgentRun { }) } - /// Advance the machine and return the next action for the driver. + /// Whether the run is waiting for its next model call to be prepared. + pub(crate) fn is_preparing_request(&self) -> bool { + matches!(self.state, RunState::PreparingRequest) + } + + /// Read the inputs for the pending model call **without advancing**. + /// + /// Pure over committed state: it consumes no turn, moves no state, and + /// poisons nothing on failure. A caller whose request preparation can fail + /// — an agent driver reading a tool registry, say — peeks first, builds the + /// request, and only then calls [`Self::commit_model_call`]. A preparation + /// failure therefore leaves the run byte-identical to before the attempt, + /// so it can be retried in this process or serialized and retried in + /// another. /// /// # Errors - /// - [`PromptError::MaxTurnsError`] when the total model-call budget is exhausted. - /// - [`PromptError::PromptCancelled`] when the machine is driven out of - /// protocol (for example, calling this while a model response is - /// pending). - pub fn next_step(&mut self) -> Result { - match std::mem::replace(&mut self.state, RunState::Failed) { - RunState::PreparingRequest => { - let Some((prompt_ref, history_for_turn)) = self.new_messages.split_last() else { - return Err(PromptError::prompt_cancelled( - self.full_history(), - "prompt loop lost its pending prompt", - )); - }; - let prompt = prompt_ref.clone(); + /// - [`PromptError::MaxTurnsError`] when the total model-call budget is + /// exhausted. The run stays intact, so raising the budget and peeking + /// again resumes it. + /// - [`PromptError::PromptCancelled`] when there is no pending prompt. + pub(crate) fn peek_model_call(&self) -> Result { + debug_assert!( + self.is_preparing_request(), + "peek_model_call is only meaningful while the run is preparing a request" + ); + let Some((prompt, history_for_turn)) = self.new_messages.split_last() else { + return Err(PromptError::prompt_cancelled( + self.full_history(), + "prompt loop lost its pending prompt", + )); + }; - if self.current_turn >= self.max_turns { - return Err(PromptError::MaxTurnsError { - max_turns: self.max_turns, - chat_history: self.full_history().into(), - prompt: prompt.into(), - }); - } + if self.current_turn >= self.max_turns { + return Err(PromptError::MaxTurnsError { + max_turns: self.max_turns, + chat_history: self.full_history().into(), + prompt: prompt.clone().into(), + }); + } - let history = - build_history_for_request(self.chat_history.as_deref(), history_for_turn); - self.current_turn += 1; - self.rollback_pending = false; - self.streamed_completion_call_recorded = false; - self.state = RunState::AwaitingModel; - Ok(AgentRunStep::CallModel { - prompt, - history, - turn: self.current_turn, - }) + Ok(ModelCallInputs { + prompt: prompt.clone(), + history: build_history_for_request(self.chat_history.as_deref(), history_for_turn), + }) + } + + /// Commit the model call previewed by [`Self::peek_model_call`], returning + /// its one-based turn index. + /// + /// This is the only place a turn is consumed. It is deliberately + /// infallible and deliberately last: everything that can fail — reading the + /// tool registry, resolving the output mode, validating the tool choice — + /// has already happened by the time it runs, so the run never advances into + /// a call that was never made. + /// + /// `advertised` records the turn's tool names on the run (see + /// [`Self::advertised_tools`]); pass `None` when the caller supplies them + /// with the [`ModelTurn`] instead. `output_tool_name` fills the run's + /// committed Tool-mode name once (#1928), pinning the mode for the rest of + /// the run. + pub(crate) fn commit_model_call( + &mut self, + advertised: Option, + output_tool_name: Option, + ) -> usize { + debug_assert!( + self.is_preparing_request(), + "commit_model_call is only valid on a peeked model call" + ); + self.set_output_tool_name(output_tool_name); + self.advertised_tools = advertised; + self.current_turn += 1; + self.rollback_pending = false; + self.streamed_completion_call_recorded = false; + self.state = RunState::AwaitingModel; + self.current_turn + } + + /// Advance the machine as far as it can go without a model call. + /// + /// Returns [`Advance::NeedsModelCall`] the moment the run wants one — + /// including when an output-mode re-prompt rolls the run back mid-advance + /// (#1928) — leaving the run in `PreparingRequest` for the caller to + /// prepare and commit. [`Self::next_step`] is this plus an immediate + /// commit, for callers with nothing that can fail in between. + pub(crate) fn advance(&mut self) -> Result { + loop { + if self.is_preparing_request() { + return Ok(Advance::NeedsModelCall); } - RunState::AwaitingAdvance(turn_state) => { - let TurnState { - message_id, - items, - has_tool_calls, - skipped, - mut internal_call_ids, - } = *turn_state; - // Tool output mode (#1928): a call to the synthetic output tool - // finalizes the run with the call's arguments as the response, - // instead of executing it as a tool. First match wins; any - // sibling tool calls in the same turn are dropped. - if has_tool_calls - && let Some(output_tool_name) = self.output_tool_name.clone() - && let Some(tool_call) = items.iter().find_map(|item| match item { - AssistantContent::ToolCall(tc) if tc.function.name == output_tool_name => { - Some(tc) - } - _ => None, - }) - { - let output_tool_calls = items - .iter() - .filter(|item| { - matches!( - item, - AssistantContent::ToolCall(tc) - if tc.function.name == output_tool_name - ) + match std::mem::replace(&mut self.state, RunState::Failed) { + // Handled above; re-entering here would advance a run the + // caller has not prepared a request for. + RunState::PreparingRequest => { + self.state = RunState::PreparingRequest; + return Ok(Advance::NeedsModelCall); + } + RunState::AwaitingAdvance(turn_state) => { + let TurnState { + message_id, + items, + has_tool_calls, + skipped, + mut internal_call_ids, + } = *turn_state; + // Tool output mode (#1928): a call to the synthetic output tool + // finalizes the run with the call's arguments as the response, + // instead of executing it as a tool. First match wins; any + // sibling tool calls in the same turn are dropped. + if has_tool_calls + && let Some(output_tool_name) = self.output_tool_name.clone() + && let Some(tool_call) = items.iter().find_map(|item| match item { + AssistantContent::ToolCall(tc) + if tc.function.name == output_tool_name => + { + Some(tc) + } + _ => None, }) - .count(); - let args = tool_call.function.arguments.clone(); - let tool_call_id = tool_call.id.clone(); - let output = json_utils::serialize_json_value(&args); - - // Validate the output against the schema's required fields and - // re-prompt while budget remains, so a model that omits fields - // gets a chance to fix it before we finalize best-effort. - let missing = self.missing_required_output_fields(&args); - if !missing.is_empty() && self.can_reprompt_for_output() { + { + let output_tool_calls = items + .iter() + .filter(|item| { + matches!( + item, + AssistantContent::ToolCall(tc) + if tc.function.name == output_tool_name + ) + }) + .count(); + let args = tool_call.function.arguments.clone(); + let tool_call_id = tool_call.id.clone(); + let output = json_utils::serialize_json_value(&args); + + // Validate the output against the schema's required fields and + // re-prompt while budget remains, so a model that omits fields + // gets a chance to fix it before we finalize best-effort. + let missing = self.missing_required_output_fields(&args); + if !missing.is_empty() && self.can_reprompt_for_output() { + self.new_messages.push(Message::Assistant { + id: message_id, + content: items.clone(), + }); + let feedback = format!( + "The `{output_tool_name}` arguments were missing required field(s): \ + {}. Call `{output_tool_name}` again with every required field.", + missing.join(", ") + ); + if let Some(user_message) = + invalid_tool_retry_user_message(&items, &tool_call_id, feedback) + { + self.new_messages.push(user_message); + } + self.reprompt_for_output(); + continue; + } + + // Finalize. The turn is persisted as the assistant's final + // *text* (keeping any reasoning, dropping every tool call) + // rather than the raw output-tool call. Otherwise the saved + // history would carry an unanswered tool_use, which providers + // reject when the conversation is replayed on a later turn. + let mut final_items: Vec = items + .iter() + .filter(|item| !matches!(item, AssistantContent::ToolCall(_))) + .cloned() + .collect(); + final_items.push(AssistantContent::text(output.clone())); self.new_messages.push(Message::Assistant { id: message_id, - content: items.clone(), + content: final_items.clone(), }); - let feedback = format!( - "The `{output_tool_name}` arguments were missing required field(s): \ - {}. Call `{output_tool_name}` again with every required field.", - missing.join(", ") - ); - if let Some(user_message) = - invalid_tool_retry_user_message(&items, &tool_call_id, feedback) - { - self.new_messages.push(user_message); - } - return self.reprompt_for_output(); + + let response = PromptResponse::new(output, self.usage) + .with_messages(self.new_messages.clone()) + .with_completion_calls(self.completion_calls.clone()) + .with_output_tool_calls(output_tool_calls) + .with_content(final_items); + self.state = RunState::Done(Box::new(response.clone())); + return Ok(Advance::Done(response)); } - // Finalize. The turn is persisted as the assistant's final - // *text* (keeping any reasoning, dropping every tool call) - // rather than the raw output-tool call. Otherwise the saved - // history would carry an unanswered tool_use, which providers - // reject when the conversation is replayed on a later turn. - let mut final_items: Vec = items - .iter() - .filter(|item| !matches!(item, AssistantContent::ToolCall(_))) - .cloned() - .collect(); - final_items.push(AssistantContent::text(output.clone())); - self.new_messages.push(Message::Assistant { - id: message_id, - content: final_items.clone(), - }); + // An empty turn is not a lost turn. Cancelling here would fail + // runs that previously succeeded: with the fabricated empty-text + // padding gone, a textless turn — a tool-call-only turn whose + // calls were all dropped, a content-filtered turn, a truncated + // stream — arrives honestly empty. `is_empty_assistant_turn` + // does the right thing: keep the turn out of history and carry on. + if !is_empty_assistant_turn(&items) { + self.new_messages.push(Message::Assistant { + id: message_id, + content: items.clone(), + }); + } - let response = PromptResponse::new(output, self.usage) - .with_messages(self.new_messages.clone()) - .with_completion_calls(self.completion_calls.clone()) - .with_output_tool_calls(output_tool_calls) - .with_content(final_items); - self.state = RunState::Done(Box::new(response.clone())); - return Ok(AgentRunStep::Done(response)); - } + if has_tool_calls { + // The model is making progress with real tools, so reset the + // output-retry budget: it is per finalization attempt, not a + // single per-run allowance an early stray turn could burn + // before the model genuinely needs to produce output (#1928). + self.output_retries = 0; + let calls: Vec = items + .iter() + .enumerate() + .filter_map(|(index, item)| match item { + AssistantContent::ToolCall(tool_call) => { + // Consume pairs positionally so duplicate + // provider IDs within one turn stay + // distinguishable. + let internal_call_id = internal_call_ids + .iter() + .position(|(id, _)| tool_call.id == id.as_str()) + .map(|pair| internal_call_ids.remove(pair).1); + Some(PendingToolCall { + tool_call: tool_call.clone(), + preresolved_result: skipped.get(&index).cloned(), + internal_call_id, + }) + } + _ => None, + }) + .collect(); + self.state = RunState::ExecutingTools(calls.clone()); + return Ok(Advance::CallTools(calls)); + } else { + // Tool output mode (#1928): the model produced a final text + // answer without calling the output tool. Re-prompt while + // budget remains so it returns structured output; the + // assistant text was already appended above, so just add the + // corrective feedback. Empty turns finalize best-effort. + // + // But if the text already *is* valid output (parses as JSON + // with every required field), accept it rather than wasting a + // turn — the model answered correctly, just via the wrong + // channel. + if let Some(output_tool_name) = self.output_tool_name.clone() + && !is_empty_assistant_turn(&items) + && self.can_reprompt_for_output() + && !self + .text_satisfies_output_schema(&assistant_text_from_choice(&items)) + { + let feedback = format!( + "Provide your final answer by calling the `{output_tool_name}` tool \ + with the structured result as its arguments, not as plain text." + ); + self.new_messages.push(Message::user(feedback)); + self.reprompt_for_output(); + continue; + } - // An empty turn is not a lost turn. Cancelling here would fail - // runs that previously succeeded: with the fabricated empty-text - // padding gone, a textless turn — a tool-call-only turn whose - // calls were all dropped, a content-filtered turn, a truncated - // stream — arrives honestly empty. `is_empty_assistant_turn` - // does the right thing: keep the turn out of history and carry on. - if !is_empty_assistant_turn(&items) { - self.new_messages.push(Message::Assistant { - id: message_id, - content: items.clone(), - }); + let response = + PromptResponse::new(assistant_text_from_choice(&items), self.usage) + .with_messages(self.new_messages.clone()) + .with_completion_calls(self.completion_calls.clone()) + .with_content(items); + self.state = RunState::Done(Box::new(response.clone())); + return Ok(Advance::Done(response)); + } } - - if has_tool_calls { - // The model is making progress with real tools, so reset the - // output-retry budget: it is per finalization attempt, not a - // single per-run allowance an early stray turn could burn - // before the model genuinely needs to produce output (#1928). - self.output_retries = 0; - let calls: Vec = items - .iter() - .enumerate() - .filter_map(|(index, item)| match item { - AssistantContent::ToolCall(tool_call) => { - // Consume pairs positionally so duplicate - // provider IDs within one turn stay - // distinguishable. - let internal_call_id = internal_call_ids - .iter() - .position(|(id, _)| tool_call.id == id.as_str()) - .map(|pair| internal_call_ids.remove(pair).1); - Some(PendingToolCall { - tool_call: tool_call.clone(), - preresolved_result: skipped.get(&index).cloned(), - internal_call_id, - }) - } - _ => None, - }) - .collect(); + RunState::ExecutingTools(calls) => { + // Idempotent, like Done: a process resuming a serialized run + // re-obtains the pending tool calls from the state itself. self.state = RunState::ExecutingTools(calls.clone()); - Ok(AgentRunStep::CallTools { calls }) - } else { - // Tool output mode (#1928): the model produced a final text - // answer without calling the output tool. Re-prompt while - // budget remains so it returns structured output; the - // assistant text was already appended above, so just add the - // corrective feedback. Empty turns finalize best-effort. - // - // But if the text already *is* valid output (parses as JSON - // with every required field), accept it rather than wasting a - // turn — the model answered correctly, just via the wrong - // channel. - if let Some(output_tool_name) = self.output_tool_name.clone() - && !is_empty_assistant_turn(&items) - && self.can_reprompt_for_output() - && !self.text_satisfies_output_schema(&assistant_text_from_choice(&items)) - { - let feedback = format!( - "Provide your final answer by calling the `{output_tool_name}` tool \ - with the structured result as its arguments, not as plain text." - ); - self.new_messages.push(Message::user(feedback)); - return self.reprompt_for_output(); - } - - let response = - PromptResponse::new(assistant_text_from_choice(&items), self.usage) - .with_messages(self.new_messages.clone()) - .with_completion_calls(self.completion_calls.clone()) - .with_content(items); - self.state = RunState::Done(Box::new(response.clone())); - Ok(AgentRunStep::Done(response)) + return Ok(Advance::CallTools(calls)); + } + RunState::Done(response) => { + let step = Advance::Done((*response).clone()); + self.state = RunState::Done(response); + return Ok(step); + } + state @ (RunState::AwaitingModel | RunState::ResolvingToolCalls(_)) => { + let reason = match &state { + RunState::AwaitingModel => { + "next_step called while a model response is pending; feed it via model_response first" + } + _ => { + "next_step called while an invalid tool-call resolution is pending; answer it via resolve_invalid_tool_call first" + } + }; + self.state = state; + return Err(self.protocol_violation(reason)); + } + RunState::Failed => { + return Err(self.protocol_violation( + "next_step called after the run already failed or was misdriven", + )); } } - RunState::ExecutingTools(calls) => { - // Idempotent, like Done: a process resuming a serialized run - // re-obtains the pending tool calls from the state itself. - let step = AgentRunStep::CallTools { - calls: calls.clone(), - }; - self.state = RunState::ExecutingTools(calls); - Ok(step) - } - RunState::Done(response) => { - let step = AgentRunStep::Done((*response).clone()); - self.state = RunState::Done(response); - Ok(step) - } - state @ (RunState::AwaitingModel | RunState::ResolvingToolCalls(_)) => { - let reason = match &state { - RunState::AwaitingModel => { - "next_step called while a model response is pending; feed it via model_response first" - } - _ => { - "next_step called while an invalid tool-call resolution is pending; answer it via resolve_invalid_tool_call first" - } - }; - self.state = state; - Err(self.protocol_violation(reason)) + } + } + + /// Advance the machine and return the next action for the driver. + /// + /// [`Self::advance`] plus an immediate commit of any model call it asks + /// for — the right entry point for a driver that builds its requests from + /// the returned inputs with nothing fallible in between. A driver whose + /// request preparation *can* fail (see [`Self::peek_model_call`]) should + /// use the two halves instead, so a failure does not consume a turn. + /// + /// # Errors + /// - [`PromptError::MaxTurnsError`] when the total model-call budget is exhausted. + /// - [`PromptError::PromptCancelled`] when the machine is driven out of + /// protocol (for example, calling this while a model response is + /// pending). + pub fn next_step(&mut self) -> Result { + match self.advance()? { + Advance::NeedsModelCall => { + let ModelCallInputs { + prompt, history, .. + } = self.peek_model_call()?; + let turn = self.commit_model_call(None, None); + Ok(AgentRunStep::CallModel { + prompt, + history, + turn, + }) } - RunState::Failed => Err(self.protocol_violation( - "next_step called after the run already failed or was misdriven", - )), + Advance::CallTools(calls) => Ok(AgentRunStep::CallTools { calls }), + Advance::Done(response) => Ok(AgentRunStep::Done(response)), } } @@ -2311,19 +2515,23 @@ mod tests { assert!(matches!(err, PromptError::PromptCancelled { .. })); } - #[test] - fn agent_run_deserializes_pre_monoid_suspended_state() { - // Pins `CompletionCall.usage`'s null tolerance on a suspended run: - // `"usage": null` (the pre-monoid Option encoding) must map to - // zero-valued usage and the run must resume. The tool calls use the - // current schema — the pre-provider-split `call_id` lift is gone - // (its ignore-the-key behavior is pinned in rig-core's message - // tests). - let fixture = r#"{"max_turns":2,"max_invalid_tool_call_retries":0,"tool_choice":null,"chat_history":null,"new_messages":[{"role":"user","content":[{"type":"text","text":"add things"}]},{"role":"assistant","id":null,"content":[{"type":"toolcall","id":"call_1","function":{"name":"add","arguments":{"x":1}},"signature":null,"additional_params":null}]}],"current_turn":1,"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15,"cached_input_tokens":0,"cache_creation_input_tokens":0,"tool_use_prompt_tokens":0,"reasoning_tokens":0},"completion_calls":[{"call_index":0,"usage":null}],"completion_call_index":1,"invalid_tool_call_retries":0,"rollback_pending":false,"streamed_completion_call_recorded":false,"state":{"ExecutingTools":[{"tool_call":{"id":"call_1","function":{"name":"add","arguments":{"x":1}},"signature":null,"additional_params":null},"preresolved_result":null,"internal_call_id":null}]}}"#; + /// A suspended run in the current serialization: tagged assistant content, + /// and `"usage": null` on the completion call pinning `CompletionCall`'s + /// null tolerance (the pre-monoid `Option` encoding must still map to + /// zero-valued usage). Fields added since are defaulted in; only the + /// version tag gates the load. + const SUSPENDED_RUN_FIXTURE: &str = r#"{"$schemaVersion":"1.0","max_turns":2,"max_invalid_tool_call_retries":0,"tool_choice":null,"chat_history":null,"new_messages":[{"role":"user","content":[{"type":"text","text":"add things"}]},{"role":"assistant","id":null,"content":[{"type":"toolcall","id":"call_1","function":{"name":"add","arguments":{"x":1}},"signature":null,"additional_params":null}]}],"current_turn":1,"usage":{"input_tokens":10,"output_tokens":5,"total_tokens":15,"cached_input_tokens":0,"cache_creation_input_tokens":0,"tool_use_prompt_tokens":0,"reasoning_tokens":0},"completion_calls":[{"call_index":0,"usage":null}],"completion_call_index":1,"invalid_tool_call_retries":0,"rollback_pending":false,"streamed_completion_call_recorded":false,"state":{"ExecutingTools":[{"tool_call":{"id":"call_1","function":{"name":"add","arguments":{"x":1}},"signature":null,"additional_params":null},"preresolved_result":null,"internal_call_id":null}]}}"#; - let mut restored: AgentRun = - serde_json::from_str(fixture).expect("old-format suspended run should deserialize"); + #[test] + fn agent_run_deserializes_suspended_state_with_defaulted_fields() { + let mut restored: AgentRun = serde_json::from_str(SUSPENDED_RUN_FIXTURE) + .expect("in-version suspended run should deserialize"); assert_eq!(restored.completion_calls()[0].usage, Usage::new()); + assert_eq!( + restored.advertised_tools(), + None, + "a run suspended by a raw hand-driver records no advertised tools" + ); let calls = expect_call_tools(&mut restored); assert_eq!(calls.len(), 1); @@ -2333,6 +2541,41 @@ mod tests { expect_call_model(&mut restored); } + /// A prose warning cannot fail a load; the version tag can. A payload + /// written by another format version is rejected outright rather than + /// silently reinterpreted — including one written before versioning + /// existed, which has no tag at all. + #[test] + fn agent_run_rejects_foreign_schema_versions() { + let untagged = SUSPENDED_RUN_FIXTURE.replace(r#""$schemaVersion":"1.0","#, ""); + let err = serde_json::from_str::(&untagged) + .expect_err("an untagged run must not deserialize"); + assert!( + err.to_string().contains("$schemaVersion"), + "error should name the missing tag, got: {err}" + ); + + let future = + SUSPENDED_RUN_FIXTURE.replace(r#""$schemaVersion":"1.0""#, r#""$schemaVersion":"2.0""#); + let err = serde_json::from_str::(&future) + .expect_err("a run from another version must not deserialize"); + let message = err.to_string(); + assert!( + message.contains("2.0") && message.contains(RUN_SCHEMA_VERSION), + "error should name both versions, got: {message}" + ); + } + + #[test] + fn serialized_run_carries_the_schema_version() { + let run = AgentRun::new("add things"); + let json: serde_json::Value = serde_json::to_value(&run).expect("run should serialize"); + assert_eq!( + json.get("$schemaVersion").and_then(|v| v.as_str()), + Some(RUN_SCHEMA_VERSION) + ); + } + #[test] fn serde_round_trip_at_exhausted_budget_preserves_boundary() { let mut run = AgentRun::new("add things").max_turns(1); diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs index e2774f2d28..a05e124329 100644 --- a/crates/rig-agent/src/agent/turn_tools.rs +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -6,10 +6,17 @@ //! together at one instant: the executable and allowed tool-name sets, the //! synthetic output-tool name, and dispatch pinned to the registry snapshot //! whose definitions the provider received. +//! +//! The type splits along the durability line: [`TurnToolNames`] is data and is +//! recorded on the [`AgentRun`](super::run::AgentRun), while the dispatch +//! target is a live object rebuilt from the registry when a run resumes +//! elsewhere. [`TurnTools`] is the two halves paired back together. use std::collections::BTreeSet; use std::sync::Arc; +use serde::{Deserialize, Serialize}; + use rig_core::message::UserContent; use super::model::ModelHandle; @@ -29,6 +36,42 @@ pub(crate) struct PreparedCompletionRequest { pub(crate) tools: TurnTools, } +/// The serializable half of [`TurnTools`]: the tool names a model call +/// advertised to the provider. +/// +/// Names are data, so they travel with the run — [`AgentRun`](super::run::AgentRun) +/// records them when a model call is committed, and they survive +/// serialization. Implementations are live objects and cannot, so the +/// dispatch target ([`ToolRegistrySnapshot`]) is rebuilt from the agent's +/// registry when a run resumes in another process. Pairing the two back +/// together yields a [`TurnTools`] in any run state, in any process, which is +/// what makes a suspended run resumable at *every* step boundary rather than +/// only while tool calls are pending. +#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct TurnToolNames { + /// The real registry tools advertised to the provider this turn. + pub executable: BTreeSet, + /// The tools the active `tool_choice` let the model call this turn, + /// including the synthetic output tool when Tool output mode is active. + pub allowed: BTreeSet, +} + +impl TurnToolNames { + /// Assemble the [`ModelTurn`] for a completion response received on the + /// turn these names were advertised for. The single construction site for + /// driver-facing turns, so the two name sets can never be transposed by a + /// caller. + pub(crate) fn model_turn(&self, response: &CompletionResponse) -> ModelTurn { + ModelTurn::new( + response.message_id.clone(), + response.choice.clone(), + response.usage, + self.executable.clone(), + self.allowed.clone(), + ) + } +} + /// One turn's advertised tool sets and their dispatch target. /// /// All four pieces were resolved together when the turn was prepared: @@ -74,6 +117,33 @@ impl std::fmt::Debug for TurnTools { } impl TurnTools { + /// Pair advertised names back with a dispatch target. + /// + /// The two halves come from different places once a run is resumed: the + /// names from the serialized [`AgentRun`](super::run::AgentRun), the + /// snapshot from the resuming process's registry. + pub(crate) fn from_parts( + snapshot: Arc, + names: TurnToolNames, + output_tool_name: Option, + ) -> Self { + Self { + snapshot, + executable_tool_names: Arc::new(names.executable), + allowed_tool_names: Arc::new(names.allowed), + output_tool_name, + } + } + + /// The advertised names, detached from the dispatch target so they can be + /// recorded on the run. + pub(crate) fn names(&self) -> TurnToolNames { + TurnToolNames { + executable: (*self.executable_tool_names).clone(), + allowed: (*self.allowed_tool_names).clone(), + } + } + /// The executable registry tools advertised to the provider this turn. pub fn executable_tool_names(&self) -> &BTreeSet { &self.executable_tool_names @@ -160,16 +230,10 @@ impl TurnTools { } /// Assemble the [`ModelTurn`] for a completion response received on this - /// prepared turn. The single construction site for driver-facing turns, - /// so the two name sets can never be transposed by a caller. + /// prepared turn. Delegates to [`TurnToolNames::model_turn`], the single + /// construction site for driver-facing turns. pub(crate) fn model_turn(&self, response: &CompletionResponse) -> ModelTurn { - ModelTurn::new( - response.message_id.clone(), - response.choice.clone(), - response.usage, - (*self.executable_tool_names).clone(), - (*self.allowed_tool_names).clone(), - ) + self.names().model_turn(response) } } From c97fc3ae47896e327da3b564fde09d346fa650ed Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 06:44:16 -0700 Subject: [PATCH 05/21] fix(agent): close the send-failure gap the resume-point claim left open Round-2 review of PR #2278 found seven defects, six of them in the round-2 redesign itself. The headline one: `7dbc0d66` documented every step of a driven run as a resume point, but only the rarer of a turn's two failures was actually recoverable. Preparation could fail without cost, because nothing advanced until a request existed. The send could not. `commit_model_call` moved the run to AwaitingModel before handing the caller the request; CompletionRequestBuilder is not Clone and send(self) consumes it; retry_model_turn accepts only AwaitingAdvance; and next_step in AwaitingModel is a protocol violation. A transient transport error - or a queued provider job that vanished - left the turn consumed and the run with no public transition out of AwaitingModel. The claim was false for the failure that actually happens. rollback_model_call adds the missing transition, on AgentRun and on AgentDriver. It refunds the turn, drops the advertised names, clears the driver's snapshot cache, and returns the run to PreparingRequest, so the next step prepares the request *again* from current configuration. Re-deriving rather than replaying is the point: a cached request carries its attempt's tool snapshot, and replaying it would advertise one set of implementations while a later turn dispatches another - the skew TurnTools exists to prevent. Tokens the provider already billed stay billed, and a streamed turn that learns its usage only after the failure can still record it, reusing the window invalid tool-call recovery already opens. This is at-least-once and says so: nothing in the run can distinguish "never arrived" from "arrived, reply lost", so bounding attempts stays with the caller who owns the transport. CompletionError::is_retryable classifies whether an attempt is worth making at all - transport failures with no status, and preserved provider statuses of 408, 409, 429 or 5xx - closing the asymmetry with tool errors, which rig has classified for some time. Two narrowings on the resume path. TurnTools::execute now refuses any name the turn did not advertise, instead of trusting its snapshot to be narrowed: in-process the two agree by construction, but a resumed turn pairs names carried on the run with a snapshot rebuilt locally, so a tool registered after suspension could be dispatched through a turn that never advertised it. And a resumed snapshot resolves the advertised names explicitly (snapshot_tool_defs_including) rather than re-running retrieval for them, so a registered *dynamic* tool the new query does not rank is no longer reported as unregistered - advice that could not be followed - nor fed to the model as not-found. The peek/commit halves are now public - peek_model_call, commit_model_call, advance, is_preparing_request, with Advance and ModelCallInputs - so any hand-driver whose preparation can fail gets the same guarantee rig's own driver has. Their preconditions become PromptError protocol violations: a debug_assert is an acceptable internal contract and an unacceptable public one. TurnToolNames is exported for real; the changelog announced it as public while the module was private and only TurnTools re-exported. Four rustdoc warnings introduced by 7dbc0d66 are gone, and the count is now a gate: 0 in rig-agent and rig-core. 510 lib tests (+15), 1239 rig-core tests, clippy -D warnings clean, fmt clean. The two dispatch fixes were each confirmed to fail before their fix rather than passing on the other's narrowing. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- MIGRATING.md | 17 ++ crates/rig-agent/CHANGELOG.md | 6 +- crates/rig-agent/src/agent/driver.rs | 345 +++++++++++++++++++++- crates/rig-agent/src/agent/mod.rs | 7 +- crates/rig-agent/src/agent/run/mod.rs | 333 +++++++++++++++++++-- crates/rig-agent/src/agent/turn_tools.rs | 79 ++++- crates/rig-agent/src/tool/server.rs | 33 +++ crates/rig-core/CHANGELOG.md | 1 + crates/rig-core/src/completion/request.rs | 108 +++++++ 9 files changed, 884 insertions(+), 45 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index 1232177fa4..fbbe89ec66 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1589,6 +1589,23 @@ in-flight suspended runs before upgrading, or discard them. In exchange, a run can now be suspended and resumed at *every* step boundary — including with a model call in flight — not only while tool calls are pending. +Because the caller owns the send, the caller also owns recovering from a failed +one: `AgentDriver::rollback_model_call()` hands a turn back when its request +could not be sent or its reply is known to be lost, refunding the turn and +returning the run to preparing so the next step yields a freshly prepared +request. It is at-least-once — a request that reached the provider and lost only +its reply will be billed twice on retry — so pair it with +`CompletionError::is_retryable()` and bound the attempts yourself. Runs driven +by `Agent::runner` are unaffected; the runner still fails the prompt. + +Two behavior changes on the driver's resume path, both narrowing what dispatch +will do: `TurnTools::execute` now refuses any name the turn did not advertise +(previously a tool registered *after* a run was suspended could be dispatched +through a turn that never advertised it), and a resumed snapshot resolves the +advertised names explicitly instead of re-running retrieval for them, so a +registered dynamic tool the new query does not rank is no longer reported as +unregistered. + An `Agent`'s default model is set at construction. Per-run overrides now go through `runner(...).using_model(...)`, `Agent::set_model`, or a `ModelSelection` hook (see the "runtime model swapping" section for the diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 2bede78126..79a4567b9d 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -36,8 +36,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `Agent::prepare_turn`, returning the new public `PreparedTurn`/`TurnTools`: the agent's baseline configuration resolved into one turn's completion request plus the turn's executable and allowed tool-name sets, the synthetic output-tool name, and tool dispatch pinned to the turn's registry snapshot — so hand-driven `AgentRun` loops (custom provider transport, durable suspend/resume) reuse the configured `Agent` instead of restating its preamble and tools. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. A prepared turn is a configuration read: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` - *(agent)* add `Agent::drive` / `Agent::drive_run`, returning the new public `AgentDriver` (with `DriveStep` and `TurnTools`): hand-drive the sans-IO `AgentRun` machine with the agent's own configuration while owning every side effect. The driver seeds the run from the agent (`default_max_turns`, `tool_choice`, output schema), yields each turn's fully configured completion request for the caller to send (or hand to a custom transport), pairs every pending tool batch with the exact registry snapshot that advertised it, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. Resuming a serialized run in a fresh process re-derives a fresh dispatch snapshot from the rebuilt agent, and surfaces missing pending tools as an error (opt out with `allow_missing_resumed_tools`). The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` -- *(agent)* `AgentDriver` is resumable at **every** step boundary, not just while tool calls are pending. The turn's advertised tool names are recorded on the run (new public `TurnToolNames`, the serializable half of `TurnTools`), so a run serialized after `DriveStep::SendRequest` — the natural suspension point for a queued or long-running provider call — can be resumed in another process and fed its reply, validated against the set the request actually carried rather than against whatever the resuming registry holds. The driver now keeps no state of its own beyond a registry-snapshot cache +- *(agent)* `AgentDriver` is resumable at **every** step boundary, not just while tool calls are pending. The turn's advertised tool names are recorded on the run (new public `TurnToolNames`, the serializable half of `TurnTools`, reachable via `AgentRun::advertised_tools()`), so a run serialized after `DriveStep::SendRequest` — the natural suspension point for a queued or long-running provider call — can be resumed in another process and fed its reply, validated against the set the request actually carried rather than against whatever the resuming registry holds. The driver now keeps no state of its own beyond a registry-snapshot cache - *(agent)* a failed model-turn preparation costs nothing: the run advances only once a request exists, so an unreachable tool server or an impossible `tool_choice` leaves the run byte-identical — same state, same turn budget — and the step can be retried in place. Previously the turn was consumed and the run was left unresumable +- *(agent)* add `AgentRun::rollback_model_call` / `AgentDriver::rollback_model_call` for the other half of that story: a request that could not be sent, or whose reply is known to be lost, hands its turn back and returns the run to preparing, so the next step yields a **freshly prepared** request rather than replaying one whose tool snapshot has gone stale. Refunds the turn, preserves usage the provider already billed (including a streamed turn that reports usage after the failure), and counts attempts via `AgentRun::model_call_rollbacks`. Deliberately at-least-once: only the caller can tell "never arrived" from "arrived, reply lost", so bounding attempts is the caller's job +- *(agent)* `AgentRun`'s peek/commit halves are public — `peek_model_call`, `commit_model_call`, `advance`, `is_preparing_request`, with `Advance` and `ModelCallInputs` — so any hand-driver whose request preparation can fail gets the same "a failure costs no turn" guarantee the driver has, not just rig's own. Their preconditions are `PromptError` protocol violations rather than debug assertions +- *(agent)* [**behavior**] `TurnTools::execute` now refuses any name the turn did not advertise, instead of trusting its snapshot to be narrowed. The two agreed by construction in-process; on a resumed turn the names come from the run and the snapshot is rebuilt locally, so a tool registered after suspension could previously be dispatched through a turn that never advertised it +- *(agent)* [**behavior**] a resumed run's dispatch snapshot resolves the turn's advertised names explicitly rather than relying on retrieval to rank them again, so a registered *dynamic* tool that the re-derived query does not return is no longer misreported as "not registered in this process" — advice that could not be followed — and no longer feeds the model a not-found result for a tool that exists - *(agent)* add `AgentDriver::request_patch` / `set_request_patch`, giving a hand-driven run the per-turn configuration the runner gets from its `CompletionCall` hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history. Because the caller owns the loop, per-turn variation needs no callback - *(agent)* add `AgentRun::tool_choice()`. A custom run's own `tool_choice` now reaches the provider when the run is hand-driven, instead of only governing the run's internal decisions while the request silently carried the agent's baseline - *(tool)* add `ToolErrorKind::NotExecutable` (`ToolExecutionError::not_executable`): the tool is advertised to the model but not executable by the dispatcher — produced when dispatching the synthetic structured-output tool, whose call carries the final structured answer diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 28bf0cf47e..466bb9498a 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -40,10 +40,21 @@ //! set the request actually carried rather than against whatever its registry //! holds now. //! -//! Preparation failures are equally survivable. Nothing advances until a -//! request exists, so a turn that fails to prepare — an unreachable tool -//! server, an impossible `tool_choice` — costs no turn from the budget and -//! leaves the run byte-identical, ready to retry. +//! A turn can fail in two places, and both are recoverable: +//! +//! - **Preparing the request.** Nothing advances until a request exists, so an +//! unreachable tool server or an impossible `tool_choice` costs no turn from +//! the budget and leaves the run byte-identical. Call [`AgentDriver::next_step`] +//! again once the cause is fixed. +//! - **Sending it.** The caller owns the send, so the caller is the only party +//! that learns it failed. [`AgentDriver::rollback_model_call`] hands the turn +//! back — refunding it and returning the run to preparing — and the next +//! `next_step` yields a freshly prepared request. This is *at-least-once*: if +//! the request reached the provider and only its reply was lost, retrying +//! bills a second completion, and only the caller can tell those apart. Use +//! [`CompletionError::is_retryable`](crate::completion::CompletionError::is_retryable) +//! to decide, and bound the attempts yourself — the driver runs no IO and +//! owns no clock. //! //! Tool *implementations* are live objects and cannot be serialized: the //! resuming process rebuilds the same `Agent` and the driver takes a fresh @@ -56,6 +67,7 @@ //! definitions alongside the serialized run; the run's own format is versioned //! by [`RUN_SCHEMA_VERSION`](super::run::RUN_SCHEMA_VERSION). +use std::collections::BTreeSet; use std::sync::Arc; use rig_core::message::UserContent; @@ -109,8 +121,9 @@ impl Agent { /// What the caller must do next to advance an [`AgentDriver`]. /// -/// Deliberately exhaustive, like [`AgentRunStep`]: a driver loop must handle -/// every step, so adding a variant is a breaking change by design. +/// Deliberately exhaustive, like [`AgentRunStep`](super::run::AgentRunStep): a +/// driver loop must handle every step, so adding a variant is a breaking change +/// by design. pub enum DriveStep { /// Send this completion request to the model — via /// [`send`](CompletionRequestBuilder::send), or @@ -292,7 +305,7 @@ impl AgentDriver { self.snapshot = Some(tools.snapshot.clone()); let turn = self .run - .commit_model_call(Some(tools.names()), tools.output_tool_name.clone()); + .commit_model_call(Some(tools.names()), tools.output_tool_name.clone())?; Ok(DriveStep::SendRequest { request: Box::new(builder), tools, @@ -330,6 +343,45 @@ impl AgentDriver { self.run.model_response(names.model_turn(response)) } + /// Hand back a model call that never produced a response, so the turn can + /// be prepared and sent again. + /// + /// [`DriveStep::SendRequest`] gives the caller a request and the caller + /// owns the send, so the caller is also the only party that learns the + /// send failed. This is how that news gets back into the run: the turn is + /// refunded and the run returns to preparing, so the next + /// [`Self::next_step`] yields a **freshly prepared** `SendRequest` — new + /// registry snapshot, new patch — rather than a replay of a request whose + /// tool snapshot has since gone stale. + /// + /// ```rust,no_run + /// # use rig_agent::agent::{AgentDriver, DriveStep}; + /// # async fn example(driver: &mut AgentDriver) -> Result<(), Box> { + /// if let DriveStep::SendRequest { request, .. } = driver.next_step().await? { + /// match request.send().await { + /// Ok(response) => { driver.model_response(&response)?; } + /// // Nothing was produced: give the turn back and try again. + /// Err(err) if err.is_retryable() => driver.rollback_model_call()?, + /// Err(err) => return Err(err.into()), + /// } + /// } + /// # Ok(()) + /// # } + /// ``` + /// + /// See [`AgentRun::rollback_model_call`] for the full semantics — in + /// particular that this is **at-least-once**: if the request reached the + /// provider and only its reply was lost, retrying bills a second + /// completion. The driver cannot tell the two apart. Bounding attempts is + /// yours to do; the driver runs no IO and owns no clock. + pub fn rollback_model_call(&mut self) -> Result<(), PromptError> { + self.run.rollback_model_call()?; + // Drop the cached dispatch target too: the retry is a new turn and + // must advertise, and dispatch through, a snapshot taken for it. + self.snapshot = None; + Ok(()) + } + /// Resolve a pending invalid tool call, exactly as /// [`AgentRun::resolve_invalid_tool_call`]. pub fn resolve_invalid_tool_call( @@ -377,7 +429,15 @@ impl AgentDriver { // rather than silently feeding not-found results to the model. The // model chose these tools from a registry this process no longer has, // and re-prompting cannot fix that. - let snapshot = Arc::new(self.fresh_snapshot().await?); + let mut snapshot = self.fresh_snapshot(&names.executable).await?; + // Narrow to what the turn advertised, mirroring what preparation does + // in-process. Without this the resumed snapshot is the whole current + // registry, and `TurnTools`' two halves — advertised names and + // dispatch target — could disagree, which is the skew the type exists + // to prevent. + snapshot.retain_names(&names.executable); + let snapshot = Arc::new(snapshot); + if !self.allow_missing_resumed_tools { let missing: Vec<&str> = calls .iter() @@ -394,10 +454,10 @@ impl AgentDriver { if !missing.is_empty() { return Err(PromptError::CompletionError(CompletionError::RequestError( format!( - "resumed run has pending tool calls {missing:?} that are not registered \ - in this process; register the tools on the agent before resuming, or \ - call `allow_missing_resumed_tools()` to dispatch anyway and feed \ - not-found results to the model" + "resumed run has pending tool calls {missing:?} that are no longer \ + registered in this process; register the tools on the agent before \ + resuming, or call `allow_missing_resumed_tools()` to dispatch anyway \ + and feed not-found results to the model" ) .into(), ))); @@ -411,8 +471,16 @@ impl AgentDriver { /// Take a registry snapshot for a run resumed in this process. /// /// The retrieval query is re-derived from the run's history, matching what - /// request preparation would have used. - async fn fresh_snapshot(&self) -> Result { + /// request preparation would have used — but retrieval alone is not enough + /// here. It selects dynamic tools by similarity, so a registered dynamic + /// tool this query does not rank would be absent, and the caller's drift + /// check could not tell that apart from a tool that really was + /// deregistered. `required` — the names the turn advertised — is resolved + /// from the registry regardless of ranking, so absence means absence. + async fn fresh_snapshot( + &self, + required: &BTreeSet, + ) -> Result { let query = self .run .full_history() @@ -421,7 +489,7 @@ impl AgentDriver { .find_map(|message| message.rag_text()); self.agent .tool_server_handle - .snapshot_tool_defs(query) + .snapshot_tool_defs_including(query, required) .await .map_err(|_| { PromptError::CompletionError(CompletionError::RequestError( @@ -985,8 +1053,253 @@ mod tests { )); } + /// The caller owns the send, so the caller is the only party that learns + /// it failed. Handing the turn back refunds it and returns the run to + /// preparing, so the run survives a failure the driver never sees. + #[tokio::test] + async fn failed_send_rolls_back_and_re_prepares() { + let model = MockCompletionModel::new([MockTurn::text("done")]); + let agent = AgentBuilder::new(model.clone()) + .default_max_turns(1) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("go"); + + let (request, _, turn) = expect_send!(driver); + assert_eq!(turn, 1); + // The send fails: drop the request without sending it. The provider + // never saw anything, so nothing was produced. + drop(request); + driver + .rollback_model_call() + .expect("a call that produced nothing can be handed back"); + assert_eq!(driver.run().turn(), 0); + assert_eq!(driver.run().model_call_rollbacks(), 1); + assert_eq!(model.request_count(), 0); + + // The budget of one is intact, so the retry can happen at all. + let (request, _, turn) = expect_send!(driver); + assert_eq!(turn, 1, "the retry takes the turn the failure did not"); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + match driver.next_step().await.expect("next_step succeeds") { + DriveStep::Done(response) => assert_eq!(response.output, "done"), + other => panic!("expected Done, got {other:?}"), + } + } + + /// The retry re-derives the request rather than replaying one. Replaying + /// would pin the failed attempt's tool snapshot, advertising one set of + /// implementations while a later turn dispatches another. + #[tokio::test] + async fn rollback_re_derives_against_the_current_registry() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("go"); + + let (_, tools, _) = expect_send!(driver); + assert!(!tools.executable_tool_names().contains("subtract")); + driver.rollback_model_call().expect("rollback succeeds"); + + agent.tool_server_handle.add_tool(MockSubtractTool).await; + let (_, tools, _) = expect_send!(driver); + assert!( + tools.executable_tool_names().contains("subtract"), + "the retry must advertise the registry as it is now: {:?}", + tools.executable_tool_names() + ); + } + + /// A run suspended mid-send is resumable, and the rollback travels with + /// it: the process that discovers the reply is lost need not be the one + /// that sent the request. + #[tokio::test] + async fn rollback_after_resume_recovers_a_lost_reply() { + let model = MockCompletionModel::new([MockTurn::text("done")]); + let agent = AgentBuilder::new(model.clone()) + .default_max_turns(1) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("go"); + let (request, _, _) = expect_send!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + drop(request); + drop(driver); + + // "Fresh process": the reply never arrived, so hand the turn back. + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = agent.drive_run(run); + driver.rollback_model_call().expect("rollback succeeds"); + assert_eq!(driver.run().turn(), 0); + + let (request, _, turn) = expect_send!(driver); + assert_eq!(turn, 1); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + } + + /// `TurnTools` promises that a name the turn did not advertise cannot + /// reach the live registry. In-process the snapshot enforces that; on a + /// resumed turn only the advertised names can, since the snapshot is + /// rebuilt here. + #[tokio::test] + async fn unadvertised_name_is_not_found_even_on_a_resumed_turn() { + let model = MockCompletionModel::new([MockTurn::tool_call( + "call_1", + "add", + json!({"x": 2, "y": 5}), + )]); + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("what is 2 + 5?"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + + // Resume against a process that has since registered another tool. + let resumed_agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(2) + .tool(MockAddTool) + .tool(MockSubtractTool) + .build(); + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = resumed_agent.drive_run(run); + let (_, tools) = expect_execute_tools!(driver); + + assert!( + !tools.executable_tool_names().contains("subtract"), + "the advertised set is the turn's, not this process's" + ); + let mut context = ToolContext::new(); + let result = tools + .execute("subtract", r#"{"x": 1, "y": 2}"#, &mut context) + .await; + assert!( + result.is_error_kind(ToolErrorKind::NotFound), + "a tool registered after the turn must not be reachable" + ); + // The advertised tool still dispatches. + let result = tools + .execute("add", r#"{"x": 2, "y": 5}"#, &mut context) + .await; + assert!(result.is_success()); + } + + /// Retrieval picks dynamic tools by similarity to the turn's query, so a + /// registered dynamic tool can be absent from a resumed snapshot purely + /// because that query did not rank it. Reporting that as "no longer + /// registered" would be false, and the advice it carries — register the + /// tool before resuming — unfollowable. + #[tokio::test] + async fn resumed_dynamic_tool_outside_retrieval_is_not_drift() { + use crate::test_utils::MockToolIndex; + use crate::tool::ToolSet; + use crate::tool::server::ToolServer; + + let model = MockCompletionModel::new([MockTurn::tool_call( + "call_1", + "subtract", + json!({"x": 5, "y": 2}), + )]); + // Retrieval finds `subtract` while the turn is prepared. + let handle = ToolServer::new() + .tool(MockAddTool) + .retrieved_tools( + 1, + MockToolIndex::new(["subtract"]), + ToolSet::from_tools(vec![MockSubtractTool]), + ) + .run(); + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool_server_handle(handle) + .build(); + + let mut driver = agent.drive("what is 5 - 2?"); + let (request, tools, _) = expect_send!(driver); + assert!(tools.executable_tool_names().contains("subtract")); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + + // Resume where retrieval ranks nothing — the tool is still registered. + let resumed_handle = ToolServer::new() + .tool(MockAddTool) + .retrieved_tools( + 1, + MockToolIndex::new(Vec::::new()), + ToolSet::from_tools(vec![MockSubtractTool]), + ) + .run(); + let resumed_agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(2) + .tool_server_handle(resumed_handle) + .build(); + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = resumed_agent.drive_run(run); + + let (calls, tools) = expect_execute_tools!(driver); + let mut context = ToolContext::new(); + let result = tools.execute_call(&calls[0], &mut context).await; + match result { + UserContent::ToolResult(result) => assert!( + !format!("{:?}", result.content).contains("not found"), + "a registered tool retrieval missed must still dispatch: {:?}", + result.content + ), + other => panic!("expected a tool result, got {other:?}"), + } + } + + /// The two halves of a resumed `TurnTools` must agree: the snapshot is + /// narrowed to what the turn advertised, not left as this process's whole + /// registry. + #[tokio::test] + async fn resumed_snapshot_equals_the_advertised_set() { + let model = MockCompletionModel::new([MockTurn::tool_call( + "call_1", + "add", + json!({"x": 2, "y": 5}), + )]); + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("what is 2 + 5?"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + + let resumed_agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(2) + .tool(MockAddTool) + .tool(MockSubtractTool) + .build(); + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let advertised = run + .advertised_tools() + .expect("the turn recorded its advertised names") + .clone(); + let mut driver = resumed_agent.drive_run(run); + let (_, tools) = expect_execute_tools!(driver); + + assert_eq!(tools.executable_tool_names(), &advertised.executable); + assert!( + !tools.executable_tool_names().contains("subtract"), + "the resuming process's extra tool is not part of this turn" + ); + } + /// Resuming is about dispatch, not about validating a request that will - /// never be built. A `tool_choice` the resuming process could not satisfy /// must not pre-empt the drift report, and must not defeat the opt-out /// that exists precisely for this situation. #[tokio::test] diff --git a/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index bf8ae0334e..b5de527be5 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -134,6 +134,9 @@ pub use prompt_request::{ TypedPromptRequest, TypedPromptResponse, }; pub use rig_core::message::Text; -pub use run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome, OutputMode, PendingToolCall}; +pub use run::{ + Advance, AgentRun, AgentRunStep, ModelCallInputs, ModelTurn, ModelTurnOutcome, OutputMode, + PendingToolCall, +}; pub use runner::AgentRunner; -pub use turn_tools::TurnTools; +pub use turn_tools::{TurnToolNames, TurnTools}; diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index fb406ebf68..ea087f0eec 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -77,6 +77,39 @@ //! # Ok(()) //! # } //! ``` +//! +//! # When building the request can fail +//! +//! [`AgentRun::next_step`] consumes a turn the moment it hands out a +//! [`AgentRunStep::CallModel`], which is right for a driver that has nothing +//! fallible to do with it. A driver that must *build* something first — read a +//! tool registry, resolve a model, validate a tool choice — uses the two halves +//! instead, so a build failure costs no turn and leaves the run intact: +//! +//! ```rust,no_run +//! use rig_agent::agent::run::{Advance, AgentRun, ModelCallInputs}; +//! +//! # async fn example(run: &mut AgentRun) -> Result<(), Box> { +//! match run.advance()? { +//! Advance::NeedsModelCall => { +//! let ModelCallInputs { prompt, history } = run.peek_model_call()?; +//! // Build the request. Failing here changes nothing about `run`: +//! // no turn spent, still `is_preparing_request()`, retry at will. +//! # let _ = (prompt, history); +//! // Only once the request exists does the run advance. +//! let _turn = run.commit_model_call(None, None)?; +//! // Send it. If the send fails or its reply is lost, hand the turn +//! // back with `run.rollback_model_call()?` and prepare again. +//! } +//! Advance::CallTools(calls) => { +//! # let _ = calls; +//! // Execute, then: run.tool_results(results)?; +//! } +//! Advance::Done(response) => println!("{}", response.output), +//! } +//! # Ok(()) +//! # } +//! ``` pub mod output_mode; pub mod streamed; @@ -388,23 +421,41 @@ pub struct AgentRun { /// advertised names arrive with the [`ModelTurn`] instead. #[serde(default)] advertised_tools: Option, + /// How many committed model calls were handed back via + /// [`AgentRun::rollback_model_call`] — sends that failed, or whose reply + /// was lost. Observability only: the run imposes no bound on it, because + /// the caller that owns the transport is the only party that can decide + /// how many attempts are reasonable. + #[serde(default)] + model_call_rollbacks: usize, state: RunState, } /// The inputs for a model call the run is ready to make, read without /// advancing anything. See [`AgentRun::peek_model_call`]. +/// +/// Deliberately exhaustive, matching [`AgentRunStep::CallModel`]'s fields: +/// destructuring it is the point, and a caller that must build a request from +/// these should be told by the compiler when there is more to build it from. #[derive(Debug, Clone)] -pub(crate) struct ModelCallInputs { - pub(crate) prompt: Message, - pub(crate) history: Vec, +pub struct ModelCallInputs { + /// The prompt message for this turn (the latest message in the run). + pub prompt: Message, + /// The chat history preceding `prompt`: the caller-provided input history + /// followed by messages accumulated by earlier turns. + pub history: Vec, } /// What [`AgentRun::advance`] resolved to. /// /// Splits "the run wants another model call" from every step that needs no -/// model call, so a caller with fallible request preparation can do that work +/// model call, so a caller whose request preparation can fail does that work /// *before* anything advances (see [`AgentRun::commit_model_call`]). -pub(crate) enum Advance { +/// +/// Deliberately exhaustive, like [`AgentRunStep`]: a driver must handle every +/// variant, so adding one is a breaking change by design. +#[derive(Debug)] +pub enum Advance { /// The run is parked in `PreparingRequest`: read the inputs with /// [`AgentRun::peek_model_call`], build the request, and commit it with /// [`AgentRun::commit_model_call`]. Nothing has advanced yet. @@ -438,6 +489,7 @@ impl AgentRun { rollback_pending: false, streamed_completion_call_recorded: false, advertised_tools: None, + model_call_rollbacks: 0, state: RunState::PreparingRequest, } } @@ -553,8 +605,13 @@ impl AgentRun { /// [`TurnTools`](crate::agent::TurnTools): pairing it with a registry /// snapshot reconstitutes the whole thing, which is what lets a driver /// resume a run suspended at *any* step rather than only while tool calls - /// are pending. - pub(crate) fn advertised_tools(&self) -> Option<&TurnToolNames> { + /// are pending. It is also the record of what the model was actually shown + /// this turn, which is worth persisting alongside the run for audit. + /// + /// `None` for a run hand-driven through [`Self::next_step`], which commits + /// no names because the driver supplies them with the [`ModelTurn`] + /// instead, and `None` before the run's first model call. + pub fn advertised_tools(&self) -> Option<&TurnToolNames> { self.advertised_tools.as_ref() } @@ -738,8 +795,10 @@ impl AgentRun { }) } - /// Whether the run is waiting for its next model call to be prepared. - pub(crate) fn is_preparing_request(&self) -> bool { + /// Whether the run is waiting for its next model call to be prepared — + /// that is, whether [`Self::peek_model_call`] and + /// [`Self::commit_model_call`] are valid right now. + pub fn is_preparing_request(&self) -> bool { matches!(self.state, RunState::PreparingRequest) } @@ -757,12 +816,16 @@ impl AgentRun { /// - [`PromptError::MaxTurnsError`] when the total model-call budget is /// exhausted. The run stays intact, so raising the budget and peeking /// again resumes it. - /// - [`PromptError::PromptCancelled`] when there is no pending prompt. - pub(crate) fn peek_model_call(&self) -> Result { - debug_assert!( - self.is_preparing_request(), - "peek_model_call is only meaningful while the run is preparing a request" - ); + /// - [`PromptError::PromptCancelled`] when there is no pending prompt, or + /// when the run is not preparing a request (check with + /// [`Self::is_preparing_request`], or reach this through + /// [`Self::advance`], which reports it). + pub fn peek_model_call(&self) -> Result { + if !self.is_preparing_request() { + return Err( + self.protocol_violation("peek_model_call called while the run wants no model call") + ); + } let Some((prompt, history_for_turn)) = self.new_messages.split_last() else { return Err(PromptError::prompt_cancelled( self.full_history(), @@ -798,22 +861,100 @@ impl AgentRun { /// with the [`ModelTurn`] instead. `output_tool_name` fills the run's /// committed Tool-mode name once (#1928), pinning the mode for the rest of /// the run. - pub(crate) fn commit_model_call( + /// + /// If the committed call then fails to reach the provider, or its reply is + /// lost, [`Self::rollback_model_call`] undoes exactly this. + /// + /// # Errors + /// [`PromptError::PromptCancelled`] when the run is not preparing a + /// request — committing a call the run never asked for would consume a + /// turn against nothing. + pub fn commit_model_call( &mut self, advertised: Option, output_tool_name: Option, - ) -> usize { - debug_assert!( - self.is_preparing_request(), - "commit_model_call is only valid on a peeked model call" - ); + ) -> Result { + if !self.is_preparing_request() { + return Err( + self.protocol_violation("commit_model_call called without a peeked model call") + ); + } self.set_output_tool_name(output_tool_name); self.advertised_tools = advertised; self.current_turn += 1; self.rollback_pending = false; self.streamed_completion_call_recorded = false; self.state = RunState::AwaitingModel; - self.current_turn + Ok(self.current_turn) + } + + /// Hand back a committed model call that never produced a response. + /// + /// For a request that could not be sent, or whose reply is known to be + /// lost: a transport error, a timeout, a cancelled or vanished provider + /// job. The turn is refunded — it produced nothing, so it costs nothing — + /// and the run returns to `PreparingRequest`, where the next + /// [`Self::peek_model_call`] / [`Self::commit_model_call`] pair (or + /// [`Self::next_step`]) prepares the request **again, from current + /// configuration**. + /// + /// Re-deriving rather than re-sending is deliberate. A request built for + /// the failed attempt carries that attempt's tool snapshot, and replaying + /// it would advertise one set of implementations while a later turn + /// dispatches another — the exact skew + /// [`TurnTools`](crate::agent::TurnTools) exists to prevent. The retry + /// therefore takes a fresh snapshot, a fresh patch, and the same history. + /// + /// # This is at-least-once + /// + /// If the request *did* reach the provider and only its reply was lost, + /// rolling back and re-sending bills a second completion. Nothing in the + /// run can distinguish "never arrived" from "arrived, reply lost" — only + /// the caller can, through provider-side idempotency or its own record of + /// what was transmitted. Roll back when you know the call produced + /// nothing; when you do not know, prefer failing the run. + /// + /// Usage accounting is preserved either way. Tokens the provider already + /// billed stay in [`Self::usage`] and [`Self::completion_calls`], and a + /// streamed turn may still record its usage after the rollback (the same + /// window invalid tool-call recovery uses, see + /// [`Self::record_streamed_completion_call`]). + /// + /// Bounding attempts is the caller's job: the machine performs no IO and + /// owns no clock, so it cannot know whether a retry is reasonable. + /// [`Self::model_call_rollbacks`] counts them for a caller that wants a + /// budget of its own, and + /// [`CompletionError::is_retryable`](crate::completion::CompletionError::is_retryable) + /// classifies whether the failure is worth another attempt at all. + /// + /// # Errors + /// [`PromptError::PromptCancelled`] when no model call is in flight. + pub fn rollback_model_call(&mut self) -> Result<(), PromptError> { + if !matches!(self.state, RunState::AwaitingModel) { + return Err(self + .protocol_violation("rollback_model_call called without a model call in flight")); + } + + // Refund the turn. `saturating_sub` is belt-and-braces: reaching + // `AwaitingModel` always went through `commit_model_call`, which + // incremented it. + self.current_turn = self.current_turn.saturating_sub(1); + self.model_call_rollbacks += 1; + // The next preparation records its own; leaving stale names here would + // let a later `model_response` validate against a turn that never ran. + self.advertised_tools = None; + // Keep the streamed-usage window open when the record has not landed + // yet: a caller that drains a broken stream for usage records it after + // this point, exactly as invalid tool-call recovery does. + self.rollback_pending = !self.streamed_completion_call_recorded; + self.state = RunState::PreparingRequest; + Ok(()) + } + + /// How many committed model calls were handed back via + /// [`Self::rollback_model_call`]. + pub fn model_call_rollbacks(&self) -> usize { + self.model_call_rollbacks } /// Advance the machine as far as it can go without a model call. @@ -823,7 +964,7 @@ impl AgentRun { /// (#1928) — leaving the run in `PreparingRequest` for the caller to /// prepare and commit. [`Self::next_step`] is this plus an immediate /// commit, for callers with nothing that can fail in between. - pub(crate) fn advance(&mut self) -> Result { + pub fn advance(&mut self) -> Result { loop { if self.is_preparing_request() { return Ok(Advance::NeedsModelCall); @@ -1048,7 +1189,7 @@ impl AgentRun { let ModelCallInputs { prompt, history, .. } = self.peek_model_call()?; - let turn = self.commit_model_call(None, None); + let turn = self.commit_model_call(None, None)?; Ok(AgentRunStep::CallModel { prompt, history, @@ -1792,6 +1933,150 @@ mod tests { } } + /// A model call that was committed and then failed to reach the provider: + /// the turn is refunded, the run is drivable again, and the retry takes + /// the turn the failure did not. + #[test] + fn rollback_model_call_refunds_the_turn_and_re_prepares() { + let mut run = AgentRun::new("add things").max_turns(1); + expect_call_model(&mut run); + assert_eq!(run.turn(), 1); + + run.rollback_model_call().expect("rollback should succeed"); + assert_eq!(run.turn(), 0, "a call that produced nothing costs nothing"); + assert_eq!(run.model_call_rollbacks(), 1); + assert!(run.is_preparing_request()); + + // The budget of one is intact, so the retry is possible at all. + let (_, _, turn) = expect_call_model(&mut run); + assert_eq!(turn, 1); + expect_continue( + run.model_response(tool_call_turn("call_1", "add")) + .expect("the retried turn is accepted"), + ); + } + + /// The advertised names belong to the call that was handed back, so they + /// must not survive it: a later `model_response` validating against a turn + /// that never ran would be validating against nothing. + #[test] + fn rollback_model_call_drops_the_advertised_tools() { + let mut run = AgentRun::new("add things").max_turns(2); + run.advance().expect("advance should succeed"); + run.peek_model_call().expect("peek should succeed"); + run.commit_model_call( + Some(TurnToolNames { + executable: tool_names(&["add"]), + allowed: tool_names(&["add"]), + }), + None, + ) + .expect("commit should succeed"); + assert!(run.advertised_tools().is_some()); + + run.rollback_model_call().expect("rollback should succeed"); + assert_eq!(run.advertised_tools(), None); + } + + /// Tokens the provider already billed stay billed. A streamed turn that + /// died mid-stream may still learn its usage afterwards, so the recording + /// window stays open across the rollback — the same window invalid + /// tool-call recovery uses. + #[test] + fn rollback_model_call_preserves_streamed_usage_accounting() { + let mut run = AgentRun::new("add things").max_turns(2); + expect_call_model(&mut run); + run.record_streamed_completion_call(Usage { + total_tokens: 11, + ..Usage::new() + }) + .expect("usage recorded before the failure"); + + run.rollback_model_call().expect("rollback should succeed"); + assert_eq!( + run.usage().total_tokens, + 11, + "billed tokens survive the rollback" + ); + assert_eq!(run.completion_calls().len(), 1); + + // A second record for the same dead turn is still rejected... + run.record_streamed_completion_call(Usage::new()) + .expect_err("the turn already recorded its usage"); + + // ...and the retried turn records its own. + expect_call_model(&mut run); + run.record_streamed_completion_call(Usage { + total_tokens: 7, + ..Usage::new() + }) + .expect("the retried turn records its own usage"); + assert_eq!(run.usage().total_tokens, 18); + assert_eq!(run.completion_calls().len(), 2); + } + + /// A stream that broke before reporting usage can still report it after + /// the rollback, which is why the window is left open when nothing was + /// recorded. + #[test] + fn rollback_model_call_leaves_the_streamed_usage_window_open() { + let mut run = AgentRun::new("add things").max_turns(2); + expect_call_model(&mut run); + run.rollback_model_call().expect("rollback should succeed"); + run.record_streamed_completion_call(Usage { + total_tokens: 3, + ..Usage::new() + }) + .expect("a drained stream may report usage after the rollback"); + assert_eq!(run.usage().total_tokens, 3); + } + + #[test] + fn rollback_model_call_requires_a_call_in_flight() { + let mut run = AgentRun::new("add things").max_turns(2); + run.rollback_model_call() + .expect_err("nothing is in flight before the first call"); + + expect_call_model(&mut run); + expect_continue( + run.model_response(tool_call_turn("call_1", "add")) + .expect("model_response should succeed"), + ); + expect_call_tools(&mut run); + run.rollback_model_call() + .expect_err("tool calls are pending, not a model call"); + } + + /// The transition is state, not driver memory, so it survives suspension. + #[test] + fn rollback_model_call_survives_a_serde_round_trip() { + let mut run = AgentRun::new("add things").max_turns(1); + expect_call_model(&mut run); + run.rollback_model_call().expect("rollback should succeed"); + + let serialized = serde_json::to_string(&run).expect("run should serialize"); + let mut restored: AgentRun = + serde_json::from_str(&serialized).expect("run should deserialize"); + assert_eq!(restored.turn(), 0); + assert_eq!(restored.model_call_rollbacks(), 1); + let (_, _, turn) = expect_call_model(&mut restored); + assert_eq!(turn, 1); + } + + /// The two halves are a public protocol, so driving them out of order is + /// an error a caller can handle — not a `debug_assert` that vanishes in + /// release builds. + #[test] + fn peek_and_commit_reject_being_driven_out_of_protocol() { + let mut run = AgentRun::new("add things").max_turns(2); + expect_call_model(&mut run); + + run.peek_model_call() + .expect_err("a model call is already in flight"); + run.commit_model_call(None, None) + .expect_err("there is no peeked call to commit"); + } + fn expect_continue(outcome: ModelTurnOutcome) -> bool { match outcome { ModelTurnOutcome::Continue { diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs index a05e124329..fe673e1a28 100644 --- a/crates/rig-agent/src/agent/turn_tools.rs +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -42,12 +42,13 @@ pub(crate) struct PreparedCompletionRequest { /// Names are data, so they travel with the run — [`AgentRun`](super::run::AgentRun) /// records them when a model call is committed, and they survive /// serialization. Implementations are live objects and cannot, so the -/// dispatch target ([`ToolRegistrySnapshot`]) is rebuilt from the agent's -/// registry when a run resumes in another process. Pairing the two back +/// dispatch target — the turn's registry snapshot — is rebuilt from the +/// agent when a run resumes in another process. Pairing the two back /// together yields a [`TurnTools`] in any run state, in any process, which is /// what makes a suspended run resumable at *every* step boundary rather than /// only while tool calls are pending. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] +#[non_exhaustive] pub struct TurnToolNames { /// The real registry tools advertised to the provider this turn. pub executable: BTreeSet, @@ -194,6 +195,16 @@ impl TurnTools { answer; the driver must consume the call's arguments instead of dispatching it." ))); } + // The advertised names are the authority, not the snapshot. In-process + // the two agree by construction (preparation narrows the snapshot to + // the advertised set), but a resumed turn pairs names carried on the + // run with a snapshot rebuilt here, so only this check makes the + // guarantee above hold on both paths. + if !self.executable_tool_names.contains(tool_name) { + return ToolResult::failed(ToolExecutionError::not_found(format!( + "`{tool_name}` was not advertised to the model on this turn" + ))); + } self.snapshot.execute(tool_name, args, context).await } @@ -323,4 +334,68 @@ mod tests { "the rejection must not leave the previous dispatch's metadata behind" ); } + + /// The advertised names are the authority, not the snapshot. Pair a + /// snapshot that *can* dispatch a tool with names that never advertised + /// it — the shape a resumed turn produces — and dispatch must still + /// refuse. + #[tokio::test] + async fn dispatch_refuses_a_name_the_turn_did_not_advertise() { + use super::{TurnToolNames, TurnTools}; + use std::collections::BTreeSet; + + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .tool(MetadataTool) + .build(); + let mut driver = agent.drive("go"); + let advertised = match driver.next_step().await.expect("prepare succeeds") { + DriveStep::SendRequest { tools, .. } => tools, + other => panic!("expected SendRequest, got {other:?}"), + }; + assert!(advertised.executable_tool_names().contains("metadata")); + + // Same dispatch target, but a turn that advertised nothing. + let disowned = TurnTools::from_parts( + advertised.snapshot.clone(), + TurnToolNames { + executable: BTreeSet::new(), + allowed: BTreeSet::new(), + }, + None, + ); + let mut context = ToolContext::new(); + let result = disowned.execute("metadata", "{}", &mut context).await; + assert!( + result.is_error_kind(ToolErrorKind::NotFound), + "a reachable implementation is still not dispatchable if the turn never advertised it" + ); + } + + /// The advertised-name gate is an `execute` early return like the + /// output-tool rejection, so it owes the same guarantee: the previous + /// dispatch's result metadata must not survive it. + #[tokio::test] + async fn unadvertised_name_rejection_clears_stale_context_metadata() { + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .tool(MetadataTool) + .build(); + let mut driver = agent.drive("go"); + let tools = match driver.next_step().await.expect("prepare succeeds") { + DriveStep::SendRequest { tools, .. } => tools, + other => panic!("expected SendRequest, got {other:?}"), + }; + + let mut context = ToolContext::new(); + let result = tools.execute("metadata", "{}", &mut context).await; + assert!(result.is_success()); + assert_eq!(context.result::(), Some(&Marker("published"))); + + let result = tools.execute("never_advertised", "{}", &mut context).await; + assert!(result.is_error_kind(ToolErrorKind::NotFound)); + assert_eq!( + context.result::(), + None, + "the rejection must not leave the previous dispatch's metadata behind" + ); + } } diff --git a/crates/rig-agent/src/tool/server.rs b/crates/rig-agent/src/tool/server.rs index d6faf61900..e26d47f236 100644 --- a/crates/rig-agent/src/tool/server.rs +++ b/crates/rig-agent/src/tool/server.rs @@ -487,6 +487,27 @@ impl ToolServerHandle { pub(crate) async fn snapshot_tool_defs( &self, prompt: Option, + ) -> Result { + self.snapshot_tool_defs_including(prompt, &BTreeSet::new()) + .await + } + + /// Resolve a snapshot that also contains `required`, whatever retrieval + /// selected. + /// + /// Retrieval picks dynamic tools by similarity to the turn's query, so a + /// registered dynamic tool can be absent from a snapshot simply because + /// this query did not rank it. That is fine when the snapshot is deciding + /// what to *advertise*, and wrong when it must *dispatch* names a previous + /// turn already advertised — a resumed run's pending calls, for instance, + /// where "absent from this snapshot" would otherwise be indistinguishable + /// from "no longer registered". Naming them explicitly resolves them from + /// the registry directly; names that really are gone stay absent, so the + /// caller's drift check still sees the truth. + pub(crate) async fn snapshot_tool_defs_including( + &self, + prompt: Option, + required: &BTreeSet, ) -> Result { let retrieval_indexes = { let state = self.0.read().await; @@ -531,6 +552,18 @@ impl ToolServerHandle { Vec::new() }; + // Append the explicitly required names after the retrieved ones: + // `snapshot_registered_tools` keeps the first declaration of a + // duplicate, so a name retrieval already selected keeps its retrieval + // ordering and a name it missed is resolved from the registry here. + let dynamic_tool_ids = if required.is_empty() { + dynamic_tool_ids + } else { + let mut ids = dynamic_tool_ids; + ids.extend(required.iter().cloned()); + ids + }; + #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] let tools = { let mut state = self.0.write().await; diff --git a/crates/rig-core/CHANGELOG.md b/crates/rig-core/CHANGELOG.md index 1f4dd8b58f..b89a0ebce5 100644 --- a/crates/rig-core/CHANGELOG.md +++ b/crates/rig-core/CHANGELOG.md @@ -56,6 +56,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- *(completion)* `CompletionError::is_retryable`: a default classification of whether re-issuing an equivalent request could plausibly succeed — transport failures with no status, and preserved provider statuses of 408, 409, 429 or 5xx. A 2xx carrying a provider-authored error envelope is deliberately not retryable, since the call completed and was rejected on its merits. Model failures previously had no equivalent to `ToolErrorKind::default_retryable`, so every caller guessed; it is a default a caller may override, not a policy - *(tool)* `ToolErrorKind::NotExecutable` and `ToolExecutionError::not_executable`: the tool is advertised to the model but not executable by the dispatcher — its call must be handled by the host. Produced by rig-agent when dispatching the synthetic structured-output tool; also the natural classification for host-executed or deferred tools. The enum is `#[non_exhaustive]`, so this is additive - *(streaming)* `wire::classify_typed_event` extends the decode-then-validate policy to typed-transport wires (bedrock, candle, gemini-grpc): modeled variants are `Known`, the SDK's non-exhaustive/unrecognized variants are `Unknown`, SDK decode errors are `Corrupt` — a typed transport earns no policy exemption - *(streaming)* `WireAdapter` gains an associated `Frame` type so typed-event wires implement the same contract over their SDK events; `classify` now takes the frame by value diff --git a/crates/rig-core/src/completion/request.rs b/crates/rig-core/src/completion/request.rs index 25aa2fd6f9..f0f20146f7 100644 --- a/crates/rig-core/src/completion/request.rs +++ b/crates/rig-core/src/completion/request.rs @@ -131,6 +131,55 @@ impl CompletionError { Self::ProviderError(error.to_string()) } } + + /// Whether re-issuing an equivalent request could plausibly succeed. + /// + /// A **default** classification, not a policy: it answers "is this failure + /// about the request or about the moment?" and leaves the decision to + /// retry — and any backoff, budget, or idempotency handling — to the + /// caller. Callers that know their provider better should override it. + /// + /// - [`HttpError`](Self::HttpError) carrying no status is a transport + /// failure: the request may never have arrived. Retryable. + /// - Any preserved provider status of `408` (request timeout), `409` + /// (conflict), `429` (too many requests), or `5xx` is retryable. These + /// are the same four classes the AI SDK treats as retryable by default. + /// - Everything else is not: a malformed request, a serialization or URL + /// error, or a response this build could not parse will fail the same + /// way every time. + /// + /// A 2xx carrying a provider-authored error envelope + /// ([`ProviderResponse`](Self::ProviderResponse)) is deliberately *not* + /// retryable — the call completed and the provider rejected it on its + /// merits. + /// + /// Note the asymmetry this removes: rig has classified tool failures by + /// retryability for some time, while model failures had no equivalent, so + /// every caller guessed. + /// + /// Useful with a hand-driven run, where the caller owns the send and must + /// decide whether a failed turn is worth handing back: + /// + /// ```rust,ignore + /// match request.send().await { + /// Ok(response) => { driver.model_response(&response)?; } + /// Err(err) if err.is_retryable() => driver.rollback_model_call()?, + /// Err(err) => return Err(err.into()), + /// } + /// ``` + pub fn is_retryable(&self) -> bool { + match self.provider_response_status() { + Some(status) => { + status == http::StatusCode::REQUEST_TIMEOUT + || status == http::StatusCode::CONFLICT + || status == http::StatusCode::TOO_MANY_REQUESTS + || status.is_server_error() + } + // No status at all: a transport failure that never reached a + // response. Anything else is a local problem with the request. + None => matches!(self, Self::HttpError(_)), + } + } } #[derive(Clone, Debug, PartialEq, Deserialize, Serialize)] @@ -1277,6 +1326,65 @@ mod tests { use super::{CompletionResponse, FinishReason, ProviderCapabilities, Usage}; use crate::message::AssistantContent; + mod retryability { + use super::super::CompletionError; + use crate::http_client; + use http::StatusCode; + + fn from_status(status: u16) -> CompletionError { + CompletionError::from_http_response( + StatusCode::from_u16(status).expect("valid status"), + "body", + ) + } + + /// The four classes that mean "the moment was wrong", not "the request + /// was wrong" — matching the AI SDK's default classification. + #[test] + fn transient_provider_statuses_are_retryable() { + for status in [408, 409, 429, 500, 502, 503, 504] { + assert!( + from_status(status).is_retryable(), + "{status} should be retryable" + ); + } + } + + #[test] + fn client_errors_are_not_retryable() { + for status in [400, 401, 403, 404, 422] { + assert!( + !from_status(status).is_retryable(), + "{status} should not be retryable" + ); + } + } + + /// A transport failure may never have reached the provider at all. + #[test] + fn transport_failures_without_a_status_are_retryable() { + let error = + CompletionError::HttpError(http_client::Error::Instance("connection reset".into())); + assert!(error.is_retryable()); + } + + /// A 2xx carrying a provider-authored error envelope completed and was + /// rejected on its merits; re-sending it changes nothing. + #[test] + fn a_provider_error_envelope_on_success_is_not_retryable() { + assert!(!from_status(200).is_retryable()); + } + + /// Local problems with the request or our parsing of the response are + /// deterministic. + #[test] + fn local_errors_are_not_retryable() { + assert!(!CompletionError::ResponseError("bad shape".into()).is_retryable()); + assert!(!CompletionError::ProviderError("opaque".into()).is_retryable()); + assert!(!CompletionError::RequestError("malformed".to_string().into()).is_retryable()); + } + } + mod message_content_validation { use super::super::CompletionRequest; use crate::message::{AssistantContent, Message, UserContent}; From 02e3324607d1a97d7557def918d89ea02ca22218 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 12:13:23 -0700 Subject: [PATCH 06/21] fix(core): stop is_retryable turning a config typo into an infinite loop Round-3 review of PR #2278 found five defects, all on the round-2 fix commit. The one that matters is `CompletionError::is_retryable`, shipped last round as "optional" and wrong in both directions. Its statusless arm keyed on the error *variant*, and `HttpError` is not a semantic class - it is a transport enum whose variants mix transient and deterministic failures. An API key read with a trailing newline reaches `bearer_auth_header`, returns `Error::InvalidHeaderValue`, carries no status, and was therefore classified retryable. A driver following the pattern in that method's own documentation - roll back on a retryable error - then re-prepares forever on a failure that can never succeed. Round 2 made the rollback cheap, which is exactly what makes this unbounded. The same arm erred the other way: `from_provider_body` produces a statusless `ProviderResponse` for Bedrock, Vertex and the gRPC Gemini client, so throttling from a non-HTTP transport was classified non-retryable and hard-failed the run. Classification now happens where the variants live. `http_client::Error::is_transient` matches exhaustively - no wildcard, so a variant added later must be classified deliberately - and excludes `Protocol`, `InvalidHeaderValue`, `NoHeaders` and `InvalidContentType` by name, the same class langgraph's `default_retry_on` and pydantic-ai's Temporal denylist exclude. `StreamEnded` and the client's own opaque `Instance` failures default to transient, because rig cannot see inside a `Box` and a dropped connection is the overwhelming case. The non-HTTP-transport gap is not papered over: statusless provider bodies stay conservatively non-retryable, and the doc names the limitation, says providers that can surface a status should use `from_http_response`, and points at the provider adapter as where this belongs. A test asserts the gap deliberately so nobody "fixes" it blind. That fix forces a correction round 2 got wrong in prose. Retryable and replay-safe are different questions, and round 2's docs let the first answer the second - `Err(err) if err.is_retryable() => driver.rollback_model_call()?` was the recommended pattern. But a stream that died *after* the request was written is retryable and not replay-safe at all: rolling back bills a second completion and repeats whatever the model already caused. `is_transient` classifies exactly those as retryable, correctly for that axis. The two questions are now separated wherever the pattern appeared - `is_retryable`, `rollback_model_call` on both types, the driver's durability section, the changelogs and MIGRATING - and the recommended pattern gained the `nothing_was_sent` term the caller alone can supply. openai-agents carries this as data (`replay_safety` on its retry advice, set from the SDK's "the request may have been accepted"); rig has no such signal, so it says so instead of implying otherwise. Three smaller fixes. `execute_call`'s pre-resolved path now clears the context's dispatch result like every other dispatch surface, so a call suppressed by invalid tool-call recovery no longer leaves the previous call's metadata readable to a loop reading it per call - the same hazard round 2 closed for `execute` and missed on its sibling. `TurnToolNames` gains a constructor: it is `#[non_exhaustive]` serialized state and `commit_model_call` is public and takes one, so the "any hand-driver, not just rig's own" claim was false. And both of rig's own `ModelCallInputs` destructurings dropped their `..`, which had been silently opting out of the compiler signal the type's own doc claims to provide. 512 rig-agent tests (+2), 1241 rig-core (+2), clippy -D warnings clean, fmt clean, 0 rustdoc warnings in both crates. Both new behavioral tests were confirmed to fail before their fix by deliberate revert. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- MIGRATING.md | 13 ++- crates/rig-agent/CHANGELOG.md | 4 +- crates/rig-agent/src/agent/driver.rs | 47 ++++---- crates/rig-agent/src/agent/run/mod.rs | 30 +++-- crates/rig-agent/src/agent/turn_tools.rs | 86 ++++++++++++++- crates/rig-core/CHANGELOG.md | 2 +- crates/rig-core/src/completion/request.rs | 129 +++++++++++++++++----- crates/rig-core/src/http_client/mod.rs | 34 ++++++ 8 files changed, 280 insertions(+), 65 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index fbbe89ec66..54bc6d60c7 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1593,9 +1593,16 @@ Because the caller owns the send, the caller also owns recovering from a failed one: `AgentDriver::rollback_model_call()` hands a turn back when its request could not be sent or its reply is known to be lost, refunding the turn and returning the run to preparing so the next step yields a freshly prepared -request. It is at-least-once — a request that reached the provider and lost only -its reply will be billed twice on retry — so pair it with -`CompletionError::is_retryable()` and bound the attempts yourself. Runs driven +request. + +Deciding to use it takes two answers, not one. `CompletionError::is_retryable()` +(new in `rig-core`) says whether a retry *could succeed* — transport failures +that could resolve on their own, and provider statuses of 408, 409, 429 or 5xx. +It does **not** say whether a retry is *safe*: a stream that died after the +request was written is retryable and may already have taken effect, so rolling +back on retryability alone bills a second completion. Only you can settle that, +through provider-side idempotency, your own record of what was transmitted, or a +transport that fails before the write. Bound the attempts yourself. Runs driven by `Agent::runner` are unaffected; the runner still fails the prompt. Two behavior changes on the driver's resume path, both narrowing what dispatch diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 79a4567b9d..02138fd280 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -38,9 +38,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `Agent::drive` / `Agent::drive_run`, returning the new public `AgentDriver` (with `DriveStep` and `TurnTools`): hand-drive the sans-IO `AgentRun` machine with the agent's own configuration while owning every side effect. The driver seeds the run from the agent (`default_max_turns`, `tool_choice`, output schema), yields each turn's fully configured completion request for the caller to send (or hand to a custom transport), pairs every pending tool batch with the exact registry snapshot that advertised it, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. Resuming a serialized run in a fresh process re-derives a fresh dispatch snapshot from the rebuilt agent, and surfaces missing pending tools as an error (opt out with `allow_missing_resumed_tools`). The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` - *(agent)* `AgentDriver` is resumable at **every** step boundary, not just while tool calls are pending. The turn's advertised tool names are recorded on the run (new public `TurnToolNames`, the serializable half of `TurnTools`, reachable via `AgentRun::advertised_tools()`), so a run serialized after `DriveStep::SendRequest` — the natural suspension point for a queued or long-running provider call — can be resumed in another process and fed its reply, validated against the set the request actually carried rather than against whatever the resuming registry holds. The driver now keeps no state of its own beyond a registry-snapshot cache - *(agent)* a failed model-turn preparation costs nothing: the run advances only once a request exists, so an unreachable tool server or an impossible `tool_choice` leaves the run byte-identical — same state, same turn budget — and the step can be retried in place. Previously the turn was consumed and the run was left unresumable -- *(agent)* add `AgentRun::rollback_model_call` / `AgentDriver::rollback_model_call` for the other half of that story: a request that could not be sent, or whose reply is known to be lost, hands its turn back and returns the run to preparing, so the next step yields a **freshly prepared** request rather than replaying one whose tool snapshot has gone stale. Refunds the turn, preserves usage the provider already billed (including a streamed turn that reports usage after the failure), and counts attempts via `AgentRun::model_call_rollbacks`. Deliberately at-least-once: only the caller can tell "never arrived" from "arrived, reply lost", so bounding attempts is the caller's job +- *(agent)* add `AgentRun::rollback_model_call` / `AgentDriver::rollback_model_call` for the other half of that story: a request that could not be sent, or whose reply is known to be lost, hands its turn back and returns the run to preparing, so the next step yields a **freshly prepared** request rather than replaying one whose tool snapshot has gone stale. Refunds the turn, preserves usage the provider already billed (including a streamed turn that reports usage after the failure), and counts attempts via `AgentRun::model_call_rollbacks`. Deciding to use it takes two answers, not one: `CompletionError::is_retryable` says whether a retry *could succeed*, and only the caller can say whether one is *safe* — a stream that died after the request was written is retryable and not replay-safe, and retrying it bills a second completion. Bounding attempts is the caller's job - *(agent)* `AgentRun`'s peek/commit halves are public — `peek_model_call`, `commit_model_call`, `advance`, `is_preparing_request`, with `Advance` and `ModelCallInputs` — so any hand-driver whose request preparation can fail gets the same "a failure costs no turn" guarantee the driver has, not just rig's own. Their preconditions are `PromptError` protocol violations rather than debug assertions - *(agent)* [**behavior**] `TurnTools::execute` now refuses any name the turn did not advertise, instead of trusting its snapshot to be narrowed. The two agreed by construction in-process; on a resumed turn the names come from the run and the snapshot is rebuilt locally, so a tool registered after suspension could previously be dispatched through a turn that never advertised it +- *(agent)* [**behavior**] `TurnTools::execute_call` clears the context's dispatch result on the pre-resolved path too, so a call suppressed by invalid tool-call recovery no longer leaves the *previous* call's result metadata readable — a loop over a turn's calls was attributing it to the suppressed one +- *(agent)* add `TurnToolNames::new`: the type is `#[non_exhaustive]` serialized state, so an external hand-driver could not build the value `AgentRun::commit_model_call` requires - *(agent)* [**behavior**] a resumed run's dispatch snapshot resolves the turn's advertised names explicitly rather than relying on retrieval to rank them again, so a registered *dynamic* tool that the re-derived query does not return is no longer misreported as "not registered in this process" — advice that could not be followed — and no longer feeds the model a not-found result for a tool that exists - *(agent)* add `AgentDriver::request_patch` / `set_request_patch`, giving a hand-driven run the per-turn configuration the runner gets from its `CompletionCall` hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history. Because the caller owns the loop, per-turn variation needs no callback - *(agent)* add `AgentRun::tool_choice()`. A custom run's own `tool_choice` now reaches the provider when the run is hand-driven, instead of only governing the run's internal decisions while the request silently carried the agent's baseline diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 466bb9498a..1df3ae729b 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -49,12 +49,13 @@ //! - **Sending it.** The caller owns the send, so the caller is the only party //! that learns it failed. [`AgentDriver::rollback_model_call`] hands the turn //! back — refunding it and returning the run to preparing — and the next -//! `next_step` yields a freshly prepared request. This is *at-least-once*: if -//! the request reached the provider and only its reply was lost, retrying -//! bills a second completion, and only the caller can tell those apart. Use +//! `next_step` yields a freshly prepared request. Deciding to use it takes +//! two answers, not one: //! [`CompletionError::is_retryable`](crate::completion::CompletionError::is_retryable) -//! to decide, and bound the attempts yourself — the driver runs no IO and -//! owns no clock. +//! says whether a retry *could succeed*, and only the caller can say whether +//! one is *safe* — a request that reached the provider and lost only its +//! reply will be billed twice. Bound the attempts yourself; the driver runs +//! no IO and owns no clock. //! //! Tool *implementations* are live objects and cannot be serialized: the //! resuming process rebuilds the same `Agent` and the driver takes a fresh @@ -281,9 +282,7 @@ impl AgentDriver { // Peek, prepare, *then* commit. Reading the inputs consumes // nothing, so everything fallible below happens while the run // is still fully intact. - let ModelCallInputs { - prompt, history, .. - } = self.run.peek_model_call()?; + let ModelCallInputs { prompt, history } = self.run.peek_model_call()?; // Pin Tool output mode once committed (#1928), mirroring the // runner: read the run's committed name into preparation, and // store the resolved name back (fill-once). @@ -354,26 +353,33 @@ impl AgentDriver { /// registry snapshot, new patch — rather than a replay of a request whose /// tool snapshot has since gone stale. /// - /// ```rust,no_run - /// # use rig_agent::agent::{AgentDriver, DriveStep}; - /// # async fn example(driver: &mut AgentDriver) -> Result<(), Box> { + /// Two questions decide whether to roll back, and the library answers only + /// the first: *could a retry succeed?* — which + /// [`CompletionError::is_retryable`](crate::completion::CompletionError::is_retryable) + /// classifies — and *is a retry safe?*, which nothing here can know. A + /// stream that died after the request was written is retryable and not + /// replay-safe: rolling back on the first question alone bills a second + /// completion and repeats whatever the model already caused. Only the + /// caller can establish the second, through provider-side idempotency, its + /// own record of what was transmitted, or a transport that fails before + /// the write. + /// + /// ```rust,ignore /// if let DriveStep::SendRequest { request, .. } = driver.next_step().await? { /// match request.send().await { /// Ok(response) => { driver.model_response(&response)?; } - /// // Nothing was produced: give the turn back and try again. - /// Err(err) if err.is_retryable() => driver.rollback_model_call()?, + /// // `nothing_was_sent` is the caller's own knowledge; the driver + /// // cannot supply it, and `is_retryable` does not answer it. + /// Err(err) if err.is_retryable() && nothing_was_sent => { + /// driver.rollback_model_call()? + /// } /// Err(err) => return Err(err.into()), /// } /// } - /// # Ok(()) - /// # } /// ``` /// - /// See [`AgentRun::rollback_model_call`] for the full semantics — in - /// particular that this is **at-least-once**: if the request reached the - /// provider and only its reply was lost, retrying bills a second - /// completion. The driver cannot tell the two apart. Bounding attempts is - /// yours to do; the driver runs no IO and owns no clock. + /// See [`AgentRun::rollback_model_call`] for the full semantics. Bounding + /// attempts is yours to do; the driver runs no IO and owns no clock. pub fn rollback_model_call(&mut self) -> Result<(), PromptError> { self.run.rollback_model_call()?; // Drop the cached dispatch target too: the retry is a new turn and @@ -1300,6 +1306,7 @@ mod tests { } /// Resuming is about dispatch, not about validating a request that will + /// never be built. A `tool_choice` the resuming process could not satisfy /// must not pre-empt the drift report, and must not defeat the opt-out /// that exists precisely for this situation. #[tokio::test] diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index ea087f0eec..6bf2744c99 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -905,14 +905,26 @@ impl AgentRun { /// [`TurnTools`](crate::agent::TurnTools) exists to prevent. The retry /// therefore takes a fresh snapshot, a fresh patch, and the same history. /// - /// # This is at-least-once + /// # Retryable is not replay-safe /// - /// If the request *did* reach the provider and only its reply was lost, - /// rolling back and re-sending bills a second completion. Nothing in the - /// run can distinguish "never arrived" from "arrived, reply lost" — only - /// the caller can, through provider-side idempotency or its own record of - /// what was transmitted. Roll back when you know the call produced - /// nothing; when you do not know, prefer failing the run. + /// Two independent questions decide whether to roll back, and only the + /// first has a library answer: + /// + /// 1. **Could a retry succeed?** A property of the failure. + /// [`CompletionError::is_retryable`](crate::completion::CompletionError::is_retryable) + /// answers it. + /// 2. **Is a retry safe?** Did the request already take effect? Nothing in + /// the run can tell "never arrived" from "arrived, reply lost", and + /// neither can `is_retryable` — a stream that died *after* the request + /// was written is highly retryable and not replay-safe at all. + /// + /// Rolling back on question 1 alone is at-least-once: a request that + /// reached the provider and lost only its reply is billed twice, and any + /// side effect the model already caused happens again. Only the caller can + /// settle question 2 — through provider-side idempotency, its own record + /// of what was transmitted, or a transport that fails before the write. + /// Roll back when you know the call produced nothing; when you do not + /// know, prefer failing the run. /// /// Usage accounting is preserved either way. Tokens the provider already /// billed stay in [`Self::usage`] and [`Self::completion_calls`], and a @@ -1186,9 +1198,7 @@ impl AgentRun { pub fn next_step(&mut self) -> Result { match self.advance()? { Advance::NeedsModelCall => { - let ModelCallInputs { - prompt, history, .. - } = self.peek_model_call()?; + let ModelCallInputs { prompt, history } = self.peek_model_call()?; let turn = self.commit_model_call(None, None)?; Ok(AgentRunStep::CallModel { prompt, diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs index fe673e1a28..91c1b71ab0 100644 --- a/crates/rig-agent/src/agent/turn_tools.rs +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -58,6 +58,27 @@ pub struct TurnToolNames { } impl TurnToolNames { + /// The names a turn advertised: the executable registry tools, and the + /// tools the active `tool_choice` allowed the model to call. + /// + /// A hand-driver that prepares its own requests commits these with + /// [`AgentRun::commit_model_call`](super::run::AgentRun::commit_model_call); + /// this is how it builds them. The type is `#[non_exhaustive]` because it + /// is serialized run state and will gain fields, which is why a struct + /// literal will not do — the constructor is the stable way in. (Contrast + /// [`ModelCallInputs`](super::run::ModelCallInputs), which is deliberately + /// exhaustive so that destructuring it keeps working: the two types answer + /// different questions and keep different answers.) + pub fn new( + executable: impl IntoIterator>, + allowed: impl IntoIterator>, + ) -> Self { + Self { + executable: executable.into_iter().map(Into::into).collect(), + allowed: allowed.into_iter().map(Into::into).collect(), + } + } + /// Assemble the [`ModelTurn`] for a completion response received on the /// turn these names were advertised for. The single construction site for /// driver-facing turns, so the two name sets can never be transposed by a @@ -212,7 +233,10 @@ impl TurnTools { /// /// Honors [`PendingToolCall::preresolved_result`] — a call suppressed by /// invalid tool-call recovery returns its pre-resolved content without - /// executing anything or touching `context`. Otherwise the call dispatches + /// executing anything. It still clears `context`'s dispatch result: the + /// suppressed call published no metadata, so leaving the *previous* call's + /// in place would attribute it to this one, and a loop over a turn's calls + /// reads that metadata per call. Otherwise the call dispatches /// through this turn's snapshot via [`execute`](Self::execute) and the /// result is assembled into the [`UserContent`] tool-result value that /// [`AgentDriver::tool_results`](super::AgentDriver::tool_results) (or @@ -227,6 +251,10 @@ impl TurnTools { context: &mut ToolContext, ) -> UserContent { if let Some(result) = &call.preresolved_result { + // Every dispatch surface clears first; this one dispatches nothing + // but is still a call in the turn's sequence, so it must not + // inherit the previous call's result metadata. + context.clear_dispatch_result(); return result.clone(); } let name = &call.tool_call.function.name; @@ -335,6 +363,62 @@ mod tests { ); } + /// A suppressed call dispatches nothing, but it is still a call in the + /// turn's sequence: a loop reading per-call metadata must not see the + /// previous call's attributed to it. + #[tokio::test] + async fn preresolved_call_clears_stale_context_metadata() { + use crate::agent::run::PendingToolCall; + use rig_core::message::{ToolCall, ToolFunction, ToolResultContent, UserContent}; + + let agent = AgentBuilder::new(MockCompletionModel::text("unused")) + .tool(MetadataTool) + .build(); + let mut driver = agent.drive("go"); + let tools = match driver.next_step().await.expect("prepare succeeds") { + DriveStep::SendRequest { tools, .. } => tools, + other => panic!("expected SendRequest, got {other:?}"), + }; + + let mut context = ToolContext::new(); + let result = tools.execute("metadata", "{}", &mut context).await; + assert!(result.is_success()); + assert_eq!(context.result::(), Some(&Marker("published"))); + + // The next call in the same turn was suppressed by invalid tool-call + // recovery, so it carries a pre-resolved result. + let suppressed = PendingToolCall { + tool_call: ToolCall::from_wire( + "call_2", + ToolFunction::new("metadata".to_string(), json!({})), + ), + preresolved_result: Some(UserContent::tool_result( + "call_2", + "metadata", + vec![ToolResultContent::text("not executed")], + )), + internal_call_id: None, + }; + let _ = tools.execute_call(&suppressed, &mut context).await; + assert_eq!( + context.result::(), + None, + "a suppressed call must not inherit the previous call's metadata" + ); + } + + /// `commit_model_call` is public and takes this type, so a hand-driver + /// outside the crate must be able to build one. `#[non_exhaustive]` blocks + /// a struct literal there, which is what the constructor is for. + #[test] + fn advertised_names_are_constructible_without_a_struct_literal() { + use super::TurnToolNames; + + let names = TurnToolNames::new(["add", "subtract"], ["add"]); + assert!(names.executable.contains("subtract")); + assert!(names.allowed.contains("add") && !names.allowed.contains("subtract")); + } + /// The advertised names are the authority, not the snapshot. Pair a /// snapshot that *can* dispatch a tool with names that never advertised /// it — the shape a resumed turn produces — and dispatch must still diff --git a/crates/rig-core/CHANGELOG.md b/crates/rig-core/CHANGELOG.md index b89a0ebce5..66c3424c19 100644 --- a/crates/rig-core/CHANGELOG.md +++ b/crates/rig-core/CHANGELOG.md @@ -56,7 +56,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -- *(completion)* `CompletionError::is_retryable`: a default classification of whether re-issuing an equivalent request could plausibly succeed — transport failures with no status, and preserved provider statuses of 408, 409, 429 or 5xx. A 2xx carrying a provider-authored error envelope is deliberately not retryable, since the call completed and was rejected on its merits. Model failures previously had no equivalent to `ToolErrorKind::default_retryable`, so every caller guessed; it is a default a caller may override, not a policy +- *(completion)* `CompletionError::is_retryable`: a conservative classification of whether re-issuing an equivalent request could plausibly succeed — preserved provider statuses of 408, 409, 429 or 5xx, plus transport failures the transport itself reports as transient (a dropped connection, a cut stream). Deterministic transport failures are excluded by name: a header value the client refused (an API key with a trailing newline, say), a request that could not be constructed, or a wrong content type will fail identically forever, and retrying them loops. A 2xx carrying a provider-authored error envelope is not retryable, since the call completed and was rejected on its merits. Two documented limits: providers on non-HTTP transports (Bedrock, Vertex, gRPC Gemini) report errors as a body with no status, so their throttling classifies as non-retryable until classification moves to the provider adapter; and this answers only whether a retry *could succeed*, never whether one is *replay-safe* - *(tool)* `ToolErrorKind::NotExecutable` and `ToolExecutionError::not_executable`: the tool is advertised to the model but not executable by the dispatcher — its call must be handled by the host. Produced by rig-agent when dispatching the synthetic structured-output tool; also the natural classification for host-executed or deferred tools. The enum is `#[non_exhaustive]`, so this is additive - *(streaming)* `wire::classify_typed_event` extends the decode-then-validate policy to typed-transport wires (bedrock, candle, gemini-grpc): modeled variants are `Known`, the SDK's non-exhaustive/unrecognized variants are `Unknown`, SDK decode errors are `Corrupt` — a typed transport earns no policy exemption - *(streaming)* `WireAdapter` gains an associated `Frame` type so typed-event wires implement the same contract over their SDK events; `classify` now takes the frame by value diff --git a/crates/rig-core/src/completion/request.rs b/crates/rig-core/src/completion/request.rs index f0f20146f7..61cdc7c9c5 100644 --- a/crates/rig-core/src/completion/request.rs +++ b/crates/rig-core/src/completion/request.rs @@ -134,36 +134,61 @@ impl CompletionError { /// Whether re-issuing an equivalent request could plausibly succeed. /// - /// A **default** classification, not a policy: it answers "is this failure - /// about the request or about the moment?" and leaves the decision to - /// retry — and any backoff, budget, or idempotency handling — to the - /// caller. Callers that know their provider better should override it. + /// A conservative **floor**, not a verdict. `true` means rig can see a + /// reason to believe a retry could succeed; `false` means it cannot see + /// one — *not* that a retry is provably pointless. It is a default a + /// caller may override, and never a policy: backoff, attempt budgets, and + /// idempotency stay with the caller. /// - /// - [`HttpError`](Self::HttpError) carrying no status is a transport - /// failure: the request may never have arrived. Retryable. - /// - Any preserved provider status of `408` (request timeout), `409` - /// (conflict), `429` (too many requests), or `5xx` is retryable. These - /// are the same four classes the AI SDK treats as retryable by default. - /// - Everything else is not: a malformed request, a serialization or URL - /// error, or a response this build could not parse will fail the same - /// way every time. + /// # What it answers /// - /// A 2xx carrying a provider-authored error envelope - /// ([`ProviderResponse`](Self::ProviderResponse)) is deliberately *not* - /// retryable — the call completed and the provider rejected it on its - /// merits. + /// - A preserved provider status of `408` (request timeout), `409` + /// (conflict), `429` (too many requests), or `5xx` — the same four + /// classes the AI SDK treats as retryable by default. + /// - A transport failure with no status, when the transport itself + /// reports one that could resolve on its own: a dropped connection or a + /// cut stream. Deterministic transport failures — a header value the + /// client refused, a request that could not be constructed, a wrong + /// content type — are **not** retryable, because retrying them is how a + /// caller ends up in a loop it cannot leave. (An API key with a trailing + /// newline is the canonical example: it fails identically forever.) + /// - A 2xx carrying a provider-authored error envelope is not retryable: + /// the call completed and the provider rejected it on its merits. /// - /// Note the asymmetry this removes: rig has classified tool failures by - /// retryability for some time, while model failures had no equivalent, so - /// every caller guessed. + /// # What it cannot answer /// - /// Useful with a hand-driven run, where the caller owns the send and must - /// decide whether a failed turn is worth handing back: + /// **Providers reached over a non-HTTP transport** — AWS Bedrock, Vertex + /// AI, the gRPC Gemini client — report errors as a body with no status + /// (see `from_provider_body`). Throttling from those providers is + /// therefore classified `false` here: rig cannot tell it from a rejection + /// without parsing an envelope it does not know. A caller that knows its + /// provider should override, and a provider that *can* surface a status + /// should route through `from_http_response` instead. Classifying at the + /// provider adapter, where the transport's own error types are visible, is + /// the real fix and is tracked separately. + /// + /// **Replay safety is a different question.** This method asks whether a + /// retry *could succeed*, not whether one is *safe*. A stream that died + /// after the request was written is retryable and may already have taken + /// effect — retrying it can bill a second completion and duplicate + /// whatever the model already caused. See + /// `AgentRun::rollback_model_call` in `rig-agent` for that axis; only the + /// caller can settle it. + /// + /// # Example + /// + /// With a hand-driven run, where the caller owns the send. Note that both + /// questions are answered before the turn is handed back: /// /// ```rust,ignore /// match request.send().await { /// Ok(response) => { driver.model_response(&response)?; } - /// Err(err) if err.is_retryable() => driver.rollback_model_call()?, + /// // `nothing_was_sent` is the caller's own knowledge — a provider + /// // idempotency key, a request log, or a transport that fails + /// // before the write. `is_retryable` cannot supply it. + /// Err(err) if err.is_retryable() && nothing_was_sent => { + /// driver.rollback_model_call()? + /// } /// Err(err) => return Err(err.into()), /// } /// ``` @@ -175,9 +200,16 @@ impl CompletionError { || status == http::StatusCode::TOO_MANY_REQUESTS || status.is_server_error() } - // No status at all: a transport failure that never reached a - // response. Anything else is a local problem with the request. - None => matches!(self, Self::HttpError(_)), + None => match self { + // The transport owns this call: it knows which of its own + // failures are deterministic. + Self::HttpError(error) => error.is_transient(), + // A provider body with no status. Rig cannot tell throttling + // from rejection without knowing the provider's envelope, so + // it does not guess. Documented above as a known gap — do not + // "fix" it by defaulting to true. + _ => false, + }, } } } @@ -1362,10 +1394,49 @@ mod tests { /// A transport failure may never have reached the provider at all. #[test] - fn transport_failures_without_a_status_are_retryable() { - let error = - CompletionError::HttpError(http_client::Error::Instance("connection reset".into())); - assert!(error.is_retryable()); + fn transient_transport_failures_are_retryable() { + for error in [ + CompletionError::HttpError(http_client::Error::Instance("connection reset".into())), + CompletionError::HttpError(http_client::Error::StreamEnded), + ] { + assert!(error.is_retryable(), "{error} should be retryable"); + } + } + + /// The failure this guards against is unbounded: a driver following + /// the documented pattern hands the turn back on every retryable + /// error, so a *deterministic* failure classified retryable is a loop + /// with no exit. An API key with a trailing newline reaches + /// `bearer_auth_header` and produces exactly this. + #[test] + fn deterministic_transport_failures_are_not_retryable() { + let invalid_header = http::HeaderValue::from_str("Bearer secret\n") + .expect_err("a header value with a newline is rejected"); + + let errors = [ + CompletionError::HttpError(http_client::Error::InvalidHeaderValue(invalid_header)), + CompletionError::HttpError(http_client::Error::NoHeaders), + CompletionError::HttpError(http_client::Error::InvalidContentType( + http::HeaderValue::from_static("text/plain"), + )), + ]; + for error in errors { + assert!( + !error.is_retryable(), + "{error} is deterministic and must not be retried" + ); + } + } + + /// Documented gap: providers on non-HTTP transports (Bedrock, Vertex, + /// gRPC Gemini) report errors as a body with no status, so rig cannot + /// tell throttling from rejection. The conservative `false` is + /// deliberate — do not "fix" it by defaulting to true; fix it at the + /// provider adapter, where the transport's own error types are visible. + #[test] + fn a_statusless_provider_body_is_not_classified() { + let error = CompletionError::from_provider_body("ThrottlingException"); + assert!(!error.is_retryable()); } /// A 2xx carrying a provider-authored error envelope completed and was diff --git a/crates/rig-core/src/http_client/mod.rs b/crates/rig-core/src/http_client/mod.rs index 02e1386472..8f6c6bce35 100644 --- a/crates/rig-core/src/http_client/mod.rs +++ b/crates/rig-core/src/http_client/mod.rs @@ -52,6 +52,40 @@ impl Error { _ => None, } } + + /// Whether this transport failure could plausibly resolve on its own. + /// + /// Answers one question only — *could a retry succeed?* — and deliberately + /// not *is a retry safe?*: a stream that died after the request was + /// written is transient and may already have taken effect. Callers pairing + /// this with a retry must settle replay safety separately. + /// + /// Deterministic failures are excluded by name. A header the client + /// refused, a request that could not be constructed, or a response whose + /// content type is wrong will fail identically on every attempt, and + /// retrying them is how a caller ends up in a loop it cannot leave. The + /// remainder — the transport's own opaque failures — default to transient, + /// because rig cannot see inside them and a connection that dropped is the + /// overwhelmingly common case. + /// + /// Matched exhaustively on purpose: a variant added later must be + /// classified deliberately rather than inherit a default. + pub(crate) fn is_transient(&self) -> bool { + match self { + // The connection or the stream failed. The request may never have + // been seen — or may have been seen and its reply lost. + Self::StreamEnded | Self::Instance(_) => true, + // Deterministic: the request, the header, our use of the client, + // or the provider's content type is wrong, and will be wrong again. + Self::Protocol(_) + | Self::InvalidHeaderValue(_) + | Self::NoHeaders + | Self::InvalidContentType(_) => false, + // Carried by status classification instead; see + // `CompletionError::is_retryable`. + Self::InvalidStatusCode(_) | Self::InvalidStatusCodeWithMessage(..) => false, + } + } } pub type Result = std::result::Result; From 344cac1297e813feaf0230b9f6b6ea6ae668aad8 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 12:55:48 -0700 Subject: [PATCH 07/21] test(agent): cassette-back the AgentDriver against real provider traffic Four review rounds and 18 findings on this PR, and the unit suite grew to 513 tests without catching several of them. The reason is structural: unit tests assert against rig's own model of a request. Round 1's finding 3 - a run's tool_choice never reaching the provider, governing only the run's internal decisions - was invisible to every test that inspected `CompletionRequest`, because the code and the test shared the same wrong model of what was being sent. `tests/README.md` already names this: "a corpus written alongside an abstraction encodes the team's model of the wire and structurally cannot falsify it". The driver is precisely that kind of component - its whole job is to build a request and pair it with the state that validates the reply - and the cassette harness matches request bodies, so a request-shape regression fails as a mock miss with a diff. Five tests, recorded against live OpenAI and replaying offline: the two-turn tool loop (both request bodies), a custom run's tool_choice on the wire, suspend-mid-model-call and resume in a fresh driver, a streamed turn driven through the driver against real SSE, and `is_retryable` against a real provider rejection. Each was falsified against the finding it guards before committing - reverting the tool_choice fix produces `body_matches: false` naming the missing field, and dropping the advertised names from `commit_model_call` fails the resume test. Round 2 of this PR shipped two tests that would have passed without the code they guarded; a test that passes both ways reads as coverage and is worse than none. Writing them turned up three things. A test that could not fail: a cassette mock miss returns 404, and so does a real provider rejection, so asserting only on status would have passed against a cassette that never matched. It now pins the provider's own error envelope via `provider_response_body()`. Any provider-error cassette test anywhere has this hazard. The driver could not drive a streamed turn at all. `AgentRun`'s streamed entry points take `&mut self` and the driver exposed only `run(&self)` and a consuming `into_run()`, so a streaming caller had to rebuild the driver - discarding the per-turn snapshot cache, which makes the driver treat a turn prepared in this very process as a resume, drift check included. That is round 4's finding 3, and it blocked the whole streaming tranche. `run_mut()` closes it, documented with the one invariant it must not break: feed responses through it, never commit or roll back a model call, or the cached dispatch target falls out of step with the turn. `advertised_tools` outlived its turn on five of six paths - round 4's finding 2. Only `rollback_model_call` cleared it; `reprompt_for_output`, `retry_model_turn`, `resolve_invalid_tool_call(Retry)`, `tool_results` and both streamed-abandon paths did not, so a run serialized while parked for a fresh call reported a turn that was over, contradicting the accessor's own doc about being an audit record. Every route back to `PreparingRequest` now goes through one `park_for_new_request` helper so the invariant cannot be half-applied. Scope is deliberately honest: of the 18 findings, this shape of suite would have caught about five. Six are documentation or API-surface defects wanting `cargo doc` and public-API gates, and seven are state-machine or local-failure defects where a unit test is the right instrument - including the one still-open finding, a deterministic transport failure classified retryable, which happens before a request exists and so leaves no traffic to record. 108 openai cassette tests pass (103 pre-existing + 5), 513 rig-agent unit tests, 1241 rig-core, clippy -D warnings clean, fmt clean, 0 rustdoc warnings. No credentials or authorization headers in the recordings. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/CHANGELOG.md | 2 + crates/rig-agent/src/agent/driver.rs | 24 ++ crates/rig-agent/src/agent/run/mod.rs | 74 +++- .../agent_driver/provider_rejection.yaml | 16 + .../agent_driver/resume_awaiting_model.yaml | 16 + .../openai/agent_driver/run_tool_choice.yaml | 16 + .../openai/agent_driver/streamed_turn.yaml | 51 +++ .../agent_driver/tool_call_round_trip.yaml | 33 ++ .../providers/openai/cassette/agent_driver.rs | 368 ++++++++++++++++++ tests/providers/openai/mod.rs | 1 + 10 files changed, 594 insertions(+), 7 deletions(-) create mode 100644 tests/cassettes/openai/agent_driver/provider_rejection.yaml create mode 100644 tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml create mode 100644 tests/cassettes/openai/agent_driver/run_tool_choice.yaml create mode 100644 tests/cassettes/openai/agent_driver/streamed_turn.yaml create mode 100644 tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml create mode 100644 tests/providers/openai/cassette/agent_driver.rs diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 02138fd280..f22b55529a 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -43,6 +43,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* [**behavior**] `TurnTools::execute` now refuses any name the turn did not advertise, instead of trusting its snapshot to be narrowed. The two agreed by construction in-process; on a resumed turn the names come from the run and the snapshot is rebuilt locally, so a tool registered after suspension could previously be dispatched through a turn that never advertised it - *(agent)* [**behavior**] `TurnTools::execute_call` clears the context's dispatch result on the pre-resolved path too, so a call suppressed by invalid tool-call recovery no longer leaves the *previous* call's result metadata readable — a loop over a turn's calls was attributing it to the suppressed one - *(agent)* add `TurnToolNames::new`: the type is `#[non_exhaustive]` serialized state, so an external hand-driver could not build the value `AgentRun::commit_model_call` requires +- *(agent)* add `AgentDriver::run_mut`, without which a streamed turn could not be driven through the driver at all — `AgentRun`'s streamed entry points take `&mut self`, so a streaming caller had to `into_run()` and rebuild the driver, discarding the per-turn snapshot cache and making the driver treat a turn prepared in the same process as a resume (drift check included). Committing or rolling back a model call through it is documented as unsupported; feeding a response is safe +- *(agent)* [**behavior**] `AgentRun::advertised_tools()` no longer outlives the turn it describes. Only `rollback_model_call` ended the turn's names; `reprompt_for_output`, `retry_model_turn`, `resolve_invalid_tool_call(Retry)`, `tool_results` and both streamed-abandon paths left them in place, so a run serialized while parked for a fresh model call reported a turn that was already over. Every route back to `PreparingRequest` now clears them - *(agent)* [**behavior**] a resumed run's dispatch snapshot resolves the turn's advertised names explicitly rather than relying on retrieval to rank them again, so a registered *dynamic* tool that the re-derived query does not return is no longer misreported as "not registered in this process" — advice that could not be followed — and no longer feeds the model a not-found result for a tool that exists - *(agent)* add `AgentDriver::request_patch` / `set_request_patch`, giving a hand-driven run the per-turn configuration the runner gets from its `CompletionCall` hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history. Because the caller owns the loop, per-turn variation needs no callback - *(agent)* add `AgentRun::tool_choice()`. A custom run's own `tool_choice` now reaches the provider when the run is hand-driven, instead of only governing the run's internal decisions while the request silently carried the agent's baseline diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 1df3ae729b..448317e629 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -255,6 +255,30 @@ impl AgentDriver { &self.run } + /// Mutable access to the run, for the entry points the driver does not + /// wrap. + /// + /// A streamed turn is fed through [`AgentRun::record_streamed_completion_call`], + /// [`AgentRun::resolve_streamed_invalid_tool_call`] and + /// [`AgentRun::streamed_turn`], all of which need `&mut AgentRun`. Driving + /// a custom streaming transport is a headline use for this type, so the + /// access has to exist; without it a streaming caller would have to + /// [`Self::into_run`], drive the turn by hand, and rebuild the driver — + /// which discards the per-turn snapshot cache and makes the driver treat a + /// turn prepared in *this* process as a resume, drift check and all. + /// + /// # Do not commit or roll back a model call through this + /// + /// Use [`Self::next_step`] and [`Self::rollback_model_call`] for those. + /// They keep the driver's cached dispatch target in step with the turn the + /// run is on; committing a turn behind the driver's back would leave the + /// previous turn's snapshot cached and dispatch this turn's calls through + /// it. Feeding a *response* — streamed or otherwise — is safe, because it + /// belongs to the turn the cache already holds. + pub fn run_mut(&mut self) -> &mut AgentRun { + &mut self.run + } + /// Consume the driver, returning the run state. pub fn into_run(self) -> AgentRun { self.run diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index 6bf2744c99..c3575ad402 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -569,6 +569,19 @@ impl AgentRun { /// which reports the pending model call without committing it. fn reprompt_for_output(&mut self) { self.output_retries += 1; + self.park_for_new_request(); + } + + /// Park the run for a fresh model call, ending the current turn. + /// + /// The turn's advertised names end with it. They describe the call the + /// model was shown, and [`Self::advertised_tools`] is documented as the + /// record of *this* turn — so a run parked here, in memory or serialized, + /// must not still report the names of a turn that finished or was + /// abandoned. Every route back to `PreparingRequest` goes through here so + /// the invariant cannot be half-applied. + fn park_for_new_request(&mut self) { + self.advertised_tools = None; self.state = RunState::PreparingRequest; } @@ -731,7 +744,7 @@ impl AgentRun { } } - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(()) } @@ -954,12 +967,11 @@ impl AgentRun { self.model_call_rollbacks += 1; // The next preparation records its own; leaving stale names here would // let a later `model_response` validate against a turn that never ran. - self.advertised_tools = None; // Keep the streamed-usage window open when the record has not landed // yet: a caller that drains a broken stream for usage records it after // this point, exactly as invalid tool-call recovery does. self.rollback_pending = !self.streamed_completion_call_recorded; - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(()) } @@ -1371,7 +1383,7 @@ impl AgentRun { )); }; self.new_messages.push(user_message); - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(ModelTurnOutcome::TurnRetried) } InvalidToolCallAction::Repair { tool_name } => { @@ -1520,7 +1532,7 @@ impl AgentRun { } self.new_messages.push(Message::User { content: results }); - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(()) } @@ -1706,7 +1718,7 @@ impl AgentRun { self.new_messages.push(assistant_message); self.new_messages.push(user_message); self.rollback_pending = true; - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(StreamedResolution::TurnAbandoned { skipped_tool_result: None, }) @@ -1759,7 +1771,7 @@ impl AgentRun { self.new_messages.push(assistant_message); self.new_messages.push(user_message); self.rollback_pending = true; - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(StreamedResolution::TurnAbandoned { skipped_tool_result: Some(Box::new(skipped_tool_result)), }) @@ -2073,6 +2085,54 @@ mod tests { assert_eq!(turn, 1); } + /// `advertised_tools` is the record of what the model was shown on the + /// *current* turn, so every route back to `PreparingRequest` must end the + /// turn's names with it — not just the rollback path. A run parked for a + /// fresh call and then serialized would otherwise report a turn that is + /// over. + #[test] + fn every_route_back_to_preparing_ends_the_turns_advertised_names() { + let advertised = TurnToolNames::new(["add"], ["add"]); + + // Normal progression: tool results end the turn. + let mut run = AgentRun::new("add things").max_turns(3); + run.advance().expect("advance"); + run.peek_model_call().expect("peek"); + run.commit_model_call(Some(advertised.clone()), None) + .expect("commit"); + expect_continue( + run.model_response(tool_call_turn("call_1", "add")) + .expect("model_response"), + ); + expect_call_tools(&mut run); + assert!(run.advertised_tools().is_some(), "the turn is still live"); + run.tool_results(vec![tool_result("call_1", "2")]) + .expect("tool_results"); + assert_eq!( + run.advertised_tools(), + None, + "tool results end the turn that advertised them" + ); + + // Hook-driven turn retry. + let mut run = AgentRun::new("add things").max_turns(3); + run.advance().expect("advance"); + run.peek_model_call().expect("peek"); + run.commit_model_call(Some(advertised.clone()), None) + .expect("commit"); + expect_continue( + run.model_response(text_turn("nope")) + .expect("model_response"), + ); + run.retry_model_turn(RetryRequest::Repeat) + .expect("retry should succeed"); + assert_eq!( + run.advertised_tools(), + None, + "a retried turn's names do not outlive it" + ); + } + /// The two halves are a public protocol, so driving them out of order is /// an error a caller can handle — not a `debug_assert` that vanishes in /// release builds. diff --git a/tests/cassettes/openai/agent_driver/provider_rejection.yaml b/tests/cassettes/openai/agent_driver/provider_rejection.yaml new file mode 100644 index 0000000000..09764851d1 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/provider_rejection.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a weather assistant.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o-this-model-does-not-exist"}' +then: + status: 404 + header: + - name: content-type + value: application/json; charset=utf-8 + body: '{"error":{"code":"model_not_found","message":"The model `gpt-4o-this-model-does-not-exist` does not exist or you do not have access to it.","param":null,"type":"invalid_request_error"}}' diff --git a/tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml b/tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml new file mode 100644 index 0000000000..66f49aaa00 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":63,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":76}}' diff --git a/tests/cassettes/openai/agent_driver/run_tool_choice.yaml b/tests/cassettes/openai/agent_driver/run_tool_choice.yaml new file mode 100644 index 0000000000..f0949fd486 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/run_tool_choice.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a weather assistant.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","tool_choice":"required","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":57,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":70}}' diff --git a/tests/cassettes/openai/agent_driver/streamed_turn.yaml b/tests/cassettes/openai/agent_driver/streamed_turn.yaml new file mode 100644 index 0000000000..c0e34bcb75 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/streamed_turn.yaml @@ -0,0 +1,51 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: text/event-stream + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true},"tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: text/event-stream; charset=utf-8 + body: |+ + data: {"choices":[{"delta":{"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"","name":"weather"},"id":"call_REDACTED_1","index":0,"type":"function"}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_1","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"{\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_2","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"city"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_3","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_4","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"Tokyo"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_5","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\"}"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_6","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{},"finish_reason":"tool_calls","index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_7","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_8","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":63,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":76}} + + data: [DONE] + +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]},{"content":"The weather in Tokyo is 22C and sunny","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"The current weather in Tokyo is 22°C and sunny.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":93,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":106}}' diff --git a/tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml b/tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml new file mode 100644 index 0000000000..d401116bc3 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":63,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":76}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]},{"content":"The weather in Tokyo is 22C and sunny","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"The current weather in Tokyo is 22°C and sunny.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":93,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":106}}' diff --git a/tests/providers/openai/cassette/agent_driver.rs b/tests/providers/openai/cassette/agent_driver.rs new file mode 100644 index 0000000000..05d9efec06 --- /dev/null +++ b/tests/providers/openai/cassette/agent_driver.rs @@ -0,0 +1,368 @@ +//! Cassette coverage for `AgentDriver` against real OpenAI traffic. +//! +//! The driver's job is to build a request and pair it with the tool state that +//! validates and dispatches the model's reply. Unit tests can check the state +//! machine, but only recorded traffic can check the *request* — and every +//! request-shape defect found in this API's review so far (a run's +//! `tool_choice` never reaching the wire, a resumed turn advertising the wrong +//! tool set) was invisible to unit tests precisely because they asserted on +//! rig's own view of the request rather than on the bytes. +//! +//! The cassette harness matches each request body against the recorded one, so +//! a request-shape regression fails as a mock miss. That is the assertion these +//! tests exist for; the response-side asserts are secondary. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use futures::StreamExt; +use rig::agent::run::StreamedTurnAssembler; +use rig::agent::{AgentRun, DriveStep}; +use rig::message::ToolChoice; +use rig::prelude::*; +use rig::providers::openai; +use rig::tool::{Tool, ToolContext}; + +use super::super::support::with_openai_completions_cassette; + +#[derive(Debug, serde::Deserialize)] +struct WeatherArgs { + city: String, +} + +#[derive(Clone)] +struct WeatherTool { + call_count: Arc, +} + +impl Tool for WeatherTool { + const NAME: &'static str = "weather"; + type Error = std::io::Error; + type Args = WeatherArgs; + type Output = String; + + fn description(&self) -> String { + "Get the current weather for a city.".to_string() + } + + fn parameters(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"] + }) + } + + fn call( + &self, + _context: &mut ToolContext, + args: Self::Args, + ) -> impl std::future::Future> + Send { + self.call_count.fetch_add(1, Ordering::SeqCst); + std::future::ready(Ok(format!("The weather in {} is 22C and sunny", args.city))) + } +} + +/// The whole hand-driven loop over real traffic: the driver builds both +/// requests, the caller sends them, and the turn that advertised the tool is +/// the turn that dispatches it. +/// +/// Locks the request bodies for both turns. The second request in particular +/// carries the assistant tool call and the tool result the driver threaded back +/// through the run — a shape no unit test observes. +#[tokio::test] +async fn drive_loop_round_trips_a_tool_call() { + with_openai_completions_cassette("agent_driver/tool_call_round_trip", |client| async move { + let calls = Arc::new(AtomicUsize::new(0)); + let agent = client + .agent(openai::GPT_4O) + .preamble("You are a weather assistant. Always use the weather tool.") + .default_max_turns(3) + .tool(WeatherTool { + call_count: calls.clone(), + }) + .build(); + + let mut driver = agent.drive("What is the weather in Tokyo?"); + + let (request, tools, turn) = match driver.next_step().await.expect("first step") { + DriveStep::SendRequest { + request, + tools, + turn, + } => (request, tools, turn), + other => panic!("expected SendRequest, got {other:?}"), + }; + assert_eq!(turn, 1); + assert!(tools.executable_tool_names().contains("weather")); + + let response = request.send().await.expect("first turn should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = match driver.next_step().await.expect("second step") { + DriveStep::ExecuteTools { calls, tools } => (calls, tools), + other => panic!("expected ExecuteTools, got {other:?}"), + }; + assert!(!pending.is_empty(), "the model should have called the tool"); + + let mut context = ToolContext::new(); + let mut results = Vec::new(); + for call in &pending { + results.push(tools.execute_call(call, &mut context).await); + } + driver.tool_results(results).expect("results accepted"); + + let request = match driver.next_step().await.expect("third step") { + DriveStep::SendRequest { request, turn, .. } => { + assert_eq!(turn, 2); + request + } + other => panic!("expected SendRequest, got {other:?}"), + }; + let response = request.send().await.expect("second turn should send"); + driver.model_response(&response).expect("turn accepted"); + + match driver.next_step().await.expect("final step") { + DriveStep::Done(response) => { + assert!( + !response.output.trim().is_empty(), + "expected a final answer" + ); + } + other => panic!("expected Done, got {other:?}"), + } + assert_eq!(calls.load(Ordering::SeqCst), 1, "the tool ran exactly once"); + }) + .await; +} + +/// A custom run is taken as-is, so its own `tool_choice` must reach the +/// provider — not merely the run's internal decisions. +/// +/// This is the finding a unit test could assert only against rig's own +/// `CompletionRequest`. Here the recorded body carries `"tool_choice"`, so a +/// driver that drops it fails as a mock miss: the request rig builds no longer +/// matches the request the provider was actually asked. +#[tokio::test] +async fn a_custom_runs_tool_choice_reaches_the_provider() { + with_openai_completions_cassette("agent_driver/run_tool_choice", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("You are a weather assistant.") + .tool(WeatherTool { + call_count: Arc::new(AtomicUsize::new(0)), + }) + .build(); + + // The choice lives on the run, not on the agent. + let run = AgentRun::new("What is the weather in Tokyo?") + .max_turns(2) + .with_tool_choice(ToolChoice::Required); + let mut driver = agent.drive_run(run); + + let request = match driver.next_step().await.expect("first step") { + DriveStep::SendRequest { request, .. } => request, + other => panic!("expected SendRequest, got {other:?}"), + }; + + let response = request.send().await.expect("request should send"); + driver.model_response(&response).expect("turn accepted"); + + // `Required` obliges a tool call, so the run must reach tools. + match driver.next_step().await.expect("second step") { + DriveStep::ExecuteTools { calls, .. } => { + assert!(!calls.is_empty(), "tool_choice=required must force a call"); + } + other => panic!("expected ExecuteTools, got {other:?}"), + } + }) + .await; +} + +/// A run suspended with a model call in flight resumes in a different driver +/// and accepts the reply. +/// +/// The request is recorded once; the resumed driver never builds one. What is +/// under test is that the turn's advertised tool names travelled with the run, +/// so the reply is validated against the set the recorded request actually +/// carried. +#[tokio::test] +async fn a_run_suspended_awaiting_the_model_resumes_and_accepts_the_reply() { + with_openai_completions_cassette("agent_driver/resume_awaiting_model", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("You are a weather assistant. Always use the weather tool.") + .default_max_turns(2) + .tool(WeatherTool { + call_count: Arc::new(AtomicUsize::new(0)), + }) + .build(); + + let mut driver = agent.drive("What is the weather in Tokyo?"); + let request = match driver.next_step().await.expect("first step") { + DriveStep::SendRequest { request, .. } => request, + other => panic!("expected SendRequest, got {other:?}"), + }; + + // Suspend with the call in flight, then send. + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + let response = request.send().await.expect("request should send"); + drop(driver); + + let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + assert!( + restored.advertised_tools().is_some(), + "the suspended run carries the turn's advertised names" + ); + let mut resumed = agent.drive_run(restored); + resumed + .model_response(&response) + .expect("a resumed run accepts the reply to its in-flight call"); + + match resumed.next_step().await.expect("second step") { + DriveStep::ExecuteTools { calls, tools } => { + assert!(!calls.is_empty()); + assert!(tools.executable_tool_names().contains("weather")); + } + other => panic!("expected ExecuteTools, got {other:?}"), + } + }) + .await; +} + +/// A hand-driven **streamed** turn goes through the driver, against real SSE +/// traffic. +/// +/// The streamed entry points live on `AgentRun` and take `&mut self`, so this +/// is only expressible because the driver hands out `run_mut()`. Without it a +/// streaming caller has to `into_run()` and rebuild the driver, which discards +/// the per-turn snapshot cache and makes the driver treat a turn prepared in +/// this very process as a resume — drift check included. The assertion that +/// matters here is the second request body: the streamed turn must thread back +/// into the run exactly as a blocking one does. +#[tokio::test] +async fn a_streamed_turn_drives_through_the_driver() { + with_openai_completions_cassette("agent_driver/streamed_turn", |client| async move { + let calls = Arc::new(AtomicUsize::new(0)); + let agent = client + .agent(openai::GPT_4O) + .preamble("You are a weather assistant. Always use the weather tool.") + .default_max_turns(3) + .tool(WeatherTool { + call_count: calls.clone(), + }) + .build(); + + let mut driver = agent.drive("What is the weather in Tokyo?"); + + let (request, tools) = match driver.next_step().await.expect("first step") { + DriveStep::SendRequest { request, tools, .. } => (request, tools), + other => panic!("expected SendRequest, got {other:?}"), + }; + + // Stream the turn and assemble it with the names this turn advertised. + let mut assembler = StreamedTurnAssembler::new( + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), + ); + let mut stream = request.stream().await.expect("stream should open"); + while let Some(item) = stream.next().await { + let item = item.expect("stream item"); + assembler.ingest(&item).expect("ingest should succeed"); + } + let final_content = stream.choice.clone(); + let streamed = assembler.finish(stream.message_id.clone(), &final_content); + + // Feed usage and the assembled turn through the driver's run. + driver + .run_mut() + .record_streamed_completion_call(stream.usage()) + .expect("usage recorded"); + driver + .run_mut() + .streamed_turn(streamed) + .expect("streamed turn accepted"); + + // The driver still owns the pairing: this turn's snapshot dispatches. + let (pending, tools) = match driver.next_step().await.expect("second step") { + DriveStep::ExecuteTools { calls, tools } => (calls, tools), + other => panic!("expected ExecuteTools, got {other:?}"), + }; + assert!(!pending.is_empty(), "the model should have called the tool"); + + let mut context = ToolContext::new(); + let mut results = Vec::new(); + for call in &pending { + results.push(tools.execute_call(call, &mut context).await); + } + driver.tool_results(results).expect("results accepted"); + + let request = match driver.next_step().await.expect("third step") { + DriveStep::SendRequest { request, .. } => request, + other => panic!("expected SendRequest, got {other:?}"), + }; + let response = request.send().await.expect("final turn should send"); + driver.model_response(&response).expect("turn accepted"); + + match driver.next_step().await.expect("final step") { + DriveStep::Done(response) => { + assert!(!response.output.trim().is_empty()); + } + other => panic!("expected Done, got {other:?}"), + } + assert_eq!(calls.load(Ordering::SeqCst), 1); + }) + .await; +} + +/// A real provider rejection is classified non-retryable, so the caller does +/// not hand the turn back and loop. +/// +/// `is_retryable` is the trigger the driver's own docs recommend for +/// `rollback_model_call`, and misclassifying a deterministic failure as +/// retryable is an unbounded loop. A synthetic error cannot falsify the +/// classification — only a real provider rejection can, which is what this +/// records. +#[tokio::test] +async fn a_provider_rejection_is_not_retryable() { + with_openai_completions_cassette("agent_driver/provider_rejection", |client| async move { + let agent = client + // A model name the provider will reject outright. + .agent("gpt-4o-this-model-does-not-exist") + .preamble("You are a weather assistant.") + .build(); + + let mut driver = agent.drive("What is the weather in Tokyo?"); + let request = match driver.next_step().await.expect("first step") { + DriveStep::SendRequest { request, .. } => request, + other => panic!("expected SendRequest, got {other:?}"), + }; + + let error = request + .send() + .await + .expect_err("an unknown model must be rejected"); + + // A cassette mock miss is *also* a 404, so asserting on the status + // alone would pass on a cassette that never matched. Pin the + // provider's own error envelope, which the harness cannot fabricate. + let body = error + .provider_response_body() + .expect("the provider's rejection body is preserved"); + assert!( + body.contains("model_not_found"), + "expected the recorded provider rejection, got: {body}" + ); + + assert!( + !error.is_retryable(), + "a provider rejection must not be classified retryable: {error}" + ); + + // The turn stays in flight: the caller decides, and here the decision + // is to fail rather than hand the turn back. + assert_eq!(driver.run().turn(), 1); + assert_eq!(driver.run().model_call_rollbacks(), 0); + }) + .await; +} diff --git a/tests/providers/openai/mod.rs b/tests/providers/openai/mod.rs index 57f38aca6a..ab9e108a12 100644 --- a/tests/providers/openai/mod.rs +++ b/tests/providers/openai/mod.rs @@ -4,6 +4,7 @@ mod regressions; mod cassette { mod agent; + mod agent_driver; mod chat_history; mod completions_api; mod document_ordering; From ccef387c34b4c7fdbea7189b4b43e2810d375a84 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 13:28:26 -0700 Subject: [PATCH 08/21] test(agent): scale the AgentDriver cassette suite to three providers 53 cassettes, 53 tests, ~60 live model calls: 25 OpenAI, 14 Anthropic, 14 Gemini. Shared driving helpers in tests/common/driver_support.rs, so each provider module carries only the client, the model id, and the wire-shape assertions. Coverage: the tool loop and parallel calls; every per-run and per-turn configuration path that reaches a request (all three tool_choice modes, and RequestPatch's preamble, tool_choice, active_tools, sampling params, extra_context and history); both resume points; rollback and re-preparation after a rejected send; a preparation failure costing no turn and no interaction; max_turns exhaustion; both output modes; and three streaming shapes. Three providers is not redundancy. The driver is provider-agnostic and the request it builds is not - tool_choice, tool declarations and the system prompt are spelled differently by each - so a per-turn patch that reaches an OpenAI request can silently fail to reach an Anthropic one. Falsification confirms it: with the run-tool_choice fix reverted, all three providers' tests fail, three independent wire encodings producing three independent mock misses. Recording against real providers found two test defects that no amount of local reasoning would have: Anthropic resolves a per-model default max_tokens, and an *unknown* model has no default - so the provider-rejection test, which used a bogus model name deliberately, failed locally with `max_tokens must be set` and never reached the wire. It was asserting on a local failure while claiming to assert on a provider one. (Incidentally a small vindication of the classification work: that local RequestError is correctly not retryable.) And a model forbidden from its only useful tool returns nothing at all. The ToolChoice::None test asserted a non-empty answer; against Anthropic the model returned content [] - honest behavior that is_empty_assistant_turn handles - because the prompt was arithmetic and the tool that would answer it had just been forbidden. The same test passed on OpenAI, whose model chose prose instead. Rewritten to assert what it exists for: the run finalizes with no tool call surviving the constraint. Exactly the doctrine's "assert on structure, not wording". 128 openai / 85 anthropic / 141 gemini cassette tests pass, 513 rig-agent unit tests, 1241 rig-core, clippy -D warnings clean, fmt clean, 0 rustdoc warnings. Secret scan clean across all 53 recordings; no authorization or x-api-key headers captured. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- tests/anthropic.rs | 2 + .../agent_driver/driver_history.yaml | 18 + .../anthropic/agent_driver/max_turns.yaml | 18 + .../agent_driver/output_mode_tool.yaml | 18 + .../agent_driver/patch_active_tools.yaml | 18 + .../agent_driver/patch_preamble.yaml | 18 + .../agent_driver/patch_tool_choice.yaml | 18 + .../agent_driver/provider_rejection.yaml | 18 + .../agent_driver/resume_awaiting_model.yaml | 18 + .../agent_driver/resume_executing_tools.yaml | 37 + .../agent_driver/run_tool_choice.yaml | 18 + .../anthropic/agent_driver/streamed_turn.yaml | 48 + .../agent_driver/tool_call_round_trip.yaml | 37 + .../agent_driver/tool_choice_none.yaml | 18 + .../agent_driver/tool_choice_specific.yaml | 18 + .../gemini/agent_driver/driver_history.yaml | 18 + .../gemini/agent_driver/max_turns.yaml | 18 + .../gemini/agent_driver/output_mode_tool.yaml | 18 + .../agent_driver/patch_active_tools.yaml | 18 + .../gemini/agent_driver/patch_preamble.yaml | 18 + .../agent_driver/patch_tool_choice.yaml | 18 + .../agent_driver/provider_rejection.yaml | 18 + .../agent_driver/resume_awaiting_model.yaml | 18 + .../agent_driver/resume_executing_tools.yaml | 37 + .../gemini/agent_driver/run_tool_choice.yaml | 18 + .../gemini/agent_driver/streamed_turn.yaml | 20 + .../agent_driver/tool_call_round_trip.yaml | 37 + .../gemini/agent_driver/tool_choice_none.yaml | 18 + .../agent_driver/tool_choice_specific.yaml | 18 + .../openai/agent_driver/driver_history.yaml | 16 + .../openai/agent_driver/max_turns.yaml | 16 + .../agent_driver/output_mode_native.yaml | 16 + .../openai/agent_driver/output_mode_tool.yaml | 16 + .../agent_driver/parallel_tool_calls.yaml | 33 + .../agent_driver/patch_active_tools.yaml | 33 + .../agent_driver/patch_extra_context.yaml | 16 + .../openai/agent_driver/patch_history.yaml | 16 + .../openai/agent_driver/patch_preamble.yaml | 16 + .../openai/agent_driver/patch_sampling.yaml | 16 + .../agent_driver/patch_tool_choice.yaml | 16 + .../openai/agent_driver/prepare_failure.yaml | 16 + .../agent_driver/provider_rejection.yaml | 2 +- .../agent_driver/resume_awaiting_model.yaml | 4 +- .../agent_driver/resume_executing_tools.yaml | 33 + .../agent_driver/resume_no_tool_leak.yaml | 16 + .../agent_driver/rollback_re_prepares.yaml | 33 + .../openai/agent_driver/run_tool_choice.yaml | 4 +- .../openai/agent_driver/streamed_patched.yaml | 42 + .../openai/agent_driver/streamed_text.yaml | 28 + .../openai/agent_driver/streamed_turn.yaml | 28 +- .../agent_driver/tool_call_round_trip.yaml | 8 +- .../openai/agent_driver/tool_choice_none.yaml | 16 + .../agent_driver/tool_choice_specific.yaml | 16 + .../agent_driver/two_tools_advertised.yaml | 33 + tests/common/driver_support.rs | 106 ++ tests/gemini.rs | 2 + tests/openai.rs | 2 + .../anthropic/cassette/agent_driver.rs | 468 ++++++++ tests/providers/anthropic/mod.rs | 1 + .../providers/gemini/cassette/agent_driver.rs | 464 ++++++++ tests/providers/gemini/mod.rs | 1 + .../providers/openai/cassette/agent_driver.rs | 1036 ++++++++++++++--- 62 files changed, 2973 insertions(+), 210 deletions(-) create mode 100644 tests/cassettes/anthropic/agent_driver/driver_history.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/max_turns.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/output_mode_tool.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/patch_active_tools.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/patch_preamble.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/patch_tool_choice.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/provider_rejection.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/resume_awaiting_model.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/resume_executing_tools.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/run_tool_choice.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/streamed_turn.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/tool_call_round_trip.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/tool_choice_none.yaml create mode 100644 tests/cassettes/anthropic/agent_driver/tool_choice_specific.yaml create mode 100644 tests/cassettes/gemini/agent_driver/driver_history.yaml create mode 100644 tests/cassettes/gemini/agent_driver/max_turns.yaml create mode 100644 tests/cassettes/gemini/agent_driver/output_mode_tool.yaml create mode 100644 tests/cassettes/gemini/agent_driver/patch_active_tools.yaml create mode 100644 tests/cassettes/gemini/agent_driver/patch_preamble.yaml create mode 100644 tests/cassettes/gemini/agent_driver/patch_tool_choice.yaml create mode 100644 tests/cassettes/gemini/agent_driver/provider_rejection.yaml create mode 100644 tests/cassettes/gemini/agent_driver/resume_awaiting_model.yaml create mode 100644 tests/cassettes/gemini/agent_driver/resume_executing_tools.yaml create mode 100644 tests/cassettes/gemini/agent_driver/run_tool_choice.yaml create mode 100644 tests/cassettes/gemini/agent_driver/streamed_turn.yaml create mode 100644 tests/cassettes/gemini/agent_driver/tool_call_round_trip.yaml create mode 100644 tests/cassettes/gemini/agent_driver/tool_choice_none.yaml create mode 100644 tests/cassettes/gemini/agent_driver/tool_choice_specific.yaml create mode 100644 tests/cassettes/openai/agent_driver/driver_history.yaml create mode 100644 tests/cassettes/openai/agent_driver/max_turns.yaml create mode 100644 tests/cassettes/openai/agent_driver/output_mode_native.yaml create mode 100644 tests/cassettes/openai/agent_driver/output_mode_tool.yaml create mode 100644 tests/cassettes/openai/agent_driver/parallel_tool_calls.yaml create mode 100644 tests/cassettes/openai/agent_driver/patch_active_tools.yaml create mode 100644 tests/cassettes/openai/agent_driver/patch_extra_context.yaml create mode 100644 tests/cassettes/openai/agent_driver/patch_history.yaml create mode 100644 tests/cassettes/openai/agent_driver/patch_preamble.yaml create mode 100644 tests/cassettes/openai/agent_driver/patch_sampling.yaml create mode 100644 tests/cassettes/openai/agent_driver/patch_tool_choice.yaml create mode 100644 tests/cassettes/openai/agent_driver/prepare_failure.yaml create mode 100644 tests/cassettes/openai/agent_driver/resume_executing_tools.yaml create mode 100644 tests/cassettes/openai/agent_driver/resume_no_tool_leak.yaml create mode 100644 tests/cassettes/openai/agent_driver/rollback_re_prepares.yaml create mode 100644 tests/cassettes/openai/agent_driver/streamed_patched.yaml create mode 100644 tests/cassettes/openai/agent_driver/streamed_text.yaml create mode 100644 tests/cassettes/openai/agent_driver/tool_choice_none.yaml create mode 100644 tests/cassettes/openai/agent_driver/tool_choice_specific.yaml create mode 100644 tests/cassettes/openai/agent_driver/two_tools_advertised.yaml create mode 100644 tests/common/driver_support.rs create mode 100644 tests/providers/anthropic/cassette/agent_driver.rs create mode 100644 tests/providers/gemini/cassette/agent_driver.rs diff --git a/tests/anthropic.rs b/tests/anthropic.rs index 6a9c6e1706..5a1a540de3 100644 --- a/tests/anthropic.rs +++ b/tests/anthropic.rs @@ -10,6 +10,8 @@ mod cassette_safety; #[path = "common/cassettes.rs"] mod cassettes; +#[path = "common/driver_support.rs"] +mod driver_support; #[path = "common/reasoning.rs"] mod reasoning; #[path = "common/support.rs"] diff --git a/tests/cassettes/anthropic/agent_driver/driver_history.yaml b/tests/cassettes/anthropic/agent_driver/driver_history.yaml new file mode 100644 index 0000000000..7a69b1820d --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/driver_history.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"My name is Ada.","type":"text"}],"role":"user"},{"content":[{"text":"Nice to meet you, Ada.","type":"text"}],"role":"assistant"},{"content":[{"text":"What is my name?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"Answer briefly.","type":"text"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"text":"Your name is Ada.","type":"text"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"end_turn","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":34,"output_tokens":8,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/max_turns.yaml b/tests/cassettes/anthropic/agent_driver/max_turns.yaml new file mode 100644 index 0000000000..52f230d586 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/max_turns.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":638,"output_tokens":69,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/output_mode_tool.yaml b/tests/cassettes/anthropic/agent_driver/output_mode_tool.yaml new file mode 100644 index 0000000000..ce6b44947c --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/output_mode_tool.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is the capital of France?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"Reply with the structured answer.\n\nWhen you have gathered enough information to answer, call the `final_result` tool exactly once with your final answer. Its arguments are the structured result and must satisfy the required schema. Do not return the final answer as plain text.","type":"text"}],"tools":[{"description":"Call this tool exactly once with your final answer when you are done. Its arguments are the structured result and must satisfy the output schema.","input_schema":{"properties":{"answer":{"type":"string"}},"required":["answer"],"type":"object"},"name":"final_result"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"answer":"The capital of France is Paris."},"name":"final_result","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":640,"output_tokens":59,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/patch_active_tools.yaml b/tests/cassettes/anthropic/agent_driver/patch_active_tools.yaml new file mode 100644 index 0000000000..52f230d586 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/patch_active_tools.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":638,"output_tokens":69,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/patch_preamble.yaml b/tests/cassettes/anthropic/agent_driver/patch_preamble.yaml new file mode 100644 index 0000000000..f06cb9fdd3 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/patch_preamble.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"Say the word banana.","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"PATCHED PREAMBLE — reply with one word.","type":"text"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"text":"Banana.","type":"text"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"end_turn","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":26,"output_tokens":6,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/patch_tool_choice.yaml b/tests/cassettes/anthropic/agent_driver/patch_tool_choice.yaml new file mode 100644 index 0000000000..d660ece459 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/patch_tool_choice.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tool_choice":{"type":"any"},"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":730,"output_tokens":53,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/provider_rejection.yaml b/tests/cassettes/anthropic/agent_driver/provider_rejection.yaml new file mode 100644 index 0000000000..dbe291bbe2 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/provider_rejection.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-this-model-does-not-exist","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}]}' +then: + status: 404 + header: + - name: content-type + value: application/json + body: '{"error":{"message":"model: claude-this-model-does-not-exist","type":"not_found_error"},"request_id":"req_REDACTED_1","type":"error"}' diff --git a/tests/cassettes/anthropic/agent_driver/resume_awaiting_model.yaml b/tests/cassettes/anthropic/agent_driver/resume_awaiting_model.yaml new file mode 100644 index 0000000000..52f230d586 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/resume_awaiting_model.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":638,"output_tokens":69,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/resume_executing_tools.yaml b/tests/cassettes/anthropic/agent_driver/resume_executing_tools.yaml new file mode 100644 index 0000000000..0b650cead4 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/resume_executing_tools.yaml @@ -0,0 +1,37 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":638,"output_tokens":69,"service_tier":"standard"}}' +--- +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"},{"content":[{"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"role":"assistant"},{"content":[{"content":[{"text":"7","type":"text"}],"tool_use_id":"toolu_REDACTED_1","type":"tool_result"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"text":"2 + 5 = **7**","type":"text"}],"id":"msg_REDACTED_2","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"end_turn","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":720,"output_tokens":13,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/run_tool_choice.yaml b/tests/cassettes/anthropic/agent_driver/run_tool_choice.yaml new file mode 100644 index 0000000000..d660ece459 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/run_tool_choice.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tool_choice":{"type":"any"},"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":730,"output_tokens":53,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/streamed_turn.yaml b/tests/cassettes/anthropic/agent_driver/streamed_turn.yaml new file mode 100644 index 0000000000..5b05465063 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/streamed_turn.yaml @@ -0,0 +1,48 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: text/event-stream + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","stream":true,"system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tool_choice":{"type":"auto"},"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: text/event-stream; charset=utf-8 + body: |+ + event: message_start + data: {"message":{"content":[],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":null,"stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":638,"output_tokens":22,"service_tier":"standard"}},"type":"message_start"} + + event: content_block_start + data: {"content_block":{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{},"name":"add","type":"tool_use"},"index":0,"type":"content_block_start"} + + event: ping + data: {"type":"ping"} + + event: content_block_delta + data: {"delta":{"partial_json":"","type":"input_json_delta"},"index":0,"type":"content_block_delta"} + + event: content_block_delta + data: {"delta":{"partial_json":"{\"x\": 2","type":"input_json_delta"},"index":0,"type":"content_block_delta"} + + event: content_block_delta + data: {"delta":{"partial_json":", \"y\":","type":"input_json_delta"},"index":0,"type":"content_block_delta"} + + event: content_block_delta + data: {"delta":{"partial_json":" 5}","type":"input_json_delta"},"index":0,"type":"content_block_delta"} + + event: content_block_stop + data: {"index":0,"type":"content_block_stop"} + + event: message_delta + data: {"delta":{"stop_details":null,"stop_reason":"tool_use","stop_sequence":null},"type":"message_delta","usage":{"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"input_tokens":638,"output_tokens":69}} + + event: message_stop + data: {"type":"message_stop"} + diff --git a/tests/cassettes/anthropic/agent_driver/tool_call_round_trip.yaml b/tests/cassettes/anthropic/agent_driver/tool_call_round_trip.yaml new file mode 100644 index 0000000000..0b650cead4 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/tool_call_round_trip.yaml @@ -0,0 +1,37 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":638,"output_tokens":69,"service_tier":"standard"}}' +--- +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"},{"content":[{"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"role":"assistant"},{"content":[{"content":[{"text":"7","type":"text"}],"tool_use_id":"toolu_REDACTED_1","type":"tool_result"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"text":"2 + 5 = **7**","type":"text"}],"id":"msg_REDACTED_2","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"end_turn","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":720,"output_tokens":13,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/tool_choice_none.yaml b/tests/cassettes/anthropic/agent_driver/tool_choice_none.yaml new file mode 100644 index 0000000000..9a2143431e --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/tool_choice_none.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"Answer in plain text.","type":"text"}],"tool_choice":{"type":"none"},"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"end_turn","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":603,"output_tokens":9,"service_tier":"standard"}}' diff --git a/tests/cassettes/anthropic/agent_driver/tool_choice_specific.yaml b/tests/cassettes/anthropic/agent_driver/tool_choice_specific.yaml new file mode 100644 index 0000000000..9dabc6d487 --- /dev/null +++ b/tests/cassettes/anthropic/agent_driver/tool_choice_specific.yaml @@ -0,0 +1,18 @@ +when: + path: /v1/messages + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + - name: anthropic-version + value: 2023-06-01 + body: '{"max_tokens":64000,"messages":[{"content":[{"text":"What is 2 + 5?","type":"text"}],"role":"user"}],"model":"claude-sonnet-4-6","system":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"tool_choice":{"name":"add","type":"tool"},"tools":[{"description":"Add x and y together","input_schema":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"},"name":"add"},{"description":"Subtract y from x (i.e.: x - y)","input_schema":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"},"name":"subtract"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"content":[{"caller":{"type":"direct"},"id":"toolu_REDACTED_1","input":{"x":2,"y":5},"name":"add","type":"tool_use"}],"id":"msg_REDACTED_1","model":"claude-sonnet-4-6","role":"assistant","stop_details":null,"stop_reason":"tool_use","stop_sequence":null,"type":"message","usage":{"cache_creation":{"ephemeral_1h_input_tokens":0,"ephemeral_5m_input_tokens":0},"cache_creation_input_tokens":0,"cache_read_input_tokens":0,"inference_geo":"global","input_tokens":824,"output_tokens":50,"service_tier":"standard"}}' diff --git a/tests/cassettes/gemini/agent_driver/driver_history.yaml b/tests/cassettes/gemini/agent_driver/driver_history.yaml new file mode 100644 index 0000000000..a581763ab1 --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/driver_history.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"My name is Ada.","thought":false}],"role":"user"},{"parts":[{"text":"Nice to meet you, Ada.","thought":false}],"role":"model"},{"parts":[{"text":"What is my name?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"Answer briefly.","thought":false}],"role":"model"},"toolConfig":null}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"text":"Your name is Ada."}],"role":"model"},"finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":5,"promptTokenCount":24,"promptTokensDetails":[{"modality":"TEXT","tokenCount":24}],"serviceTier":"standard","thoughtsTokenCount":35,"totalTokenCount":64}}' diff --git a/tests/cassettes/gemini/agent_driver/max_turns.yaml b/tests/cassettes/gemini/agent_driver/max_turns.yaml new file mode 100644 index 0000000000..6251bedd9d --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/max_turns.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":103,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"serviceTier":"standard","thoughtsTokenCount":47,"totalTokenCount":168}}' diff --git a/tests/cassettes/gemini/agent_driver/output_mode_tool.yaml b/tests/cassettes/gemini/agent_driver/output_mode_tool.yaml new file mode 100644 index 0000000000..8b215a907b --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/output_mode_tool.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is the capital of France?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"Reply with the structured answer.\n\nWhen you have gathered enough information to answer, call the `final_result` tool exactly once with your final answer. Its arguments are the structured result and must satisfy the required schema. Do not return the final answer as plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Call this tool exactly once with your final answer when you are done. Its arguments are the structured result and must satisfy the output schema.","name":"final_result","parameters":{"properties":{"answer":{"type":"string"}},"required":["answer"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"answer":"Paris"},"name":"final_result"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":15,"promptTokenCount":122,"promptTokensDetails":[{"modality":"TEXT","tokenCount":122}],"serviceTier":"standard","thoughtsTokenCount":69,"totalTokenCount":206}}' diff --git a/tests/cassettes/gemini/agent_driver/patch_active_tools.yaml b/tests/cassettes/gemini/agent_driver/patch_active_tools.yaml new file mode 100644 index 0000000000..aeb5f633a8 --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/patch_active_tools.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":103,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"serviceTier":"standard","thoughtsTokenCount":54,"totalTokenCount":175}}' diff --git a/tests/cassettes/gemini/agent_driver/patch_preamble.yaml b/tests/cassettes/gemini/agent_driver/patch_preamble.yaml new file mode 100644 index 0000000000..7156fd4d42 --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/patch_preamble.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"Say the word banana.","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"PATCHED PREAMBLE — reply with one word.","thought":false}],"role":"model"},"toolConfig":null}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"text":"Banana"}],"role":"model"},"finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":1,"promptTokenCount":18,"promptTokensDetails":[{"modality":"TEXT","tokenCount":18}],"serviceTier":"standard","thoughtsTokenCount":342,"totalTokenCount":361}}' diff --git a/tests/cassettes/gemini/agent_driver/patch_tool_choice.yaml b/tests/cassettes/gemini/agent_driver/patch_tool_choice.yaml new file mode 100644 index 0000000000..85d51f589b --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/patch_tool_choice.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":{"functionCallingConfig":{"mode":"ANY"}},"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":103,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"serviceTier":"standard","thoughtsTokenCount":61,"totalTokenCount":182}}' diff --git a/tests/cassettes/gemini/agent_driver/provider_rejection.yaml b/tests/cassettes/gemini/agent_driver/provider_rejection.yaml new file mode 100644 index 0000000000..82ea577c78 --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/provider_rejection.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-this-model-does-not-exist:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null}' +then: + status: 404 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"error":{"code":404,"message":"models/gemini-this-model-does-not-exist is not found for API version v1beta, or is not supported for generateContent. Call ModelService.ListModels to see the list of available models and their supported methods.","status":"NOT_FOUND"}}' diff --git a/tests/cassettes/gemini/agent_driver/resume_awaiting_model.yaml b/tests/cassettes/gemini/agent_driver/resume_awaiting_model.yaml new file mode 100644 index 0000000000..aeb5f633a8 --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/resume_awaiting_model.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":103,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"serviceTier":"standard","thoughtsTokenCount":54,"totalTokenCount":175}}' diff --git a/tests/cassettes/gemini/agent_driver/resume_executing_tools.yaml b/tests/cassettes/gemini/agent_driver/resume_executing_tools.yaml new file mode 100644 index 0000000000..4f6a0896ab --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/resume_executing_tools.yaml @@ -0,0 +1,37 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":103,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"serviceTier":"standard","thoughtsTokenCount":79,"totalTokenCount":200}}' +--- +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"},{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thought":false,"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},{"parts":[{"functionResponse":{"name":"add","response":{"result":7}},"thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"text":"7"}],"role":"model"},"finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_2","usageMetadata":{"candidatesTokenCount":1,"promptTokenCount":134,"promptTokensDetails":[{"modality":"TEXT","tokenCount":134}],"serviceTier":"standard","totalTokenCount":135}}' diff --git a/tests/cassettes/gemini/agent_driver/run_tool_choice.yaml b/tests/cassettes/gemini/agent_driver/run_tool_choice.yaml new file mode 100644 index 0000000000..3c1c29722b --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/run_tool_choice.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":{"functionCallingConfig":{"mode":"ANY"}},"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":103,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"serviceTier":"standard","thoughtsTokenCount":47,"totalTokenCount":168}}' diff --git a/tests/cassettes/gemini/agent_driver/streamed_turn.yaml b/tests/cassettes/gemini/agent_driver/streamed_turn.yaml new file mode 100644 index 0000000000..9caf9533df --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/streamed_turn.yaml @@ -0,0 +1,20 @@ +when: + path: /v1beta/models/gemini-2.5-flash:streamGenerateContent + method: POST + query_param: + - name: alt + value: sse + - name: key + value: '[REDACTED]' + header: + - name: accept + value: text/event-stream + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: text/event-stream + body: "data: {\"candidates\":[{\"content\":{\"parts\":[{\"functionCall\":{\"args\":{\"x\":2,\"y\":5},\"name\":\"add\"},\"thoughtSignature\":\"signature_REDACTED_1\"}],\"role\":\"model\"},\"finishMessage\":\"Model generated function call(s).\",\"finishReason\":\"STOP\",\"index\":0}],\"modelVersion\":\"gemini-2.5-flash\",\"responseId\":\"id_REDACTED_1\",\"usageMetadata\":{\"candidatesTokenCount\":18,\"promptTokenCount\":103,\"promptTokensDetails\":[{\"modality\":\"TEXT\",\"tokenCount\":103}],\"serviceTier\":\"standard\",\"thoughtsTokenCount\":51,\"totalTokenCount\":172}}\r\n\r\n" diff --git a/tests/cassettes/gemini/agent_driver/tool_call_round_trip.yaml b/tests/cassettes/gemini/agent_driver/tool_call_round_trip.yaml new file mode 100644 index 0000000000..5c39c5dd4a --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/tool_call_round_trip.yaml @@ -0,0 +1,37 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":103,"promptTokensDetails":[{"modality":"TEXT","tokenCount":103}],"serviceTier":"standard","thoughtsTokenCount":47,"totalTokenCount":168}}' +--- +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"},{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thought":false,"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},{"parts":[{"functionResponse":{"name":"add","response":{"result":7}},"thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":null,"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"text":"7"}],"role":"model"},"finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_2","usageMetadata":{"candidatesTokenCount":1,"promptTokenCount":134,"promptTokensDetails":[{"modality":"TEXT","tokenCount":134}],"serviceTier":"standard","totalTokenCount":135}}' diff --git a/tests/cassettes/gemini/agent_driver/tool_choice_none.yaml b/tests/cassettes/gemini/agent_driver/tool_choice_none.yaml new file mode 100644 index 0000000000..81bec5bad0 --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/tool_choice_none.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"Answer in plain text.","thought":false}],"role":"model"},"toolConfig":{"functionCallingConfig":{"mode":"NONE"}},"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"text":"7"}],"role":"model"},"finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":1,"promptTokenCount":69,"promptTokensDetails":[{"modality":"TEXT","tokenCount":69}],"serviceTier":"standard","thoughtsTokenCount":49,"totalTokenCount":119}}' diff --git a/tests/cassettes/gemini/agent_driver/tool_choice_specific.yaml b/tests/cassettes/gemini/agent_driver/tool_choice_specific.yaml new file mode 100644 index 0000000000..df80f6740f --- /dev/null +++ b/tests/cassettes/gemini/agent_driver/tool_choice_specific.yaml @@ -0,0 +1,18 @@ +when: + path: /v1beta/models/gemini-2.5-flash:generateContent + method: POST + query_param: + - name: key + value: '[REDACTED]' + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"contents":[{"parts":[{"text":"What is 2 + 5?","thought":false}],"role":"user"}],"generationConfig":null,"safetySettings":null,"systemInstruction":{"parts":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","thought":false}],"role":"model"},"toolConfig":{"functionCallingConfig":{"allowed_function_names":["add"],"mode":"ANY"}},"tools":[{"codeExecution":null,"functionDeclarations":[{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}}]}]}' +then: + status: 200 + header: + - name: content-type + value: application/json; charset=UTF-8 + body: '{"candidates":[{"content":{"parts":[{"functionCall":{"args":{"x":2,"y":5},"name":"add"},"thoughtSignature":"signature_REDACTED_1"}],"role":"model"},"finishMessage":"Model generated function call(s).","finishReason":"STOP","index":0}],"modelVersion":"gemini-2.5-flash","responseId":"id_REDACTED_1","usageMetadata":{"candidatesTokenCount":18,"promptTokenCount":164,"promptTokensDetails":[{"modality":"TEXT","tokenCount":164}],"serviceTier":"standard","thoughtsTokenCount":56,"totalTokenCount":238}}' diff --git a/tests/cassettes/openai/agent_driver/driver_history.yaml b/tests/cassettes/openai/agent_driver/driver_history.yaml new file mode 100644 index 0000000000..80f25eb1e2 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/driver_history.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Answer briefly.","type":"text"}],"role":"system"},{"content":"My name is Ada.","role":"user"},{"content":[{"text":"Nice to meet you, Ada.","type":"text"}],"role":"assistant"},{"content":"What is my name?","role":"user"}],"model":"gpt-4o"}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"Your name is Ada.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":5,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":39,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":44}}' diff --git a/tests/cassettes/openai/agent_driver/max_turns.yaml b/tests/cassettes/openai/agent_driver/max_turns.yaml new file mode 100644 index 0000000000..6fbb9ee313 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/max_turns.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' diff --git a/tests/cassettes/openai/agent_driver/output_mode_native.yaml b/tests/cassettes/openai/agent_driver/output_mode_native.yaml new file mode 100644 index 0000000000..a9353fe327 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/output_mode_native.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Reply with the structured answer.","type":"text"}],"role":"system"},{"content":"What is the capital of France?","role":"user"}],"model":"gpt-4o","response_format":{"json_schema":{"name":"response_schema","schema":{"additionalProperties":false,"properties":{"answer":{"type":"string"}},"required":["answer"],"type":"object"},"strict":true},"type":"json_schema"}}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"{\"answer\":\"The capital of France is Paris.\"}","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":11,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":50,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":61}}' diff --git a/tests/cassettes/openai/agent_driver/output_mode_tool.yaml b/tests/cassettes/openai/agent_driver/output_mode_tool.yaml new file mode 100644 index 0000000000..f5d6ba04cd --- /dev/null +++ b/tests/cassettes/openai/agent_driver/output_mode_tool.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Reply with the structured answer.\n\nWhen you have gathered enough information to answer, call the `final_result` tool exactly once with your final answer. Its arguments are the structured result and must satisfy the required schema. Do not return the final answer as plain text.","type":"text"}],"role":"system"},{"content":"What is the capital of France?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Call this tool exactly once with your final answer when you are done. Its arguments are the structured result and must satisfy the output schema.","name":"final_result","parameters":{"properties":{"answer":{"type":"string"}},"required":["answer"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"answer\":\"The capital of France is Paris.\"}","name":"final_result"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":20,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":123,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":143}}' diff --git a/tests/cassettes/openai/agent_driver/parallel_tool_calls.yaml b/tests/cassettes/openai/agent_driver/parallel_tool_calls.yaml new file mode 100644 index 0000000000..dc7c1aa494 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/parallel_tool_calls.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"Compute 2 + 5 and 9 - 3. Use the tools for both.","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\": 2, \"y\": 5}","name":"add"},"id":"call_REDACTED_1","type":"function"},{"function":{"arguments":"{\"x\": 9, \"y\": 3}","name":"subtract"},"id":"call_REDACTED_2","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":50,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":159,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":209}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"Compute 2 + 5 and 9 - 3. Use the tools for both.","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"},{"function":{"arguments":"{\"x\":9,\"y\":3}","name":"subtract"},"id":"call_REDACTED_2","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"},{"content":"6","role":"tool","tool_call_id":"call_REDACTED_2"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7 and 6","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":5,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":225,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":230}}' diff --git a/tests/cassettes/openai/agent_driver/patch_active_tools.yaml b/tests/cassettes/openai/agent_driver/patch_active_tools.yaml new file mode 100644 index 0000000000..95ca7fdfb0 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/patch_active_tools.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":131,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":133}}' diff --git a/tests/cassettes/openai/agent_driver/patch_extra_context.yaml b/tests/cassettes/openai/agent_driver/patch_extra_context.yaml new file mode 100644 index 0000000000..44af771aa7 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/patch_extra_context.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Answer using the provided context only.","type":"text"}],"role":"system"},{"content":"\nThe launch code is banana.\n\n","role":"user"},{"content":"What is the launch code?","role":"user"}],"model":"gpt-4o"}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"I''m sorry, I can''t assist with that request.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":10,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":45,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":55}}' diff --git a/tests/cassettes/openai/agent_driver/patch_history.yaml b/tests/cassettes/openai/agent_driver/patch_history.yaml new file mode 100644 index 0000000000..16972fb8d5 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/patch_history.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Answer briefly.","type":"text"}],"role":"system"},{"content":"Remember this: the code word is banana.","role":"user"},{"content":[{"text":"Noted.","type":"text"}],"role":"assistant"},{"content":"What did I just say?","role":"user"}],"model":"gpt-4o"}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"You mentioned that the code word is banana.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":9,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":40,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":49}}' diff --git a/tests/cassettes/openai/agent_driver/patch_preamble.yaml b/tests/cassettes/openai/agent_driver/patch_preamble.yaml new file mode 100644 index 0000000000..e7d24e8c83 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/patch_preamble.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"PATCHED PREAMBLE — reply with one word.","type":"text"}],"role":"system"},{"content":"Say the word banana.","role":"user"}],"model":"gpt-4o"}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"Banana.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":3,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":27,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":30}}' diff --git a/tests/cassettes/openai/agent_driver/patch_sampling.yaml b/tests/cassettes/openai/agent_driver/patch_sampling.yaml new file mode 100644 index 0000000000..ebacfac886 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/patch_sampling.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"max_tokens":16,"messages":[{"content":[{"text":"Reply with one word.","type":"text"}],"role":"system"},{"content":"Say the word banana.","role":"user"}],"model":"gpt-4o","temperature":0.0}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"Banana.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":3,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":21,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":24}}' diff --git a/tests/cassettes/openai/agent_driver/patch_tool_choice.yaml b/tests/cassettes/openai/agent_driver/patch_tool_choice.yaml new file mode 100644 index 0000000000..28769034d7 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/patch_tool_choice.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tool_choice":"required","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' diff --git a/tests/cassettes/openai/agent_driver/prepare_failure.yaml b/tests/cassettes/openai/agent_driver/prepare_failure.yaml new file mode 100644 index 0000000000..6fbb9ee313 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/prepare_failure.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' diff --git a/tests/cassettes/openai/agent_driver/provider_rejection.yaml b/tests/cassettes/openai/agent_driver/provider_rejection.yaml index 09764851d1..543c27cafb 100644 --- a/tests/cassettes/openai/agent_driver/provider_rejection.yaml +++ b/tests/cassettes/openai/agent_driver/provider_rejection.yaml @@ -7,7 +7,7 @@ when: value: '*/*' - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a weather assistant.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o-this-model-does-not-exist"}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o-this-model-does-not-exist"}' then: status: 404 header: diff --git a/tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml b/tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml index 66f49aaa00..6fbb9ee313 100644 --- a/tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml +++ b/tests/cassettes/openai/agent_driver/resume_awaiting_model.yaml @@ -7,10 +7,10 @@ when: value: '*/*' - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' then: status: 200 header: - name: content-type value: application/json - body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":63,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":76}}' + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' diff --git a/tests/cassettes/openai/agent_driver/resume_executing_tools.yaml b/tests/cassettes/openai/agent_driver/resume_executing_tools.yaml new file mode 100644 index 0000000000..95ca7fdfb0 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/resume_executing_tools.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":131,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":133}}' diff --git a/tests/cassettes/openai/agent_driver/resume_no_tool_leak.yaml b/tests/cassettes/openai/agent_driver/resume_no_tool_leak.yaml new file mode 100644 index 0000000000..6fbb9ee313 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/resume_no_tool_leak.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' diff --git a/tests/cassettes/openai/agent_driver/rollback_re_prepares.yaml b/tests/cassettes/openai/agent_driver/rollback_re_prepares.yaml new file mode 100644 index 0000000000..95ca7fdfb0 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/rollback_re_prepares.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":131,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":133}}' diff --git a/tests/cassettes/openai/agent_driver/run_tool_choice.yaml b/tests/cassettes/openai/agent_driver/run_tool_choice.yaml index f0949fd486..28769034d7 100644 --- a/tests/cassettes/openai/agent_driver/run_tool_choice.yaml +++ b/tests/cassettes/openai/agent_driver/run_tool_choice.yaml @@ -7,10 +7,10 @@ when: value: '*/*' - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a weather assistant.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","tool_choice":"required","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tool_choice":"required","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' then: status: 200 header: - name: content-type value: application/json - body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":57,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":70}}' + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' diff --git a/tests/cassettes/openai/agent_driver/streamed_patched.yaml b/tests/cassettes/openai/agent_driver/streamed_patched.yaml new file mode 100644 index 0000000000..4bc46b6c3b --- /dev/null +++ b/tests/cassettes/openai/agent_driver/streamed_patched.yaml @@ -0,0 +1,42 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: text/event-stream + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true},"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: text/event-stream; charset=utf-8 + body: |+ + data: {"choices":[{"delta":{"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"","name":"add"},"id":"call_REDACTED_1","index":0,"type":"function"}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_1","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"{\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_2","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"x"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_3","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_4","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"2"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_5","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":",\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_6","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"y"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_7","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_8","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"5"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_9","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"}"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_10","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{},"finish_reason":"tool_calls","index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_11","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_12","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}} + + data: [DONE] + diff --git a/tests/cassettes/openai/agent_driver/streamed_text.yaml b/tests/cassettes/openai/agent_driver/streamed_text.yaml new file mode 100644 index 0000000000..116888c2a8 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/streamed_text.yaml @@ -0,0 +1,28 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: text/event-stream + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Reply with one short sentence.","type":"text"}],"role":"system"},{"content":"Say hello.","role":"user"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}' +then: + status: 200 + header: + - name: content-type + value: text/event-stream; charset=utf-8 + body: |+ + data: {"choices":[{"delta":{"content":"","refusal":null,"role":"assistant"},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_1","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"content":"Hello"},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_2","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"content":"!"},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_3","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{},"finish_reason":"stop","index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_4","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":20,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":22}} + + data: [DONE] + diff --git a/tests/cassettes/openai/agent_driver/streamed_turn.yaml b/tests/cassettes/openai/agent_driver/streamed_turn.yaml index c0e34bcb75..b22d6b0ed0 100644 --- a/tests/cassettes/openai/agent_driver/streamed_turn.yaml +++ b/tests/cassettes/openai/agent_driver/streamed_turn.yaml @@ -7,28 +7,36 @@ when: value: text/event-stream - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true},"tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true},"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"type":"number"},"y":{"type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' then: status: 200 header: - name: content-type value: text/event-stream; charset=utf-8 body: |+ - data: {"choices":[{"delta":{"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"","name":"weather"},"id":"call_REDACTED_1","index":0,"type":"function"}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_1","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + data: {"choices":[{"delta":{"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"","name":"add"},"id":"call_REDACTED_1","index":0,"type":"function"}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_1","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"{\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_2","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} - data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"city"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_3","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"x"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_3","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} - data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_4","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_4","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} - data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"Tokyo"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_5","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"2"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_5","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} - data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\"}"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_6","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":",\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_6","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} - data: {"choices":[{"delta":{},"finish_reason":"tool_calls","index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_7","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"y"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_7","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} - data: {"choices":[],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_8","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":63,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":76}} + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_8","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"5"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_9","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"}"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_10","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{},"finish_reason":"tool_calls","index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_11","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_12","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":93,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":110}} data: [DONE] @@ -42,10 +50,10 @@ when: value: '*/*' - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]},{"content":"The weather in Tokyo is 22C and sunny","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"type":"number"},"y":{"type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' then: status: 200 header: - name: content-type value: application/json - body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"The current weather in Tokyo is 22°C and sunny.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":93,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":106}}' + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":118,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":120}}' diff --git a/tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml b/tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml index d401116bc3..8d492204c6 100644 --- a/tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml +++ b/tests/cassettes/openai/agent_driver/tool_call_round_trip.yaml @@ -7,13 +7,13 @@ when: value: '*/*' - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"type":"number"},"y":{"type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' then: status: 200 header: - name: content-type value: application/json - body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":63,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":76}}' + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":93,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":110}}' --- when: path: /v1/chat/completions @@ -24,10 +24,10 @@ when: value: '*/*' - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a weather assistant. Always use the weather tool.","type":"text"}],"role":"system"},{"content":"What is the weather in Tokyo?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"city\":\"Tokyo\"}","name":"weather"},"id":"call_REDACTED_1","type":"function"}]},{"content":"The weather in Tokyo is 22C and sunny","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Get the current weather for a city.","name":"weather","parameters":{"properties":{"city":{"type":"string"}},"required":["city"],"type":"object"}},"type":"function"}]}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"type":"number"},"y":{"type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' then: status: 200 header: - name: content-type value: application/json - body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"The current weather in Tokyo is 22°C and sunny.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":13,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":93,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":106}}' + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":118,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":120}}' diff --git a/tests/cassettes/openai/agent_driver/tool_choice_none.yaml b/tests/cassettes/openai/agent_driver/tool_choice_none.yaml new file mode 100644 index 0000000000..9164f2d73d --- /dev/null +++ b/tests/cassettes/openai/agent_driver/tool_choice_none.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tool_choice":"none","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"2 + 5 is 7.","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":8,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":73,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":81}}' diff --git a/tests/cassettes/openai/agent_driver/tool_choice_specific.yaml b/tests/cassettes/openai/agent_driver/tool_choice_specific.yaml new file mode 100644 index 0000000000..1904e46098 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/tool_choice_specific.yaml @@ -0,0 +1,16 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tool_choice":{"function":{"name":"add"},"type":"function"},"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":9,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":156,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":165}}' diff --git a/tests/cassettes/openai/agent_driver/two_tools_advertised.yaml b/tests/cassettes/openai/agent_driver/two_tools_advertised.yaml new file mode 100644 index 0000000000..84efd7a056 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/two_tools_advertised.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":148,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":165}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":173,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":175}}' diff --git a/tests/common/driver_support.rs b/tests/common/driver_support.rs new file mode 100644 index 0000000000..0a03efa38d --- /dev/null +++ b/tests/common/driver_support.rs @@ -0,0 +1,106 @@ +//! Shared helpers for `AgentDriver` cassette suites. +//! +//! Every provider's driver suite drives the same protocol — prepare, send, +//! feed, dispatch, repeat — so the loop lives here and the per-provider tests +//! keep only what is provider-specific: the client, the model id, and the +//! assertions about the request the provider actually received. +//! +//! The assertions these helpers make are deliberately structural (a tool call +//! happened, the tool ran, the final answer is non-empty). Model wording varies +//! between recordings; request *shape* does not, and that is what the cassette +//! harness pins for us. + +#![allow(dead_code)] + +use rig::agent::{AgentDriver, DriveStep, PendingToolCall, TurnTools}; +use rig::completion::PromptError; +use rig::tool::ToolContext; + +/// Preamble that reliably drives a tool call on every provider tested. +pub(crate) const FORCE_TOOLS_PREAMBLE: &str = "You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text."; + +/// A prompt whose only sensible answer is one `add` call. +pub(crate) const ADD_PROMPT: &str = "What is 2 + 5?"; + +/// Advance the driver, expecting a request to send. +pub(crate) async fn expect_send( + driver: &mut AgentDriver, +) -> ( + Box>, + TurnTools, + usize, +) { + match driver.next_step().await.expect("next_step should succeed") { + DriveStep::SendRequest { + request, + tools, + turn, + } => (request, tools, turn), + other => panic!("expected SendRequest, got {other:?}"), + } +} + +/// Advance the driver, expecting pending tool calls. +pub(crate) async fn expect_execute_tools( + driver: &mut AgentDriver, +) -> (Vec, TurnTools) { + match driver.next_step().await.expect("next_step should succeed") { + DriveStep::ExecuteTools { calls, tools } => (calls, tools), + other => panic!("expected ExecuteTools, got {other:?}"), + } +} + +/// Advance the driver, expecting the run to be finished. +pub(crate) async fn expect_done(driver: &mut AgentDriver) -> rig::agent::PromptResponse { + match driver.next_step().await.expect("next_step should succeed") { + DriveStep::Done(response) => response, + other => panic!("expected Done, got {other:?}"), + } +} + +/// Dispatch every pending call through the turn that advertised it, and feed +/// the results back. +pub(crate) async fn dispatch_and_feed( + driver: &mut AgentDriver, + calls: &[PendingToolCall], + tools: &TurnTools, +) { + let mut context = ToolContext::new(); + let mut results = Vec::new(); + for call in calls { + results.push(tools.execute_call(call, &mut context).await); + } + driver + .tool_results(results) + .expect("tool results should be accepted"); +} + +/// Drive a blocking run to completion, returning the final response. +/// +/// The loop a caller writes by hand, in the shape the driver's own docs use. +pub(crate) async fn drive_to_completion( + driver: &mut AgentDriver, +) -> Result { + loop { + match driver.next_step().await? { + DriveStep::SendRequest { request, .. } => { + let response = request.send().await.map_err(PromptError::CompletionError)?; + driver.model_response(&response)?; + } + DriveStep::ExecuteTools { calls, tools } => { + let mut context = ToolContext::new(); + let mut results = Vec::new(); + for call in &calls { + results.push(tools.execute_call(call, &mut context).await); + } + driver.tool_results(results)?; + } + DriveStep::Done(response) => return Ok(response), + } + } +} + +// Deliberately no "inspect the built request" helper. `build()` consumes the +// builder, so a test that inspects a request cannot also send it — and the +// recorded cassette body is the stronger assertion anyway: it is what the +// provider received, not what rig believed it was sending. diff --git a/tests/gemini.rs b/tests/gemini.rs index 699804ca64..1f23a25ecf 100644 --- a/tests/gemini.rs +++ b/tests/gemini.rs @@ -10,6 +10,8 @@ mod cassette_safety; #[path = "common/cassettes.rs"] mod cassettes; +#[path = "common/driver_support.rs"] +mod driver_support; #[path = "common/reasoning.rs"] mod reasoning; #[path = "common/support.rs"] diff --git a/tests/openai.rs b/tests/openai.rs index 1433ecd419..87c4f7649e 100644 --- a/tests/openai.rs +++ b/tests/openai.rs @@ -10,6 +10,8 @@ mod cassette_safety; #[path = "common/cassettes.rs"] mod cassettes; +#[path = "common/driver_support.rs"] +mod driver_support; #[path = "common/reasoning.rs"] mod reasoning; #[path = "common/support.rs"] diff --git a/tests/providers/anthropic/cassette/agent_driver.rs b/tests/providers/anthropic/cassette/agent_driver.rs new file mode 100644 index 0000000000..1d42b52731 --- /dev/null +++ b/tests/providers/anthropic/cassette/agent_driver.rs @@ -0,0 +1,468 @@ +//! Cassette coverage for `AgentDriver` against real Anthropic traffic. +//! +//! The driver is provider-agnostic; the request it builds is not. Anthropic +//! spells `tool_choice`, tool declarations and the system prompt differently +//! from OpenAI, so the same driver invariants need their own recordings — a +//! per-turn patch that reaches an OpenAI request could silently fail to reach +//! an Anthropic one, and only a recorded body would show it. +//! +//! See `tests/providers/openai/cassette/agent_driver.rs` for the full +//! rationale; this module carries the provider-portable core of that suite. + +use futures::StreamExt; +use rig::agent::run::{OutputMode, StreamedTurnAssembler}; +use rig::agent::{AgentRun, RequestPatch}; +use rig::completion::PromptError; +use rig::message::{Message, ToolChoice}; +use rig::prelude::*; +use rig::providers::anthropic; + +use super::super::support::with_anthropic_cassette; +use crate::driver_support::{ + ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, + expect_execute_tools, expect_send, +}; +use crate::support::{Adder, Subtract}; + +const MODEL: &str = anthropic::completion::CLAUDE_SONNET_4_6; + +/// The whole hand-driven loop over real Anthropic traffic, locking both +/// request bodies. +#[tokio::test] +async fn drive_loop_round_trips_a_tool_call() { + with_anthropic_cassette("agent_driver/tool_call_round_trip", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 1); + assert!(tools.executable_tool_names().contains("add")); + + let response = request.send().await.expect("first turn should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty(), "the model should have called the tool"); + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let response = drive_to_completion(&mut driver) + .await + .expect("run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// A custom run's `tool_choice` must reach Anthropic, whose wire spelling +/// differs from OpenAI's. +#[tokio::test] +async fn a_custom_runs_tool_choice_reaches_the_provider() { + with_anthropic_cassette("agent_driver/run_tool_choice", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::Required); + let mut driver = agent.drive_run(run); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, _) = expect_execute_tools(&mut driver).await; + assert!( + !pending.is_empty(), + "tool_choice=required must force a call" + ); + }) + .await; +} + +/// `ToolChoice::None` forbids tool use for the turn: the run finalizes +/// without ever reaching a tool step. +/// +/// Deliberately *not* asserting a non-empty answer. Recording this against +/// Anthropic showed the model returning a genuinely empty turn — forbidden +/// from the only tool that would answer an arithmetic prompt, it emitted no +/// content at all. That is honest provider behavior and `is_empty_assistant_turn` +/// handles it; the invariant under test is that no tool call happened, not +/// that the model had something to say. +#[tokio::test] +async fn tool_choice_none_forbids_tools_on_the_wire() { + with_anthropic_cassette("agent_driver/tool_choice_none", |client| async move { + let agent = client + .agent(MODEL) + .preamble("Answer in plain text.") + .tool(Adder) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::None); + let mut driver = agent.drive_run(run); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!( + tools.allowed_tool_names().is_empty(), + "ToolChoice::None allows nothing to be called" + ); + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + // Straight to Done: a tool step here would mean the constraint did not + // reach the provider. + let response = expect_done(&mut driver).await; + assert!( + response + .content + .iter() + .all(|item| !matches!(item, rig::message::AssistantContent::ToolCall(_))), + "no tool call may survive ToolChoice::None" + ); + }) + .await; +} + +/// `ToolChoice::Specific` names one tool on the Anthropic wire. +#[tokio::test] +async fn tool_choice_specific_names_the_tool_on_the_wire() { + with_anthropic_cassette("agent_driver/tool_choice_specific", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .tool(Subtract) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::Specific { + function_names: vec!["add".to_string()], + }); + let mut driver = agent.drive_run(run); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.allowed_tool_names().contains("add")); + assert!(!tools.allowed_tool_names().contains("subtract")); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, _) = expect_execute_tools(&mut driver).await; + assert_eq!(pending[0].tool_call.function.name, "add"); + }) + .await; +} + +/// A patched preamble replaces the agent's system prompt for the turn. +#[tokio::test] +async fn a_patched_preamble_replaces_the_agents_on_the_wire() { + with_anthropic_cassette("agent_driver/patch_preamble", |client| async move { + let agent = client + .agent(MODEL) + .preamble("BASELINE PREAMBLE — must not appear in the request") + .build(); + + let mut driver = agent + .drive("Say the word banana.") + .request_patch(RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word.")); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// `active_tools` narrows the tools Anthropic is shown. +#[tokio::test] +async fn a_patched_active_tools_narrows_the_advertised_set() { + with_anthropic_cassette("agent_driver/patch_active_tools", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .tool(Subtract) + .build(); + + let mut driver = agent + .drive(ADD_PROMPT) + .request_patch(RequestPatch::new().active_tools(["add"])); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.executable_tool_names().contains("add")); + assert!(!tools.executable_tool_names().contains("subtract")); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + }) + .await; +} + +/// An explicit patch outranks the run's own choice. +#[tokio::test] +async fn a_patched_tool_choice_outranks_the_runs() { + with_anthropic_cassette("agent_driver/patch_tool_choice", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::None); + let mut driver = agent + .drive_run(run) + .request_patch(RequestPatch::new().tool_choice(ToolChoice::Required)); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.allowed_tool_names().contains("add")); + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, _) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty()); + }) + .await; +} + +/// A driver-level history leads the Anthropic request. +#[tokio::test] +async fn driver_history_leads_the_request() { + with_anthropic_cassette("agent_driver/driver_history", |client| async move { + let agent = client.agent(MODEL).preamble("Answer briefly.").build(); + + let mut driver = agent.drive("What is my name?").history(vec![ + Message::user("My name is Ada."), + Message::assistant("Nice to meet you, Ada."), + ]); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// A run suspended mid-model-call resumes in a fresh driver. +#[tokio::test] +async fn a_run_suspended_awaiting_the_model_resumes_and_accepts_the_reply() { + with_anthropic_cassette("agent_driver/resume_awaiting_model", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(2) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + let response = request.send().await.expect("should send"); + drop(driver); + + let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + assert!(restored.advertised_tools().is_some()); + let mut resumed = agent.drive_run(restored); + resumed + .model_response(&response) + .expect("a resumed run accepts the reply to its in-flight call"); + + let (pending, tools) = expect_execute_tools(&mut resumed).await; + assert!(!pending.is_empty()); + assert!(tools.executable_tool_names().contains("add")); + }) + .await; +} + +/// A run suspended with tool calls pending resumes and completes; the second +/// request is built by the resumed driver. +#[tokio::test] +async fn a_run_suspended_executing_tools_resumes_and_completes() { + with_anthropic_cassette("agent_driver/resume_executing_tools", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools(&mut driver).await; + + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + drop(driver); + + let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut resumed = agent.drive_run(restored); + let (pending, tools) = expect_execute_tools(&mut resumed).await; + dispatch_and_feed(&mut resumed, &pending, &tools).await; + + let response = drive_to_completion(&mut resumed) + .await + .expect("resumed run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// Tool output mode against Anthropic: the synthetic output tool is advertised +/// and allowed, never executable, and the run finalizes on its call. +#[tokio::test] +async fn tool_output_mode_finalizes_via_the_output_tool() { + with_anthropic_cassette("agent_driver/output_mode_tool", |client| async move { + let agent = client + .agent(MODEL) + .preamble("Reply with the structured answer.") + .output_schema_raw( + serde_json::from_value(serde_json::json!({ + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"] + })) + .expect("valid schema"), + ) + .output_mode(OutputMode::Tool) + .build(); + + let mut driver = agent.drive("What is the capital of France?"); + let (request, tools, _) = expect_send(&mut driver).await; + let output_tool = tools + .output_tool_name() + .expect("Tool mode advertises an output tool") + .to_owned(); + assert!(tools.allowed_tool_names().contains(&output_tool)); + assert!(!tools.executable_tool_names().contains(&output_tool)); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(response.output.contains("answer")); + }) + .await; +} + +/// A streamed turn driven through the driver against real Anthropic SSE. +#[tokio::test] +async fn a_streamed_turn_drives_through_the_driver() { + with_anthropic_cassette("agent_driver/streamed_turn", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send(&mut driver).await; + + let mut assembler = StreamedTurnAssembler::new( + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), + ); + let mut stream = request.stream().await.expect("stream should open"); + while let Some(item) = stream.next().await { + assembler + .ingest(&item.expect("stream item")) + .expect("ingest should succeed"); + } + let final_content = stream.choice.clone(); + let streamed = assembler.finish(stream.message_id.clone(), &final_content); + + driver + .run_mut() + .record_streamed_completion_call(stream.usage()) + .expect("usage recorded"); + driver + .run_mut() + .streamed_turn(streamed) + .expect("streamed turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty()); + dispatch_and_feed(&mut driver, &pending, &tools).await; + }) + .await; +} + +/// A rejected send can be handed back and the turn prepared again. +#[tokio::test] +async fn a_provider_rejection_is_not_retryable() { + with_anthropic_cassette("agent_driver/provider_rejection", |client| async move { + // Anthropic requires `max_tokens`; without it the request fails + // locally and never reaches the provider, which is a different + // assertion than the one this test is making. + let agent = client + .agent("claude-this-model-does-not-exist") + .preamble(FORCE_TOOLS_PREAMBLE) + .max_tokens(64) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + let error = request + .send() + .await + .expect_err("an unknown model must be rejected"); + + // Pin the provider's own envelope: a cassette mock miss is a 404 too. + let body = error + .provider_response_body() + .expect("the provider's rejection body is preserved"); + assert!( + body.contains("not_found") || body.contains("model"), + "expected the recorded provider rejection, got: {body}" + ); + assert!( + !error.is_retryable(), + "a rejection is not retryable: {error}" + ); + }) + .await; +} + +/// Exhausting the model-call budget stops the run before a second request. +#[tokio::test] +async fn max_turns_exhaustion_stops_before_a_second_send() { + with_anthropic_cassette("agent_driver/max_turns", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT).max_turns(1); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let error = driver + .next_step() + .await + .expect_err("the budget of one is spent"); + assert!(matches!(error, PromptError::MaxTurnsError { .. })); + }) + .await; +} diff --git a/tests/providers/anthropic/mod.rs b/tests/providers/anthropic/mod.rs index 19bc3123f2..584e813ee4 100644 --- a/tests/providers/anthropic/mod.rs +++ b/tests/providers/anthropic/mod.rs @@ -2,6 +2,7 @@ mod support; mod cassette { mod agent; + mod agent_driver; mod default_max_turns; mod document_file_id; mod empty_end_turn; diff --git a/tests/providers/gemini/cassette/agent_driver.rs b/tests/providers/gemini/cassette/agent_driver.rs new file mode 100644 index 0000000000..ebfc7aa902 --- /dev/null +++ b/tests/providers/gemini/cassette/agent_driver.rs @@ -0,0 +1,464 @@ +//! Cassette coverage for `AgentDriver` against real Gemini traffic. +//! +//! The driver is provider-agnostic; the request it builds is not. Gemini +//! spells `tool_choice`, tool declarations and the system prompt differently +//! from OpenAI, so the same driver invariants need their own recordings — a +//! per-turn patch that reaches an OpenAI request could silently fail to reach +//! a Gemini one, and only a recorded body would show it. +//! +//! See `tests/providers/openai/cassette/agent_driver.rs` for the full +//! rationale; this module carries the provider-portable core of that suite. + +use futures::StreamExt; +use rig::agent::run::{OutputMode, StreamedTurnAssembler}; +use rig::agent::{AgentRun, RequestPatch}; +use rig::completion::PromptError; +use rig::message::{Message, ToolChoice}; +use rig::prelude::*; +use rig::providers::gemini; + +use super::super::support::with_gemini_cassette; +use crate::driver_support::{ + ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, + expect_execute_tools, expect_send, +}; +use crate::support::{Adder, Subtract}; + +const MODEL: &str = gemini::completion::GEMINI_2_5_FLASH; + +/// The whole hand-driven loop over real Gemini traffic, locking both +/// request bodies. +#[tokio::test] +async fn drive_loop_round_trips_a_tool_call() { + with_gemini_cassette("agent_driver/tool_call_round_trip", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 1); + assert!(tools.executable_tool_names().contains("add")); + + let response = request.send().await.expect("first turn should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty(), "the model should have called the tool"); + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let response = drive_to_completion(&mut driver) + .await + .expect("run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// A custom run's `tool_choice` must reach Gemini, whose wire spelling +/// differs from OpenAI's. +#[tokio::test] +async fn a_custom_runs_tool_choice_reaches_the_provider() { + with_gemini_cassette("agent_driver/run_tool_choice", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::Required); + let mut driver = agent.drive_run(run); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, _) = expect_execute_tools(&mut driver).await; + assert!( + !pending.is_empty(), + "tool_choice=required must force a call" + ); + }) + .await; +} + +/// `ToolChoice::None` forbids tool use for the turn: the run finalizes +/// without ever reaching a tool step. +/// +/// Deliberately *not* asserting a non-empty answer. Recording this against +/// Gemini showed the model returning a genuinely empty turn — forbidden +/// from the only tool that would answer an arithmetic prompt, it emitted no +/// content at all. That is honest provider behavior and `is_empty_assistant_turn` +/// handles it; the invariant under test is that no tool call happened, not +/// that the model had something to say. +#[tokio::test] +async fn tool_choice_none_forbids_tools_on_the_wire() { + with_gemini_cassette("agent_driver/tool_choice_none", |client| async move { + let agent = client + .agent(MODEL) + .preamble("Answer in plain text.") + .tool(Adder) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::None); + let mut driver = agent.drive_run(run); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!( + tools.allowed_tool_names().is_empty(), + "ToolChoice::None allows nothing to be called" + ); + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + // Straight to Done: a tool step here would mean the constraint did not + // reach the provider. + let response = expect_done(&mut driver).await; + assert!( + response + .content + .iter() + .all(|item| !matches!(item, rig::message::AssistantContent::ToolCall(_))), + "no tool call may survive ToolChoice::None" + ); + }) + .await; +} + +/// `ToolChoice::Specific` names one tool on the Gemini wire. +#[tokio::test] +async fn tool_choice_specific_names_the_tool_on_the_wire() { + with_gemini_cassette("agent_driver/tool_choice_specific", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .tool(Subtract) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::Specific { + function_names: vec!["add".to_string()], + }); + let mut driver = agent.drive_run(run); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.allowed_tool_names().contains("add")); + assert!(!tools.allowed_tool_names().contains("subtract")); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, _) = expect_execute_tools(&mut driver).await; + assert_eq!(pending[0].tool_call.function.name, "add"); + }) + .await; +} + +/// A patched preamble replaces the agent's system prompt for the turn. +#[tokio::test] +async fn a_patched_preamble_replaces_the_agents_on_the_wire() { + with_gemini_cassette("agent_driver/patch_preamble", |client| async move { + let agent = client + .agent(MODEL) + .preamble("BASELINE PREAMBLE — must not appear in the request") + .build(); + + let mut driver = agent + .drive("Say the word banana.") + .request_patch(RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word.")); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// `active_tools` narrows the tools Gemini is shown. +#[tokio::test] +async fn a_patched_active_tools_narrows_the_advertised_set() { + with_gemini_cassette("agent_driver/patch_active_tools", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .tool(Subtract) + .build(); + + let mut driver = agent + .drive(ADD_PROMPT) + .request_patch(RequestPatch::new().active_tools(["add"])); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.executable_tool_names().contains("add")); + assert!(!tools.executable_tool_names().contains("subtract")); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + }) + .await; +} + +/// An explicit patch outranks the run's own choice. +#[tokio::test] +async fn a_patched_tool_choice_outranks_the_runs() { + with_gemini_cassette("agent_driver/patch_tool_choice", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::None); + let mut driver = agent + .drive_run(run) + .request_patch(RequestPatch::new().tool_choice(ToolChoice::Required)); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.allowed_tool_names().contains("add")); + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, _) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty()); + }) + .await; +} + +/// A driver-level history leads the Gemini request. +#[tokio::test] +async fn driver_history_leads_the_request() { + with_gemini_cassette("agent_driver/driver_history", |client| async move { + let agent = client.agent(MODEL).preamble("Answer briefly.").build(); + + let mut driver = agent.drive("What is my name?").history(vec![ + Message::user("My name is Ada."), + Message::assistant("Nice to meet you, Ada."), + ]); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// A run suspended mid-model-call resumes in a fresh driver. +#[tokio::test] +async fn a_run_suspended_awaiting_the_model_resumes_and_accepts_the_reply() { + with_gemini_cassette("agent_driver/resume_awaiting_model", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(2) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + let response = request.send().await.expect("should send"); + drop(driver); + + let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + assert!(restored.advertised_tools().is_some()); + let mut resumed = agent.drive_run(restored); + resumed + .model_response(&response) + .expect("a resumed run accepts the reply to its in-flight call"); + + let (pending, tools) = expect_execute_tools(&mut resumed).await; + assert!(!pending.is_empty()); + assert!(tools.executable_tool_names().contains("add")); + }) + .await; +} + +/// A run suspended with tool calls pending resumes and completes; the second +/// request is built by the resumed driver. +#[tokio::test] +async fn a_run_suspended_executing_tools_resumes_and_completes() { + with_gemini_cassette("agent_driver/resume_executing_tools", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools(&mut driver).await; + + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + drop(driver); + + let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut resumed = agent.drive_run(restored); + let (pending, tools) = expect_execute_tools(&mut resumed).await; + dispatch_and_feed(&mut resumed, &pending, &tools).await; + + let response = drive_to_completion(&mut resumed) + .await + .expect("resumed run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// Tool output mode against Gemini: the synthetic output tool is advertised +/// and allowed, never executable, and the run finalizes on its call. +#[tokio::test] +async fn tool_output_mode_finalizes_via_the_output_tool() { + with_gemini_cassette("agent_driver/output_mode_tool", |client| async move { + let agent = client + .agent(MODEL) + .preamble("Reply with the structured answer.") + .output_schema_raw( + serde_json::from_value(serde_json::json!({ + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"] + })) + .expect("valid schema"), + ) + .output_mode(OutputMode::Tool) + .build(); + + let mut driver = agent.drive("What is the capital of France?"); + let (request, tools, _) = expect_send(&mut driver).await; + let output_tool = tools + .output_tool_name() + .expect("Tool mode advertises an output tool") + .to_owned(); + assert!(tools.allowed_tool_names().contains(&output_tool)); + assert!(!tools.executable_tool_names().contains(&output_tool)); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(response.output.contains("answer")); + }) + .await; +} + +/// A streamed turn driven through the driver against real Gemini SSE. +#[tokio::test] +async fn a_streamed_turn_drives_through_the_driver() { + with_gemini_cassette("agent_driver/streamed_turn", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send(&mut driver).await; + + let mut assembler = StreamedTurnAssembler::new( + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), + ); + let mut stream = request.stream().await.expect("stream should open"); + while let Some(item) = stream.next().await { + assembler + .ingest(&item.expect("stream item")) + .expect("ingest should succeed"); + } + let final_content = stream.choice.clone(); + let streamed = assembler.finish(stream.message_id.clone(), &final_content); + + driver + .run_mut() + .record_streamed_completion_call(stream.usage()) + .expect("usage recorded"); + driver + .run_mut() + .streamed_turn(streamed) + .expect("streamed turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty()); + dispatch_and_feed(&mut driver, &pending, &tools).await; + }) + .await; +} + +/// A rejected send can be handed back and the turn prepared again. +#[tokio::test] +async fn a_provider_rejection_is_not_retryable() { + with_gemini_cassette("agent_driver/provider_rejection", |client| async move { + let agent = client + .agent("gemini-this-model-does-not-exist") + .preamble(FORCE_TOOLS_PREAMBLE) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + let error = request + .send() + .await + .expect_err("an unknown model must be rejected"); + + // Pin the provider's own envelope: a cassette mock miss is a 404 too. + let body = error + .provider_response_body() + .expect("the provider's rejection body is preserved"); + assert!( + body.contains("not found") || body.contains("NOT_FOUND") || body.contains("model"), + "expected the recorded provider rejection, got: {body}" + ); + assert!( + !error.is_retryable(), + "a rejection is not retryable: {error}" + ); + }) + .await; +} + +/// Exhausting the model-call budget stops the run before a second request. +#[tokio::test] +async fn max_turns_exhaustion_stops_before_a_second_send() { + with_gemini_cassette("agent_driver/max_turns", |client| async move { + let agent = client + .agent(MODEL) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT).max_turns(1); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let error = driver + .next_step() + .await + .expect_err("the budget of one is spent"); + assert!(matches!(error, PromptError::MaxTurnsError { .. })); + }) + .await; +} diff --git a/tests/providers/gemini/mod.rs b/tests/providers/gemini/mod.rs index ed96f161d8..ae04e1ce61 100644 --- a/tests/providers/gemini/mod.rs +++ b/tests/providers/gemini/mod.rs @@ -5,6 +5,7 @@ mod tools_support; mod cassette { mod agent; + mod agent_driver; mod agent_run_recovery; mod agent_run_resume; mod agent_run_stepping; diff --git a/tests/providers/openai/cassette/agent_driver.rs b/tests/providers/openai/cassette/agent_driver.rs index 05d9efec06..2a45227ecb 100644 --- a/tests/providers/openai/cassette/agent_driver.rs +++ b/tests/providers/openai/cassette/agent_driver.rs @@ -2,54 +2,64 @@ //! //! The driver's job is to build a request and pair it with the tool state that //! validates and dispatches the model's reply. Unit tests can check the state -//! machine, but only recorded traffic can check the *request* — and every -//! request-shape defect found in this API's review so far (a run's -//! `tool_choice` never reaching the wire, a resumed turn advertising the wrong -//! tool set) was invisible to unit tests precisely because they asserted on -//! rig's own view of the request rather than on the bytes. +//! machine; only recorded traffic can check the *request*. Every request-shape +//! defect found in this API's review (a run's `tool_choice` never reaching the +//! wire, a resumed turn advertising the wrong tool set) was invisible to unit +//! tests precisely because those tests asserted on rig's own view of the +//! request rather than on the bytes. //! -//! The cassette harness matches each request body against the recorded one, so -//! a request-shape regression fails as a mock miss. That is the assertion these -//! tests exist for; the response-side asserts are secondary. +//! The harness matches each request body against the recorded one, so a +//! request-shape regression fails as a mock miss with a body diff. **That is +//! the assertion these tests exist for**; the response-side asserts are +//! secondary and deliberately structural, since model wording varies between +//! recordings. use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use futures::StreamExt; -use rig::agent::run::StreamedTurnAssembler; -use rig::agent::{AgentRun, DriveStep}; -use rig::message::ToolChoice; +use rig::agent::run::{OutputMode, StreamedTurnAssembler}; +use rig::agent::{AgentRun, RequestPatch}; +use rig::completion::PromptError; +use rig::message::{Message, ToolChoice}; use rig::prelude::*; use rig::providers::openai; use rig::tool::{Tool, ToolContext}; use super::super::support::with_openai_completions_cassette; +use crate::driver_support::{ + ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, + expect_execute_tools, expect_send, +}; +use crate::support::{Adder, Subtract}; -#[derive(Debug, serde::Deserialize)] -struct WeatherArgs { - city: String, +/// A counting `add` so tests can assert the tool ran exactly once. +#[derive(Clone)] +struct CountingAdder { + calls: Arc, } -#[derive(Clone)] -struct WeatherTool { - call_count: Arc, +#[derive(serde::Deserialize)] +struct AddArgs { + x: i32, + y: i32, } -impl Tool for WeatherTool { - const NAME: &'static str = "weather"; +impl Tool for CountingAdder { + const NAME: &'static str = "add"; type Error = std::io::Error; - type Args = WeatherArgs; - type Output = String; + type Args = AddArgs; + type Output = i32; fn description(&self) -> String { - "Get the current weather for a city.".to_string() + "Add x and y together".to_string() } fn parameters(&self) -> serde_json::Value { serde_json::json!({ "type": "object", - "properties": { "city": { "type": "string" } }, - "required": ["city"] + "properties": { "x": { "type": "number" }, "y": { "type": "number" } }, + "required": ["x", "y"] }) } @@ -58,209 +68,828 @@ impl Tool for WeatherTool { _context: &mut ToolContext, args: Self::Args, ) -> impl std::future::Future> + Send { - self.call_count.fetch_add(1, Ordering::SeqCst); - std::future::ready(Ok(format!("The weather in {} is 22C and sunny", args.city))) + self.calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Ok(args.x + args.y)) } } +// ── Tranche 1: the loop ────────────────────────────────────────────────── + /// The whole hand-driven loop over real traffic: the driver builds both /// requests, the caller sends them, and the turn that advertised the tool is /// the turn that dispatches it. /// -/// Locks the request bodies for both turns. The second request in particular -/// carries the assistant tool call and the tool result the driver threaded back -/// through the run — a shape no unit test observes. +/// Locks both request bodies. The second in particular carries the assistant +/// tool call and the tool result the driver threaded back through the run — a +/// shape no unit test observes. #[tokio::test] async fn drive_loop_round_trips_a_tool_call() { with_openai_completions_cassette("agent_driver/tool_call_round_trip", |client| async move { let calls = Arc::new(AtomicUsize::new(0)); let agent = client .agent(openai::GPT_4O) - .preamble("You are a weather assistant. Always use the weather tool.") + .preamble(FORCE_TOOLS_PREAMBLE) .default_max_turns(3) - .tool(WeatherTool { - call_count: calls.clone(), + .tool(CountingAdder { + calls: calls.clone(), }) .build(); - let mut driver = agent.drive("What is the weather in Tokyo?"); + let mut driver = agent.drive(ADD_PROMPT); - let (request, tools, turn) = match driver.next_step().await.expect("first step") { - DriveStep::SendRequest { - request, - tools, - turn, - } => (request, tools, turn), - other => panic!("expected SendRequest, got {other:?}"), - }; + let (request, tools, turn) = expect_send(&mut driver).await; assert_eq!(turn, 1); - assert!(tools.executable_tool_names().contains("weather")); + assert!(tools.executable_tool_names().contains("add")); let response = request.send().await.expect("first turn should send"); driver.model_response(&response).expect("turn accepted"); - let (pending, tools) = match driver.next_step().await.expect("second step") { - DriveStep::ExecuteTools { calls, tools } => (calls, tools), - other => panic!("expected ExecuteTools, got {other:?}"), - }; + let (pending, tools) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty(), "the model should have called the tool"); + dispatch_and_feed(&mut driver, &pending, &tools).await; - let mut context = ToolContext::new(); - let mut results = Vec::new(); - for call in &pending { - results.push(tools.execute_call(call, &mut context).await); - } - driver.tool_results(results).expect("results accepted"); - - let request = match driver.next_step().await.expect("third step") { - DriveStep::SendRequest { request, turn, .. } => { - assert_eq!(turn, 2); - request - } - other => panic!("expected SendRequest, got {other:?}"), - }; + let (request, _, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 2); let response = request.send().await.expect("second turn should send"); driver.model_response(&response).expect("turn accepted"); - match driver.next_step().await.expect("final step") { - DriveStep::Done(response) => { - assert!( - !response.output.trim().is_empty(), - "expected a final answer" - ); - } - other => panic!("expected Done, got {other:?}"), - } + let response = expect_done(&mut driver).await; + assert!( + !response.output.trim().is_empty(), + "expected a final answer" + ); assert_eq!(calls.load(Ordering::SeqCst), 1, "the tool ran exactly once"); }) .await; } +/// Two tools advertised, one called: the request carries both, so a narrowing +/// regression shows up as a body diff. +#[tokio::test] +async fn both_registered_tools_are_advertised() { + with_openai_completions_cassette("agent_driver/two_tools_advertised", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .tool(Subtract) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.executable_tool_names().contains("add")); + assert!(tools.executable_tool_names().contains("subtract")); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let response = drive_to_completion(&mut driver) + .await + .expect("run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// A prompt needing two independent tool calls: every pending call dispatches +/// through the same advertising turn, and the request that follows carries all +/// their results. +#[tokio::test] +async fn parallel_tool_calls_all_dispatch_through_one_turn() { + with_openai_completions_cassette("agent_driver/parallel_tool_calls", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(4) + .tool(Adder) + .tool(Subtract) + .build(); + + let mut driver = agent.drive("Compute 2 + 5 and 9 - 3. Use the tools for both."); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty()); + // Whatever the model asked for, every call dispatches through the turn + // that advertised it. + for call in &pending { + assert!( + tools + .executable_tool_names() + .contains(&call.tool_call.function.name), + "the model called a tool this turn never advertised" + ); + } + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let response = drive_to_completion(&mut driver) + .await + .expect("run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +// ── Tranche 2: per-run and per-turn configuration on the wire ──────────── + /// A custom run is taken as-is, so its own `tool_choice` must reach the /// provider — not merely the run's internal decisions. /// -/// This is the finding a unit test could assert only against rig's own -/// `CompletionRequest`. Here the recorded body carries `"tool_choice"`, so a -/// driver that drops it fails as a mock miss: the request rig builds no longer -/// matches the request the provider was actually asked. +/// The finding a unit test could only assert against rig's own +/// `CompletionRequest`. Reverting the fix makes this fail as a mock miss whose +/// diff names the missing `tool_choice`. #[tokio::test] async fn a_custom_runs_tool_choice_reaches_the_provider() { with_openai_completions_cassette("agent_driver/run_tool_choice", |client| async move { let agent = client .agent(openai::GPT_4O) - .preamble("You are a weather assistant.") - .tool(WeatherTool { - call_count: Arc::new(AtomicUsize::new(0)), - }) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) .build(); - // The choice lives on the run, not on the agent. - let run = AgentRun::new("What is the weather in Tokyo?") + let run = AgentRun::new(ADD_PROMPT) .max_turns(2) .with_tool_choice(ToolChoice::Required); let mut driver = agent.drive_run(run); - let request = match driver.next_step().await.expect("first step") { - DriveStep::SendRequest { request, .. } => request, - other => panic!("expected SendRequest, got {other:?}"), + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, _) = expect_execute_tools(&mut driver).await; + assert!( + !pending.is_empty(), + "tool_choice=required must force a call" + ); + }) + .await; +} + +/// `ToolChoice::None` forbids tool use for the turn: the request carries the +/// constraint and the advertised-but-not-allowed set is empty. +#[tokio::test] +async fn tool_choice_none_forbids_tools_on_the_wire() { + with_openai_completions_cassette("agent_driver/tool_choice_none", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Answer in plain text.") + .tool(Adder) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::None); + let mut driver = agent.drive_run(run); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!( + tools.allowed_tool_names().is_empty(), + "ToolChoice::None allows nothing to be called" + ); + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// `ToolChoice::Specific` names one tool; the request must carry that name. +#[tokio::test] +async fn tool_choice_specific_names_the_tool_on_the_wire() { + with_openai_completions_cassette("agent_driver/tool_choice_specific", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .tool(Subtract) + .build(); + + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::Specific { + function_names: vec!["add".to_string()], + }); + let mut driver = agent.drive_run(run); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.allowed_tool_names().contains("add")); + assert!( + !tools.allowed_tool_names().contains("subtract"), + "Specific narrows what the model may call" + ); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, _) = expect_execute_tools(&mut driver).await; + assert_eq!(pending[0].tool_call.function.name, "add"); + }) + .await; +} + +/// The driver runs no hooks, so `RequestPatch` is the seam through which a +/// hand-driven turn gets per-turn configuration. Nothing else exercises it, so +/// each field that reaches the request gets a recorded body. +#[tokio::test] +async fn a_patched_preamble_replaces_the_agents_on_the_wire() { + with_openai_completions_cassette("agent_driver/patch_preamble", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("BASELINE PREAMBLE — must not appear in the request") + .build(); + + let mut driver = agent + .drive("Say the word banana.") + .request_patch(RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word.")); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// An explicit patch outranks the run's own choice. +#[tokio::test] +async fn a_patched_tool_choice_outranks_the_runs() { + with_openai_completions_cassette("agent_driver/patch_tool_choice", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + // The run says None; the patch says Required. The patch wins, so the + // recorded request carries `required`. + let run = AgentRun::new(ADD_PROMPT) + .max_turns(2) + .with_tool_choice(ToolChoice::None); + let mut driver = agent + .drive_run(run) + .request_patch(RequestPatch::new().tool_choice(ToolChoice::Required)); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!( + tools.allowed_tool_names().contains("add"), + "the patch's Required must govern, not the run's None" + ); + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, _) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty()); + }) + .await; +} + +/// `active_tools` narrows the advertised set for the turn, so the request's +/// `tools` array shrinks. +#[tokio::test] +async fn a_patched_active_tools_narrows_the_advertised_set() { + with_openai_completions_cassette("agent_driver/patch_active_tools", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .tool(Subtract) + .build(); + + let mut driver = agent + .drive(ADD_PROMPT) + .request_patch(RequestPatch::new().active_tools(["add"])); + + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.executable_tool_names().contains("add")); + assert!( + !tools.executable_tool_names().contains("subtract"), + "active_tools must narrow the advertised set: {:?}", + tools.executable_tool_names() + ); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + let response = drive_to_completion(&mut driver) + .await + .expect("run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// Sampling parameters are per-turn request fields; the recorded body pins +/// them. +#[tokio::test] +async fn patched_sampling_parameters_reach_the_request() { + with_openai_completions_cassette("agent_driver/patch_sampling", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Reply with one word.") + .temperature(0.9) + .build(); + + let mut driver = agent + .drive("Say the word banana.") + .request_patch(RequestPatch::new().temperature(0.0).max_tokens(16)); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_done(&mut driver).await; + }) + .await; +} + +/// Extra context documents are appended to the turn's request. +#[tokio::test] +async fn patched_extra_context_reaches_the_request() { + with_openai_completions_cassette("agent_driver/patch_extra_context", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Answer using the provided context only.") + .build(); + + let document = rig::completion::Document { + id: "note-1".to_string(), + text: "The launch code is banana.".to_string(), + additional_props: Default::default(), }; + let mut driver = agent + .drive("What is the launch code?") + .request_patch(RequestPatch::new().extra_context(vec![document])); - let response = request.send().await.expect("request should send"); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} - // `Required` obliges a tool call, so the run must reach tools. - match driver.next_step().await.expect("second step") { - DriveStep::ExecuteTools { calls, .. } => { - assert!(!calls.is_empty(), "tool_choice=required must force a call"); - } - other => panic!("expected ExecuteTools, got {other:?}"), - } +/// A patched history replaces the run's for the turn. +#[tokio::test] +async fn a_patched_history_replaces_the_runs_for_the_turn() { + with_openai_completions_cassette("agent_driver/patch_history", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Answer briefly.") + .build(); + + let mut driver = + agent + .drive("What did I just say?") + .request_patch(RequestPatch::new().history(vec![ + Message::user("Remember this: the code word is banana."), + Message::assistant("Noted."), + ])); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// A driver-level history seeds the run and leads the request. +#[tokio::test] +async fn driver_history_leads_the_request() { + with_openai_completions_cassette("agent_driver/driver_history", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Answer briefly.") + .build(); + + let mut driver = agent.drive("What is my name?").history(vec![ + Message::user("My name is Ada."), + Message::assistant("Nice to meet you, Ada."), + ]); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); }) .await; } +// ── Tranche 3: suspend, resume, drift ──────────────────────────────────── + /// A run suspended with a model call in flight resumes in a different driver /// and accepts the reply. /// -/// The request is recorded once; the resumed driver never builds one. What is -/// under test is that the turn's advertised tool names travelled with the run, -/// so the reply is validated against the set the recorded request actually -/// carried. +/// The request is recorded once; the resumed driver never builds one. Under +/// test is that the turn's advertised names travelled with the run, so the +/// reply is validated against the set the recorded request actually carried. #[tokio::test] async fn a_run_suspended_awaiting_the_model_resumes_and_accepts_the_reply() { with_openai_completions_cassette("agent_driver/resume_awaiting_model", |client| async move { let agent = client .agent(openai::GPT_4O) - .preamble("You are a weather assistant. Always use the weather tool.") + .preamble(FORCE_TOOLS_PREAMBLE) .default_max_turns(2) - .tool(WeatherTool { - call_count: Arc::new(AtomicUsize::new(0)), - }) + .tool(Adder) .build(); - let mut driver = agent.drive("What is the weather in Tokyo?"); - let request = match driver.next_step().await.expect("first step") { - DriveStep::SendRequest { request, .. } => request, - other => panic!("expected SendRequest, got {other:?}"), - }; + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; - // Suspend with the call in flight, then send. let serialized = serde_json::to_string(driver.run()).expect("run serializes"); - let response = request.send().await.expect("request should send"); + let response = request.send().await.expect("should send"); drop(driver); let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); assert!( restored.advertised_tools().is_some(), - "the suspended run carries the turn's advertised names" + "a suspended run carries the turn's advertised names" ); let mut resumed = agent.drive_run(restored); resumed .model_response(&response) .expect("a resumed run accepts the reply to its in-flight call"); - match resumed.next_step().await.expect("second step") { - DriveStep::ExecuteTools { calls, tools } => { - assert!(!calls.is_empty()); - assert!(tools.executable_tool_names().contains("weather")); - } - other => panic!("expected ExecuteTools, got {other:?}"), - } + let (pending, tools) = expect_execute_tools(&mut resumed).await; + assert!(!pending.is_empty()); + assert!(tools.executable_tool_names().contains("add")); + }) + .await; +} + +/// A run suspended with tool calls pending resumes and completes, and the +/// second request is built by the *resumed* driver — so its recorded body is +/// the assertion that resume rebuilds the request faithfully. +#[tokio::test] +async fn a_run_suspended_executing_tools_resumes_and_completes() { + with_openai_completions_cassette("agent_driver/resume_executing_tools", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools(&mut driver).await; + + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + drop(driver); + + let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut resumed = agent.drive_run(restored); + let (pending, tools) = expect_execute_tools(&mut resumed).await; + dispatch_and_feed(&mut resumed, &pending, &tools).await; + + let response = drive_to_completion(&mut resumed) + .await + .expect("resumed run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// The resumed turn advertises what *that turn* advertised, not what the +/// resuming process happens to have registered. +/// +/// The resuming agent registers an extra tool. If the resumed dispatch target +/// were the current registry rather than the turn's recorded set, the extra +/// tool would leak into a turn that never advertised it. +#[tokio::test] +async fn a_resumed_turn_advertises_its_own_tools_not_the_processs() { + with_openai_completions_cassette("agent_driver/resume_no_tool_leak", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let _ = expect_execute_tools(&mut driver).await; + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + let advertised = driver + .run() + .advertised_tools() + .expect("the turn recorded its names") + .clone(); + drop(driver); + + // The resuming process has since registered `subtract`. + let resumed_agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .tool(Subtract) + .build(); + let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut resumed = resumed_agent.drive_run(restored); + + let (pending, tools) = expect_execute_tools(&mut resumed).await; + assert_eq!( + tools.executable_tool_names(), + &advertised.executable, + "the resumed turn's tool set is the turn's, not the process's" + ); + assert!( + !tools.executable_tool_names().contains("subtract"), + "a tool registered after suspension must not join this turn" + ); + dispatch_and_feed(&mut resumed, &pending, &tools).await; }) .await; } -/// A hand-driven **streamed** turn goes through the driver, against real SSE -/// traffic. +// ── Tranche 4: failure and recovery ────────────────────────────────────── + +/// A real provider rejection is classified non-retryable, so the caller does +/// not hand the turn back and loop. +/// +/// A synthetic error cannot falsify the classification; only a real provider +/// rejection can. +#[tokio::test] +async fn a_provider_rejection_is_not_retryable() { + with_openai_completions_cassette("agent_driver/provider_rejection", |client| async move { + let agent = client + .agent("gpt-4o-this-model-does-not-exist") + .preamble(FORCE_TOOLS_PREAMBLE) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, _) = expect_send(&mut driver).await; + + let error = request + .send() + .await + .expect_err("an unknown model must be rejected"); + + // A cassette mock miss is *also* a 404, so asserting on the status + // alone would pass against a cassette that never matched. Pin the + // provider's own envelope, which the harness cannot fabricate. + let body = error + .provider_response_body() + .expect("the provider's rejection body is preserved"); + assert!( + body.contains("model_not_found"), + "expected the recorded provider rejection, got: {body}" + ); + assert!( + !error.is_retryable(), + "a provider rejection must not be classified retryable: {error}" + ); + + assert_eq!(driver.run().turn(), 1, "the turn is still in flight"); + assert_eq!(driver.run().model_call_rollbacks(), 0); + }) + .await; +} + +/// A rejected send can be handed back and the turn prepared again. +/// +/// Two recorded interactions: the rejection, then a fresh request. The second +/// body is the assertion — the retry must be a *new build*, not a replay of +/// the request that failed. +#[tokio::test] +async fn a_rejected_send_rolls_back_and_re_prepares() { + with_openai_completions_cassette("agent_driver/rollback_re_prepares", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(Adder) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, _, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 1); + + // Drop the request unsent: the provider never saw anything. + drop(request); + driver + .rollback_model_call() + .expect("a call that produced nothing can be handed back"); + assert_eq!(driver.run().turn(), 0, "the turn is refunded"); + assert_eq!(driver.run().model_call_rollbacks(), 1); + + // The retry takes the turn the failure did not, and its request is + // freshly built. + let (request, _, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 1); + let response = request.send().await.expect("the retry should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + let response = drive_to_completion(&mut driver) + .await + .expect("run should finish"); + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// A preparation failure costs no turn and consumes no interaction. +/// +/// The cassette's contribution is the *negative*: exactly one interaction is +/// recorded, and it is consumed only by the retry. A driver that advanced +/// before preparing would burn the turn and the cassette would still have an +/// unconsumed interaction at teardown. +#[tokio::test] +async fn a_preparation_failure_costs_no_turn_and_no_interaction() { + with_openai_completions_cassette("agent_driver/prepare_failure", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(2) + .tool(Adder) + .build(); + + // An `active_tools` allow-list naming a tool this turn does not have: + // preparation fails locally, with no provider round trip. + let mut driver = agent + .drive(ADD_PROMPT) + .request_patch(RequestPatch::new().active_tools(["nonexistent_tool"])); + + let error = driver + .next_step() + .await + .expect_err("active_tools naming an unavailable tool must fail at prepare time"); + assert!(matches!(error, PromptError::CompletionError(_))); + assert_eq!( + driver.run().turn(), + 0, + "a request that never left the process must not consume a turn" + ); + + // Fix the cause and drive the very same step again. + driver.set_request_patch(RequestPatch::new()); + let (request, tools, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 1, "the retry takes the turn the failure did not"); + assert!(tools.executable_tool_names().contains("add")); + let response = request.send().await.expect("the retry should send"); + driver.model_response(&response).expect("turn accepted"); + }) + .await; +} + +/// Exhausting the model-call budget stops the run before it builds a second +/// request — so the cassette holds exactly one interaction. +#[tokio::test] +async fn max_turns_exhaustion_stops_before_a_second_send() { + with_openai_completions_cassette("agent_driver/max_turns", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .tool(Adder) + .build(); + + // One model call only; the tool call cannot be answered. + let mut driver = agent.drive(ADD_PROMPT).max_turns(1); + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let error = driver + .next_step() + .await + .expect_err("the budget of one is spent"); + assert!( + matches!(error, PromptError::MaxTurnsError { .. }), + "expected MaxTurnsError, got {error:?}" + ); + }) + .await; +} + +// ── Tranche 5: output modes ────────────────────────────────────────────── + +/// Tool output mode advertises a synthetic output tool alongside the real +/// ones, and the run finalizes on its call rather than dispatching it. +#[tokio::test] +async fn tool_output_mode_finalizes_via_the_output_tool() { + with_openai_completions_cassette("agent_driver/output_mode_tool", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Reply with the structured answer.") + .output_schema_raw( + serde_json::from_value(serde_json::json!({ + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"] + })) + .expect("valid schema"), + ) + .output_mode(OutputMode::Tool) + .build(); + + let mut driver = agent.drive("What is the capital of France?"); + let (request, tools, _) = expect_send(&mut driver).await; + let output_tool = tools + .output_tool_name() + .expect("Tool mode advertises an output tool") + .to_owned(); + assert!( + tools.allowed_tool_names().contains(&output_tool), + "the output tool is allowed" + ); + assert!( + !tools.executable_tool_names().contains(&output_tool), + "the output tool is never executable" + ); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + + // The output-tool call is intercepted by the run, never surfaced. + let response = expect_done(&mut driver).await; + assert!( + response.output.contains("answer"), + "expected structured output, got {}", + response.output + ); + }) + .await; +} + +/// Native output mode sets the provider's own structured-output constraint, so +/// the request carries a response format rather than a synthetic tool. +#[tokio::test] +async fn native_output_mode_uses_the_provider_constraint() { + with_openai_completions_cassette("agent_driver/output_mode_native", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Reply with the structured answer.") + .output_schema_raw( + serde_json::from_value(serde_json::json!({ + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"], + "additionalProperties": false + })) + .expect("valid schema"), + ) + .output_mode(OutputMode::Native) + .build(); + + let mut driver = agent.drive("What is the capital of France?"); + let (request, tools, _) = expect_send(&mut driver).await; + assert!( + tools.output_tool_name().is_none(), + "Native mode advertises no synthetic tool" + ); + + let response = request.send().await.expect("should send"); + driver.model_response(&response).expect("turn accepted"); + let response = expect_done(&mut driver).await; + assert!(response.output.contains("answer")); + }) + .await; +} + +// ── Tranche 6: streaming ───────────────────────────────────────────────── + +/// A hand-driven **streamed** turn goes through the driver, against real SSE. /// /// The streamed entry points live on `AgentRun` and take `&mut self`, so this /// is only expressible because the driver hands out `run_mut()`. Without it a /// streaming caller has to `into_run()` and rebuild the driver, which discards /// the per-turn snapshot cache and makes the driver treat a turn prepared in -/// this very process as a resume — drift check included. The assertion that -/// matters here is the second request body: the streamed turn must thread back -/// into the run exactly as a blocking one does. +/// this very process as a resume — drift check included. #[tokio::test] async fn a_streamed_turn_drives_through_the_driver() { with_openai_completions_cassette("agent_driver/streamed_turn", |client| async move { let calls = Arc::new(AtomicUsize::new(0)); let agent = client .agent(openai::GPT_4O) - .preamble("You are a weather assistant. Always use the weather tool.") + .preamble(FORCE_TOOLS_PREAMBLE) .default_max_turns(3) - .tool(WeatherTool { - call_count: calls.clone(), + .tool(CountingAdder { + calls: calls.clone(), }) .build(); - let mut driver = agent.drive("What is the weather in Tokyo?"); - - let (request, tools) = match driver.next_step().await.expect("first step") { - DriveStep::SendRequest { request, tools, .. } => (request, tools), - other => panic!("expected SendRequest, got {other:?}"), - }; + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send(&mut driver).await; - // Stream the turn and assemble it with the names this turn advertised. let mut assembler = StreamedTurnAssembler::new( tools.executable_tool_names().clone(), tools.allowed_tool_names().clone(), @@ -273,7 +902,6 @@ async fn a_streamed_turn_drives_through_the_driver() { let final_content = stream.choice.clone(); let streamed = assembler.finish(stream.message_id.clone(), &final_content); - // Feed usage and the assembled turn through the driver's run. driver .run_mut() .record_streamed_completion_call(stream.usage()) @@ -283,86 +911,112 @@ async fn a_streamed_turn_drives_through_the_driver() { .streamed_turn(streamed) .expect("streamed turn accepted"); - // The driver still owns the pairing: this turn's snapshot dispatches. - let (pending, tools) = match driver.next_step().await.expect("second step") { - DriveStep::ExecuteTools { calls, tools } => (calls, tools), - other => panic!("expected ExecuteTools, got {other:?}"), - }; + let (pending, tools) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty(), "the model should have called the tool"); + dispatch_and_feed(&mut driver, &pending, &tools).await; - let mut context = ToolContext::new(); - let mut results = Vec::new(); - for call in &pending { - results.push(tools.execute_call(call, &mut context).await); - } - driver.tool_results(results).expect("results accepted"); + let response = drive_to_completion(&mut driver) + .await + .expect("run should finish"); + assert!(!response.output.trim().is_empty()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + }) + .await; +} - let request = match driver.next_step().await.expect("third step") { - DriveStep::SendRequest { request, .. } => request, - other => panic!("expected SendRequest, got {other:?}"), - }; - let response = request.send().await.expect("final turn should send"); - driver.model_response(&response).expect("turn accepted"); +/// A streamed text-only turn: no tool calls, so the run finalizes straight +/// from the assembled stream. +#[tokio::test] +async fn a_streamed_text_turn_finalizes_the_run() { + with_openai_completions_cassette("agent_driver/streamed_text", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Reply with one short sentence.") + .build(); + + let mut driver = agent.drive("Say hello."); + let (request, tools, _) = expect_send(&mut driver).await; - match driver.next_step().await.expect("final step") { - DriveStep::Done(response) => { - assert!(!response.output.trim().is_empty()); - } - other => panic!("expected Done, got {other:?}"), + let mut assembler = StreamedTurnAssembler::new( + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), + ); + let mut stream = request.stream().await.expect("stream should open"); + while let Some(item) = stream.next().await { + assembler + .ingest(&item.expect("stream item")) + .expect("ingest should succeed"); } - assert_eq!(calls.load(Ordering::SeqCst), 1); + let final_content = stream.choice.clone(); + let streamed = assembler.finish(stream.message_id.clone(), &final_content); + + driver + .run_mut() + .record_streamed_completion_call(stream.usage()) + .expect("usage recorded"); + driver + .run_mut() + .streamed_turn(streamed) + .expect("streamed turn accepted"); + + let response = expect_done(&mut driver).await; + assert!(!response.output.trim().is_empty()); + assert_eq!( + driver.run().completion_calls().len(), + 1, + "a streamed turn records exactly one completion call" + ); }) .await; } -/// A real provider rejection is classified non-retryable, so the caller does -/// not hand the turn back and loop. -/// -/// `is_retryable` is the trigger the driver's own docs recommend for -/// `rollback_model_call`, and misclassifying a deterministic failure as -/// retryable is an unbounded loop. A synthetic error cannot falsify the -/// classification — only a real provider rejection can, which is what this -/// records. +/// A streamed turn under a `RequestPatch`: the patch reaches the streaming +/// request exactly as it reaches a blocking one. #[tokio::test] -async fn a_provider_rejection_is_not_retryable() { - with_openai_completions_cassette("agent_driver/provider_rejection", |client| async move { +async fn a_streamed_turn_honors_the_request_patch() { + with_openai_completions_cassette("agent_driver/streamed_patched", |client| async move { let agent = client - // A model name the provider will reject outright. - .agent("gpt-4o-this-model-does-not-exist") - .preamble("You are a weather assistant.") + .agent(openai::GPT_4O) + .preamble("BASELINE — must not appear") + .default_max_turns(2) + .tool(Adder) + .tool(Subtract) .build(); - let mut driver = agent.drive("What is the weather in Tokyo?"); - let request = match driver.next_step().await.expect("first step") { - DriveStep::SendRequest { request, .. } => request, - other => panic!("expected SendRequest, got {other:?}"), - }; + let mut driver = agent.drive(ADD_PROMPT).request_patch( + RequestPatch::new() + .preamble(FORCE_TOOLS_PREAMBLE) + .active_tools(["add"]), + ); - let error = request - .send() - .await - .expect_err("an unknown model must be rejected"); + let (request, tools, _) = expect_send(&mut driver).await; + assert!(!tools.executable_tool_names().contains("subtract")); - // A cassette mock miss is *also* a 404, so asserting on the status - // alone would pass on a cassette that never matched. Pin the - // provider's own error envelope, which the harness cannot fabricate. - let body = error - .provider_response_body() - .expect("the provider's rejection body is preserved"); - assert!( - body.contains("model_not_found"), - "expected the recorded provider rejection, got: {body}" + let mut assembler = StreamedTurnAssembler::new( + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), ); + let mut stream = request.stream().await.expect("stream should open"); + while let Some(item) = stream.next().await { + assembler + .ingest(&item.expect("stream item")) + .expect("ingest should succeed"); + } + let final_content = stream.choice.clone(); + let streamed = assembler.finish(stream.message_id.clone(), &final_content); - assert!( - !error.is_retryable(), - "a provider rejection must not be classified retryable: {error}" - ); + driver + .run_mut() + .record_streamed_completion_call(stream.usage()) + .expect("usage recorded"); + driver + .run_mut() + .streamed_turn(streamed) + .expect("streamed turn accepted"); - // The turn stays in flight: the caller decides, and here the decision - // is to fail rather than hand the turn back. - assert_eq!(driver.run().turn(), 1); - assert_eq!(driver.run().model_call_rollbacks(), 0); + let (pending, tools) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty()); + dispatch_and_feed(&mut driver, &pending, &tools).await; }) .await; } From 67ce85414507ec7d26d071d003fc733210496d6a Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 14:30:52 -0700 Subject: [PATCH 09/21] fix(core): make the cut-stream retryability claim true instead of asserted Round-5 review found six defects. Two were the same one seen from both sides, and together they mean the retryability work shipped in rounds 3 and 4 did not cover the failure it was written for. `from_stream_transport` folded every statusless transport error into `ProviderError(String)` - rig's unclassifiable bucket - and every streaming provider routes mid-stream failures through it: anthropic, openai, gemini's two APIs, cohere, the openai-compatible shim. So `is_retryable` answered false for a cut stream on every provider, while its own docs, the changelog and MIGRATING all said otherwise, and `rollback_model_call`'s guidance was built on it. The test that claimed to cover this hand-built `HttpError(StreamEnded)`, which nothing produces. In the SSE layer `StreamEnded` is the *normal* end-of-stream sentinel that every provider matches and breaks on; its only error-producing sites are the wasm stub clients, where it means the transport cannot send at all. So the arm classifying it transient existed only to satisfy that test. Round 2's rule - falsify against the finding - passes here and is not enough; the tightened rule is that a test asserting a provider-facing failure must reach it through a path production uses, and must say in its doc comment what that path is. The fix is codex's: a semantic variant survives stringification where a generic one does not. `CodexErrorDetails::Stream(String)` carries the same lossy payload and stays classifiable because the variant does. `CompletionError::StreamInterrupted` is that, gated on the transport's own `is_transient` - the first draft routed *every* statusless failure there and two pre-existing tests caught it, since `InvalidContentType` reaches the same call sites and is deterministic. Deterministic statusless failures keep `ProviderError`; those tests pass unchanged. `Instance` is now classified by what it wraps rather than assumed transient - round 4's still-open finding. It is the client's catch-all and holds both a dropped connection and a request the client refused to build, and a `base_url` missing its scheme loops any caller driving off `is_retryable`. The `source()` chain is walked for a `reqwest::Error` that `is_builder()` or `is_redirect()`, following openai-agents' `iter_error_chain`; anything unrecognised keeps the transient default. Four smaller fixes. `commit_model_call` now enforces `max_turns`: the check lived only in `peek_model_call` while commit - public, and documented as "the only place a turn is consumed" - incremented unconditionally, so a hand-driver that peeked once and re-committed drove calls forever. `AgentDriver::max_invalid_tool_call_retries` exists, since `drive()` seeded zero and the documented `Retry` resolution could never succeed. `ModelTurnOutcome` is `#[must_use]`, which immediately caught every drop site including this PR's own test helper - a hallucinated tool name used to surface two steps later as an unrelated protocol violation. And resume no longer runs a retrieval query whose result is discarded in full by the narrowing that follows it, which also removes a resume-time failure when the index is down. One honest limit recorded rather than papered over: the cassette derived for `StreamInterrupted` cannot produce it. A mock that stops writing closes cleanly, so the SSE layer sees end-of-stream, not a transport failure; only a connection severed mid-frame does, and no mock can do that. The test was kept and inverted to pin the clean close and say why the classification stays unit-tested. 518 rig-agent, 1251 rig-core, 129/86/141 cassette, clippy -D warnings clean, fmt clean, 0 rustdoc warnings. Every new behavioral test was falsified against its fix by deliberate revert. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/CHANGELOG.md | 4 + crates/rig-agent/src/agent/driver.rs | 159 ++++++++++++---- crates/rig-agent/src/agent/run/mod.rs | 82 +++++++-- crates/rig-core/CHANGELOG.md | 3 + crates/rig-core/src/completion/request.rs | 170 ++++++++++++++++-- crates/rig-core/src/http_client/mod.rs | 46 ++++- .../agent_driver/streamed_interrupted.yaml | 27 +++ tests/common/driver_support.rs | 35 +++- .../anthropic/cassette/agent_driver.rs | 28 ++- .../providers/gemini/cassette/agent_driver.rs | 28 ++- .../providers/openai/cassette/agent_driver.rs | 119 +++++++++--- 11 files changed, 582 insertions(+), 119 deletions(-) create mode 100644 tests/cassettes/openai/agent_driver/streamed_interrupted.yaml diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index f22b55529a..3b43bc66c6 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -43,6 +43,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* [**behavior**] `TurnTools::execute` now refuses any name the turn did not advertise, instead of trusting its snapshot to be narrowed. The two agreed by construction in-process; on a resumed turn the names come from the run and the snapshot is rebuilt locally, so a tool registered after suspension could previously be dispatched through a turn that never advertised it - *(agent)* [**behavior**] `TurnTools::execute_call` clears the context's dispatch result on the pre-resolved path too, so a call suppressed by invalid tool-call recovery no longer leaves the *previous* call's result metadata readable — a loop over a turn's calls was attributing it to the suppressed one - *(agent)* add `TurnToolNames::new`: the type is `#[non_exhaustive]` serialized state, so an external hand-driver could not build the value `AgentRun::commit_model_call` requires +- *(agent)* add `AgentDriver::max_invalid_tool_call_retries`, mirroring the runner's. `Agent::drive` seeds the budget at zero and the driver exposed no override, so answering a `NeedsResolution` with `InvalidToolCallAction::Retry` — the path `AgentDriver::resolve_invalid_tool_call`'s own docs point at — could never succeed on a driver built the documented way +- *(agent)* [**breaking**] `ModelTurnOutcome` is `#[must_use]`. It carries `NeedsResolution`, which must be answered before the run may advance; dropping it let a hallucinated tool name surface two steps later as an unrelated "next_step called while an invalid tool-call resolution is pending". Callers writing `driver.model_response(&r)?;` now get a warning +- *(agent)* [**behavior**] `AgentRun::commit_model_call` enforces the model-call budget. The check previously lived only in `peek_model_call`, while `commit_model_call` — public, and documented as "the only place a turn is consumed" — incremented unconditionally, so a hand-driver that peeked once and re-committed drove model calls forever against a budget never consulted +- *(agent)* [**behavior**] resuming a run no longer runs a retrieval query. The resumed snapshot is narrowed to the turn's advertised names immediately after it is taken, so the query's contribution was discarded in full while costing a vector search — and made resume *fail* when the index was unavailable, even when every pending call was a static tool - *(agent)* add `AgentDriver::run_mut`, without which a streamed turn could not be driven through the driver at all — `AgentRun`'s streamed entry points take `&mut self`, so a streaming caller had to `into_run()` and rebuild the driver, discarding the per-turn snapshot cache and making the driver treat a turn prepared in the same process as a resume (drift check included). Committing or rolling back a model call through it is documented as unsupported; feeding a response is safe - *(agent)* [**behavior**] `AgentRun::advertised_tools()` no longer outlives the turn it describes. Only `rollback_model_call` ended the turn's names; `reprompt_for_output`, `retry_model_turn`, `resolve_invalid_tool_call(Retry)`, `tool_results` and both streamed-abandon paths left them in place, so a run serialized while parked for a fresh model call reported a turn that was already over. Every route back to `PreparingRequest` now clears them - *(agent)* [**behavior**] a resumed run's dispatch snapshot resolves the turn's advertised names explicitly rather than relying on retrieval to rank them again, so a registered *dynamic* tool that the re-derived query does not return is no longer misreported as "not registered in this process" — advice that could not be followed — and no longer feeds the model a not-found result for a tool that exists diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 448317e629..27283ba470 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -216,6 +216,19 @@ impl AgentDriver { self } + /// Set the retry budget for [`InvalidToolCallAction::Retry`] resolutions, + /// mirroring [`AgentRunner::max_invalid_tool_call_retries`](super::AgentRunner::max_invalid_tool_call_retries). + /// + /// [`Agent::drive`] seeds this at **zero**, so answering a + /// [`ModelTurnOutcome::NeedsResolution`] with `Retry` fails until you raise + /// it — the run has no budget to spend and reports the invalid call + /// instead. Invalid tool-call retries also consume the total model-call + /// budget, so raise [`Self::max_turns`] alongside it. + pub fn max_invalid_tool_call_retries(mut self, retries: usize) -> Self { + self.run = self.run.max_invalid_tool_call_retries(retries); + self + } + /// Set the per-turn request configuration for the turns this driver /// prepares. /// @@ -498,28 +511,28 @@ impl AgentDriver { Ok(TurnTools::from_parts(snapshot, names, output_tool_name)) } - /// Take a registry snapshot for a run resumed in this process. + /// Take a registry snapshot for a run resumed in this process, containing + /// exactly the names the turn advertised and still has. /// - /// The retrieval query is re-derived from the run's history, matching what - /// request preparation would have used — but retrieval alone is not enough - /// here. It selects dynamic tools by similarity, so a registered dynamic - /// tool this query does not rank would be absent, and the caller's drift - /// check could not tell that apart from a tool that really was - /// deregistered. `required` — the names the turn advertised — is resolved - /// from the registry regardless of ranking, so absence means absence. + /// **No retrieval query is run**, deliberately. Retrieval selects dynamic + /// tools by similarity, and this snapshot is narrowed to `required` + /// immediately afterwards, so every retrieved name that is not in + /// `required` is discarded and every name in `required` is resolved from + /// the registry by name regardless of ranking. A query would therefore + /// contribute nothing to the result while costing a vector search — and, + /// worse, would make resuming a run *fail* when the index is unavailable, + /// even if every pending call is a static tool. + /// + /// Resolving `required` by name is also what makes the caller's drift check + /// meaningful: absence from this snapshot means the tool is gone, not + /// merely that a query did not rank it. async fn fresh_snapshot( &self, required: &BTreeSet, ) -> Result { - let query = self - .run - .full_history() - .iter() - .rev() - .find_map(|message| message.rag_text()); self.agent .tool_server_handle - .snapshot_tool_defs_including(query, required) + .snapshot_tool_defs_including(None, required) .await .map_err(|_| { PromptError::CompletionError(CompletionError::RequestError( @@ -568,6 +581,18 @@ mod tests { })) } + /// Assert a model turn was accepted outright. + /// + /// `ModelTurnOutcome` is `#[must_use]`: `NeedsResolution` has to be + /// answered before the run may advance, so these tests name an unexpected + /// one rather than dropping it and hitting a protocol violation later. + fn expect_continue(outcome: ModelTurnOutcome) { + match outcome { + ModelTurnOutcome::Continue { .. } => {} + other => panic!("expected the turn to be accepted, got {other:?}"), + } + } + /// Expect the next step to be `SendRequest`, panicking otherwise. macro_rules! expect_send { ($driver:expr) => { @@ -613,7 +638,7 @@ mod tests { assert!(tools.executable_tool_names().contains("add")); assert_eq!(tools.executable_tool_names(), tools.allowed_tool_names()); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); // Mutate the live registry AFTER the turn advertised its tools. agent.tool_server_handle.remove_tool("add").await; @@ -646,7 +671,7 @@ mod tests { let (request, _, turn) = expect_send!(driver); assert_eq!(turn, 2); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); match driver.next_step().await.expect("next_step succeeds") { DriveStep::Done(response) => assert_eq!(response.output, "3"), other => panic!("expected Done, got {other:?}"), @@ -717,7 +742,7 @@ mod tests { let mut driver = agent.drive("compute"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); match driver.next_step().await.expect("next_step succeeds") { DriveStep::Done(response) => { assert!( @@ -751,7 +776,7 @@ mod tests { .expect("Tool mode commits a name on turn 1") .to_owned(); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn processed"); + expect_continue(driver.model_response(&response).expect("turn processed")); // Change the tool set between turns: retire every executable tool. agent.tool_server_handle.remove_tool("add").await; @@ -783,7 +808,7 @@ mod tests { let mut driver = agent.drive("what is 2 + 5?"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); let _ = expect_execute_tools!(driver); let serialized = serde_json::to_string(driver.run()).expect("run serializes"); drop(driver); @@ -805,7 +830,7 @@ mod tests { driver.tool_results(results).expect("results accepted"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); match driver.next_step().await.expect("next_step succeeds") { DriveStep::Done(response) => assert_eq!(response.output, "7"), other => panic!("expected Done, got {other:?}"), @@ -829,7 +854,7 @@ mod tests { let mut driver = agent.drive("what is 2 + 5?"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); let _ = expect_execute_tools!(driver); let serialized = serde_json::to_string(driver.run()).expect("run serializes"); @@ -979,7 +1004,7 @@ mod tests { assert_eq!(turn, 1, "the retry takes the turn the failure did not"); assert!(tools.executable_tool_names().contains("add")); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); } /// The natural suspension point for a caller that owns the transport is @@ -1012,9 +1037,11 @@ mod tests { .build(); let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); let mut driver = agent.drive_run(run); - driver - .model_response(&response) - .expect("a run resumed mid-model-call accepts its reply"); + expect_continue( + driver + .model_response(&response) + .expect("a run resumed mid-model-call accepts its reply"), + ); let (calls, tools) = expect_execute_tools!(driver); let mut context = ToolContext::new(); @@ -1025,7 +1052,7 @@ mod tests { driver.tool_results(results).expect("results accepted"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); match driver.next_step().await.expect("next_step succeeds") { DriveStep::Done(response) => assert_eq!(response.output, "7"), other => panic!("expected Done, got {other:?}"), @@ -1111,7 +1138,7 @@ mod tests { let (request, _, turn) = expect_send!(driver); assert_eq!(turn, 1, "the retry takes the turn the failure did not"); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); match driver.next_step().await.expect("next_step succeeds") { DriveStep::Done(response) => assert_eq!(response.output, "done"), other => panic!("expected Done, got {other:?}"), @@ -1167,7 +1194,7 @@ mod tests { let (request, _, turn) = expect_send!(driver); assert_eq!(turn, 1); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); } /// `TurnTools` promises that a name the turn did not advertise cannot @@ -1188,7 +1215,7 @@ mod tests { let mut driver = agent.drive("what is 2 + 5?"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); let _ = expect_execute_tools!(driver); let serialized = serde_json::to_string(driver.run()).expect("run serializes"); @@ -1221,6 +1248,72 @@ mod tests { assert!(result.is_success()); } + /// `Retry` needs a budget, and `drive()` seeds zero — so the driver must + /// expose the setter the runner has, or the documented resolution path is + /// unreachable for every driver built the documented way. + #[tokio::test] + async fn a_retry_resolution_needs_a_budget_the_driver_can_set() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("call_1", "nonexistent", json!({})), + MockTurn::text("ok"), + ]); + let agent = AgentBuilder::new(model) + .default_max_turns(3) + .tool(MockAddTool) + .build(); + + // Seeded at zero: the retry has nothing to spend. + let mut driver = agent.drive("go"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + let outcome = driver + .model_response(&response) + .expect("the invalid call needs resolution"); + assert!(matches!(outcome, ModelTurnOutcome::NeedsResolution(_))); + driver + .resolve_invalid_tool_call(InvalidToolCallAction::Retry { + feedback: "use a registered tool".to_string(), + }) + .expect_err("a retry budget of zero cannot retry"); + } + + /// With a budget, the same resolution is accepted and the run re-prepares. + #[tokio::test] + async fn a_retry_resolution_succeeds_once_the_budget_is_set() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("call_1", "nonexistent", json!({})), + MockTurn::text("recovered"), + ]); + let agent = AgentBuilder::new(model) + .default_max_turns(3) + .tool(MockAddTool) + .build(); + + let mut driver = agent.drive("go").max_invalid_tool_call_retries(1); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + let outcome = driver + .model_response(&response) + .expect("the invalid call needs resolution"); + assert!(matches!(outcome, ModelTurnOutcome::NeedsResolution(_))); + + let outcome = driver + .resolve_invalid_tool_call(InvalidToolCallAction::Retry { + feedback: "use a registered tool".to_string(), + }) + .expect("a retry budget of one accepts the retry"); + assert!(matches!(outcome, ModelTurnOutcome::TurnRetried)); + + // The retry re-prepares and the run completes. + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + expect_continue(driver.model_response(&response).expect("turn accepted")); + match driver.next_step().await.expect("next_step succeeds") { + DriveStep::Done(response) => assert_eq!(response.output, "recovered"), + other => panic!("expected Done, got {other:?}"), + } + } + /// Retrieval picks dynamic tools by similarity to the turn's query, so a /// registered dynamic tool can be absent from a resumed snapshot purely /// because that query did not rank it. Reporting that as "no longer @@ -1255,7 +1348,7 @@ mod tests { let (request, tools, _) = expect_send!(driver); assert!(tools.executable_tool_names().contains("subtract")); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); let _ = expect_execute_tools!(driver); let serialized = serde_json::to_string(driver.run()).expect("run serializes"); @@ -1305,7 +1398,7 @@ mod tests { let mut driver = agent.drive("what is 2 + 5?"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); let _ = expect_execute_tools!(driver); let serialized = serde_json::to_string(driver.run()).expect("run serializes"); @@ -1348,7 +1441,7 @@ mod tests { let mut driver = agent.drive("what is 2 + 5?"); let (request, _, _) = expect_send!(driver); let response = request.send().await.expect("scripted turn"); - driver.model_response(&response).expect("turn accepted"); + expect_continue(driver.model_response(&response).expect("turn accepted")); let _ = expect_execute_tools!(driver); let serialized = serde_json::to_string(driver.run()).expect("run serializes"); diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index c3575ad402..a3066c7470 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -288,6 +288,10 @@ impl ModelTurn { /// Deliberately exhaustive: a driver must handle every outcome, so adding a /// variant is a breaking change by design. #[derive(Debug)] +#[must_use = "a model turn can need resolution before the run may advance; \ + answer `NeedsResolution` via `resolve_invalid_tool_call` — dropping \ + the outcome surfaces the same problem two steps later as an \ + unrelated protocol violation"] pub enum ModelTurnOutcome { /// The turn was accepted. Unless `response_hook_suppressed` is set, the /// driver should run its completion-response hook now, then call @@ -846,13 +850,7 @@ impl AgentRun { )); }; - if self.current_turn >= self.max_turns { - return Err(PromptError::MaxTurnsError { - max_turns: self.max_turns, - chat_history: self.full_history().into(), - prompt: prompt.clone().into(), - }); - } + self.check_turn_budget()?; Ok(ModelCallInputs { prompt: prompt.clone(), @@ -860,6 +858,31 @@ impl AgentRun { }) } + /// Whether another model call fits in the run's budget. + /// + /// Checked in **two** places on purpose. [`Self::peek_model_call`] runs it + /// so a caller learns the budget is spent before building a request it + /// cannot send; [`Self::commit_model_call`] runs it because that is where + /// the turn is actually consumed, and the two halves are public — nothing + /// obliges a caller to peek before every commit, and a caller that peeks + /// once and re-commits (after a rollback, say) would otherwise drive model + /// calls forever against a budget never consulted. + fn check_turn_budget(&self) -> Result<(), PromptError> { + if self.current_turn < self.max_turns { + return Ok(()); + } + let prompt = self + .new_messages + .last() + .cloned() + .unwrap_or_else(|| Message::user("")); + Err(PromptError::MaxTurnsError { + max_turns: self.max_turns, + chat_history: self.full_history().into(), + prompt: prompt.into(), + }) + } + /// Commit the model call previewed by [`Self::peek_model_call`], returning /// its one-based turn index. /// @@ -879,9 +902,13 @@ impl AgentRun { /// lost, [`Self::rollback_model_call`] undoes exactly this. /// /// # Errors - /// [`PromptError::PromptCancelled`] when the run is not preparing a - /// request — committing a call the run never asked for would consume a - /// turn against nothing. + /// - [`PromptError::PromptCancelled`] when the run is not preparing a + /// request — committing a call the run never asked for would consume a + /// turn against nothing. + /// - [`PromptError::MaxTurnsError`] when the model-call budget is spent. + /// Checked here as well as in [`Self::peek_model_call`], because this is + /// where the turn is actually consumed and a caller is not obliged to + /// peek before every commit. pub fn commit_model_call( &mut self, advertised: Option, @@ -892,6 +919,7 @@ impl AgentRun { self.protocol_violation("commit_model_call called without a peeked model call") ); } + self.check_turn_budget()?; self.set_output_tool_name(output_tool_name); self.advertised_tools = advertised; self.current_turn += 1; @@ -2133,6 +2161,40 @@ mod tests { ); } + /// The budget is checked where it is *spent*, not only where it is + /// previewed. Both halves are public, and nothing obliges a caller to peek + /// before every commit — a caller that peeks once and re-commits (after a + /// rollback, say) would otherwise drive model calls forever. + #[test] + fn commit_model_call_enforces_the_turn_budget_without_a_peek() { + let mut run = AgentRun::new("add things").max_turns(1); + + // Spend the budget through the normal path. + expect_call_model(&mut run); + run.rollback_model_call().expect("rollback should succeed"); + expect_call_model(&mut run); + assert_eq!(run.turn(), 1); + expect_continue( + run.model_response(tool_call_turn("call_1", "add")) + .expect("model_response should succeed"), + ); + expect_call_tools(&mut run); + run.tool_results(vec![tool_result("call_1", "2")]) + .expect("tool_results should succeed"); + + // The run wants another call and the budget is spent. Committing + // without peeking must not get one. + assert!(run.is_preparing_request()); + let err = run + .commit_model_call(None, None) + .expect_err("the budget is spent"); + assert!( + matches!(err, PromptError::MaxTurnsError { .. }), + "expected MaxTurnsError, got {err:?}" + ); + assert_eq!(run.turn(), 1, "a refused commit consumes nothing"); + } + /// The two halves are a public protocol, so driving them out of order is /// an error a caller can handle — not a `debug_assert` that vanishes in /// release builds. diff --git a/crates/rig-core/CHANGELOG.md b/crates/rig-core/CHANGELOG.md index 66c3424c19..6b11f65874 100644 --- a/crates/rig-core/CHANGELOG.md +++ b/crates/rig-core/CHANGELOG.md @@ -56,6 +56,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- *(completion)* `CompletionError::StreamInterrupted`: the response stream failed after the provider accepted the request. Previously every streaming provider folded these into `ProviderError(String)` — rig's unclassifiable bucket — which meant `is_retryable` answered `false` for the one failure the hand-driven rollback story is built around, on every provider, while the docs claimed otherwise. The string payload is unchanged; only the variant carrying it is, and that is what makes the failure classifiable at all. Retryable and deliberately **not** replay-safe: the request reached the provider, so a retry may bill a second completion. The enum is `#[non_exhaustive]`, so this is additive +- *(completion)* [**behavior**] `is_retryable` no longer treats `http_client::Error::StreamEnded` as transient. In the SSE layer it is the *normal* end-of-stream sentinel that every provider matches and breaks on, never surfacing as an error; its only error-producing sites are the wasm stub clients, where it means the transport cannot send at all. Neither is a retryable failure +- *(completion)* [**behavior**] `is_retryable` classifies `http_client::Error::Instance` by what it wraps rather than assuming it transient. `Instance` is the client's catch-all and carries both a dropped connection *and* a request the client refused to build — a `base_url` missing its scheme, say, which fails identically forever and loops any caller driving off this method. The wrapped error's `source()` chain is walked for a `reqwest::Error` that `is_builder()` or `is_redirect()`; anything unrecognised keeps the transient default - *(completion)* `CompletionError::is_retryable`: a conservative classification of whether re-issuing an equivalent request could plausibly succeed — preserved provider statuses of 408, 409, 429 or 5xx, plus transport failures the transport itself reports as transient (a dropped connection, a cut stream). Deterministic transport failures are excluded by name: a header value the client refused (an API key with a trailing newline, say), a request that could not be constructed, or a wrong content type will fail identically forever, and retrying them loops. A 2xx carrying a provider-authored error envelope is not retryable, since the call completed and was rejected on its merits. Two documented limits: providers on non-HTTP transports (Bedrock, Vertex, gRPC Gemini) report errors as a body with no status, so their throttling classifies as non-retryable until classification moves to the provider adapter; and this answers only whether a retry *could succeed*, never whether one is *replay-safe* - *(tool)* `ToolErrorKind::NotExecutable` and `ToolExecutionError::not_executable`: the tool is advertised to the model but not executable by the dispatcher — its call must be handled by the host. Produced by rig-agent when dispatching the synthetic structured-output tool; also the natural classification for host-executed or deferred tools. The enum is `#[non_exhaustive]`, so this is additive - *(streaming)* `wire::classify_typed_event` extends the decode-then-validate policy to typed-transport wires (bedrock, candle, gemini-grpc): modeled variants are `Known`, the SDK's non-exhaustive/unrecognized variants are `Unknown`, SDK decode errors are `Corrupt` — a typed transport earns no policy exemption diff --git a/crates/rig-core/src/completion/request.rs b/crates/rig-core/src/completion/request.rs index 61cdc7c9c5..9359553bd1 100644 --- a/crates/rig-core/src/completion/request.rs +++ b/crates/rig-core/src/completion/request.rs @@ -114,6 +114,27 @@ pub enum CompletionError { /// Raw error response preserved from the completion model provider #[error("ProviderResponseError: {0}")] ProviderResponse(provider_response::ProviderResponseError), + + /// The response stream failed after the request was accepted and before + /// the provider finished emitting it. + /// + /// Distinct from [`ProviderError`](Self::ProviderError) — rig's + /// unclassifiable bucket — because the *cause* is known even when the + /// provider's own error type is not: the connection or the stream broke + /// mid-response. Keeping that as its own variant is what lets + /// [`is_retryable`](Self::is_retryable) classify it at all; folding it into + /// a display string would destroy the classification while preserving only + /// the message. + /// + /// **Retryable, and not replay-safe.** The request reached the provider — + /// that is what "after the request was accepted" means — so a retry may + /// bill a second completion and repeat any side effect the model already + /// caused. This is the canonical instance of the two questions + /// `is_retryable` deliberately does not conflate; see its docs, and + /// `AgentRun::rollback_model_call` in `rig-agent` for the axis it cannot + /// answer for you. + #[error("StreamInterrupted: {0}")] + StreamInterrupted(String), } crate::provider_response::impl_provider_response_helpers!(CompletionError); @@ -122,11 +143,24 @@ impl CompletionError { /// Maps an SSE transport error into a completion error without flattening HTTP failures. /// /// Non-success HTTP responses remain [`CompletionError::HttpError`] so provider response - /// helpers can read status and body. Other transport failures keep the existing - /// [`CompletionError::ProviderError`] display string behavior. + /// helpers can read status and body. A statusless failure the transport + /// reports as transient — the connection or the stream broke — becomes + /// [`CompletionError::StreamInterrupted`], which is what lets + /// [`is_retryable`](Self::is_retryable) classify it; every mid-stream + /// failure used to fold into [`CompletionError::ProviderError`], the + /// unclassifiable bucket, so a cut stream answered `false` on every + /// provider while the docs claimed otherwise. + /// + /// Deterministic statusless failures — a wrong content type, a request that + /// could not be built — keep the existing `ProviderError` behavior. They + /// arrive here from the same call sites but will fail identically on every + /// attempt, and widening `StreamInterrupted` to cover them would trade one + /// misclassification for another. pub(crate) fn from_stream_transport(error: http_client::Error) -> Self { if error.non_success_status().is_some() { Self::HttpError(error) + } else if error.is_transient() { + Self::StreamInterrupted(error.to_string()) } else { Self::ProviderError(error.to_string()) } @@ -145,13 +179,17 @@ impl CompletionError { /// - A preserved provider status of `408` (request timeout), `409` /// (conflict), `429` (too many requests), or `5xx` — the same four /// classes the AI SDK treats as retryable by default. + /// - [`StreamInterrupted`](Self::StreamInterrupted): the response stream + /// broke after the provider accepted the request. Retryable — and see + /// the replay-safety note below, because this is the case where the two + /// questions diverge most sharply. /// - A transport failure with no status, when the transport itself - /// reports one that could resolve on its own: a dropped connection or a - /// cut stream. Deterministic transport failures — a header value the - /// client refused, a request that could not be constructed, a wrong - /// content type — are **not** retryable, because retrying them is how a - /// caller ends up in a loop it cannot leave. (An API key with a trailing - /// newline is the canonical example: it fails identically forever.) + /// reports one that could resolve on its own — a dropped connection. + /// Deterministic transport failures — a header value the client refused, + /// a request that could not be constructed, a wrong content type — are + /// **not** retryable, because retrying them is how a caller ends up in a + /// loop it cannot leave. (An API key with a trailing newline is the + /// canonical example: it fails identically forever.) /// - A 2xx carrying a provider-authored error envelope is not retryable: /// the call completed and the provider rejected it on its merits. /// @@ -204,6 +242,10 @@ impl CompletionError { // The transport owns this call: it knows which of its own // failures are deterministic. Self::HttpError(error) => error.is_transient(), + // The stream broke after the provider accepted the request. + // Retryable, and the one case where "could a retry succeed" + // and "is a retry safe" give different answers. + Self::StreamInterrupted(_) => true, // A provider body with no status. Rig cannot tell throttling // from rejection without knowing the provider's envelope, so // it does not guess. Documented above as a known gap — do not @@ -1395,12 +1437,112 @@ mod tests { /// A transport failure may never have reached the provider at all. #[test] fn transient_transport_failures_are_retryable() { - for error in [ - CompletionError::HttpError(http_client::Error::Instance("connection reset".into())), - CompletionError::HttpError(http_client::Error::StreamEnded), - ] { - assert!(error.is_retryable(), "{error} should be retryable"); - } + // Produced by `instance_error` on every blocking send path when the + // client's own request fails — the connection dropped. + let error = + CompletionError::HttpError(http_client::Error::Instance("connection reset".into())); + assert!(error.is_retryable(), "{error} should be retryable"); + } + + /// Produced by `from_stream_transport`, which every streaming provider + /// routes mid-stream failures through — anthropic, openai, gemini + /// (both APIs), cohere, and the openai-compatible shim. + /// + /// This case previously classified `false`: the statusless branch + /// folded into `ProviderError`, rig's unclassifiable bucket, so the + /// documented "a cut stream is retryable" guarantee was unreachable on + /// every provider. The test that claimed to cover it constructed + /// `HttpError(StreamEnded)`, a shape no production path produces. + #[test] + fn an_interrupted_stream_is_retryable() { + let error = CompletionError::from_stream_transport(http_client::Error::Instance( + "connection reset by peer".into(), + )); + assert!( + matches!(error, CompletionError::StreamInterrupted(_)), + "a statusless stream failure must keep its classification, got {error:?}" + ); + assert!(error.is_retryable()); + } + + /// `StreamInterrupted` covers the *transient* statusless failures and + /// no more. A deterministic one arriving through the same call site — + /// a provider that answered with the wrong content type — keeps the + /// unclassifiable `ProviderError` behavior, because it will fail + /// identically on every attempt and widening the retryable variant to + /// cover it would trade one misclassification for another. + #[test] + fn a_deterministic_stream_failure_is_not_interrupted() { + let error = + CompletionError::from_stream_transport(http_client::Error::InvalidContentType( + http::HeaderValue::from_static("application/json"), + )); + assert!( + matches!(error, CompletionError::ProviderError(_)), + "a deterministic stream failure must not become StreamInterrupted, got {error:?}" + ); + assert!(!error.is_retryable()); + } + + /// A stream failure that *does* carry a status keeps going through the + /// status classification, so a mid-stream 400 stays non-retryable. + #[test] + fn a_stream_failure_with_a_status_is_classified_by_status() { + let rejected = CompletionError::from_stream_transport( + http_client::Error::InvalidStatusCodeWithMessage( + http::StatusCode::BAD_REQUEST, + "bad request".to_string(), + ), + ); + assert!(!rejected.is_retryable()); + + let throttled = CompletionError::from_stream_transport( + http_client::Error::InvalidStatusCodeWithMessage( + http::StatusCode::TOO_MANY_REQUESTS, + "slow down".to_string(), + ), + ); + assert!(throttled.is_retryable()); + } + + /// The new variant must not widen the unclassifiable bucket: a plain + /// `ProviderError` stays non-retryable. + #[test] + fn a_provider_error_remains_unclassifiable() { + assert!(!CompletionError::ProviderError("opaque".into()).is_retryable()); + } + + /// `StreamEnded` is the SSE layer's normal end-of-stream sentinel — + /// every provider matches it and breaks rather than surfacing it — and + /// its only error-producing sites are the wasm stub clients, where it + /// means the transport cannot send at all. Neither is transient. + #[test] + fn the_end_of_stream_sentinel_is_not_a_retryable_failure() { + let error = CompletionError::HttpError(http_client::Error::StreamEnded); + assert!(!error.is_retryable()); + } + + /// `Instance` carries both classes. A request the client refused to + /// build — a `base_url` missing its scheme is the canonical case — + /// fails identically forever, and retrying it is an unbounded loop. + #[tokio::test] + async fn a_client_side_build_failure_is_not_retryable() { + let reqwest_error = reqwest::Client::new() + .get("api.openai.com/v1/chat/completions") // no scheme + .send() + .await + .expect_err("reqwest refuses a schemeless URL"); + assert!( + reqwest_error.is_builder(), + "expected a builder-kind error, got {reqwest_error:?}" + ); + + let error = + CompletionError::HttpError(http_client::Error::Instance(Box::new(reqwest_error))); + assert!( + !error.is_retryable(), + "a deterministic client error must not be retried: {error}" + ); } /// The failure this guards against is unbounded: a driver following diff --git a/crates/rig-core/src/http_client/mod.rs b/crates/rig-core/src/http_client/mod.rs index 8f6c6bce35..24a54ef9f2 100644 --- a/crates/rig-core/src/http_client/mod.rs +++ b/crates/rig-core/src/http_client/mod.rs @@ -72,15 +72,21 @@ impl Error { /// classified deliberately rather than inherit a default. pub(crate) fn is_transient(&self) -> bool { match self { - // The connection or the stream failed. The request may never have - // been seen — or may have been seen and its reply lost. - Self::StreamEnded | Self::Instance(_) => true, + // The client's own failure. Usually a connection that dropped, but + // not always — classify by what it actually wraps. + Self::Instance(error) => !is_deterministic_client_error(error.as_ref()), // Deterministic: the request, the header, our use of the client, // or the provider's content type is wrong, and will be wrong again. Self::Protocol(_) | Self::InvalidHeaderValue(_) | Self::NoHeaders | Self::InvalidContentType(_) => false, + // Not a failure worth retrying in either of its two roles: in the + // streaming layer it is the *normal* end-of-stream sentinel, which + // every provider matches and breaks on rather than surfacing as an + // error, and its only error-producing sites are the wasm stub + // clients, where it means "this transport cannot send at all". + Self::StreamEnded => false, // Carried by status classification instead; see // `CompletionError::is_retryable`. Self::InvalidStatusCode(_) | Self::InvalidStatusCodeWithMessage(..) => false, @@ -88,6 +94,40 @@ impl Error { } } +/// Whether a boxed client error is one that will fail identically on every +/// attempt. +/// +/// [`Error::Instance`] is the catch-all for the underlying HTTP client, and the +/// client puts both classes in it: a connection reset (transient) and a request +/// that could not be built or a URL it refused to parse (deterministic — a +/// `base_url` missing its scheme fails forever, and retrying it is an unbounded +/// loop for any caller driving off [`CompletionError::is_retryable`]). +/// +/// The classification is not on the box's own type but on what it wraps, so the +/// chain is walked to the bottom — a client error is frequently a wrapper. An +/// unrecognised error defaults to transient: rig cannot see inside it, and a +/// dropped connection is overwhelmingly the common case. +#[cfg(not(target_family = "wasm"))] +fn is_deterministic_client_error(error: &(dyn std::error::Error + 'static)) -> bool { + let mut current = Some(error); + while let Some(error) = current { + if let Some(reqwest_error) = error.downcast_ref::() + && (reqwest_error.is_builder() || reqwest_error.is_redirect()) + { + return true; + } + current = error.source(); + } + false +} + +/// wasm has no reqwest error to inspect, so nothing is recognised as +/// deterministic and everything keeps the transient default. +#[cfg(target_family = "wasm")] +fn is_deterministic_client_error(_error: &(dyn std::error::Error + 'static)) -> bool { + false +} + pub type Result = std::result::Result; #[cfg(not(target_family = "wasm"))] diff --git a/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml b/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml new file mode 100644 index 0000000000..873a18abae --- /dev/null +++ b/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml @@ -0,0 +1,27 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: text/event-stream + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true},"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"type":"number"},"y":{"type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: text/event-stream; charset=utf-8 + body: |+ + data: {"choices":[{"delta":{"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"","name":"add"},"id":"call_REDACTED_1","index":0,"type":"function"}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_1","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"{\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_2","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"x"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_3","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"\":"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_4","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"2"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_5","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + + data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":",\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_6","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} diff --git a/tests/common/driver_support.rs b/tests/common/driver_support.rs index 0a03efa38d..3820626ba1 100644 --- a/tests/common/driver_support.rs +++ b/tests/common/driver_support.rs @@ -12,8 +12,10 @@ #![allow(dead_code)] -use rig::agent::{AgentDriver, DriveStep, PendingToolCall, TurnTools}; -use rig::completion::PromptError; +use rig::agent::{ + AgentDriver, DriveStep, InvalidToolCallAction, ModelTurnOutcome, PendingToolCall, TurnTools, +}; +use rig::completion::{CompletionResponse, PromptError}; use rig::tool::ToolContext; /// Preamble that reliably drives a tool call on every provider tested. @@ -75,9 +77,31 @@ pub(crate) async fn dispatch_and_feed( .expect("tool results should be accepted"); } +/// Feed a model response and assert the turn was accepted outright. +/// +/// `ModelTurnOutcome` is `#[must_use]` because `NeedsResolution` must be +/// answered before the run may advance. These suites drive well-behaved turns, +/// so an unexpected resolution request is a test failure worth naming — not a +/// value to drop on the floor, which would resurface two steps later as an +/// unrelated protocol violation. +pub(crate) fn expect_turn_accepted(driver: &mut AgentDriver, response: &CompletionResponse) { + match driver + .model_response(response) + .expect("the model turn should be accepted") + { + ModelTurnOutcome::Continue { .. } => {} + other => panic!("expected the turn to be accepted outright, got {other:?}"), + } +} + /// Drive a blocking run to completion, returning the final response. /// -/// The loop a caller writes by hand, in the shape the driver's own docs use. +/// The loop a caller writes by hand, in the shape the driver's own docs use — +/// including the arm most hand-written loops forget. A model that hallucinates +/// a tool name yields `NeedsResolution`, and the run cannot advance until it is +/// answered; this helper answers `Fail`, the correct default for a caller with +/// no recovery policy, so the run surfaces the invalid call rather than a +/// protocol violation about it. pub(crate) async fn drive_to_completion( driver: &mut AgentDriver, ) -> Result { @@ -85,7 +109,10 @@ pub(crate) async fn drive_to_completion( match driver.next_step().await? { DriveStep::SendRequest { request, .. } => { let response = request.send().await.map_err(PromptError::CompletionError)?; - driver.model_response(&response)?; + let mut outcome = driver.model_response(&response)?; + while let ModelTurnOutcome::NeedsResolution(_) = outcome { + outcome = driver.resolve_invalid_tool_call(InvalidToolCallAction::Fail)?; + } } DriveStep::ExecuteTools { calls, tools } => { let mut context = ToolContext::new(); diff --git a/tests/providers/anthropic/cassette/agent_driver.rs b/tests/providers/anthropic/cassette/agent_driver.rs index 1d42b52731..94ce33badd 100644 --- a/tests/providers/anthropic/cassette/agent_driver.rs +++ b/tests/providers/anthropic/cassette/agent_driver.rs @@ -20,7 +20,7 @@ use rig::providers::anthropic; use super::super::support::with_anthropic_cassette; use crate::driver_support::{ ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, - expect_execute_tools, expect_send, + expect_execute_tools, expect_send, expect_turn_accepted, }; use crate::support::{Adder, Subtract}; @@ -44,7 +44,7 @@ async fn drive_loop_round_trips_a_tool_call() { assert!(tools.executable_tool_names().contains("add")); let response = request.send().await.expect("first turn should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty(), "the model should have called the tool"); @@ -76,7 +76,7 @@ async fn a_custom_runs_tool_choice_reaches_the_provider() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert!( @@ -116,7 +116,7 @@ async fn tool_choice_none_forbids_tools_on_the_wire() { "ToolChoice::None allows nothing to be called" ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); // Straight to Done: a tool step here would mean the constraint did not // reach the provider. @@ -155,7 +155,7 @@ async fn tool_choice_specific_names_the_tool_on_the_wire() { assert!(!tools.allowed_tool_names().contains("subtract")); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert_eq!(pending[0].tool_call.function.name, "add"); }) @@ -177,7 +177,7 @@ async fn a_patched_preamble_replaces_the_agents_on_the_wire() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -205,7 +205,7 @@ async fn a_patched_active_tools_narrows_the_advertised_set() { assert!(!tools.executable_tool_names().contains("subtract")); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; }) @@ -232,7 +232,7 @@ async fn a_patched_tool_choice_outranks_the_runs() { let (request, tools, _) = expect_send(&mut driver).await; assert!(tools.allowed_tool_names().contains("add")); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty()); }) @@ -252,7 +252,7 @@ async fn driver_history_leads_the_request() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -280,9 +280,7 @@ async fn a_run_suspended_awaiting_the_model_resumes_and_accepts_the_reply() { let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); assert!(restored.advertised_tools().is_some()); let mut resumed = agent.drive_run(restored); - resumed - .model_response(&response) - .expect("a resumed run accepts the reply to its in-flight call"); + expect_turn_accepted(&mut resumed, &response); let (pending, tools) = expect_execute_tools(&mut resumed).await; assert!(!pending.is_empty()); @@ -306,7 +304,7 @@ async fn a_run_suspended_executing_tools_resumes_and_completes() { let mut driver = agent.drive(ADD_PROMPT); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let _ = expect_execute_tools(&mut driver).await; let serialized = serde_json::to_string(driver.run()).expect("run serializes"); @@ -354,7 +352,7 @@ async fn tool_output_mode_finalizes_via_the_output_tool() { assert!(!tools.executable_tool_names().contains(&output_tool)); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(response.output.contains("answer")); }) @@ -453,7 +451,7 @@ async fn max_turns_exhaustion_stops_before_a_second_send() { let mut driver = agent.drive(ADD_PROMPT).max_turns(1); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; diff --git a/tests/providers/gemini/cassette/agent_driver.rs b/tests/providers/gemini/cassette/agent_driver.rs index ebfc7aa902..15ee183a5e 100644 --- a/tests/providers/gemini/cassette/agent_driver.rs +++ b/tests/providers/gemini/cassette/agent_driver.rs @@ -20,7 +20,7 @@ use rig::providers::gemini; use super::super::support::with_gemini_cassette; use crate::driver_support::{ ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, - expect_execute_tools, expect_send, + expect_execute_tools, expect_send, expect_turn_accepted, }; use crate::support::{Adder, Subtract}; @@ -44,7 +44,7 @@ async fn drive_loop_round_trips_a_tool_call() { assert!(tools.executable_tool_names().contains("add")); let response = request.send().await.expect("first turn should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty(), "the model should have called the tool"); @@ -76,7 +76,7 @@ async fn a_custom_runs_tool_choice_reaches_the_provider() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert!( @@ -116,7 +116,7 @@ async fn tool_choice_none_forbids_tools_on_the_wire() { "ToolChoice::None allows nothing to be called" ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); // Straight to Done: a tool step here would mean the constraint did not // reach the provider. @@ -155,7 +155,7 @@ async fn tool_choice_specific_names_the_tool_on_the_wire() { assert!(!tools.allowed_tool_names().contains("subtract")); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert_eq!(pending[0].tool_call.function.name, "add"); }) @@ -177,7 +177,7 @@ async fn a_patched_preamble_replaces_the_agents_on_the_wire() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -205,7 +205,7 @@ async fn a_patched_active_tools_narrows_the_advertised_set() { assert!(!tools.executable_tool_names().contains("subtract")); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; }) @@ -232,7 +232,7 @@ async fn a_patched_tool_choice_outranks_the_runs() { let (request, tools, _) = expect_send(&mut driver).await; assert!(tools.allowed_tool_names().contains("add")); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty()); }) @@ -252,7 +252,7 @@ async fn driver_history_leads_the_request() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -280,9 +280,7 @@ async fn a_run_suspended_awaiting_the_model_resumes_and_accepts_the_reply() { let restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); assert!(restored.advertised_tools().is_some()); let mut resumed = agent.drive_run(restored); - resumed - .model_response(&response) - .expect("a resumed run accepts the reply to its in-flight call"); + expect_turn_accepted(&mut resumed, &response); let (pending, tools) = expect_execute_tools(&mut resumed).await; assert!(!pending.is_empty()); @@ -306,7 +304,7 @@ async fn a_run_suspended_executing_tools_resumes_and_completes() { let mut driver = agent.drive(ADD_PROMPT); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let _ = expect_execute_tools(&mut driver).await; let serialized = serde_json::to_string(driver.run()).expect("run serializes"); @@ -354,7 +352,7 @@ async fn tool_output_mode_finalizes_via_the_output_tool() { assert!(!tools.executable_tool_names().contains(&output_tool)); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(response.output.contains("answer")); }) @@ -449,7 +447,7 @@ async fn max_turns_exhaustion_stops_before_a_second_send() { let mut driver = agent.drive(ADD_PROMPT).max_turns(1); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; diff --git a/tests/providers/openai/cassette/agent_driver.rs b/tests/providers/openai/cassette/agent_driver.rs index 2a45227ecb..90528eedd6 100644 --- a/tests/providers/openai/cassette/agent_driver.rs +++ b/tests/providers/openai/cassette/agent_driver.rs @@ -29,7 +29,7 @@ use rig::tool::{Tool, ToolContext}; use super::super::support::with_openai_completions_cassette; use crate::driver_support::{ ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, - expect_execute_tools, expect_send, + expect_execute_tools, expect_send, expect_turn_accepted, }; use crate::support::{Adder, Subtract}; @@ -102,7 +102,7 @@ async fn drive_loop_round_trips_a_tool_call() { assert!(tools.executable_tool_names().contains("add")); let response = request.send().await.expect("first turn should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty(), "the model should have called the tool"); @@ -111,7 +111,7 @@ async fn drive_loop_round_trips_a_tool_call() { let (request, _, turn) = expect_send(&mut driver).await; assert_eq!(turn, 2); let response = request.send().await.expect("second turn should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!( @@ -142,7 +142,7 @@ async fn both_registered_tools_are_advertised() { assert!(tools.executable_tool_names().contains("subtract")); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; @@ -171,7 +171,7 @@ async fn parallel_tool_calls_all_dispatch_through_one_turn() { let mut driver = agent.drive("Compute 2 + 5 and 9 - 3. Use the tools for both."); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty()); @@ -219,7 +219,7 @@ async fn a_custom_runs_tool_choice_reaches_the_provider() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert!( @@ -252,7 +252,7 @@ async fn tool_choice_none_forbids_tools_on_the_wire() { "ToolChoice::None allows nothing to be called" ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -285,7 +285,7 @@ async fn tool_choice_specific_names_the_tool_on_the_wire() { ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert_eq!(pending[0].tool_call.function.name, "add"); }) @@ -309,7 +309,7 @@ async fn a_patched_preamble_replaces_the_agents_on_the_wire() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -341,7 +341,7 @@ async fn a_patched_tool_choice_outranks_the_runs() { "the patch's Required must govern, not the run's None" ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, _) = expect_execute_tools(&mut driver).await; assert!(!pending.is_empty()); }) @@ -374,7 +374,7 @@ async fn a_patched_active_tools_narrows_the_advertised_set() { ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; let response = drive_to_completion(&mut driver) @@ -402,7 +402,7 @@ async fn patched_sampling_parameters_reach_the_request() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let _ = expect_done(&mut driver).await; }) .await; @@ -428,7 +428,7 @@ async fn patched_extra_context_reaches_the_request() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -454,7 +454,7 @@ async fn a_patched_history_replaces_the_runs_for_the_turn() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -477,7 +477,7 @@ async fn driver_history_leads_the_request() { let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) @@ -515,9 +515,7 @@ async fn a_run_suspended_awaiting_the_model_resumes_and_accepts_the_reply() { "a suspended run carries the turn's advertised names" ); let mut resumed = agent.drive_run(restored); - resumed - .model_response(&response) - .expect("a resumed run accepts the reply to its in-flight call"); + expect_turn_accepted(&mut resumed, &response); let (pending, tools) = expect_execute_tools(&mut resumed).await; assert!(!pending.is_empty()); @@ -542,7 +540,7 @@ async fn a_run_suspended_executing_tools_resumes_and_completes() { let mut driver = agent.drive(ADD_PROMPT); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let _ = expect_execute_tools(&mut driver).await; let serialized = serde_json::to_string(driver.run()).expect("run serializes"); @@ -580,7 +578,7 @@ async fn a_resumed_turn_advertises_its_own_tools_not_the_processs() { let mut driver = agent.drive(ADD_PROMPT); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let _ = expect_execute_tools(&mut driver).await; let serialized = serde_json::to_string(driver.run()).expect("run serializes"); let advertised = driver @@ -692,7 +690,7 @@ async fn a_rejected_send_rolls_back_and_re_prepares() { let (request, _, turn) = expect_send(&mut driver).await; assert_eq!(turn, 1); let response = request.send().await.expect("the retry should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; @@ -743,7 +741,7 @@ async fn a_preparation_failure_costs_no_turn_and_no_interaction() { assert_eq!(turn, 1, "the retry takes the turn the failure did not"); assert!(tools.executable_tool_names().contains("add")); let response = request.send().await.expect("the retry should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); }) .await; } @@ -763,7 +761,7 @@ async fn max_turns_exhaustion_stops_before_a_second_send() { let mut driver = agent.drive(ADD_PROMPT).max_turns(1); let (request, _, _) = expect_send(&mut driver).await; let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; @@ -817,7 +815,7 @@ async fn tool_output_mode_finalizes_via_the_output_tool() { ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); // The output-tool call is intercepted by the run, never surfaced. let response = expect_done(&mut driver).await; @@ -858,7 +856,7 @@ async fn native_output_mode_uses_the_provider_constraint() { ); let response = request.send().await.expect("should send"); - driver.model_response(&response).expect("turn accepted"); + expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; assert!(response.output.contains("answer")); }) @@ -1020,3 +1018,74 @@ async fn a_streamed_turn_honors_the_request_patch() { }) .await; } + +/// A stream that stops mid-tool-arguments, derived from +/// `streamed_turn.yaml` by deleting every event after a partial `arguments` +/// delta — no `finish_reason` chunk, no usage chunk, no `[DONE]`. Removing +/// recorded data is a scrub; nothing here was invented. +/// +/// This exists to pin what a *truncated* stream actually does, which is not +/// the same thing as an *interrupted* one: a mock server that stops writing +/// closes the connection cleanly, so the SSE layer sees end-of-stream rather +/// than a transport failure. See the assertions for what that means for the +/// assembled turn. +#[tokio::test] +async fn a_truncated_stream_ends_the_turn_without_a_transport_error() { + with_openai_completions_cassette("agent_driver/streamed_interrupted", |client| async move { + // `CountingAdder`, not `Adder`: this cassette is derived from + // `streamed_turn.yaml`, and the request body it matches carries that + // tool's schema. + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(CountingAdder { + calls: Arc::new(AtomicUsize::new(0)), + }) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send(&mut driver).await; + + let mut assembler = StreamedTurnAssembler::new( + tools.executable_tool_names().clone(), + tools.allowed_tool_names().clone(), + ); + let mut stream = request.stream().await.expect("stream should open"); + let mut stream_error = None; + while let Some(item) = stream.next().await { + match item { + Ok(item) => { + assembler.ingest(&item).expect("ingest should succeed"); + } + Err(err) => { + stream_error = Some(err); + break; + } + } + } + + // Observed, not assumed: a cleanly-closed truncated body is not a + // transport failure. `CompletionError::StreamInterrupted` needs a + // connection severed mid-frame, which no mock can produce — so its + // classification stays unit-tested in rig-core. + assert!( + stream_error.is_none(), + "a truncated body closes cleanly, it does not error: {stream_error:?}" + ); + + // The turn assembles from what arrived: a tool call whose arguments + // never completed. + let final_content = stream.choice.clone(); + let streamed = assembler.finish(stream.message_id.clone(), &final_content); + driver + .run_mut() + .record_streamed_completion_call(stream.usage()) + .expect("usage recorded"); + driver + .run_mut() + .streamed_turn(streamed) + .expect("the truncated turn is still a turn"); + }) + .await; +} From f6b5c304ef3598fed6d6b3eb1cf87fc2bab9c8f1 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 14:52:54 -0700 Subject: [PATCH 10/21] fix(tests): restore the truncated stream cassette to scrubbed form CI's `stable / test` failed on `cassettes_do_not_contain_obvious_secrets`, which is an idempotence check: the scrubber must be a no-op on a committed cassette. `streamed_interrupted.yaml` was hand-derived from `streamed_turn.yaml` by deleting the events after a partial arguments delta, and the derivation stripped the trailing newline the `body: |+` block scalar preserves - so re-scrubbing produced a different file and the guard fired. No secret was ever present; the scan reports any non-canonical form through the same assertion. Rebuilt keeping the blank line that terminates the last retained SSE event, which is part of the block scalar's value. The guard was right to fire, and worth noting for the next hand-derived cassette: a scrub has to leave the file in exactly the form the recorder would have written, not merely remove data. `cargo nextest run --locked --features bedrock` - the CI invocation - passes 3558/3558. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- tests/cassettes/openai/agent_driver/streamed_interrupted.yaml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml b/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml index 873a18abae..44284b7e7f 100644 --- a/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml +++ b/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml @@ -25,3 +25,4 @@ then: data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":"2"},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_5","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} data: {"choices":[{"delta":{"tool_calls":[{"function":{"arguments":",\""},"index":0}]},"finish_reason":null,"index":0,"logprobs":null}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","obfuscation":"obfuscation_REDACTED_6","object":"chat.completion.chunk","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":null} + From 54a6348cae7672c9d77a79d8a3511691276676b0 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 15:17:46 -0700 Subject: [PATCH 11/21] docs: write the release notes for readers of the release, not the branch Round-6 review found no correctness defects and two documentation ones, both from writing notes as if for a reviewer of this branch rather than a user of the release. Five `[**behavior**]` entries in rig-agent's `### Added` block amended APIs that the same block introduces. `TurnTools` was added by the first bullet and its `execute` behavior "changed" four bullets later; `commit_model_call` was made public in one entry and "previously" checked the budget only in `peek_model_call` in another; `advertised_tools()` and the resume path likewise had a "previously" no user can have experienced. rig-core was starker still - two behavior bullets amending `is_retryable` sat directly above the bullet that adds it. Someone upgrading from 0.40 would audit their code for a dispatch change that cannot reach them, or conclude `advertised_tools()` existed and their persisted state needed review. Each behavior note is folded into the entry that introduces its API, so every bullet now describes the thing as shipped. Two entries were sorted rather than merged: `ModelTurnOutcome` predates this release, so `#[must_use]` is a genuine breaking change to an existing type and moves to `### Changed`, and `ToolErrorKind::NotExecutable` stays its own addition. MIGRATING's resume-path paragraph drops its "previously" for the same reason. `DriveStep::SendRequest`'s `tools` field was documented as "informational", which is true for a blocking send and false for a streamed one: `StreamedTurnAssembler::new` needs the turn's executable and allowed name sets, and this field is the only place a driven turn can get them. Every streamed cassette test in this PR uses it that way while the module example and both rewritten examples destructure `{ request, .. }`, so a caller following the docs into a streaming loop had nowhere to go. The field doc now says which send needs it and names the assembler. 3558/3558 on the CI invocation, 0 rustdoc warnings, fmt clean. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- MIGRATING.md | 14 +++++------ crates/rig-agent/CHANGELOG.md | 35 +++++++++++++--------------- crates/rig-agent/src/agent/driver.rs | 19 +++++++++++++-- crates/rig-core/CHANGELOG.md | 4 +--- 4 files changed, 41 insertions(+), 31 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index 54bc6d60c7..8a8316b13b 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1605,13 +1605,13 @@ through provider-side idempotency, your own record of what was transmitted, or a transport that fails before the write. Bound the attempts yourself. Runs driven by `Agent::runner` are unaffected; the runner still fails the prompt. -Two behavior changes on the driver's resume path, both narrowing what dispatch -will do: `TurnTools::execute` now refuses any name the turn did not advertise -(previously a tool registered *after* a run was suspended could be dispatched -through a turn that never advertised it), and a resumed snapshot resolves the -advertised names explicitly instead of re-running retrieval for them, so a -registered dynamic tool the new query does not rank is no longer reported as -unregistered. +Two things to know about dispatch on a resumed run, since both differ from what +you might assume: `TurnTools::execute` refuses any name the turn did not +advertise, so a tool registered *after* a run was suspended cannot be dispatched +through a turn that never offered it to the model; and the resumed snapshot +resolves the advertised names against the registry directly rather than +re-running retrieval, so a registered dynamic tool is found whether or not a +fresh query would rank it, and absence really does mean the tool is gone. An `Agent`'s default model is set at construction. Per-run overrides now go through `runner(...).using_model(...)`, `Agent::set_model`, or a diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 3b43bc66c6..788d71c9a8 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -17,6 +17,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* [**breaking**] `PromptResponse` JSON serialized before the `content` field existed no longer deserializes — the missing-`content` reconstruction (and the serde shadow repr that carried it) is deleted and `content` is a required field; the JSON wire shape is unchanged +- *(agent)* [**breaking**] `ModelTurnOutcome` is `#[must_use]`. It carries `NeedsResolution`, which must be answered via `resolve_invalid_tool_call` before the run may advance; dropping it lets a hallucinated tool name surface two steps later as an unrelated "next_step called while an invalid tool-call resolution is pending". Code writing `run.model_response(turn)?;` now warns + - *(agent)* [**breaking**] the serialized `AgentRun` carries a `$schemaVersion` tag (`rig_agent::agent::run::RUN_SCHEMA_VERSION`, currently `"1.0"`) and a build reads only the version it writes. A payload from another version — including one written before versioning existed, which has no tag — fails deserialization with a named error instead of being silently reinterpreted. Runs suspended by an earlier build cannot be resumed by this one; drain them before upgrading - *(agent)* [**breaking**] message content is a plain `Vec`, following rig-core's removal of the non-empty container — every content list the agent constructs, inspects, or hands to hooks is a `Vec`, and the `OneOrMany` re-export is gone from this crate @@ -33,26 +35,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `AgentHook::on_reasoning_delta` with the Rig stream correlator, optional provider reasoning id, current fragment, and per-part aggregate; reasoning hooks share the existing observation-interest and stop-before-yield semantics used by other streaming deltas -- *(agent)* add `Agent::prepare_turn`, returning the new public `PreparedTurn`/`TurnTools`: the agent's baseline configuration resolved into one turn's completion request plus the turn's executable and allowed tool-name sets, the synthetic output-tool name, and tool dispatch pinned to the turn's registry snapshot — so hand-driven `AgentRun` loops (custom provider transport, durable suspend/resume) reuse the configured `Agent` instead of restating its preamble and tools. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. A prepared turn is a configuration read: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` - -- *(agent)* add `Agent::drive` / `Agent::drive_run`, returning the new public `AgentDriver` (with `DriveStep` and `TurnTools`): hand-drive the sans-IO `AgentRun` machine with the agent's own configuration while owning every side effect. The driver seeds the run from the agent (`default_max_turns`, `tool_choice`, output schema), yields each turn's fully configured completion request for the caller to send (or hand to a custom transport), pairs every pending tool batch with the exact registry snapshot that advertised it, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip. Resuming a serialized run in a fresh process re-derives a fresh dispatch snapshot from the rebuilt agent, and surfaces missing pending tools as an error (opt out with `allow_missing_resumed_tools`). The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` -- *(agent)* `AgentDriver` is resumable at **every** step boundary, not just while tool calls are pending. The turn's advertised tool names are recorded on the run (new public `TurnToolNames`, the serializable half of `TurnTools`, reachable via `AgentRun::advertised_tools()`), so a run serialized after `DriveStep::SendRequest` — the natural suspension point for a queued or long-running provider call — can be resumed in another process and fed its reply, validated against the set the request actually carried rather than against whatever the resuming registry holds. The driver now keeps no state of its own beyond a registry-snapshot cache -- *(agent)* a failed model-turn preparation costs nothing: the run advances only once a request exists, so an unreachable tool server or an impossible `tool_choice` leaves the run byte-identical — same state, same turn budget — and the step can be retried in place. Previously the turn was consumed and the run was left unresumable -- *(agent)* add `AgentRun::rollback_model_call` / `AgentDriver::rollback_model_call` for the other half of that story: a request that could not be sent, or whose reply is known to be lost, hands its turn back and returns the run to preparing, so the next step yields a **freshly prepared** request rather than replaying one whose tool snapshot has gone stale. Refunds the turn, preserves usage the provider already billed (including a streamed turn that reports usage after the failure), and counts attempts via `AgentRun::model_call_rollbacks`. Deciding to use it takes two answers, not one: `CompletionError::is_retryable` says whether a retry *could succeed*, and only the caller can say whether one is *safe* — a stream that died after the request was written is retryable and not replay-safe, and retrying it bills a second completion. Bounding attempts is the caller's job -- *(agent)* `AgentRun`'s peek/commit halves are public — `peek_model_call`, `commit_model_call`, `advance`, `is_preparing_request`, with `Advance` and `ModelCallInputs` — so any hand-driver whose request preparation can fail gets the same "a failure costs no turn" guarantee the driver has, not just rig's own. Their preconditions are `PromptError` protocol violations rather than debug assertions -- *(agent)* [**behavior**] `TurnTools::execute` now refuses any name the turn did not advertise, instead of trusting its snapshot to be narrowed. The two agreed by construction in-process; on a resumed turn the names come from the run and the snapshot is rebuilt locally, so a tool registered after suspension could previously be dispatched through a turn that never advertised it -- *(agent)* [**behavior**] `TurnTools::execute_call` clears the context's dispatch result on the pre-resolved path too, so a call suppressed by invalid tool-call recovery no longer leaves the *previous* call's result metadata readable — a loop over a turn's calls was attributing it to the suppressed one -- *(agent)* add `TurnToolNames::new`: the type is `#[non_exhaustive]` serialized state, so an external hand-driver could not build the value `AgentRun::commit_model_call` requires -- *(agent)* add `AgentDriver::max_invalid_tool_call_retries`, mirroring the runner's. `Agent::drive` seeds the budget at zero and the driver exposed no override, so answering a `NeedsResolution` with `InvalidToolCallAction::Retry` — the path `AgentDriver::resolve_invalid_tool_call`'s own docs point at — could never succeed on a driver built the documented way -- *(agent)* [**breaking**] `ModelTurnOutcome` is `#[must_use]`. It carries `NeedsResolution`, which must be answered before the run may advance; dropping it let a hallucinated tool name surface two steps later as an unrelated "next_step called while an invalid tool-call resolution is pending". Callers writing `driver.model_response(&r)?;` now get a warning -- *(agent)* [**behavior**] `AgentRun::commit_model_call` enforces the model-call budget. The check previously lived only in `peek_model_call`, while `commit_model_call` — public, and documented as "the only place a turn is consumed" — incremented unconditionally, so a hand-driver that peeked once and re-committed drove model calls forever against a budget never consulted -- *(agent)* [**behavior**] resuming a run no longer runs a retrieval query. The resumed snapshot is narrowed to the turn's advertised names immediately after it is taken, so the query's contribution was discarded in full while costing a vector search — and made resume *fail* when the index was unavailable, even when every pending call was a static tool -- *(agent)* add `AgentDriver::run_mut`, without which a streamed turn could not be driven through the driver at all — `AgentRun`'s streamed entry points take `&mut self`, so a streaming caller had to `into_run()` and rebuild the driver, discarding the per-turn snapshot cache and making the driver treat a turn prepared in the same process as a resume (drift check included). Committing or rolling back a model call through it is documented as unsupported; feeding a response is safe -- *(agent)* [**behavior**] `AgentRun::advertised_tools()` no longer outlives the turn it describes. Only `rollback_model_call` ended the turn's names; `reprompt_for_output`, `retry_model_turn`, `resolve_invalid_tool_call(Retry)`, `tool_results` and both streamed-abandon paths left them in place, so a run serialized while parked for a fresh model call reported a turn that was already over. Every route back to `PreparingRequest` now clears them -- *(agent)* [**behavior**] a resumed run's dispatch snapshot resolves the turn's advertised names explicitly rather than relying on retrieval to rank them again, so a registered *dynamic* tool that the re-derived query does not return is no longer misreported as "not registered in this process" — advice that could not be followed — and no longer feeds the model a not-found result for a tool that exists -- *(agent)* add `AgentDriver::request_patch` / `set_request_patch`, giving a hand-driven run the per-turn configuration the runner gets from its `CompletionCall` hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history. Because the caller owns the loop, per-turn variation needs no callback -- *(agent)* add `AgentRun::tool_choice()`. A custom run's own `tool_choice` now reaches the provider when the run is hand-driven, instead of only governing the run's internal decisions while the request silently carried the agent's baseline +- *(agent)* add `Agent::drive` / `Agent::drive_run`, returning the new public `AgentDriver` (with `DriveStep`, `TurnTools` and `TurnToolNames`): hand-drive the sans-IO `AgentRun` machine with the agent's own configuration while owning every side effect. The driver seeds the run from the agent (`default_max_turns`, `tool_choice`, output schema), yields each turn's fully configured completion request for the caller to send (or hand to a custom transport), pairs every pending tool batch with the exact registry snapshot that advertised it, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. Dispatch is gated on the turn's advertised names, so a tool registered after a turn was prepared — or after a run was suspended — cannot be dispatched through a turn that never advertised it. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip, and such a failure costs nothing: the run advances only once a request exists, so it keeps its state and its turn budget and the step can be retried in place. The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` + +- *(agent)* a driven run is resumable at **every** step boundary, not only while tool calls are pending. The turn's advertised tool names travel with the run — `TurnToolNames`, the serializable half of `TurnTools`, readable via `AgentRun::advertised_tools()` and constructible with `TurnToolNames::new` — so a run serialized after `DriveStep::SendRequest`, the natural suspension point for a queued or long-running provider call, resumes in another process and accepts its reply, validated against the set the request actually carried rather than whatever the resuming registry holds. The names describe the turn in flight and are cleared on every route back to preparing, so a run parked for a fresh call never reports a turn that is over. Resuming rebuilds the dispatch target by resolving those names against the registry directly, without re-running retrieval: a registered *dynamic* tool the turn used is found whether or not a fresh query would rank it, absence means the tool is genuinely gone (surfaced as an error, opt out with `allow_missing_resumed_tools`), and resume does not fail when a vector index is unavailable + +- *(agent)* `AgentRun`'s peek/commit halves are public — `peek_model_call`, `commit_model_call`, `advance`, `is_preparing_request`, with `Advance` and `ModelCallInputs` — so any hand-driver whose request preparation can fail gets the same "a failure costs no turn" guarantee the driver has, not just rig's own. Preconditions are `PromptError` protocol violations rather than debug assertions, and the model-call budget is enforced at `commit_model_call`, where the turn is actually spent, as well as at `peek_model_call` + +- *(agent)* add `AgentRun::rollback_model_call` / `AgentDriver::rollback_model_call`: a request that could not be sent, or whose reply is known to be lost, hands its turn back and returns the run to preparing, so the next step yields a **freshly prepared** request rather than replaying one whose tool snapshot has gone stale. Refunds the turn, preserves usage the provider already billed (including a streamed turn that reports usage after the failure), and counts attempts via `AgentRun::model_call_rollbacks`. Deciding to use it takes two answers, not one: `CompletionError::is_retryable` says whether a retry *could succeed*, and only the caller can say whether one is *safe* — a stream that died after the request was written is retryable and not replay-safe, and retrying it bills a second completion. Bounding attempts is the caller's job + +- *(agent)* add `AgentDriver::request_patch` / `set_request_patch`, giving a hand-driven run the per-turn configuration the runner gets from its `CompletionCall` hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history. Because the caller owns the loop, per-turn variation needs no callback. `AgentRun::tool_choice()` is public and a custom run's own choice reaches the provider, rather than governing only the run's internal decisions while the request carries the agent's baseline + +- *(agent)* add `AgentDriver::run_mut`, for the `AgentRun` entry points the driver does not wrap — the streamed ones (`record_streamed_completion_call`, `resolve_streamed_invalid_tool_call`, `streamed_turn`) all take `&mut self`, and driving a custom streaming transport is a headline use for this type. Committing or rolling back a model call through it is documented as unsupported, because that would leave the driver's cached dispatch target on the previous turn; feeding a response is safe + +- *(agent)* add `AgentDriver::max_invalid_tool_call_retries`, mirroring the runner's. `Agent::drive` seeds the budget at zero, so answering a `NeedsResolution` with `InvalidToolCallAction::Retry` needs this raised first - *(tool)* add `ToolErrorKind::NotExecutable` (`ToolExecutionError::not_executable`): the tool is advertised to the model but not executable by the dispatcher — produced when dispatching the synthetic structured-output tool, whose call carries the final structured answer + - *(agent)* add opaque, cloneable `ModelHandle` values with by-value `ProviderCapabilities` snapshots, plus default replacement, per-run default override (`using_model`), and hook-driven per-call selection via `AgentHook::on_model_select` - *(agent)* add run-local extractor default-model overrides used across retries diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 27283ba470..9b7c332cc0 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -143,9 +143,24 @@ pub enum DriveStep { /// `.record_content_telemetry(false)` on it to opt a hand-driven turn /// out. request: Box>, - /// The turn's advertised tool sets — informational here (the driver - /// assembles the model turn itself); the same value arrives on the + /// The turn's advertised tool sets. + /// + /// For a **blocking** send this is informational — the driver builds + /// the model turn itself when you hand the response to + /// [`AgentDriver::model_response`] — and the same value arrives on the /// following [`ExecuteTools`](Self::ExecuteTools) step for dispatch. + /// + /// For a **streamed** send it is required, and this is the only place + /// a driven turn can get it: + /// [`StreamedTurnAssembler::new`](super::run::StreamedTurnAssembler::new) + /// takes the turn's executable and allowed name sets, so build the + /// assembler from + /// [`executable_tool_names`](TurnTools::executable_tool_names) and + /// [`allowed_tool_names`](TurnTools::allowed_tool_names) here, then + /// feed the assembled turn through + /// [`AgentDriver::run_mut`]. Destructuring this step as + /// `SendRequest { request, .. }` is fine for a blocking loop and will + /// leave a streaming one with no way to validate the model's calls. tools: TurnTools, /// One-based index of this model call within the run. turn: usize, diff --git a/crates/rig-core/CHANGELOG.md b/crates/rig-core/CHANGELOG.md index 6b11f65874..34b64a5ae0 100644 --- a/crates/rig-core/CHANGELOG.md +++ b/crates/rig-core/CHANGELOG.md @@ -57,9 +57,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - *(completion)* `CompletionError::StreamInterrupted`: the response stream failed after the provider accepted the request. Previously every streaming provider folded these into `ProviderError(String)` — rig's unclassifiable bucket — which meant `is_retryable` answered `false` for the one failure the hand-driven rollback story is built around, on every provider, while the docs claimed otherwise. The string payload is unchanged; only the variant carrying it is, and that is what makes the failure classifiable at all. Retryable and deliberately **not** replay-safe: the request reached the provider, so a retry may bill a second completion. The enum is `#[non_exhaustive]`, so this is additive -- *(completion)* [**behavior**] `is_retryable` no longer treats `http_client::Error::StreamEnded` as transient. In the SSE layer it is the *normal* end-of-stream sentinel that every provider matches and breaks on, never surfacing as an error; its only error-producing sites are the wasm stub clients, where it means the transport cannot send at all. Neither is a retryable failure -- *(completion)* [**behavior**] `is_retryable` classifies `http_client::Error::Instance` by what it wraps rather than assuming it transient. `Instance` is the client's catch-all and carries both a dropped connection *and* a request the client refused to build — a `base_url` missing its scheme, say, which fails identically forever and loops any caller driving off this method. The wrapped error's `source()` chain is walked for a `reqwest::Error` that `is_builder()` or `is_redirect()`; anything unrecognised keeps the transient default -- *(completion)* `CompletionError::is_retryable`: a conservative classification of whether re-issuing an equivalent request could plausibly succeed — preserved provider statuses of 408, 409, 429 or 5xx, plus transport failures the transport itself reports as transient (a dropped connection, a cut stream). Deterministic transport failures are excluded by name: a header value the client refused (an API key with a trailing newline, say), a request that could not be constructed, or a wrong content type will fail identically forever, and retrying them loops. A 2xx carrying a provider-authored error envelope is not retryable, since the call completed and was rejected on its merits. Two documented limits: providers on non-HTTP transports (Bedrock, Vertex, gRPC Gemini) report errors as a body with no status, so their throttling classifies as non-retryable until classification moves to the provider adapter; and this answers only whether a retry *could succeed*, never whether one is *replay-safe* +- *(completion)* `CompletionError::is_retryable`: a conservative classification of whether re-issuing an equivalent request could plausibly succeed — preserved provider statuses of 408, 409, 429 or 5xx, a `StreamInterrupted` response stream, and transport failures the transport itself reports as transient. Deterministic failures are excluded by name: a header value the client refused (an API key with a trailing newline, say), a request the client could not build (a `base_url` missing its scheme), a wrong content type — all fail identically forever, and retrying them loops. `Instance`, the HTTP client's catch-all, is classified by walking what it wraps rather than assumed transient, and `StreamEnded` is not treated as a failure at all: in the SSE layer it is the normal end-of-stream sentinel. A 2xx carrying a provider-authored error envelope is not retryable, since the call completed and was rejected on its merits. Two documented limits: providers on non-HTTP transports (Bedrock, Vertex, gRPC Gemini) report errors as a body with no status, so their throttling classifies as non-retryable until classification moves to the provider adapter; and this answers only whether a retry *could succeed*, never whether one is *replay-safe* - *(tool)* `ToolErrorKind::NotExecutable` and `ToolExecutionError::not_executable`: the tool is advertised to the model but not executable by the dispatcher — its call must be handled by the host. Produced by rig-agent when dispatching the synthetic structured-output tool; also the natural classification for host-executed or deferred tools. The enum is `#[non_exhaustive]`, so this is additive - *(streaming)* `wire::classify_typed_event` extends the decode-then-validate policy to typed-transport wires (bedrock, candle, gemini-grpc): modeled variants are `Known`, the SDK's non-exhaustive/unrecognized variants are `Unknown`, SDK decode errors are `Corrupt` — a typed transport earns no policy exemption - *(streaming)* `WireAdapter` gains an associated `Frame` type so typed-event wires implement the same contract over their SDK events; `classify` now takes the frame by value From 76ab0c922212c926bcf794a84600abe6e3c68798 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 16:03:31 -0700 Subject: [PATCH 12/21] fix(agent): answer "what was this turn allowed to do" from the turn A per-turn RequestPatch may override the run's tool choice, and until now nothing recorded that it had. `commit_model_call` stored the advertised tool names and the output-tool name; four sites - invalid-tool-call hook context and the Skip rejection, on the unary and streamed paths - read `AgentRun::tool_choice`, the run's baseline. So the state machine could disagree with the request that actually went out, in both directions: a Skip permitted under a baseline of Required when the request carried None, or a hook told the choice was None when the request required a tool. This is reachable on main today through a CompletionCall hook - RequestPatch::tool_choice predates the driver, preparation already prefers it, and the runner seeds the run's choice from the agent. The driver's patch seam adds a second door to the same room, which is why it is worth fixing here rather than filing. PreparedTurnMetadata folds the advertised names together with the effective tool choice and the output tool, and is what a commit records. The resolved choice comes back from preparation itself (`PreparedCompletionRequest::tool_choice`) rather than being re-derived at the call site: preparation is where the baseline and the patch are reconciled, so it is the only place that can answer without repeating the merge rule. `effective_tool_choice()` prefers the committed turn and falls back to the baseline for runs hand-driven through `next_step`, which commit no metadata because their names arrive with the ModelTurn. Being committed state, it survives serialization: a run suspended with a resolution pending resumes answering the way the process that sent the request would have. Four tests, three of them falsified against the fix by reverting the four read sites; the fourth pins the baseline fallback. 522 rig-agent tests, 54 driver cassettes across three providers replay unchanged, clippy and fmt clean, 0 rustdoc warnings. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/CHANGELOG.md | 2 + crates/rig-agent/src/agent/completion.rs | 3 + crates/rig-agent/src/agent/driver.rs | 7 +- crates/rig-agent/src/agent/mod.rs | 2 +- crates/rig-agent/src/agent/run/mod.rs | 224 ++++++++++++++++++----- crates/rig-agent/src/agent/turn_tools.rs | 58 +++++- 6 files changed, 244 insertions(+), 52 deletions(-) diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 788d71c9a8..311196f069 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -39,6 +39,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* a driven run is resumable at **every** step boundary, not only while tool calls are pending. The turn's advertised tool names travel with the run — `TurnToolNames`, the serializable half of `TurnTools`, readable via `AgentRun::advertised_tools()` and constructible with `TurnToolNames::new` — so a run serialized after `DriveStep::SendRequest`, the natural suspension point for a queued or long-running provider call, resumes in another process and accepts its reply, validated against the set the request actually carried rather than whatever the resuming registry holds. The names describe the turn in flight and are cleared on every route back to preparing, so a run parked for a fresh call never reports a turn that is over. Resuming rebuilds the dispatch target by resolving those names against the registry directly, without re-running retrieval: a registered *dynamic* tool the turn used is found whether or not a fresh query would rank it, absence means the tool is genuinely gone (surfaced as an error, opt out with `allow_missing_resumed_tools`), and resume does not fail when a vector index is unavailable +- *(agent)* a committed model call records what the turn *resolved to*, not just what it advertised: `PreparedTurnMetadata` (tool names, the tool choice the request actually carried, the synthetic output tool) is committed atomically by `AgentRun::commit_model_call` and readable via `AgentRun::prepared_turn()`. Invalid-tool-call context and the `Skip` rejection now answer from the turn rather than the run's baseline, so a per-turn `RequestPatch` that overrides `tool_choice` — from a `CompletionCall` hook or from a hand-driven turn — can no longer leave the state machine disagreeing with the request that went out. `AgentRun::advertised_tools()` remains as shorthand for the names + - *(agent)* `AgentRun`'s peek/commit halves are public — `peek_model_call`, `commit_model_call`, `advance`, `is_preparing_request`, with `Advance` and `ModelCallInputs` — so any hand-driver whose request preparation can fail gets the same "a failure costs no turn" guarantee the driver has, not just rig's own. Preconditions are `PromptError` protocol violations rather than debug assertions, and the model-call budget is enforced at `commit_model_call`, where the turn is actually spent, as well as at `peek_model_call` - *(agent)* add `AgentRun::rollback_model_call` / `AgentDriver::rollback_model_call`: a request that could not be sent, or whose reply is known to be lost, hands its turn back and returns the run to preparing, so the next step yields a **freshly prepared** request rather than replaying one whose tool snapshot has gone stale. Refunds the turn, preserves usage the provider already billed (including a streamed turn that reports usage after the failure), and counts attempts via `AgentRun::model_call_rollbacks`. Deciding to use it takes two answers, not one: `CompletionError::is_retryable` says whether a retry *could succeed*, and only the caller can say whether one is *safe* — a stream that died after the request was written is retryable and not replay-safe, and retrying it bills a second completion. Bounding attempts is the caller's job diff --git a/crates/rig-agent/src/agent/completion.rs b/crates/rig-agent/src/agent/completion.rs index a1c656cef6..869bb6a7d0 100644 --- a/crates/rig-agent/src/agent/completion.rs +++ b/crates/rig-agent/src/agent/completion.rs @@ -576,6 +576,9 @@ pub(crate) async fn build_prepared_completion_request( Ok(PreparedCompletionRequest { builder: completion_request, + // The reconciled baseline-plus-patch choice, so no caller has to + // repeat the merge rule to learn what the request carries. + tool_choice: tool_choice.cloned(), tools: TurnTools { snapshot: Arc::new(tool_snapshot), executable_tool_names: Arc::new(executable_tool_names), diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 9b7c332cc0..7a2ab33caa 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -352,11 +352,10 @@ impl AgentDriver { .await .map_err(PromptError::CompletionError)?; - let PreparedCompletionRequest { builder, tools } = prepared; + let metadata = prepared.turn_metadata(); + let PreparedCompletionRequest { builder, tools, .. } = prepared; self.snapshot = Some(tools.snapshot.clone()); - let turn = self - .run - .commit_model_call(Some(tools.names()), tools.output_tool_name.clone())?; + let turn = self.run.commit_model_call(Some(metadata))?; Ok(DriveStep::SendRequest { request: Box::new(builder), tools, diff --git a/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index b5de527be5..5a424eb356 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -139,4 +139,4 @@ pub use run::{ PendingToolCall, }; pub use runner::AgentRunner; -pub use turn_tools::{TurnToolNames, TurnTools}; +pub use turn_tools::{PreparedTurnMetadata, TurnToolNames, TurnTools}; diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index a3066c7470..35778f0c75 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -31,7 +31,7 @@ //! //! | version | change | //! | --- | --- | -//! | `1.0` | Initial versioned format. Records the turn's advertised tool names so a run suspended mid-model-call can be resumed. | +//! | `1.0` | Initial versioned format. Records what each model call resolved to — advertised tool names, effective tool choice, output tool — so a run suspended mid-model-call can be resumed and answers about the turn come from the turn. | //! //! `AgentRun` deliberately contains no model, tool registry, memory backend, or //! hook stack. Hand-driving it is a low-level provider integration: the caller @@ -97,7 +97,7 @@ //! // no turn spent, still `is_preparing_request()`, retry at will. //! # let _ = (prompt, history); //! // Only once the request exists does the run advance. -//! let _turn = run.commit_model_call(None, None)?; +//! let _turn = run.commit_model_call(None)?; //! // Send it. If the send fails or its reply is lost, hand the turn //! // back with `run.rollback_model_call()?` and prepare again. //! } @@ -131,7 +131,7 @@ use crate::{ assistant_text_from_choice, build_full_history, build_history_for_request, invalid_tool_retry_user_message, is_empty_assistant_turn, tool_result_message, }, - agent::turn_tools::TurnToolNames, + agent::turn_tools::{PreparedTurnMetadata, TurnToolNames}, completion::{Message, PromptError, Usage}, json_utils, }; @@ -415,16 +415,22 @@ pub struct AgentRun { /// [`AgentRunStep::CallModel`] is emitted. #[serde(default)] streamed_completion_call_recorded: bool, - /// Tool names advertised on the most recently committed model call. + /// What the most recently committed model call resolved to: its advertised + /// tool names, the tool choice the request actually carried, and the + /// synthetic output tool. /// /// Recorded by [`AgentRun::commit_model_call`] and kept for the whole turn /// — through validation, resolution and tool execution — so a driver can /// rebuild the turn's [`TurnTools`](crate::agent::TurnTools) in *any* - /// state, including in a process that only just deserialized the run. + /// state, including in a process that only just deserialized the run, and + /// so every question about what the turn was allowed to do is answered by + /// the turn rather than by the run's baseline. + /// /// `None` for runs hand-driven through [`AgentRun::next_step`], where the - /// advertised names arrive with the [`ModelTurn`] instead. + /// advertised names arrive with the [`ModelTurn`] instead; those runs fall + /// back to the baseline, which is all they ever had. #[serde(default)] - advertised_tools: Option, + prepared_turn: Option, /// How many committed model calls were handed back via /// [`AgentRun::rollback_model_call`] — sends that failed, or whose reply /// was lost. Observability only: the run imposes no bound on it, because @@ -492,7 +498,7 @@ impl AgentRun { invalid_tool_call_retries: 0, rollback_pending: false, streamed_completion_call_recorded: false, - advertised_tools: None, + prepared_turn: None, model_call_rollbacks: 0, state: RunState::PreparingRequest, } @@ -578,14 +584,13 @@ impl AgentRun { /// Park the run for a fresh model call, ending the current turn. /// - /// The turn's advertised names end with it. They describe the call the - /// model was shown, and [`Self::advertised_tools`] is documented as the - /// record of *this* turn — so a run parked here, in memory or serialized, - /// must not still report the names of a turn that finished or was + /// The turn's prepared metadata ends with it. It describes the call the + /// model was shown — see [`Self::prepared_turn`] — so a run parked here, in + /// memory or serialized, must not still report a turn that finished or was /// abandoned. Every route back to `PreparingRequest` goes through here so /// the invariant cannot be half-applied. fn park_for_new_request(&mut self) { - self.advertised_tools = None; + self.prepared_turn = None; self.state = RunState::PreparingRequest; } @@ -615,21 +620,43 @@ impl AgentRun { self.tool_choice.as_ref() } - /// The tool names advertised on the current turn's model call, when the - /// call was committed with them (see [`Self::commit_model_call`]). + /// What the current turn's model call resolved to, when the call was + /// committed with it (see [`Self::commit_model_call`]). /// - /// This is the durable half of the turn's - /// [`TurnTools`](crate::agent::TurnTools): pairing it with a registry - /// snapshot reconstitutes the whole thing, which is what lets a driver + /// The durable half of the turn: pairing its tool names with a registry + /// snapshot reconstitutes the turn's + /// [`TurnTools`](crate::agent::TurnTools), which is what lets a driver /// resume a run suspended at *any* step rather than only while tool calls /// are pending. It is also the record of what the model was actually shown - /// this turn, which is worth persisting alongside the run for audit. + /// and actually allowed to do this turn, which is worth persisting + /// alongside the run for audit. /// /// `None` for a run hand-driven through [`Self::next_step`], which commits - /// no names because the driver supplies them with the [`ModelTurn`] - /// instead, and `None` before the run's first model call. + /// no metadata because the driver supplies the names with the + /// [`ModelTurn`] instead, and `None` before the run's first model call. + pub fn prepared_turn(&self) -> Option<&PreparedTurnMetadata> { + self.prepared_turn.as_ref() + } + + /// The tool names the current turn advertised. Shorthand for + /// [`Self::prepared_turn`]'s `tools`. pub fn advertised_tools(&self) -> Option<&TurnToolNames> { - self.advertised_tools.as_ref() + self.prepared_turn.as_ref().map(|turn| &turn.tools) + } + + /// The tool choice that governs the current turn: the one the request + /// actually carried when a turn is in flight, and the run's baseline + /// otherwise. + /// + /// This is the answer every question about *what this turn was allowed to + /// do* needs. A per-turn patch may have overridden the baseline, and + /// reading [`Self::tool_choice`] instead would make the machine disagree + /// with the request that went out. + fn effective_tool_choice(&self) -> Option<&ToolChoice> { + self.prepared_turn + .as_ref() + .map(|turn| turn.tool_choice.as_ref()) + .unwrap_or(self.tool_choice.as_ref()) } /// Set the synthetic output-tool name for Tool output mode (see #1928). @@ -806,7 +833,7 @@ impl AgentRun { )), available_tools: resolving.executable_tool_names.iter().cloned().collect(), allowed_tools: resolving.allowed_tool_names.iter().cloned().collect(), - tool_choice: self.tool_choice.clone(), + tool_choice: self.effective_tool_choice().cloned(), chat_history: self.diagnostic_history(resolving), is_streaming: false, }) @@ -892,11 +919,13 @@ impl AgentRun { /// has already happened by the time it runs, so the run never advances into /// a call that was never made. /// - /// `advertised` records the turn's tool names on the run (see - /// [`Self::advertised_tools`]); pass `None` when the caller supplies them - /// with the [`ModelTurn`] instead. `output_tool_name` fills the run's - /// committed Tool-mode name once (#1928), pinning the mode for the rest of - /// the run. + /// `prepared` records what the turn resolved to (see + /// [`Self::prepared_turn`]): its advertised tool names, the tool choice the + /// request actually carried, and the synthetic output tool. Pass `None` + /// when the caller supplies the names with the [`ModelTurn`] instead — a + /// run driven that way answers from its baseline, which is all it has. The + /// metadata's `output_tool_name` fills the run's committed Tool-mode name + /// once (#1928), pinning the mode for the rest of the run. /// /// If the committed call then fails to reach the provider, or its reply is /// lost, [`Self::rollback_model_call`] undoes exactly this. @@ -911,8 +940,7 @@ impl AgentRun { /// peek before every commit. pub fn commit_model_call( &mut self, - advertised: Option, - output_tool_name: Option, + prepared: Option, ) -> Result { if !self.is_preparing_request() { return Err( @@ -920,8 +948,12 @@ impl AgentRun { ); } self.check_turn_budget()?; - self.set_output_tool_name(output_tool_name); - self.advertised_tools = advertised; + self.set_output_tool_name( + prepared + .as_ref() + .and_then(|turn| turn.output_tool_name.clone()), + ); + self.prepared_turn = prepared; self.current_turn += 1; self.rollback_pending = false; self.streamed_completion_call_recorded = false; @@ -1239,7 +1271,7 @@ impl AgentRun { match self.advance()? { Advance::NeedsModelCall => { let ModelCallInputs { prompt, history } = self.peek_model_call()?; - let turn = self.commit_model_call(None, None)?; + let turn = self.commit_model_call(None)?; Ok(AgentRunStep::CallModel { prompt, history, @@ -1437,7 +1469,7 @@ impl AgentRun { Err(PromptError::prompt_cancelled(diagnostic_history, reason)) } InvalidToolCallAction::Skip { reason } => { - if matches!(self.tool_choice, Some(ToolChoice::None)) { + if matches!(self.effective_tool_choice(), Some(ToolChoice::None)) { return Err(unknown_tool_call_error( tool_call.function.name, executable_tool_names, @@ -1680,7 +1712,7 @@ impl AgentRun { args: invalid.args.clone(), available_tools: invalid.executable_tool_names.iter().cloned().collect(), allowed_tools: invalid.allowed_tool_names.iter().cloned().collect(), - tool_choice: self.tool_choice.clone(), + tool_choice: self.effective_tool_choice().cloned(), chat_history: self .streamed_diagnostic_history(partial, Some(invalid.tool_call.clone())), is_streaming: true, @@ -1768,7 +1800,7 @@ impl AgentRun { Err(PromptError::prompt_cancelled(diagnostic_history, reason)) } InvalidToolCallAction::Skip { reason } => { - if matches!(self.tool_choice, Some(ToolChoice::None)) { + if matches!(self.effective_tool_choice(), Some(ToolChoice::None)) { self.state = RunState::Failed; return Err(unknown_tool_call_error( invalid.tool_call.function.name.clone(), @@ -2014,13 +2046,10 @@ mod tests { let mut run = AgentRun::new("add things").max_turns(2); run.advance().expect("advance should succeed"); run.peek_model_call().expect("peek should succeed"); - run.commit_model_call( - Some(TurnToolNames { - executable: tool_names(&["add"]), - allowed: tool_names(&["add"]), - }), + run.commit_model_call(Some(PreparedTurnMetadata::new( + TurnToolNames::new(["add"], ["add"]), None, - ) + ))) .expect("commit should succeed"); assert!(run.advertised_tools().is_some()); @@ -2126,7 +2155,7 @@ mod tests { let mut run = AgentRun::new("add things").max_turns(3); run.advance().expect("advance"); run.peek_model_call().expect("peek"); - run.commit_model_call(Some(advertised.clone()), None) + run.commit_model_call(Some(PreparedTurnMetadata::new(advertised.clone(), None))) .expect("commit"); expect_continue( run.model_response(tool_call_turn("call_1", "add")) @@ -2146,7 +2175,7 @@ mod tests { let mut run = AgentRun::new("add things").max_turns(3); run.advance().expect("advance"); run.peek_model_call().expect("peek"); - run.commit_model_call(Some(advertised.clone()), None) + run.commit_model_call(Some(PreparedTurnMetadata::new(advertised.clone(), None))) .expect("commit"); expect_continue( run.model_response(text_turn("nope")) @@ -2161,6 +2190,109 @@ mod tests { ); } + /// Drive to an invalid tool call on a run whose committed turn resolved to + /// `committed`, over a baseline of `baseline`. + fn run_with_divergent_choice( + baseline: Option, + committed: Option, + ) -> AgentRun { + let mut run = AgentRun::new("go").max_turns(3); + if let Some(baseline) = baseline { + run = run.with_tool_choice(baseline); + } + run.advance().expect("advance"); + run.peek_model_call().expect("peek"); + run.commit_model_call(Some(PreparedTurnMetadata::new( + TurnToolNames::new(["add"], ["add"]), + committed, + ))) + .expect("commit"); + run + } + + /// A per-turn patch may override the run's baseline choice, and everything + /// asking "what was this turn allowed to do" must read what the request + /// actually carried. Reading the baseline instead lets a `Skip` through + /// that the wire never justified. + #[test] + fn skip_is_rejected_by_the_turns_choice_not_the_runs_baseline() { + // Baseline permits tools; the turn that went out forbade them. + let mut run = run_with_divergent_choice(Some(ToolChoice::Required), Some(ToolChoice::None)); + let outcome = run + .model_response(tool_call_turn("call_1", "nonexistent")) + .expect("the invalid call needs resolution"); + assert!(matches!(outcome, ModelTurnOutcome::NeedsResolution(_))); + + run.resolve_invalid_tool_call(InvalidToolCallAction::Skip { + reason: "skip it".to_string(), + }) + .expect_err("the turn forbade tools, so a skipped call cannot be justified"); + } + + /// And the other direction: the hook context must report the choice the + /// request carried, not the baseline it overrode. + #[test] + fn invalid_call_context_reports_the_turns_choice() { + let mut run = run_with_divergent_choice(Some(ToolChoice::None), Some(ToolChoice::Required)); + let outcome = run + .model_response(tool_call_turn("call_1", "nonexistent")) + .expect("the invalid call needs resolution"); + let ModelTurnOutcome::NeedsResolution(context) = outcome else { + panic!("expected NeedsResolution"); + }; + assert_eq!( + context.tool_choice, + Some(ToolChoice::Required), + "the context must report what the request carried" + ); + } + + /// The effective choice is committed state, so it survives a suspension + /// taken while a resolution is pending. + #[test] + fn the_turns_choice_survives_serialize_and_resume() { + let mut run = run_with_divergent_choice(Some(ToolChoice::Required), Some(ToolChoice::None)); + let outcome = run + .model_response(tool_call_turn("call_1", "nonexistent")) + .expect("the invalid call needs resolution"); + assert!(matches!(outcome, ModelTurnOutcome::NeedsResolution(_))); + + let serialized = serde_json::to_string(&run).expect("run serializes"); + let mut restored: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + assert_eq!( + restored + .prepared_turn() + .expect("the turn's metadata survives") + .tool_choice, + Some(ToolChoice::None) + ); + restored + .resolve_invalid_tool_call(InvalidToolCallAction::Skip { + reason: "skip it".to_string(), + }) + .expect_err("the resumed run answers with the turn's choice, not the baseline"); + } + + /// A run driven through `next_step` commits no metadata — the names arrive + /// with the `ModelTurn` — so it answers from its baseline, which is all it + /// ever had. + #[test] + fn a_run_without_committed_metadata_falls_back_to_its_baseline() { + let mut run = AgentRun::new("go") + .max_turns(3) + .with_tool_choice(ToolChoice::None); + expect_call_model(&mut run); + assert_eq!(run.prepared_turn(), None); + + let outcome = run + .model_response(tool_call_turn("call_1", "nonexistent")) + .expect("the invalid call needs resolution"); + let ModelTurnOutcome::NeedsResolution(context) = outcome else { + panic!("expected NeedsResolution"); + }; + assert_eq!(context.tool_choice, Some(ToolChoice::None)); + } + /// The budget is checked where it is *spent*, not only where it is /// previewed. Both halves are public, and nothing obliges a caller to peek /// before every commit — a caller that peeks once and re-commits (after a @@ -2186,7 +2318,7 @@ mod tests { // without peeking must not get one. assert!(run.is_preparing_request()); let err = run - .commit_model_call(None, None) + .commit_model_call(None) .expect_err("the budget is spent"); assert!( matches!(err, PromptError::MaxTurnsError { .. }), @@ -2205,7 +2337,7 @@ mod tests { run.peek_model_call() .expect_err("a model call is already in flight"); - run.commit_model_call(None, None) + run.commit_model_call(None) .expect_err("there is no peeked call to commit"); } diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs index 91c1b71ab0..5caf1035ca 100644 --- a/crates/rig-agent/src/agent/turn_tools.rs +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -17,7 +17,7 @@ use std::sync::Arc; use serde::{Deserialize, Serialize}; -use rig_core::message::UserContent; +use rig_core::message::{ToolChoice, UserContent}; use super::model::ModelHandle; use super::run::{ModelTurn, PendingToolCall}; @@ -34,6 +34,19 @@ pub(crate) struct PreparedCompletionRequest { pub(crate) builder: CompletionRequestBuilder, /// The turn's tool sets and dispatch target. pub(crate) tools: TurnTools, + /// The tool choice the built request actually carries, after any per-turn + /// patch. Returned rather than re-derived by callers: preparation is where + /// the baseline and the patch are reconciled, so it is the only place that + /// can answer without repeating the merge rule. + pub(crate) tool_choice: Option, +} + +impl PreparedCompletionRequest { + /// What this turn resolved to, ready to commit on the run. + pub(crate) fn turn_metadata(&self) -> PreparedTurnMetadata { + PreparedTurnMetadata::new(self.tools.names(), self.tool_choice.clone()) + .with_output_tool_name(self.tools.output_tool_name.clone()) + } } /// The serializable half of [`TurnTools`]: the tool names a model call @@ -94,6 +107,49 @@ impl TurnToolNames { } } +/// What a prepared turn resolved to, recorded when its model call is committed. +/// +/// The turn's own answers, not the run's baseline. A per-turn +/// [`RequestPatch`](crate::agent::RequestPatch) may override the tool choice — +/// from a `CompletionCall` hook under the runner, or from the preparation +/// callback under [`AgentDriver`](super::AgentDriver) — and everything +/// downstream that reasons about *what this turn was allowed to do* must read +/// what was actually sent. Reading the run's baseline instead makes the state +/// machine disagree with the wire: a `Skip` resolution permitted under a +/// baseline of `Required` when the request that went out carried `None`, or an +/// invalid-tool-call hook told the choice was `None` when the request required +/// a tool. +/// +/// Serialized with the run, so a resumed turn answers those questions the same +/// way the process that sent it would have. +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[non_exhaustive] +pub struct PreparedTurnMetadata { + /// The tool names this turn advertised. + pub tools: TurnToolNames, + /// The tool choice the request actually carried, after any per-turn patch. + pub tool_choice: Option, + /// The synthetic structured-output tool advertised this turn, if any. + pub output_tool_name: Option, +} + +impl PreparedTurnMetadata { + /// The metadata for a turn prepared with these tools and choice. + pub fn new(tools: TurnToolNames, tool_choice: Option) -> Self { + Self { + tools, + tool_choice, + output_tool_name: None, + } + } + + /// Record the synthetic structured-output tool advertised this turn. + pub fn with_output_tool_name(mut self, name: Option) -> Self { + self.output_tool_name = name; + self + } +} + /// One turn's advertised tool sets and their dispatch target. /// /// All four pieces were resolved together when the turn was prepared: From 5a1d1774b2ea4f83cfd0f8175a0d8af8ef041691 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 16:09:29 -0700 Subject: [PATCH 13/21] refactor!(agent): per-turn configuration is an input, not driver state `AgentDriver` held a sticky `RequestPatch` that `drive_run` silently reset to empty, while the module docs claimed "the serializable state is *all* of the state: the driver holds nothing it could lose". Both cannot be true. A multi-turn run resumed after a patched turn advertised different tools, sent different history, or used a different preamble on the next request - and no resume test combined resume with a patch, which is why six review rounds missed it. The fix is to stop holding it. `next_step_with` takes a preparation callback that receives the prompt, the history and the prospective turn index, and returns the turn's patch and optionally a model. It runs before anything advances, so a caller decision that fails costs no turn - the same guarantee preparation already had. `next_step` remains, and delegates with an empty preparation. This is the AI SDK's `prepareStep` with the callback made async, and it is what the runner already does internally: compute per-turn config from hooks, pass it to preparation. Policy for a turn that has not happened yet is not run state, so there is nothing to serialize and nothing to lose. `TurnPreparation` carries a model as well as a patch, so folding the runner's model selection into the same seam is a call-site change rather than a signature change. One cassette was re-recorded deliberately. `patch_active_tools` encoded the sticky behavior - both recorded turns carried the narrowed tool set, because the patch leaked into the second. Under the new contract only the turn given the patch is narrowed, so the test now drives both turns and asserts the second advertises the full set again. Both request bodies are the assertion: a patch that outlived its turn would narrow the second body and fail as a mock miss. That test is now the regression guard for the defect this commit fixes. 522 rig-agent tests, 54 driver cassettes across three providers, clippy -D warnings clean, fmt clean, 0 rustdoc warnings. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/CHANGELOG.md | 2 +- crates/rig-agent/src/agent/driver.rs | 211 ++++++++++++------ crates/rig-agent/src/agent/mod.rs | 2 +- .../agent_driver/patch_active_tools.yaml | 4 +- tests/common/driver_support.rs | 29 ++- .../anthropic/cassette/agent_driver.rs | 29 +-- .../providers/gemini/cassette/agent_driver.rs | 29 +-- .../providers/openai/cassette/agent_driver.rs | 125 +++++++---- 8 files changed, 286 insertions(+), 145 deletions(-) diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 311196f069..54eb2521bf 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -45,7 +45,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `AgentRun::rollback_model_call` / `AgentDriver::rollback_model_call`: a request that could not be sent, or whose reply is known to be lost, hands its turn back and returns the run to preparing, so the next step yields a **freshly prepared** request rather than replaying one whose tool snapshot has gone stale. Refunds the turn, preserves usage the provider already billed (including a streamed turn that reports usage after the failure), and counts attempts via `AgentRun::model_call_rollbacks`. Deciding to use it takes two answers, not one: `CompletionError::is_retryable` says whether a retry *could succeed*, and only the caller can say whether one is *safe* — a stream that died after the request was written is retryable and not replay-safe, and retrying it bills a second completion. Bounding attempts is the caller's job -- *(agent)* add `AgentDriver::request_patch` / `set_request_patch`, giving a hand-driven run the per-turn configuration the runner gets from its `CompletionCall` hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history. Because the caller owns the loop, per-turn variation needs no callback. `AgentRun::tool_choice()` is public and a custom run's own choice reaches the provider, rather than governing only the run's internal decisions while the request carries the agent's baseline +- *(agent)* add `AgentDriver::next_step_with`, taking a per-turn preparation callback (`TurnPreparation`, `TurnPreparationContext`) that supplies the `RequestPatch` and optionally a model for the turn about to be prepared. This is how a hand-driven run gets what the runner gets from its `CompletionCall` and model-selection hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history — and the callback runs *before* anything advances, so a decision that fails costs no turn. Per-turn configuration is an input, never driver state: a driver holding it would resume a serialized run with configuration the suspending process never recorded, and nothing on the run would say so. `AgentRun::tool_choice()` is public and a custom run's own choice reaches the provider, rather than governing only the run's internal decisions while the request carries the agent's baseline - *(agent)* add `AgentDriver::run_mut`, for the `AgentRun` entry points the driver does not wrap — the streamed ones (`record_streamed_completion_call`, `resolve_streamed_invalid_tool_call`, `streamed_turn`) all take `&mut self`, and driving a custom streaming transport is a headline use for this type. Committing or rolling back a model call through it is documented as unsupported, because that would leave the driver's cached dispatch target on the previous turn; feeding a response is safe diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 7a2ab33caa..9f9dfd739f 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -23,10 +23,14 @@ //! side effect stays with the caller. //! //! It owns that pairing without *holding* it. Everything durable lives on the -//! [`AgentRun`] — including the turn's advertised tool names — and the driver -//! keeps only a cache of the live registry snapshot, which cannot be -//! serialized in any design. That is what makes the durability guarantees -//! below hold at every step rather than at one of them. +//! [`AgentRun`] — including what each committed turn resolved to — and the +//! driver's only other field is a cache of the live registry snapshot, which +//! cannot be serialized in any design. Per-turn configuration is an input to +//! [`AgentDriver::next_step_with`], never a field: policy for a turn that has +//! not happened yet is not run state, and a driver that stored it would resume +//! a serialized run with configuration the suspending process never recorded. +//! That is what makes the durability guarantees below hold at every step +//! rather than at one of them. //! //! # Durability //! @@ -84,6 +88,7 @@ use crate::completion::{ CompletionError, CompletionRequestBuilder, CompletionResponse, Message, PromptError, }; use crate::tool::server::ToolRegistrySnapshot; +use rig_core::wasm_compat::WasmBoxedFuture; impl Agent { /// Hand-drive this agent: build a driver whose run is seeded from the @@ -114,7 +119,6 @@ impl Agent { agent: self.clone(), run, snapshot: None, - request_patch: RequestPatch::new(), allow_missing_resumed_tools: false, } } @@ -134,11 +138,12 @@ pub enum DriveStep { /// The fully configured request: the agent's preamble (with any /// output-mode augmentation), static context, model parameters, /// `tool_choice`, and this turn's tool definitions — with the - /// driver's [`RequestPatch`](AgentDriver::request_patch) applied over - /// that baseline, which is how a hand-driven turn gets the per-turn - /// preamble, `tool_choice`, or `active_tools` narrowing the runner - /// gets from its `CompletionCall` hooks. No hooks run on this path; - /// the patch is the seam. The request still honors the agent's + /// [`RequestPatch`] from [`AgentDriver::next_step_with`]'s callback + /// applied over that baseline, which is how a hand-driven turn gets + /// the per-turn preamble, `tool_choice`, or `active_tools` narrowing + /// the runner gets from its `CompletionCall` hooks. No hooks run on + /// this path; the patch is the seam. The request still honors the + /// agent's /// `record_telemetry_content` for provider-level spans; call /// `.record_content_telemetry(false)` on it to opt a hand-driven turn /// out. @@ -198,6 +203,49 @@ impl std::fmt::Debug for DriveStep { } } +/// What a caller may decide about the turn that is about to be prepared. +/// +/// Handed to the callback of [`AgentDriver::next_step_with`] before anything +/// advances, so a decision that fails costs no turn. +#[non_exhaustive] +pub struct TurnPreparationContext<'a> { + /// The prompt this turn will send. + pub prompt: &'a Message, + /// The history preceding it. + pub history: &'a [Message], + /// One-based index this call *would* take, once committed. + pub turn: usize, +} + +/// A caller's decisions for one turn. +/// +/// Per-turn configuration is an **input to preparation**, never state the +/// driver holds. Future policy is not run state: storing it would make a +/// resumed run silently prepare its next request differently from the process +/// that suspended it, and there would be nothing on the run to say so. +#[derive(Debug, Default)] +#[non_exhaustive] +pub struct TurnPreparation { + /// Per-turn overrides layered over the agent's baseline. Each set field + /// replaces the configured value for this turn; unset fields inherit it. + pub patch: RequestPatch, + /// The model to use for this turn. `None` uses the agent's. + pub model: Option, +} + +impl TurnPreparation { + /// Prepare this turn with `patch` layered over the agent's baseline. + pub fn with_patch(patch: RequestPatch) -> Self { + Self { patch, model: None } + } + + /// Use `model` for this turn instead of the agent's. + pub fn using_model(mut self, model: ModelHandle) -> Self { + self.model = Some(model); + self + } +} + /// Hand-drives one [`AgentRun`] with one [`Agent`]'s configuration. Built by /// [`Agent::drive`] / [`Agent::drive_run`]; see the [module docs](self) for /// the driving protocol and the boundary with [`AgentRunner`](super::AgentRunner). @@ -212,9 +260,6 @@ pub struct AgentDriver { /// through the exact implementations the provider was shown, and is /// rebuilt on demand when a resumed run reaches its pending tool calls. snapshot: Option>, - /// Per-turn request configuration applied to every turn this driver - /// prepares. See [`Self::request_patch`]. - request_patch: RequestPatch, allow_missing_resumed_tools: bool, } @@ -244,29 +289,6 @@ impl AgentDriver { self } - /// Set the per-turn request configuration for the turns this driver - /// prepares. - /// - /// The driver runs no hooks, so this is how a hand-driven run gets what - /// the runner gets from its `CompletionCall` hooks: a per-turn preamble, - /// sampling parameters, `tool_choice`, `active_tools` narrowing, extra - /// context, or a substituted history. Each set field replaces the agent's - /// configured value for the turn; unset fields inherit it. - /// - /// The patch applies to every turn this driver prepares. Because the - /// caller owns the loop, per-turn variation needs no callback — call - /// [`Self::set_request_patch`] between steps. - pub fn request_patch(mut self, patch: RequestPatch) -> Self { - self.request_patch = patch; - self - } - - /// Replace the per-turn request configuration in place, so a driving loop - /// can vary it from turn to turn. See [`Self::request_patch`]. - pub fn set_request_patch(&mut self, patch: RequestPatch) { - self.request_patch = patch; - } - /// Opt out of the resumed-run drift check: dispatch pending calls whose /// tools are missing from this process's registry anyway, feeding the /// resulting not-found errors to the model instead of surfacing the drift @@ -329,19 +351,75 @@ impl AgentDriver { /// retried once the cause is fixed (a tool server that was briefly /// unreachable, say), here or in another process. pub async fn next_step(&mut self) -> Result { + self.next_step_with(|_| Box::pin(async { Ok(TurnPreparation::default()) })) + .await + } + + /// Advance to the next step, deciding this turn's configuration first. + /// + /// The callback runs **before anything advances**, is handed the prompt, + /// history and prospective turn index, and returns the turn's + /// [`RequestPatch`] and optionally a model. It is the hand-driven + /// equivalent of the runner's `CompletionCall` and model-selection hooks, + /// and it is where a caller does per-turn work that can fail: a callback + /// that returns `Err` costs no turn, exactly like a preparation failure, + /// because the commit is still ahead of it. + /// + /// Per-turn configuration is an input, never driver state. A driver that + /// stored it would resume a serialized run with configuration the + /// suspending process never recorded, and nothing on the run would say so. + /// + /// ```rust,ignore + /// let step = driver + /// .next_step_with(|ctx| { + /// Box::pin(async move { + /// Ok(TurnPreparation::with_patch( + /// RequestPatch::new().active_tools(tools_for(ctx.turn)), + /// )) + /// }) + /// }) + /// .await?; + /// ``` + pub async fn next_step_with(&mut self, prepare: F) -> Result + where + F: for<'a> FnOnce( + TurnPreparationContext<'a>, + ) -> WasmBoxedFuture<'a, Result>, + { match self.run.advance()? { Advance::NeedsModelCall => { - // Peek, prepare, *then* commit. Reading the inputs consumes - // nothing, so everything fallible below happens while the run - // is still fully intact. + // Peek, decide, prepare, *then* commit. Reading the inputs + // consumes nothing, so everything fallible below — including + // the caller's own callback — happens while the run is still + // fully intact. let ModelCallInputs { prompt, history } = self.run.peek_model_call()?; + let turn = self.run.turn() + 1; + let preparation = prepare(TurnPreparationContext { + prompt: &prompt, + history: &history, + turn, + }) + .await?; + + // The run's own choice is the baseline for a hand-driven turn: + // a custom run handed to `drive_run` is taken as-is, so its + // choice must reach the provider. An explicit patch outranks + // it, exactly as a per-turn patch outranks the agent's. + let mut patch = preparation.patch; + if patch.tool_choice.is_none() { + patch.tool_choice = self.run.tool_choice().cloned(); + } + // Pin Tool output mode once committed (#1928), mirroring the // runner: read the run's committed name into preparation, and // store the resolved name back (fill-once). let committed = self.run.output_tool_name().map(str::to_owned); - let patch = self.effective_request_patch(); + let mut baseline = TurnBaseline::from_agent(&self.agent); + if let Some(model) = preparation.model.as_ref() { + baseline.model = model; + } let prepared = build_prepared_completion_request( - TurnBaseline::from_agent(&self.agent), + baseline, TurnRequest { prompt, chat_history: &history, @@ -554,22 +632,6 @@ impl AgentDriver { )) }) } - - /// The patch actually applied to the next turn's request. - /// - /// The run's own `tool_choice` is the driver's baseline — a run built with - /// [`AgentRun::with_tool_choice`] and handed to - /// [`Agent::drive_run`](super::Agent::drive_run) is taken as-is, so its - /// choice must reach the provider and not merely the run's internal - /// decisions. An explicit [`Self::request_patch`] outranks it, exactly as - /// a per-turn patch outranks the agent's baseline everywhere else. - fn effective_request_patch(&self) -> RequestPatch { - let mut patch = self.request_patch.clone(); - if patch.tool_choice.is_none() { - patch.tool_choice = self.run.tool_choice().cloned(); - } - patch - } } #[cfg(test)] @@ -607,6 +669,24 @@ mod tests { } } + /// Expect the next step to be `SendRequest` with a per-turn preparation. + macro_rules! expect_send_with { + ($driver:expr, $patch:expr) => { + match $driver + .next_step_with(|_| Box::pin(async { Ok(TurnPreparation::with_patch($patch)) })) + .await + .expect("next_step_with succeeds") + { + DriveStep::SendRequest { + request, + tools, + turn, + } => (request, tools, turn), + other => panic!("expected SendRequest, got {other:?}"), + } + }; + } + /// Expect the next step to be `SendRequest`, panicking otherwise. macro_rules! expect_send { ($driver:expr) => { @@ -1100,16 +1180,15 @@ mod tests { .tool(MockAddTool) .tool(MockSubtractTool) .build(); - let mut driver = agent - .drive_run(AgentRun::new("go").with_tool_choice(ToolChoice::None)) - .request_patch( - RequestPatch::new() - .preamble("patched preamble") - .tool_choice(ToolChoice::Required) - .active_tools(["add"]), - ); + let mut driver = agent.drive_run(AgentRun::new("go").with_tool_choice(ToolChoice::None)); - let (request, tools, _) = expect_send!(driver); + let (request, tools, _) = expect_send_with!( + driver, + RequestPatch::new() + .preamble("patched preamble") + .tool_choice(ToolChoice::Required) + .active_tools(["add"]) + ); assert!( tools.executable_tool_names().contains("add") && !tools.executable_tool_names().contains("subtract"), diff --git a/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index 5a424eb356..8de4457ef0 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -116,7 +116,7 @@ pub(crate) const UNKNOWN_AGENT_NAME: &str = "Unnamed Agent"; pub use builder::{AgentBuilder, NoToolConfig, WithBuilderTools, WithToolServerHandle}; pub use completion::Agent; -pub use driver::{AgentDriver, DriveStep}; +pub use driver::{AgentDriver, DriveStep, TurnPreparation, TurnPreparationContext}; pub use hook::CompletionCall as CompletionCallEvent; pub use hook::{ AgentHook, CompletionCallAction, CompletionResponse as CompletionResponseEvent, HookContext, diff --git a/tests/cassettes/openai/agent_driver/patch_active_tools.yaml b/tests/cassettes/openai/agent_driver/patch_active_tools.yaml index 95ca7fdfb0..6704875b3c 100644 --- a/tests/cassettes/openai/agent_driver/patch_active_tools.yaml +++ b/tests/cassettes/openai/agent_driver/patch_active_tools.yaml @@ -24,10 +24,10 @@ when: value: '*/*' - name: content-type value: application/json - body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' then: status: 200 header: - name: content-type value: application/json - body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":131,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":133}}' + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_2","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":173,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":175}}' diff --git a/tests/common/driver_support.rs b/tests/common/driver_support.rs index 3820626ba1..4b50c4fe6c 100644 --- a/tests/common/driver_support.rs +++ b/tests/common/driver_support.rs @@ -13,7 +13,8 @@ #![allow(dead_code)] use rig::agent::{ - AgentDriver, DriveStep, InvalidToolCallAction, ModelTurnOutcome, PendingToolCall, TurnTools, + AgentDriver, DriveStep, InvalidToolCallAction, ModelTurnOutcome, PendingToolCall, RequestPatch, + TurnPreparation, TurnTools, }; use rig::completion::{CompletionResponse, PromptError}; use rig::tool::ToolContext; @@ -24,6 +25,32 @@ pub(crate) const FORCE_TOOLS_PREAMBLE: &str = "You are a calculator assistant. Y /// A prompt whose only sensible answer is one `add` call. pub(crate) const ADD_PROMPT: &str = "What is 2 + 5?"; +/// Advance the driver with a per-turn patch, expecting a request to send. +/// +/// Per-turn configuration is an input to preparation, so a test that wants a +/// patched turn supplies it here rather than configuring the driver. +pub(crate) async fn expect_send_patched( + driver: &mut AgentDriver, + patch: RequestPatch, +) -> ( + Box>, + TurnTools, + usize, +) { + match driver + .next_step_with(|_| Box::pin(async { Ok(TurnPreparation::with_patch(patch)) })) + .await + .expect("next_step_with should succeed") + { + DriveStep::SendRequest { + request, + tools, + turn, + } => (request, tools, turn), + other => panic!("expected SendRequest, got {other:?}"), + } +} + /// Advance the driver, expecting a request to send. pub(crate) async fn expect_send( driver: &mut AgentDriver, diff --git a/tests/providers/anthropic/cassette/agent_driver.rs b/tests/providers/anthropic/cassette/agent_driver.rs index 94ce33badd..dc425fd46b 100644 --- a/tests/providers/anthropic/cassette/agent_driver.rs +++ b/tests/providers/anthropic/cassette/agent_driver.rs @@ -20,7 +20,7 @@ use rig::providers::anthropic; use super::super::support::with_anthropic_cassette; use crate::driver_support::{ ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, - expect_execute_tools, expect_send, expect_turn_accepted, + expect_execute_tools, expect_send, expect_send_patched, expect_turn_accepted, }; use crate::support::{Adder, Subtract}; @@ -171,11 +171,13 @@ async fn a_patched_preamble_replaces_the_agents_on_the_wire() { .preamble("BASELINE PREAMBLE — must not appear in the request") .build(); - let mut driver = agent - .drive("Say the word banana.") - .request_patch(RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word.")); + let mut driver = agent.drive("Say the word banana."); - let (request, _, _) = expect_send(&mut driver).await; + let (request, _, _) = expect_send_patched( + &mut driver, + RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word."), + ) + .await; let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; @@ -196,11 +198,10 @@ async fn a_patched_active_tools_narrows_the_advertised_set() { .tool(Subtract) .build(); - let mut driver = agent - .drive(ADD_PROMPT) - .request_patch(RequestPatch::new().active_tools(["add"])); + let mut driver = agent.drive(ADD_PROMPT); - let (request, tools, _) = expect_send(&mut driver).await; + let (request, tools, _) = + expect_send_patched(&mut driver, RequestPatch::new().active_tools(["add"])).await; assert!(tools.executable_tool_names().contains("add")); assert!(!tools.executable_tool_names().contains("subtract")); @@ -225,11 +226,13 @@ async fn a_patched_tool_choice_outranks_the_runs() { let run = AgentRun::new(ADD_PROMPT) .max_turns(2) .with_tool_choice(ToolChoice::None); - let mut driver = agent - .drive_run(run) - .request_patch(RequestPatch::new().tool_choice(ToolChoice::Required)); + let mut driver = agent.drive_run(run); - let (request, tools, _) = expect_send(&mut driver).await; + let (request, tools, _) = expect_send_patched( + &mut driver, + RequestPatch::new().tool_choice(ToolChoice::Required), + ) + .await; assert!(tools.allowed_tool_names().contains("add")); let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); diff --git a/tests/providers/gemini/cassette/agent_driver.rs b/tests/providers/gemini/cassette/agent_driver.rs index 15ee183a5e..11025f5bf7 100644 --- a/tests/providers/gemini/cassette/agent_driver.rs +++ b/tests/providers/gemini/cassette/agent_driver.rs @@ -20,7 +20,7 @@ use rig::providers::gemini; use super::super::support::with_gemini_cassette; use crate::driver_support::{ ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, - expect_execute_tools, expect_send, expect_turn_accepted, + expect_execute_tools, expect_send, expect_send_patched, expect_turn_accepted, }; use crate::support::{Adder, Subtract}; @@ -171,11 +171,13 @@ async fn a_patched_preamble_replaces_the_agents_on_the_wire() { .preamble("BASELINE PREAMBLE — must not appear in the request") .build(); - let mut driver = agent - .drive("Say the word banana.") - .request_patch(RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word.")); + let mut driver = agent.drive("Say the word banana."); - let (request, _, _) = expect_send(&mut driver).await; + let (request, _, _) = expect_send_patched( + &mut driver, + RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word."), + ) + .await; let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; @@ -196,11 +198,10 @@ async fn a_patched_active_tools_narrows_the_advertised_set() { .tool(Subtract) .build(); - let mut driver = agent - .drive(ADD_PROMPT) - .request_patch(RequestPatch::new().active_tools(["add"])); + let mut driver = agent.drive(ADD_PROMPT); - let (request, tools, _) = expect_send(&mut driver).await; + let (request, tools, _) = + expect_send_patched(&mut driver, RequestPatch::new().active_tools(["add"])).await; assert!(tools.executable_tool_names().contains("add")); assert!(!tools.executable_tool_names().contains("subtract")); @@ -225,11 +226,13 @@ async fn a_patched_tool_choice_outranks_the_runs() { let run = AgentRun::new(ADD_PROMPT) .max_turns(2) .with_tool_choice(ToolChoice::None); - let mut driver = agent - .drive_run(run) - .request_patch(RequestPatch::new().tool_choice(ToolChoice::Required)); + let mut driver = agent.drive_run(run); - let (request, tools, _) = expect_send(&mut driver).await; + let (request, tools, _) = expect_send_patched( + &mut driver, + RequestPatch::new().tool_choice(ToolChoice::Required), + ) + .await; assert!(tools.allowed_tool_names().contains("add")); let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); diff --git a/tests/providers/openai/cassette/agent_driver.rs b/tests/providers/openai/cassette/agent_driver.rs index 90528eedd6..fe816b85d5 100644 --- a/tests/providers/openai/cassette/agent_driver.rs +++ b/tests/providers/openai/cassette/agent_driver.rs @@ -19,7 +19,7 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use futures::StreamExt; use rig::agent::run::{OutputMode, StreamedTurnAssembler}; -use rig::agent::{AgentRun, RequestPatch}; +use rig::agent::{AgentRun, RequestPatch, TurnPreparation}; use rig::completion::PromptError; use rig::message::{Message, ToolChoice}; use rig::prelude::*; @@ -29,7 +29,7 @@ use rig::tool::{Tool, ToolContext}; use super::super::support::with_openai_completions_cassette; use crate::driver_support::{ ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_done, - expect_execute_tools, expect_send, expect_turn_accepted, + expect_execute_tools, expect_send, expect_send_patched, expect_turn_accepted, }; use crate::support::{Adder, Subtract}; @@ -303,11 +303,13 @@ async fn a_patched_preamble_replaces_the_agents_on_the_wire() { .preamble("BASELINE PREAMBLE — must not appear in the request") .build(); - let mut driver = agent - .drive("Say the word banana.") - .request_patch(RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word.")); + let mut driver = agent.drive("Say the word banana."); - let (request, _, _) = expect_send(&mut driver).await; + let (request, _, _) = expect_send_patched( + &mut driver, + RequestPatch::new().preamble("PATCHED PREAMBLE — reply with one word."), + ) + .await; let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; @@ -331,11 +333,13 @@ async fn a_patched_tool_choice_outranks_the_runs() { let run = AgentRun::new(ADD_PROMPT) .max_turns(2) .with_tool_choice(ToolChoice::None); - let mut driver = agent - .drive_run(run) - .request_patch(RequestPatch::new().tool_choice(ToolChoice::Required)); + let mut driver = agent.drive_run(run); - let (request, tools, _) = expect_send(&mut driver).await; + let (request, tools, _) = expect_send_patched( + &mut driver, + RequestPatch::new().tool_choice(ToolChoice::Required), + ) + .await; assert!( tools.allowed_tool_names().contains("add"), "the patch's Required must govern, not the run's None" @@ -348,10 +352,15 @@ async fn a_patched_tool_choice_outranks_the_runs() { .await; } -/// `active_tools` narrows the advertised set for the turn, so the request's -/// `tools` array shrinks. +/// `active_tools` narrows the advertised set **for the turn it is given to**, +/// and for no other. +/// +/// Both request bodies are the assertion: turn one carries only `add`, turn two +/// carries `add` and `subtract` again. A patch that leaked into the following +/// turn — which is what a driver holding it as state does — would narrow the +/// second body too and fail as a mock miss. #[tokio::test] -async fn a_patched_active_tools_narrows_the_advertised_set() { +async fn a_patched_active_tools_narrows_only_the_turn_it_is_given_to() { with_openai_completions_cassette("agent_driver/patch_active_tools", |client| async move { let agent = client .agent(openai::GPT_4O) @@ -361,11 +370,10 @@ async fn a_patched_active_tools_narrows_the_advertised_set() { .tool(Subtract) .build(); - let mut driver = agent - .drive(ADD_PROMPT) - .request_patch(RequestPatch::new().active_tools(["add"])); + let mut driver = agent.drive(ADD_PROMPT); - let (request, tools, _) = expect_send(&mut driver).await; + let (request, tools, _) = + expect_send_patched(&mut driver, RequestPatch::new().active_tools(["add"])).await; assert!(tools.executable_tool_names().contains("add")); assert!( !tools.executable_tool_names().contains("subtract"), @@ -377,9 +385,19 @@ async fn a_patched_active_tools_narrows_the_advertised_set() { expect_turn_accepted(&mut driver, &response); let (pending, tools) = expect_execute_tools(&mut driver).await; dispatch_and_feed(&mut driver, &pending, &tools).await; - let response = drive_to_completion(&mut driver) - .await - .expect("run should finish"); + + // No patch this time. The narrowing does not persist: per-turn + // configuration was an input to the previous preparation, not state. + let (request, tools, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 2); + assert!( + tools.executable_tool_names().contains("subtract"), + "the narrowing must not outlive the turn it was given to: {:?}", + tools.executable_tool_names() + ); + let response = request.send().await.expect("should send"); + expect_turn_accepted(&mut driver, &response); + let response = expect_done(&mut driver).await; assert!(!response.output.trim().is_empty()); }) .await; @@ -396,11 +414,13 @@ async fn patched_sampling_parameters_reach_the_request() { .temperature(0.9) .build(); - let mut driver = agent - .drive("Say the word banana.") - .request_patch(RequestPatch::new().temperature(0.0).max_tokens(16)); + let mut driver = agent.drive("Say the word banana."); - let (request, _, _) = expect_send(&mut driver).await; + let (request, _, _) = expect_send_patched( + &mut driver, + RequestPatch::new().temperature(0.0).max_tokens(16), + ) + .await; let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); let _ = expect_done(&mut driver).await; @@ -422,11 +442,13 @@ async fn patched_extra_context_reaches_the_request() { text: "The launch code is banana.".to_string(), additional_props: Default::default(), }; - let mut driver = agent - .drive("What is the launch code?") - .request_patch(RequestPatch::new().extra_context(vec![document])); + let mut driver = agent.drive("What is the launch code?"); - let (request, _, _) = expect_send(&mut driver).await; + let (request, _, _) = expect_send_patched( + &mut driver, + RequestPatch::new().extra_context(vec![document]), + ) + .await; let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; @@ -444,15 +466,16 @@ async fn a_patched_history_replaces_the_runs_for_the_turn() { .preamble("Answer briefly.") .build(); - let mut driver = - agent - .drive("What did I just say?") - .request_patch(RequestPatch::new().history(vec![ - Message::user("Remember this: the code word is banana."), - Message::assistant("Noted."), - ])); + let mut driver = agent.drive("What did I just say?"); - let (request, _, _) = expect_send(&mut driver).await; + let (request, _, _) = expect_send_patched( + &mut driver, + RequestPatch::new().history(vec![ + Message::user("Remember this: the code word is banana."), + Message::assistant("Noted."), + ]), + ) + .await; let response = request.send().await.expect("should send"); expect_turn_accepted(&mut driver, &response); let response = expect_done(&mut driver).await; @@ -718,14 +741,19 @@ async fn a_preparation_failure_costs_no_turn_and_no_interaction() { .tool(Adder) .build(); - // An `active_tools` allow-list naming a tool this turn does not have: - // preparation fails locally, with no provider round trip. - let mut driver = agent - .drive(ADD_PROMPT) - .request_patch(RequestPatch::new().active_tools(["nonexistent_tool"])); + let mut driver = agent.drive(ADD_PROMPT); + // An `active_tools` allow-list naming a tool this turn does not have: + // preparation fails locally, with no provider round trip. The patch is + // an input to *this* preparation, so the failure is scoped to it. let error = driver - .next_step() + .next_step_with(|_| { + Box::pin(async { + Ok(TurnPreparation::with_patch( + RequestPatch::new().active_tools(["nonexistent_tool"]), + )) + }) + }) .await .expect_err("active_tools naming an unavailable tool must fail at prepare time"); assert!(matches!(error, PromptError::CompletionError(_))); @@ -735,8 +763,8 @@ async fn a_preparation_failure_costs_no_turn_and_no_interaction() { "a request that never left the process must not consume a turn" ); - // Fix the cause and drive the very same step again. - driver.set_request_patch(RequestPatch::new()); + // The next turn simply does not pass that patch. Nothing to reset: + // per-turn configuration was never stored. let (request, tools, turn) = expect_send(&mut driver).await; assert_eq!(turn, 1, "the retry takes the turn the failure did not"); assert!(tools.executable_tool_names().contains("add")); @@ -981,13 +1009,14 @@ async fn a_streamed_turn_honors_the_request_patch() { .tool(Subtract) .build(); - let mut driver = agent.drive(ADD_PROMPT).request_patch( + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send_patched( + &mut driver, RequestPatch::new() .preamble(FORCE_TOOLS_PREAMBLE) .active_tools(["add"]), - ); - - let (request, tools, _) = expect_send(&mut driver).await; + ) + .await; assert!(!tools.executable_tool_names().contains("subtract")); let mut assembler = StreamedTurnAssembler::new( From 2322d392b98d0dab867c801dbd9bd0aa704035c2 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 16:11:05 -0700 Subject: [PATCH 14/21] refactor!(agent): streamed turns enter through the driver, not around it `run_mut()` existed because `AgentRun`'s streamed entry points take `&mut self` and the driver wrapped none of them. Its own documentation carried a "do not commit or roll back through this" clause, which is the tell: a safe API whose primary streaming path needs an escape hatch around its invariants is unfinished, and every streamed cassette test in this PR used it. `record_stream_usage`, `accept_streamed_turn` and `resolve_streamed_invalid_tool_call` replace it. Streaming is now a mode of the same object rather than a second protocol - pydantic-ai's `ModelRequestNode` carries a `stream()` method for the same reason - and the turn stays paired with the snapshot that prepared it without asking the caller to be careful. The nine streamed cassette tests across three providers are the proof the API is sufficient: they were the only callers, and they replay unchanged through it. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/CHANGELOG.md | 2 +- crates/rig-agent/src/agent/driver.rs | 66 ++++++++++++------- .../anthropic/cassette/agent_driver.rs | 6 +- .../providers/gemini/cassette/agent_driver.rs | 6 +- .../providers/openai/cassette/agent_driver.rs | 36 +++++----- 5 files changed, 64 insertions(+), 52 deletions(-) diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 54eb2521bf..24892b0f41 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -47,7 +47,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `AgentDriver::next_step_with`, taking a per-turn preparation callback (`TurnPreparation`, `TurnPreparationContext`) that supplies the `RequestPatch` and optionally a model for the turn about to be prepared. This is how a hand-driven run gets what the runner gets from its `CompletionCall` and model-selection hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` narrowing, extra context, substituted history — and the callback runs *before* anything advances, so a decision that fails costs no turn. Per-turn configuration is an input, never driver state: a driver holding it would resume a serialized run with configuration the suspending process never recorded, and nothing on the run would say so. `AgentRun::tool_choice()` is public and a custom run's own choice reaches the provider, rather than governing only the run's internal decisions while the request carries the agent's baseline -- *(agent)* add `AgentDriver::run_mut`, for the `AgentRun` entry points the driver does not wrap — the streamed ones (`record_streamed_completion_call`, `resolve_streamed_invalid_tool_call`, `streamed_turn`) all take `&mut self`, and driving a custom streaming transport is a headline use for this type. Committing or rolling back a model call through it is documented as unsupported, because that would leave the driver's cached dispatch target on the previous turn; feeding a response is safe +- *(agent)* add driver-owned streamed ingress — `AgentDriver::record_stream_usage`, `accept_streamed_turn`, `resolve_streamed_invalid_tool_call` — so a streamed turn enters through the same object that prepared it and stays paired with its dispatch snapshot. Driving a custom streaming transport is a headline use for this type, and it previously required reaching around the driver into the run; the alternative, rebuilding the driver from `into_run()`, discards the per-turn snapshot cache and makes the driver treat a turn prepared in the same process as a resume, drift check included - *(agent)* add `AgentDriver::max_invalid_tool_call_retries`, mirroring the runner's. `Agent::drive` seeds the budget at zero, so answering a `NeedsResolution` with `InvalidToolCallAction::Retry` needs this raised first - *(tool)* add `ToolErrorKind::NotExecutable` (`ToolExecutionError::not_executable`): the tool is advertised to the model but not executable by the dispatcher — produced when dispatching the synthetic structured-output tool, whose call carries the final structured answer diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 9f9dfd739f..c39283342c 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -79,13 +79,17 @@ use rig_core::message::UserContent; use super::completion::{Agent, TurnBaseline, TurnRequest, build_prepared_completion_request}; use super::model::ModelHandle; -use super::run::{Advance, AgentRun, ModelCallInputs, ModelTurnOutcome, PendingToolCall}; +use super::run::{ + Advance, AgentRun, ModelCallInputs, ModelTurnOutcome, PartialStreamedTurn, PendingToolCall, + StreamedInvalidToolCall, StreamedResolution, StreamedTurn, +}; use super::runner::build_agent_run; use super::turn_tools::{PreparedCompletionRequest, TurnTools}; use crate::agent::hook::{InvalidToolCallAction, RequestPatch}; +use crate::agent::prompt_request::CompletionCall; use crate::agent::prompt_request::PromptResponse; use crate::completion::{ - CompletionError, CompletionRequestBuilder, CompletionResponse, Message, PromptError, + CompletionError, CompletionRequestBuilder, CompletionResponse, Message, PromptError, Usage, }; use crate::tool::server::ToolRegistrySnapshot; use rig_core::wasm_compat::WasmBoxedFuture; @@ -163,7 +167,7 @@ pub enum DriveStep { /// [`executable_tool_names`](TurnTools::executable_tool_names) and /// [`allowed_tool_names`](TurnTools::allowed_tool_names) here, then /// feed the assembled turn through - /// [`AgentDriver::run_mut`]. Destructuring this step as + /// [`AgentDriver::accept_streamed_turn`]. Destructuring this step as /// `SendRequest { request, .. }` is fine for a blocking loop and will /// leave a streaming one with no way to validate the model's calls. tools: TurnTools, @@ -305,28 +309,46 @@ impl AgentDriver { &self.run } - /// Mutable access to the run, for the entry points the driver does not - /// wrap. + /// Record one provider completion call for a streamed turn. /// - /// A streamed turn is fed through [`AgentRun::record_streamed_completion_call`], - /// [`AgentRun::resolve_streamed_invalid_tool_call`] and - /// [`AgentRun::streamed_turn`], all of which need `&mut AgentRun`. Driving - /// a custom streaming transport is a headline use for this type, so the - /// access has to exist; without it a streaming caller would have to - /// [`Self::into_run`], drive the turn by hand, and rebuild the driver — - /// which discards the per-turn snapshot cache and makes the driver treat a - /// turn prepared in *this* process as a resume, drift check and all. + /// A streamed turn learns its usage from the provider's final stream event, + /// which arrives separately from the assembled turn — including for turns + /// abandoned by invalid tool-call recovery, where the stream is drained for + /// usage after the rollback. Exactly as + /// [`AgentRun::record_streamed_completion_call`], but through the driver so + /// the turn stays paired with the snapshot that prepared it. + pub fn record_stream_usage(&mut self, usage: Usage) -> Result { + self.run.record_streamed_completion_call(usage) + } + + /// Feed the assembled streamed turn for the pending + /// [`DriveStep::SendRequest`]. /// - /// # Do not commit or roll back a model call through this + /// The streamed counterpart of [`Self::model_response`]: the run then + /// proceeds exactly as it would for a blocking turn, and the next + /// [`Self::next_step`] yields `ExecuteTools` or `Done` paired with this + /// turn's dispatch snapshot. /// - /// Use [`Self::next_step`] and [`Self::rollback_model_call`] for those. - /// They keep the driver's cached dispatch target in step with the turn the - /// run is on; committing a turn behind the driver's back would leave the - /// previous turn's snapshot cached and dispatch this turn's calls through - /// it. Feeding a *response* — streamed or otherwise — is safe, because it - /// belongs to the turn the cache already holds. - pub fn run_mut(&mut self) -> &mut AgentRun { - &mut self.run + /// Build the turn with a [`StreamedTurnAssembler`](super::run::StreamedTurnAssembler) + /// constructed from the [`TurnTools`] the matching `SendRequest` carried; + /// see that module's docs for the full streaming protocol. + pub fn accept_streamed_turn(&mut self, turn: StreamedTurn) -> Result<(), PromptError> { + self.run.streamed_turn(turn) + } + + /// Resolve an invalid tool call surfaced mid-stream, exactly as + /// [`AgentRun::resolve_streamed_invalid_tool_call`]. + /// + /// Answering while the stream is still open is the point: a doomed turn can + /// be abandoned without paying for the rest of the provider's output. + pub fn resolve_streamed_invalid_tool_call( + &mut self, + partial: &PartialStreamedTurn, + invalid: &StreamedInvalidToolCall, + action: InvalidToolCallAction, + ) -> Result { + self.run + .resolve_streamed_invalid_tool_call(partial, invalid, action) } /// Consume the driver, returning the run state. diff --git a/tests/providers/anthropic/cassette/agent_driver.rs b/tests/providers/anthropic/cassette/agent_driver.rs index dc425fd46b..3531ce9e82 100644 --- a/tests/providers/anthropic/cassette/agent_driver.rs +++ b/tests/providers/anthropic/cassette/agent_driver.rs @@ -390,12 +390,10 @@ async fn a_streamed_turn_drives_through_the_driver() { let streamed = assembler.finish(stream.message_id.clone(), &final_content); driver - .run_mut() - .record_streamed_completion_call(stream.usage()) + .record_stream_usage(stream.usage()) .expect("usage recorded"); driver - .run_mut() - .streamed_turn(streamed) + .accept_streamed_turn(streamed) .expect("streamed turn accepted"); let (pending, tools) = expect_execute_tools(&mut driver).await; diff --git a/tests/providers/gemini/cassette/agent_driver.rs b/tests/providers/gemini/cassette/agent_driver.rs index 11025f5bf7..ad4241a9d2 100644 --- a/tests/providers/gemini/cassette/agent_driver.rs +++ b/tests/providers/gemini/cassette/agent_driver.rs @@ -390,12 +390,10 @@ async fn a_streamed_turn_drives_through_the_driver() { let streamed = assembler.finish(stream.message_id.clone(), &final_content); driver - .run_mut() - .record_streamed_completion_call(stream.usage()) + .record_stream_usage(stream.usage()) .expect("usage recorded"); driver - .run_mut() - .streamed_turn(streamed) + .accept_streamed_turn(streamed) .expect("streamed turn accepted"); let (pending, tools) = expect_execute_tools(&mut driver).await; diff --git a/tests/providers/openai/cassette/agent_driver.rs b/tests/providers/openai/cassette/agent_driver.rs index fe816b85d5..726366d75a 100644 --- a/tests/providers/openai/cassette/agent_driver.rs +++ b/tests/providers/openai/cassette/agent_driver.rs @@ -895,11 +895,13 @@ async fn native_output_mode_uses_the_provider_constraint() { /// A hand-driven **streamed** turn goes through the driver, against real SSE. /// -/// The streamed entry points live on `AgentRun` and take `&mut self`, so this -/// is only expressible because the driver hands out `run_mut()`. Without it a -/// streaming caller has to `into_run()` and rebuild the driver, which discards -/// the per-turn snapshot cache and makes the driver treat a turn prepared in -/// this very process as a resume — drift check included. +/// The streamed turn enters through the driver's own ingress +/// (`record_stream_usage` / `accept_streamed_turn`), so the turn stays paired +/// with the snapshot that prepared it. Before those existed a streaming caller +/// had to reach around the driver into the run, and the alternative — rebuild +/// the driver from `into_run()` — discards the per-turn snapshot cache and +/// makes the driver treat a turn prepared in this very process as a resume, +/// drift check included. #[tokio::test] async fn a_streamed_turn_drives_through_the_driver() { with_openai_completions_cassette("agent_driver/streamed_turn", |client| async move { @@ -929,12 +931,10 @@ async fn a_streamed_turn_drives_through_the_driver() { let streamed = assembler.finish(stream.message_id.clone(), &final_content); driver - .run_mut() - .record_streamed_completion_call(stream.usage()) + .record_stream_usage(stream.usage()) .expect("usage recorded"); driver - .run_mut() - .streamed_turn(streamed) + .accept_streamed_turn(streamed) .expect("streamed turn accepted"); let (pending, tools) = expect_execute_tools(&mut driver).await; @@ -977,12 +977,10 @@ async fn a_streamed_text_turn_finalizes_the_run() { let streamed = assembler.finish(stream.message_id.clone(), &final_content); driver - .run_mut() - .record_streamed_completion_call(stream.usage()) + .record_stream_usage(stream.usage()) .expect("usage recorded"); driver - .run_mut() - .streamed_turn(streamed) + .accept_streamed_turn(streamed) .expect("streamed turn accepted"); let response = expect_done(&mut driver).await; @@ -1033,12 +1031,10 @@ async fn a_streamed_turn_honors_the_request_patch() { let streamed = assembler.finish(stream.message_id.clone(), &final_content); driver - .run_mut() - .record_streamed_completion_call(stream.usage()) + .record_stream_usage(stream.usage()) .expect("usage recorded"); driver - .run_mut() - .streamed_turn(streamed) + .accept_streamed_turn(streamed) .expect("streamed turn accepted"); let (pending, tools) = expect_execute_tools(&mut driver).await; @@ -1108,12 +1104,10 @@ async fn a_truncated_stream_ends_the_turn_without_a_transport_error() { let final_content = stream.choice.clone(); let streamed = assembler.finish(stream.message_id.clone(), &final_content); driver - .run_mut() - .record_streamed_completion_call(stream.usage()) + .record_stream_usage(stream.usage()) .expect("usage recorded"); driver - .run_mut() - .streamed_turn(streamed) + .accept_streamed_turn(streamed) .expect("the truncated turn is still a turn"); }) .await; From 4685f1503663522036fafa74747639734ec9e541 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 16:18:57 -0700 Subject: [PATCH 15/21] test(agent): pin runner/driver parity so divergence fails a test `AgentRun` has two configured coordinators, and six review rounds kept finding places where they had drifted. Consolidating them is the fix; this is the measurement, and it comes first so the refactor has a safety net rather than a hope. Four scenarios, each recorded twice into one cassette - once through `AgentRunner`, once through `AgentDriver`. The harness matches request bodies, so a coordinator that builds a different request for the same configuration fails as a mock miss with a body diff naming the field. The first recording establishes something worth writing down: diffing the runner half of each cassette against the driver half, the request bodies are identical byte for byte, except the provider-assigned tool-call id - which differs between any two live runs and is renumbered by the scrubber. The two coordinators already agree completely on *what* they send. Where they differ is *when* they commit. The runner spends the turn before running its completion-call hooks, its model selection and its request preparation, each of which can terminate the run; the driver prepares first and commits last. That changes no request, which is exactly why it survived review, and why the wire tests are paired here with unit tests for the boundary: a preparation callback returning `Err` leaves `run.turn()` at zero and reaches no provider, and the same step succeeds once the decision does. Two scenarios needed rethinking on contact with a live provider. `ToolChoice::Required` forbids the model from ever answering in text, so a run under it always ends by exhausting its budget - on both coordinators. That is a parity claim too, and a sharper one than "both finish": same configuration, same first request, same terminal error, same accounting. 3568/3568 on the CI invocation, 133 openai cassette tests, 524 rig-agent, 1251 rig-core, clippy and fmt clean, 0 rustdoc warnings. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/src/agent/driver.rs | 91 +++++++ .../openai/coordinator_parity/custom_run.yaml | 33 +++ .../openai/coordinator_parity/plain_turn.yaml | 33 +++ .../coordinator_parity/tool_choice.yaml | 33 +++ .../coordinator_parity/tool_round_trip.yaml | 67 +++++ .../openai/cassette/coordinator_parity.rs | 255 ++++++++++++++++++ tests/providers/openai/mod.rs | 1 + 7 files changed, 513 insertions(+) create mode 100644 tests/cassettes/openai/coordinator_parity/custom_run.yaml create mode 100644 tests/cassettes/openai/coordinator_parity/plain_turn.yaml create mode 100644 tests/cassettes/openai/coordinator_parity/tool_choice.yaml create mode 100644 tests/cassettes/openai/coordinator_parity/tool_round_trip.yaml create mode 100644 tests/providers/openai/cassette/coordinator_parity.rs diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index c39283342c..15b243bc35 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -1363,6 +1363,97 @@ mod tests { assert!(result.is_success()); } + /// A caller decision that fails costs no turn. + /// + /// The commit is behind the preparation callback, so a callback returning + /// `Err` — the driver's equivalent of a `CompletionCall` hook terminating + /// the run, or a model selection refusing — leaves the run exactly as it + /// was. The runner's loop commits *before* its hooks and its selection + /// run, and this is the boundary that must not drift back. + #[tokio::test] + async fn a_failed_preparation_decision_costs_no_turn() { + let model = MockCompletionModel::new([MockTurn::text("done")]); + let agent = AgentBuilder::new(model.clone()) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("go"); + + let error = driver + .next_step_with(|_| { + Box::pin(async { + Err(PromptError::CompletionError(CompletionError::RequestError( + "a hook said stop".into(), + ))) + }) + }) + .await + .expect_err("the callback refused the turn"); + assert!(matches!(error, PromptError::CompletionError(_))); + assert_eq!( + driver.run().turn(), + 0, + "a decision that never produced a request must not consume a turn" + ); + assert_eq!(model.request_count(), 0, "and must not reach the provider"); + + // The very same step succeeds once the decision does. + let (request, _, turn) = expect_send!(driver); + assert_eq!(turn, 1, "the retry takes the turn the refusal did not"); + let response = request.send().await.expect("scripted turn"); + expect_continue(driver.model_response(&response).expect("turn accepted")); + } + + /// The callback sees the turn it is deciding for, before that turn exists + /// on the run. + #[tokio::test] + async fn the_preparation_callback_sees_the_prospective_turn() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("call_1", "add", json!({"x": 1, "y": 2})), + MockTurn::text("3"), + ]); + let agent = AgentBuilder::new(model) + .default_max_turns(3) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("add 1 and 2"); + + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + for _ in 0..1 { + let seen = seen.clone(); + let (request, tools, turn) = match driver + .next_step_with(move |ctx| { + seen.lock().expect("lock").push(ctx.turn); + Box::pin(async { Ok(TurnPreparation::default()) }) + }) + .await + .expect("next_step_with succeeds") + { + DriveStep::SendRequest { + request, + tools, + turn, + } => (request, tools, turn), + other => panic!("expected SendRequest, got {other:?}"), + }; + assert_eq!(turn, 1); + let response = request.send().await.expect("scripted turn"); + expect_continue(driver.model_response(&response).expect("turn accepted")); + let (calls, _) = expect_execute_tools!(driver); + let mut context = ToolContext::new(); + let mut results = Vec::new(); + for call in &calls { + results.push(tools.execute_call(call, &mut context).await); + } + driver.tool_results(results).expect("results accepted"); + } + assert_eq!( + *seen.lock().expect("lock"), + vec![1], + "the callback is told the turn it is preparing, before it is committed" + ); + } + /// `Retry` needs a budget, and `drive()` seeds zero — so the driver must /// expose the setter the runner has, or the documented resolution path is /// unreachable for every driver built the documented way. diff --git a/tests/cassettes/openai/coordinator_parity/custom_run.yaml b/tests/cassettes/openai/coordinator_parity/custom_run.yaml new file mode 100644 index 0000000000..482f14fa27 --- /dev/null +++ b/tests/cassettes/openai/coordinator_parity/custom_run.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","temperature":0.0,"tool_choice":"required","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","temperature":0.0,"tool_choice":"required","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_2","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' diff --git a/tests/cassettes/openai/coordinator_parity/plain_turn.yaml b/tests/cassettes/openai/coordinator_parity/plain_turn.yaml new file mode 100644 index 0000000000..43f233ed8c --- /dev/null +++ b/tests/cassettes/openai/coordinator_parity/plain_turn.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Reply with one short sentence.","type":"text"}],"role":"system"},{"content":"Say hello.","role":"user"}],"model":"gpt-4o","temperature":0.0}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"Hello!","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":20,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":22}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"Reply with one short sentence.","type":"text"}],"role":"system"},{"content":"Say hello.","role":"user"}],"model":"gpt-4o","temperature":0.0}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"Hello!","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":20,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":22}}' diff --git a/tests/cassettes/openai/coordinator_parity/tool_choice.yaml b/tests/cassettes/openai/coordinator_parity/tool_choice.yaml new file mode 100644 index 0000000000..bca11e3e1f --- /dev/null +++ b/tests/cassettes/openai/coordinator_parity/tool_choice.yaml @@ -0,0 +1,33 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","temperature":0.0,"tool_choice":"required","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":148,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":165}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","temperature":0.0,"tool_choice":"required","tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"},{"function":{"description":"Subtract y from x (i.e.: x - y)","name":"subtract","parameters":{"properties":{"x":{"description":"The number to subtract from","type":"number"},"y":{"description":"The number to subtract","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_2","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":148,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":165}}' diff --git a/tests/cassettes/openai/coordinator_parity/tool_round_trip.yaml b/tests/cassettes/openai/coordinator_parity/tool_round_trip.yaml new file mode 100644 index 0000000000..8311bd59c6 --- /dev/null +++ b/tests/cassettes/openai/coordinator_parity/tool_round_trip.yaml @@ -0,0 +1,67 @@ +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","temperature":0.0,"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_1","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_1","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_1"}],"model":"gpt-4o","temperature":0.0,"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_2","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":131,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":133}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"}],"model":"gpt-4o","temperature":0.0,"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"tool_calls","index":0,"logprobs":null,"message":{"annotations":[],"content":null,"refusal":null,"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_2","type":"function"}]}}],"created":0,"id":"chatcmpl-REDACTED_3","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":17,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":106,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":123}}' +--- +when: + path: /v1/chat/completions + method: POST + query_param: [] + header: + - name: accept + value: '*/*' + - name: content-type + value: application/json + body: '{"messages":[{"content":[{"text":"You are a calculator assistant. You MUST use the provided tools for every arithmetic operation instead of computing it yourself. After you have the tool results, reply with the final numeric answer in plain text.","type":"text"}],"role":"system"},{"content":"What is 2 + 5?","role":"user"},{"role":"assistant","tool_calls":[{"function":{"arguments":"{\"x\":2,\"y\":5}","name":"add"},"id":"call_REDACTED_2","type":"function"}]},{"content":"7","role":"tool","tool_call_id":"call_REDACTED_2"}],"model":"gpt-4o","temperature":0.0,"tools":[{"function":{"description":"Add x and y together","name":"add","parameters":{"properties":{"x":{"description":"The first number to add","type":"number"},"y":{"description":"The second number to add","type":"number"}},"required":["x","y"],"type":"object"}},"type":"function"}]}' +then: + status: 200 + header: + - name: content-type + value: application/json + body: '{"choices":[{"finish_reason":"stop","index":0,"logprobs":null,"message":{"annotations":[],"content":"7","refusal":null,"role":"assistant"}}],"created":0,"id":"chatcmpl-REDACTED_4","model":"gpt-4o-2024-08-06","object":"chat.completion","service_tier":"default","system_fingerprint":"fp_REDACTED_1","usage":{"completion_tokens":2,"completion_tokens_details":{"accepted_prediction_tokens":0,"audio_tokens":0,"reasoning_tokens":0,"rejected_prediction_tokens":0},"prompt_tokens":131,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":133}}' diff --git a/tests/providers/openai/cassette/coordinator_parity.rs b/tests/providers/openai/cassette/coordinator_parity.rs new file mode 100644 index 0000000000..af2e43478d --- /dev/null +++ b/tests/providers/openai/cassette/coordinator_parity.rs @@ -0,0 +1,255 @@ +//! Runner/driver parity, as a test rather than a review finding. +//! +//! `AgentRun` has two configured coordinators: `AgentRunner`'s internal loop and +//! the public `AgentDriver`. Both prepare requests, pair turns with tool-registry +//! snapshots and manage structured-output metadata. Nothing structural stops +//! them drifting, and review has repeatedly found places where they had. +//! +//! **One cassette, two coordinators.** Each scenario is recorded twice — once +//! per coordinator — and the harness matches request bodies. If the runner and +//! the driver build different requests for the same configuration, the second +//! pass fails as a mock miss with a body diff naming the field. That is the +//! assertion; the state comparisons below are secondary and catch divergence +//! the wire cannot show, like a differently committed turn budget. +//! +//! When this suite fails after a change to either coordinator, the question is +//! not "which assertion do I relax" but "which of the two did I mean to +//! change". +//! +//! # What the first recording established +//! +//! Diffing the runner half of each cassette against the driver half: the +//! request bodies are **identical**, byte for byte, except for the +//! provider-assigned tool-call id — which differs between any two live runs and +//! is renumbered by the scrubber. So as of this suite's first recording the two +//! coordinators agree completely on *what* they send. +//! +//! Where they still differ is *when* they commit a turn: the runner spends the +//! turn before running its completion-call hooks, its model selection and its +//! request preparation, each of which can terminate the run, while the driver +//! prepares first and commits last. That divergence is invisible here, because +//! it changes no request — which is exactly why it survived several reviews, +//! and why these tests pin the wire while the commit-boundary unit tests pin +//! the rest. + +use rig::agent::AgentRun; +use rig::completion::PromptError; +use rig::message::ToolChoice; +use rig::prelude::*; +use rig::providers::openai; + +use super::super::support::with_openai_completions_cassette; +use crate::driver_support::{ + ADD_PROMPT, FORCE_TOOLS_PREAMBLE, dispatch_and_feed, drive_to_completion, expect_execute_tools, + expect_send, expect_turn_accepted, +}; +use crate::support::{Adder, Subtract}; + +/// A plain single-turn run: no tools, no patch. The narrowest possible parity +/// claim, and the one that fails first if either coordinator changes how it +/// builds a baseline request. +#[tokio::test] +async fn a_plain_turn_is_identical_through_both_coordinators() { + with_openai_completions_cassette("coordinator_parity/plain_turn", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("Reply with one short sentence.") + .temperature(0.0) + .build(); + + // Runner. + let runner_response = agent + .runner("Say hello.") + .run() + .await + .expect("the runner should finish"); + + // Driver. Same agent, same prompt, no per-turn configuration — so the + // request must be byte-identical, which the cassette enforces. + let mut driver = agent.drive("Say hello."); + let driver_response = drive_to_completion(&mut driver) + .await + .expect("the driver should finish"); + + assert_eq!( + driver.run().turn(), + 1, + "one model call, same as the runner's single completion call" + ); + assert_eq!( + runner_response.completion_calls.len(), + driver.run().completion_calls().len(), + "both coordinators account one completion call per model call" + ); + assert!(!runner_response.output.trim().is_empty()); + assert!(!driver_response.output.trim().is_empty()); + }) + .await; +} + +/// A two-turn tool round trip. The second request is the interesting one: it +/// carries the assistant tool call and the tool result each coordinator +/// threaded back through the run, so a divergence in how either writes history +/// shows up as a body diff. +#[tokio::test] +async fn a_tool_round_trip_is_identical_through_both_coordinators() { + with_openai_completions_cassette("coordinator_parity/tool_round_trip", |client| async move { + let build = |client: &openai::CompletionsClient| { + client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .temperature(0.0) + .default_max_turns(3) + .tool(Adder) + .build() + }; + + let runner_response = build(&client) + .runner(ADD_PROMPT) + .run() + .await + .expect("the runner should finish"); + + let agent = build(&client); + let mut driver = agent.drive(ADD_PROMPT); + let driver_response = drive_to_completion(&mut driver) + .await + .expect("the driver should finish"); + + assert_eq!( + runner_response.completion_calls.len(), + driver.run().completion_calls().len(), + "both coordinators spend the same number of model calls on the same work" + ); + assert_eq!(driver.run().turn(), 2); + assert!(!runner_response.output.trim().is_empty()); + assert!(!driver_response.output.trim().is_empty()); + }) + .await; +} + +/// `tool_choice` on the agent must reach the provider identically from both. +/// +/// `Required` forbids the model from ever answering in text, so a run under it +/// always ends by exhausting its budget — on **both** coordinators, which is +/// the parity claim. Same configuration, same first request body (the cassette +/// enforces it), same terminal condition, same accounting. +#[tokio::test] +async fn a_required_tool_choice_is_identical_through_both_coordinators() { + with_openai_completions_cassette("coordinator_parity/tool_choice", |client| async move { + let build = |client: &openai::CompletionsClient| { + client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .temperature(0.0) + .tool_choice(ToolChoice::Required) + .tool(Adder) + .tool(Subtract) + .build() + }; + + let runner_error = build(&client) + .runner(ADD_PROMPT) + .max_turns(1) + .run() + .await + .expect_err("Required with a budget of one cannot finish"); + assert!( + matches!( + runner_error, + PromptError::MaxTurnsError { max_turns: 1, .. } + ), + "expected the runner to exhaust its budget, got {runner_error:?}" + ); + + let agent = build(&client); + let mut driver = agent.drive(ADD_PROMPT).max_turns(1); + let (request, tools, _) = expect_send(&mut driver).await; + assert!(tools.allowed_tool_names().contains("add")); + let response = request.send().await.expect("should send"); + expect_turn_accepted(&mut driver, &response); + + // The turn recorded what it resolved to, and it is the agent's choice — + // the same one the runner sent. + let prepared = driver + .run() + .prepared_turn() + .expect("the committed turn records its metadata") + .clone(); + assert_eq!(prepared.tool_choice, Some(ToolChoice::Required)); + + let (pending, tools) = expect_execute_tools(&mut driver).await; + assert!(!pending.is_empty(), "Required must force a tool call"); + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let driver_error = driver + .next_step() + .await + .expect_err("the driver exhausts the same budget"); + assert!( + matches!( + driver_error, + PromptError::MaxTurnsError { max_turns: 1, .. } + ), + "expected the driver to exhaust its budget, got {driver_error:?}" + ); + assert_eq!( + driver.run().completion_calls().len(), + 1, + "one model call spent, same as the runner" + ); + }) + .await; +} + +/// A custom `AgentRun` driven by hand reaches the same request the runner +/// builds from the equivalent agent configuration. +/// +/// The run carries the choice rather than the agent, which is the path that +/// used to drop it silently. Recorded for each, so a regression there is a mock +/// miss rather than a behavioral difference nobody notices. +#[tokio::test] +async fn a_custom_run_matches_the_runners_equivalent_configuration() { + with_openai_completions_cassette("coordinator_parity/custom_run", |client| async move { + // Runner: the choice comes from the agent. + let runner_error = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .temperature(0.0) + .tool_choice(ToolChoice::Required) + .tool(Adder) + .build() + .runner(ADD_PROMPT) + .max_turns(1) + .run() + .await + .expect_err("Required with a budget of one cannot finish"); + assert!(matches!( + runner_error, + PromptError::MaxTurnsError { max_turns: 1, .. } + )); + + // Driver: an equivalent run carrying the choice itself, on an agent + // that has none. + let bare = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .temperature(0.0) + .tool(Adder) + .build(); + let run = AgentRun::new(ADD_PROMPT) + .max_turns(1) + .with_tool_choice(ToolChoice::Required); + let mut driver = bare.drive_run(run); + + let (request, _, _) = expect_send(&mut driver).await; + let response = request.send().await.expect("should send"); + expect_turn_accepted(&mut driver, &response); + let (pending, _) = expect_execute_tools(&mut driver).await; + assert!( + !pending.is_empty(), + "the run's own Required must have reached the provider" + ); + }) + .await; +} diff --git a/tests/providers/openai/mod.rs b/tests/providers/openai/mod.rs index ab9e108a12..e23bb1cbbc 100644 --- a/tests/providers/openai/mod.rs +++ b/tests/providers/openai/mod.rs @@ -7,6 +7,7 @@ mod cassette { mod agent_driver; mod chat_history; mod completions_api; + mod coordinator_parity; mod document_ordering; mod extractor; mod extractor_usage; From 95a4eba4e601414fae7669ae044658b70f93dea8 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 16:28:04 -0700 Subject: [PATCH 16/21] fix(agent): the runner spends a turn only once its request exists `drive_agent` called `run.next_step()`, which commits, and then ran three things that can terminate the run: completion-call hooks, model selection and request preparation. A stop or a failure in any of them consumed a turn against a call that never happened - and `AgentRun::next_step`'s own doc says it is for a caller with nothing fallible in between. The loop now uses the peek/commit halves `AgentDriver` already used and this PR made public: `advance` reports that the run wants a model call without spending it, `peek_model_call` reads the inputs, and `commit_model_call` runs once the request exists - recording what the turn resolved to, which also retires the separate `set_output_tool_name` bookkeeping. This is not observable through `AgentRunner`'s API: every stop path ends the run, so no caller can see the turn count afterwards. It matters for two reasons. The two coordinators now agree on when a turn is spent, which is the invariant the parity suite exists to hold. And anything that later resumes or retries a runner-driven run - the direction the ECS work points - would otherwise have inherited a discrepancy that is invisible until it isn't. Deliberately *not* the whole consolidation. `drive_agent` still owns its own `AgentRun` and its own pending tool snapshot rather than driving an `AgentDriver`, because `TurnSource::run_model_turn` and `run_tool_calls` take `&mut AgentRun` and moving them onto the driver means changing the trait and both implementations - 650 lines of stream-macro code whose borrow structure deserves its own change. The commit boundary was the correctness half and it is done; the deduplication half is a follow-up with a clear path, and the parity suite is now in place to guard it. 3568/3568 on the CI invocation. Every existing runner and streaming test passes unchanged, which is the bar for a refactor of this loop. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/CHANGELOG.md | 2 + .../src/agent/prompt_request/streaming.rs | 42 +++++++++++++++---- 2 files changed, 37 insertions(+), 7 deletions(-) diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 24892b0f41..1144a8e25e 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -39,6 +39,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* a driven run is resumable at **every** step boundary, not only while tool calls are pending. The turn's advertised tool names travel with the run — `TurnToolNames`, the serializable half of `TurnTools`, readable via `AgentRun::advertised_tools()` and constructible with `TurnToolNames::new` — so a run serialized after `DriveStep::SendRequest`, the natural suspension point for a queued or long-running provider call, resumes in another process and accepts its reply, validated against the set the request actually carried rather than whatever the resuming registry holds. The names describe the turn in flight and are cleared on every route back to preparing, so a run parked for a fresh call never reports a turn that is over. Resuming rebuilds the dispatch target by resolving those names against the registry directly, without re-running retrieval: a registered *dynamic* tool the turn used is found whether or not a fresh query would rank it, absence means the tool is genuinely gone (surfaced as an error, opt out with `allow_missing_resumed_tools`), and resume does not fail when a vector index is unavailable +- *(agent)* [**behavior**] the runner spends a turn only once its request exists. Its loop advanced the run first and then ran completion-call hooks, model selection and request preparation — three steps that can terminate the run — so a hook stop, a selection stop or a preparation failure consumed a turn against a call that never happened. It now advances with `AgentRun::advance`, peeks the inputs, and commits after preparation succeeds, which is the boundary `AgentDriver` already used. Not observable through `AgentRunner`'s own API, since every stop path ends the run either way; it matters because the two coordinators now agree on when a turn is spent, and because anything that later resumes or retries a runner-driven run would have inherited the discrepancy + - *(agent)* a committed model call records what the turn *resolved to*, not just what it advertised: `PreparedTurnMetadata` (tool names, the tool choice the request actually carried, the synthetic output tool) is committed atomically by `AgentRun::commit_model_call` and readable via `AgentRun::prepared_turn()`. Invalid-tool-call context and the `Skip` rejection now answer from the turn rather than the run's baseline, so a per-turn `RequestPatch` that overrides `tool_choice` — from a `CompletionCall` hook or from a hand-driven turn — can no longer leave the state machine disagreeing with the request that went out. `AgentRun::advertised_tools()` remains as shorthand for the names - *(agent)* `AgentRun`'s peek/commit halves are public — `peek_model_call`, `commit_model_call`, `advance`, `is_preparing_request`, with `Advance` and `ModelCallInputs` — so any hand-driver whose request preparation can fail gets the same "a failure costs no turn" guarantee the driver has, not just rig's own. Preconditions are `PromptError` protocol violations rather than debug assertions, and the model-call budget is enforced at `commit_model_call`, where the turn is actually spent, as well as at `peek_model_call` diff --git a/crates/rig-agent/src/agent/prompt_request/streaming.rs b/crates/rig-agent/src/agent/prompt_request/streaming.rs index fb86b66916..c0acb51ce9 100644 --- a/crates/rig-agent/src/agent/prompt_request/streaming.rs +++ b/crates/rig-agent/src/agent/prompt_request/streaming.rs @@ -12,7 +12,7 @@ use crate::{ }, agent::prompt_request::{assistant_text_from_choice, is_empty_assistant_turn}, agent::run::{ - AgentRun, AgentRunStep, PendingToolCall, + AgentRun, PendingToolCall, streamed::{StreamedResolution, StreamedTurnAssembler, StreamedTurnEvent}, }, agent::runner::{ @@ -31,7 +31,11 @@ use tracing_futures::Instrument; use super::{CompletionCall, PromptResponse, forward_prompt_setters}; use crate::{ - agent::{Agent, model::ModelHandle}, + agent::{ + Agent, + model::ModelHandle, + run::{Advance, ModelCallInputs}, + }, completion::{CompletionError, PromptError}, }; use rig_core::message::{Message, Text}; @@ -490,7 +494,12 @@ where let mut previous_model: Option = None; 'outer: loop { - let step = match run.next_step() { + // `advance` never commits a model call: it reports that the run + // wants one and leaves the turn unspent. Everything fallible in the + // arm below — completion-call hooks, model selection, request + // preparation — therefore runs *before* the commit, so a stop or a + // failure costs no turn. `next_step` would have spent it first. + let step = match run.advance() { Ok(step) => step, Err(err) => { store_error_usage(&runner, &run); @@ -500,7 +509,16 @@ where }; match step { - AgentRunStep::CallModel { prompt, history, turn } => { + Advance::NeedsModelCall => { + let ModelCallInputs { prompt, history } = match run.peek_model_call() { + Ok(inputs) => inputs, + Err(err) => { + store_error_usage(&runner, &run); + yield Err(Box::new(err).into()); + break 'outer; + } + }; + let turn = run.turn() + 1; drop(pending_tool_snapshot.take()); if runner.max_turns > 1 { tracing::info!("Current conversation Turns: {}/{}", turn, runner.max_turns); @@ -591,7 +609,16 @@ where break 'outer; } }; - run.set_output_tool_name(prepared.tools.output_tool_name.clone()); + // The request exists, so the turn is real: commit it with + // what it resolved to. Everything that could have stopped + // the run is behind us, and nothing after this point can + // leave a turn spent against no request. + let turn_metadata = prepared.turn_metadata(); + if let Err(err) = run.commit_model_call(Some(turn_metadata)) { + store_error_usage(&runner, &run); + yield Err(Box::new(err).into()); + break 'outer; + } let turn_tool_snapshot = prepared.tools.snapshot.clone(); if runner.record_telemetry_content { let input_messages = prepared.builder.messages_for_telemetry(); @@ -634,7 +661,7 @@ where } pending_tool_snapshot = Some(turn_tool_snapshot); } - AgentRunStep::CallTools { calls } => { + Advance::CallTools(calls) => { let Some(tool_snapshot) = pending_tool_snapshot.take() else { store_error_usage(&runner, &run); yield Err(StreamingError::Completion(CompletionError::ResponseError( @@ -667,7 +694,7 @@ where break 'outer; } } - AgentRunStep::Done(response) => { + Advance::Done(response) => { // Run-completion marker, unifying the blocking and streaming // drivers' run-finished logs into one shared event. tracing::info!( @@ -1697,6 +1724,7 @@ mod migrated_tests { use crate::agent::AgentBuilder; use crate::agent::hook::{AgentHook, HookContext}; use crate::agent::prompt_request::{TOOL_NOT_EXECUTED_DUE_TO_INVALID_PEER, tool_result_output}; + use crate::agent::run::AgentRunStep; use crate::client::AgentClientExt; use crate::completion::{CompletionRequest, Prompt, PromptError, ToolDefinition, Usage}; use crate::streaming::{StreamingPrompt, ToolCallDeltaContent}; From e25a9e3df6e815e783400e812fac0c0eae837287 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 16:46:16 -0700 Subject: [PATCH 17/21] fix(agent): name the prompt-cache cost of per-turn tool narrowing Rebasing onto main brought in `cassette_cache_prefix`, a guard that fails when a recorded conversation moves the wire prefix providers cache on. It immediately caught `patch_active_tools`, and it was right to. That cassette narrows the advertised tools on turn 1 and re-advertises the full set on turn 2 - which is the assertion, since a patch that outlived its turn would keep the second turn narrowed. But the tools array is part of the cached prefix, so growing it between turns busts the cache. The guard's own message names dynamic tool disclosure as a legitimate case for the exemption list, so the cassette is added there with that reason. The more useful half is what it says about the feature. `RequestPatch::active_tools` now documents the cost where a caller will see it: the saving from advertising fewer tools is paid back, and then some, on every later turn of the run, so narrow for a reason - a turn that must not reach a destructive tool - rather than to trim tokens. 3588/3588 on the CI invocation after the rebase onto 3de43b96. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- crates/rig-agent/src/agent/hook.rs | 9 +++++++++ tests/cassette_cache_prefix.rs | 10 ++++++++++ 2 files changed, 19 insertions(+) diff --git a/crates/rig-agent/src/agent/hook.rs b/crates/rig-agent/src/agent/hook.rs index d56151c9e3..a4d1020552 100644 --- a/crates/rig-agent/src/agent/hook.rs +++ b/crates/rig-agent/src/agent/hook.rs @@ -743,6 +743,15 @@ impl RequestPatch { } /// Sets the allow-list used to narrow the tools advertised for this turn. + /// + /// # This costs prompt-cache hits + /// + /// The tools array is part of the wire prefix providers cache on. Narrowing + /// it on one turn and not the next changes that prefix, so the following + /// turn cannot reuse the cached one — the saving from advertising fewer + /// tools is paid back, and then some, on every later turn of the run. + /// Narrow for a reason (a turn that must not reach a destructive tool), + /// not to trim tokens. pub fn active_tools(mut self, values: I) -> Self where I: IntoIterator, diff --git a/tests/cassette_cache_prefix.rs b/tests/cassette_cache_prefix.rs index e8be98848f..13b48b8a9f 100644 --- a/tests/cassette_cache_prefix.rs +++ b/tests/cassette_cache_prefix.rs @@ -44,6 +44,16 @@ const MOVES_CACHE_PREFIX: &[(&str, &str)] = &[ "streaming twin of request_overridden_by_hook_blocking; same deliberate \ hook rewrite", ), + ( + "openai/agent_driver/patch_active_tools.yaml", + "the scenario under test is per-turn `active_tools` narrowing: turn 1 is \ + given a patch restricting the advertised set to `add`, turn 2 is given \ + none and re-advertises `add` and `subtract`. The tools array growing \ + between turns is the assertion — a patch that outlived its turn would \ + keep the second turn narrowed — so the moving prefix is the recorded \ + behavior, not a defect in it. It is also the real cost of the feature, \ + and `RequestPatch::active_tools` says so", + ), ( "openai/streaming_grammar/three_turn_tool_session.yaml", "turns 2 and 3 are hand-built `completion_request`s that deliberately omit \ From 2cc23dd0f7ef3ebe9c0200f6cb1f37f65f9136e2 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 17:01:50 -0700 Subject: [PATCH 18/21] docs: point the migration guide at the API that exists MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `5a1d1774` deleted `AgentDriver::request_patch` / `set_request_patch` in favour of `next_step_with`, updated the changelog, and left MIGRATING.md telling users to call the removed methods. That paragraph is the only user-facing prose describing how a hand-driven turn gets a preamble, `tool_choice` or `active_tools` override, so a reader following it hits `no method named request_patch` with no pointer to the mechanism that replaced it. Rewritten around `next_step_with`, with a worked snippet, and saying why it is a callback rather than a setter: the callback runs before the turn is committed, so a decision that fails costs no turn, and configuration for a turn that has not happened yet is not run state to be serialized. `TurnRequest::patch` carried the same stale intra-doc link. It never failed `cargo doc` because `agent::completion` is private and the type is `pub(crate)`, so rustdoc never resolved it — which is why the doc gate stayed green while the link rotted. Claude-Session: https://claude.ai/code/session_0119GvrZzdUge8Q11SxHx4PA --- MIGRATING.md | 37 +++++++++++++++++++----- crates/rig-agent/src/agent/completion.rs | 9 +++--- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index 8a8316b13b..39ddb94a7b 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1572,14 +1572,35 @@ not an execution path — hooks, memory, retrieval policy, and telemetry still run only under `Agent::runner`. See `examples/agent_run_stepping` and `examples/agent_with_durable_approval`. -A hand-driven turn gets per-turn configuration through -`AgentDriver::request_patch(RequestPatch)` (or `set_request_patch` between -steps) — the same `RequestPatch` the runner merges from its `CompletionCall` -hooks, covering the per-turn preamble, sampling parameters, `tool_choice`, -`active_tools` narrowing, extra context, and substituted history. A custom -run's own `AgentRun::with_tool_choice` now also reaches the provider; before, -it governed only the run's internal decisions while the request carried the -agent's baseline. +A hand-driven turn gets per-turn configuration from +`AgentDriver::next_step_with`, whose callback receives the prompt, the history +and the turn about to be prepared, and returns a `TurnPreparation` carrying the +turn's `RequestPatch` and optionally a model: + +```rust +let step = driver + .next_step_with(|ctx| { + Box::pin(async move { + Ok(TurnPreparation::with_patch( + RequestPatch::new().preamble(preamble_for(ctx.turn)), + )) + }) + }) + .await?; +``` + +That is the same `RequestPatch` the runner merges from its `CompletionCall` +hooks — per-turn preamble, sampling parameters, `tool_choice`, `active_tools` +narrowing, extra context, substituted history — and the callback runs *before* +the turn is committed, so a decision that fails costs no turn. It is an input +rather than driver state on purpose: configuration for a turn that has not +happened yet is not run state, and a driver holding it would resume a +serialized run with settings the suspending process never recorded. `next_step()` +remains for turns that need none. + +A custom run's own `AgentRun::with_tool_choice` also reaches the provider; +before, it governed only the run's internal decisions while the request carried +the agent's baseline. **Serialized `AgentRun` state does not survive the upgrade.** Payloads now carry a `$schemaVersion` tag and a build reads only the version it writes, so a diff --git a/crates/rig-agent/src/agent/completion.rs b/crates/rig-agent/src/agent/completion.rs index 869bb6a7d0..63593c021e 100644 --- a/crates/rig-agent/src/agent/completion.rs +++ b/crates/rig-agent/src/agent/completion.rs @@ -261,10 +261,11 @@ pub(crate) struct TurnRequest<'a> { /// The run's already-committed Tool-mode output name, re-advertised so the /// mode cannot flip or re-pick a name mid-run (#1928). pub(crate) committed_output_tool: Option<&'a str>, - /// Per-turn overrides — from `CompletionCall` hooks in the runner, from - /// [`AgentDriver::request_patch`](crate::agent::AgentDriver::request_patch) - /// when hand-driven. The single seam through which a turn diverges from - /// the baseline. + /// Per-turn overrides — merged from `CompletionCall` hooks in the runner, + /// returned by + /// [`AgentDriver::next_step_with`](crate::agent::AgentDriver::next_step_with)'s + /// preparation callback when hand-driven. The single seam through which a + /// turn diverges from the baseline. pub(crate) patch: Option<&'a RequestPatch>, } From 96ba87ff5e908a33ce0b9827b8ef1e822859a667 Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 18:13:13 -0700 Subject: [PATCH 19/21] fix(agent): the request decides what a streamed turn may call, not its caller MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `AgentDriver::model_response` built its `ModelTurn` from the metadata committed when the request was built, but `AgentRun::streamed_turn` validated against the sets a caller handed `StreamedTurnAssembler::new` — two same-typed `BTreeSet` arguments, transposable and widenable. A turn whose effective choice was `Specific(["add"])` therefore accepted and dispatched a streamed `subtract` if the assembler was built wide, against a policy the provider was told to enforce. `PreparedTurnMetadata` exists to answer *what was this turn allowed to do*; it was consulted on one ingress and ignored on the other. Add `effective_tool_names`, the name-set counterpart of the existing `effective_tool_choice`, and read it wherever a turn's tool sets are validated or reported: streamed ingestion, the mid-stream invalid-call hook context and its `Repair` check, and blocking ingestion too, so one sentence describes both paths. Runs hand-driven through `AgentRun` commit no metadata and keep validating against their carried sets. Then close the constructor a caller could get wrong: `StreamedTurnAssembler::new` now takes the paired `TurnToolNames`, and `TurnTools::streamed_turn_assembler` — the streaming counterpart of `TurnToolNames::model_turn` — is how a driven turn builds one, so the driver path never spells the sets at all. --- crates/rig-agent/src/agent/driver.rs | 189 ++++++++++++++++-- .../src/agent/prompt_request/streaming.rs | 7 +- crates/rig-agent/src/agent/run/mod.rs | 78 +++++++- crates/rig-agent/src/agent/run/streamed.rs | 29 ++- crates/rig-agent/src/agent/turn_tools.rs | 16 +- .../anthropic/cassette/agent_driver.rs | 7 +- .../providers/gemini/cassette/agent_driver.rs | 7 +- .../gemini/cassette/agent_run_streamed.rs | 8 +- .../providers/openai/cassette/agent_driver.rs | 22 +- 9 files changed, 295 insertions(+), 68 deletions(-) diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 15b243bc35..77538891f1 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -160,16 +160,13 @@ pub enum DriveStep { /// following [`ExecuteTools`](Self::ExecuteTools) step for dispatch. /// /// For a **streamed** send it is required, and this is the only place - /// a driven turn can get it: - /// [`StreamedTurnAssembler::new`](super::run::StreamedTurnAssembler::new) - /// takes the turn's executable and allowed name sets, so build the - /// assembler from - /// [`executable_tool_names`](TurnTools::executable_tool_names) and - /// [`allowed_tool_names`](TurnTools::allowed_tool_names) here, then - /// feed the assembled turn through + /// a driven turn can get it: call + /// [`TurnTools::streamed_turn_assembler`] here to assemble the + /// provider's stream, then feed the assembled turn through /// [`AgentDriver::accept_streamed_turn`]. Destructuring this step as /// `SendRequest { request, .. }` is fine for a blocking loop and will - /// leave a streaming one with no way to validate the model's calls. + /// leave a streaming one with no way to catch an invalid tool call + /// while the stream is still open. tools: TurnTools, /// One-based index of this model call within the run. turn: usize, @@ -329,9 +326,17 @@ impl AgentDriver { /// [`Self::next_step`] yields `ExecuteTools` or `Done` paired with this /// turn's dispatch snapshot. /// - /// Build the turn with a [`StreamedTurnAssembler`](super::run::StreamedTurnAssembler) - /// constructed from the [`TurnTools`] the matching `SendRequest` carried; - /// see that module's docs for the full streaming protocol. + /// Build the turn with the + /// [`StreamedTurnAssembler`](super::run::StreamedTurnAssembler) from the + /// matching `SendRequest`'s [`TurnTools::streamed_turn_assembler`]; see + /// that module's docs for the full streaming protocol. + /// + /// Tool calls are validated against what the *request* advertised — the + /// metadata committed when this turn was prepared, not the names the + /// assembler was built with. The two agree whenever the assembler came + /// from its own turn's `TurnTools`; when they do not, the request wins, + /// so no assembler a caller mis-built can get a tool dispatched that this + /// turn's `tool_choice` forbade. pub fn accept_streamed_turn(&mut self, turn: StreamedTurn) -> Result<(), PromptError> { self.run.streamed_turn(turn) } @@ -660,11 +665,13 @@ impl AgentDriver { mod tests { use super::*; use crate::agent::AgentBuilder; - use crate::agent::run::OutputMode; - use crate::completion::Message; + use crate::agent::TurnToolNames; + use crate::agent::run::{AgentRunStep, OutputMode, StreamedTurnAssembler, StreamedTurnEvent}; + use crate::completion::{AssistantContent, Message, Usage}; + use crate::streaming::StreamedAssistantContent; use crate::test_utils::{MockAddTool, MockCompletionModel, MockSubtractTool, MockTurn}; use crate::tool::{ToolContext, ToolErrorKind}; - use rig_core::message::ToolChoice; + use rig_core::message::{ToolCall, ToolCallId, ToolChoice, ToolFunction}; use serde_json::json; fn schema(value: serde_json::Value) -> schemars::Schema { @@ -1680,4 +1687,158 @@ mod tests { "the opt-out must dispatch and produce a tool result" ); } + + // ── Streamed ingress authority ────────────────────────────────────── + // + // `StreamedTurnAssembler` is built by the caller, so the sets it validates + // against are the caller's. These pin that a mis-built one cannot widen + // what a *driven* turn is allowed to do: the run answers from the metadata + // committed when its request was built. The bad assemblers below are + // constructed deliberately — a caller can construct them, which is the + // whole point — but only by going around + // `TurnTools::streamed_turn_assembler`. + + /// Drive one streamed turn whose assembler was built with `names`, + /// streaming a single call to `tool_name`, and hand the result to the + /// driver. + async fn stream_one_call( + driver: &mut AgentDriver, + names: &TurnToolNames, + tool_name: &str, + ) -> Result<(), PromptError> { + let (_, _, _) = expect_send_with!( + driver, + RequestPatch::new().tool_choice(ToolChoice::Specific { + function_names: vec!["add".to_string()] + }) + ); + let mut assembler = StreamedTurnAssembler::new(names); + let tool_call = ToolCall::new( + ToolCallId::new("call_1").expect("non-empty id"), + ToolFunction::new(tool_name.to_string(), json!({"x": 1, "y": 2})), + ); + let events = assembler + .ingest(&StreamedAssistantContent::ToolCall { + tool_call: tool_call.clone(), + internal_call_id: "internal_1".to_string(), + }) + .expect("ingest succeeds"); + // A widened assembler raises nothing mid-stream — that is exactly the + // early exit the caller forfeits, and why the run must still refuse. + assert!( + !events + .iter() + .any(|event| matches!(event, StreamedTurnEvent::InvalidToolCall(_))), + "the mis-built assembler was supposed to wave this call through" + ); + driver + .record_stream_usage(Usage::new()) + .expect("usage recorded"); + let turn = assembler.finish(None, &[AssistantContent::ToolCall(tool_call)]); + driver.accept_streamed_turn(turn) + } + + fn driven_to_add_only() -> Agent { + AgentBuilder::new(MockCompletionModel::text("unused")) + .default_max_turns(2) + .tool(MockAddTool) + .tool(MockSubtractTool) + .build() + } + + /// Criterion: the request's `tool_choice` narrowed the turn to `add`, so a + /// streamed `subtract` is refused however wide the assembler's set was. + #[tokio::test] + async fn a_widened_assembler_cannot_widen_a_driven_streamed_turn() { + let agent = driven_to_add_only(); + let mut driver = agent.drive("go"); + let widened = TurnToolNames::new(["add", "subtract"], ["add", "subtract"]); + + let err = stream_one_call(&mut driver, &widened, "subtract") + .await + .expect_err("the request forbade `subtract`, so the run must refuse it"); + let message = err.to_string(); + assert!( + message.contains("subtract"), + "the refusal must name the call: {message}" + ); + } + + /// Criterion: the same, for the transposition the old two-set constructor + /// invited. Here `allowed` receives the executable set, which under a + /// `Specific` choice is the wider of the two. + #[tokio::test] + async fn a_transposed_assembler_cannot_widen_a_driven_streamed_turn() { + let agent = driven_to_add_only(); + let mut driver = agent.drive("go"); + // What a caller writing `new(allowed, executable)` would have got. + let transposed = TurnToolNames::new(["add"], ["add", "subtract"]); + + stream_one_call(&mut driver, &transposed, "subtract") + .await + .expect_err("the committed metadata, not the assembler, decides"); + } + + /// Criterion: the authority rule does not reject what the turn *did* + /// advertise — the narrowed call still goes through. + #[tokio::test] + async fn the_committed_tool_still_passes_the_streamed_ingress() { + let agent = driven_to_add_only(); + let mut driver = agent.drive("go"); + let widened = TurnToolNames::new(["add", "subtract"], ["add", "subtract"]); + + stream_one_call(&mut driver, &widened, "add") + .await + .expect("`add` was advertised, so it is accepted"); + let (calls, _) = expect_execute_tools!(driver); + assert_eq!(calls[0].tool_call.function.name, "add"); + } + + /// Criterion: the fallback arm. A run hand-driven through `AgentRun` + /// commits no metadata and has nothing but its carried sets, so those + /// still govern — otherwise the rule would silently allow *everything* on + /// the raw path. + #[tokio::test] + async fn a_raw_run_still_validates_against_its_carried_sets() { + let mut run = AgentRun::new("go"); + let AgentRunStep::CallModel { .. } = run.next_step().expect("first step") else { + panic!("expected a model call"); + }; + + let mut assembler = StreamedTurnAssembler::new(&TurnToolNames::new(["add"], ["add"])); + let tool_call = ToolCall::new( + ToolCallId::new("call_1").expect("non-empty id"), + ToolFunction::new("add".to_string(), json!({"x": 1, "y": 2})), + ); + assembler + .ingest(&StreamedAssistantContent::ToolCall { + tool_call: tool_call.clone(), + internal_call_id: "internal_1".to_string(), + }) + .expect("ingest succeeds"); + let turn = assembler.finish(None, &[AssistantContent::ToolCall(tool_call)]); + run.streamed_turn(turn) + .expect("the carried sets allowed `add`"); + + // And the negative: a call the carried sets never allowed. + let mut run = AgentRun::new("go"); + let AgentRunStep::CallModel { .. } = run.next_step().expect("first step") else { + panic!("expected a model call"); + }; + let tool_call = ToolCall::new( + ToolCallId::new("call_1").expect("non-empty id"), + ToolFunction::new("subtract".to_string(), json!({})), + ); + // Built by hand rather than assembled: the assembler's own check would + // have caught this mid-stream, and the run's is what is under test. + let turn = StreamedTurn { + message_id: None, + choice: vec![AssistantContent::ToolCall(tool_call)], + executable_tool_names: ["add".to_string()].into_iter().collect(), + allowed_tool_names: ["add".to_string()].into_iter().collect(), + internal_call_ids: Vec::new(), + }; + run.streamed_turn(turn) + .expect_err("the carried sets never allowed `subtract`"); + } } diff --git a/crates/rig-agent/src/agent/prompt_request/streaming.rs b/crates/rig-agent/src/agent/prompt_request/streaming.rs index c0acb51ce9..1485a174fa 100644 --- a/crates/rig-agent/src/agent/prompt_request/streaming.rs +++ b/crates/rig-agent/src/agent/prompt_request/streaming.rs @@ -13,7 +13,7 @@ use crate::{ agent::prompt_request::{assistant_text_from_choice, is_empty_assistant_turn}, agent::run::{ AgentRun, PendingToolCall, - streamed::{StreamedResolution, StreamedTurnAssembler, StreamedTurnEvent}, + streamed::{StreamedResolution, StreamedTurnEvent}, }, agent::runner::{ AgentRunner, CompletionCallOutcome, ModelTurnDecision, ToolExecution, acquire_agent_span, @@ -1113,10 +1113,7 @@ impl TurnSource for StreamingTurnSource { // `ModelTurnFinished` event carries the turn's usage. let mut last_usage = crate::completion::Usage::new(); - let mut assembler = StreamedTurnAssembler::new( - (*prepared.tools.executable_tool_names).clone(), - (*prepared.tools.allowed_tool_names).clone(), - ); + let mut assembler = prepared.tools.streamed_turn_assembler(); let mut completion_call_emitted = false; let mut turn_abandoned = false; let mut provider_final_seen = false; diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index 35778f0c75..ab9e6b4e6a 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -659,6 +659,32 @@ impl AgentRun { .unwrap_or(self.tool_choice.as_ref()) } + /// The tool names that govern the current turn: the committed metadata's + /// when the turn was prepared with it, and `carried` otherwise. + /// + /// The name-set counterpart of [`Self::effective_tool_choice`], and the + /// same rule. A prepared turn's request went out with a resolved tool set + /// — narrowed by `active_tools`, restricted by a `Specific` choice, widened + /// by the synthetic output tool — and *what this turn was allowed to do* is + /// that set, not one an ingesting caller supplies alongside the response. + /// The blocking path could not disagree (the driver builds the + /// [`ModelTurn`] from the committed names itself), but the streamed path + /// carries sets a caller handed the assembler, and a widened or transposed + /// pair there would get the model a tool the request forbade. + /// + /// `carried` is the fallback for a run hand-driven through + /// [`Self::next_step`], which commits no metadata and has only what its + /// driver supplies. + fn effective_tool_names<'a>( + &'a self, + carried: (&'a BTreeSet, &'a BTreeSet), + ) -> (&'a BTreeSet, &'a BTreeSet) { + match self.prepared_turn.as_ref() { + Some(prepared) => (&prepared.tools.executable, &prepared.tools.allowed), + None => carried, + } + } + /// Set the synthetic output-tool name for Tool output mode (see #1928). /// When a model turn calls this tool, the run finalizes with the call's /// arguments (serialized JSON) as the response. @@ -1307,13 +1333,23 @@ impl AgentRun { .iter() .any(|item| matches!(item, AssistantContent::ToolCall(_))); + // Same rule as the streamed path: a prepared turn answers from what it + // committed. `AgentDriver` builds this `ModelTurn` from those very + // names, so the two agree by construction there; applying the rule + // here anyway means one sentence describes both ingress paths, and no + // caller of `AgentRun::model_response` can widen a driven turn either. + let (executable_tool_names, allowed_tool_names) = + self.effective_tool_names((&turn.executable_tool_names, &turn.allowed_tool_names)); + let (executable_tool_names, allowed_tool_names) = + (executable_tool_names.clone(), allowed_tool_names.clone()); + self.state = RunState::ResolvingToolCalls(Box::new(ResolvingState { message_id: turn.message_id, original_choice: turn.choice, items, next_index: 0, - executable_tool_names: turn.executable_tool_names, - allowed_tool_names: turn.allowed_tool_names, + executable_tool_names, + allowed_tool_names, skipped: BTreeMap::new(), recovered: false, any_skipped: false, @@ -1705,13 +1741,19 @@ impl AgentRun { partial: &PartialStreamedTurn, invalid: &StreamedInvalidToolCall, ) -> InvalidToolCallContext { + // A hook is told what the *request* advertised, for the same reason it + // is told the effective tool choice below: the two must describe one + // request, or a repair suggested from the hook's view of the tool set + // is rejected by the run's. + let (executable_tool_names, allowed_tool_names) = self + .effective_tool_names((&invalid.executable_tool_names, &invalid.allowed_tool_names)); InvalidToolCallContext { tool_name: invalid.tool_call.function.name.clone(), tool_call_id: Some(invalid.tool_call.id.as_str().to_owned()), internal_call_id: Some(invalid.internal_call_id.clone()), args: invalid.args.clone(), - available_tools: invalid.executable_tool_names.iter().cloned().collect(), - allowed_tools: invalid.allowed_tool_names.iter().cloned().collect(), + available_tools: executable_tool_names.iter().cloned().collect(), + allowed_tools: allowed_tool_names.iter().cloned().collect(), tool_choice: self.effective_tool_choice().cloned(), chat_history: self .streamed_diagnostic_history(partial, Some(invalid.tool_call.clone())), @@ -1740,9 +1782,15 @@ impl AgentRun { let diagnostic_history = self.streamed_diagnostic_history(partial, Some(invalid.tool_call.clone())); - let executable_tool_names: Vec = - invalid.executable_tool_names.iter().cloned().collect(); - let allowed_tool_names: Vec = invalid.allowed_tool_names.iter().cloned().collect(); + // What the request advertised, not what the assembler was told — the + // `Repair` arm below validates the replacement name against + // `allowed`, and accepting one the request never advertised would + // dispatch a tool the provider was told did not exist this turn. + let (executable, allowed) = self + .effective_tool_names((&invalid.executable_tool_names, &invalid.allowed_tool_names)); + let allowed_names = allowed.clone(); + let executable_tool_names: Vec = executable.iter().cloned().collect(); + let allowed_tool_names: Vec = allowed.iter().cloned().collect(); match action { InvalidToolCallAction::Fail => { @@ -1784,7 +1832,7 @@ impl AgentRun { }) } InvalidToolCallAction::Repair { tool_name } => { - if !invalid.allowed_tool_names.contains(&tool_name) { + if !allowed_names.contains(&tool_name) { self.state = RunState::Failed; return Err(unknown_tool_call_error( tool_name, @@ -1868,11 +1916,19 @@ impl AgentRun { .iter() .any(|item| matches!(item, AssistantContent::ToolCall(_))); + // The turn's own answer, not the assembler's. `StreamedTurn` carries + // the sets a caller handed `StreamedTurnAssembler::new`, and a driven + // turn's request went out under the sets committed with it. + let (executable_tool_names, allowed_tool_names) = + self.effective_tool_names((&turn.executable_tool_names, &turn.allowed_tool_names)); + let (executable_tool_names, allowed_tool_names) = + (executable_tool_names.clone(), allowed_tool_names.clone()); + for item in &turn.choice { let AssistantContent::ToolCall(tool_call) = item else { continue; }; - if !turn.allowed_tool_names.contains(&tool_call.function.name) { + if !allowed_tool_names.contains(&tool_call.function.name) { let mut diagnostic_messages = self.new_messages.clone(); if !is_empty_assistant_turn(&turn.choice) { diagnostic_messages.push(Message::Assistant { @@ -1885,8 +1941,8 @@ impl AgentRun { self.state = RunState::Failed; return Err(unknown_tool_call_error( tool_call.function.name.clone(), - turn.executable_tool_names.iter().cloned().collect(), - turn.allowed_tool_names.iter().cloned().collect(), + executable_tool_names.iter().cloned().collect(), + allowed_tool_names.iter().cloned().collect(), diagnostic_history, )); } diff --git a/crates/rig-agent/src/agent/run/streamed.rs b/crates/rig-agent/src/agent/run/streamed.rs index c8e4d044b2..a53734c703 100644 --- a/crates/rig-agent/src/agent/run/streamed.rs +++ b/crates/rig-agent/src/agent/run/streamed.rs @@ -43,6 +43,7 @@ use rig_core::message::{ use crate::{ agent::prompt_request::{TOOL_NOT_EXECUTED_DUE_TO_INVALID_PEER, tool_result_message}, + agent::turn_tools::TurnToolNames, completion::{CompletionError, Message, Usage}, json_utils, streaming::{StreamedAssistantContent, ToolCallDeltaContent}, @@ -400,15 +401,27 @@ impl Drop for ExclusionCount { } impl StreamedTurnAssembler { - /// Create an assembler for one streamed turn with the tool names + /// Create an assembler for one streamed turn from the tool names /// advertised to the provider for that turn. - pub fn new( - executable_tool_names: BTreeSet, - allowed_tool_names: BTreeSet, - ) -> Self { + /// + /// Takes the paired [`TurnToolNames`] rather than two same-typed sets, so + /// the executable and allowed sets cannot be transposed here — the same + /// reason [`TurnToolNames::model_turn`] is the only construction site for + /// driver-facing blocking turns. Under [`AgentDriver`](crate::agent::AgentDriver) + /// prefer [`TurnTools::streamed_turn_assembler`](crate::agent::TurnTools::streamed_turn_assembler), + /// which takes the names straight off the turn the matching + /// `SendRequest` carried. + /// + /// These names govern *mid-stream* validation only. What a turn is + /// finally allowed to do is answered by the run itself, from the metadata + /// committed when the request was built — see + /// [`AgentRun::streamed_turn`](super::AgentRun::streamed_turn) — so a + /// widened set here cannot get a forbidden call dispatched. It can only + /// cost the early exit this type exists to provide. + pub fn new(names: &TurnToolNames) -> Self { Self { - executable_tool_names, - allowed_tool_names, + executable_tool_names: names.executable.clone(), + allowed_tool_names: names.allowed.clone(), text: String::new(), saw_text: false, reasoning_parts: Vec::new(), @@ -907,7 +920,7 @@ mod tests { } fn assembler() -> StreamedTurnAssembler { - StreamedTurnAssembler::new(tool_names(&["add"]), tool_names(&["add"])) + StreamedTurnAssembler::new(&TurnToolNames::new(["add"], ["add"])) } fn text_item(text: &str) -> StreamedAssistantContent { diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs index 5caf1035ca..1c1cbe6a42 100644 --- a/crates/rig-agent/src/agent/turn_tools.rs +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -20,7 +20,7 @@ use serde::{Deserialize, Serialize}; use rig_core::message::{ToolChoice, UserContent}; use super::model::ModelHandle; -use super::run::{ModelTurn, PendingToolCall}; +use super::run::{ModelTurn, PendingToolCall, StreamedTurnAssembler}; use crate::completion::{CompletionRequestBuilder, CompletionResponse}; use crate::tool::server::ToolRegistrySnapshot; use crate::tool::{ToolContext, ToolExecutionError, ToolResult}; @@ -330,6 +330,20 @@ impl TurnTools { pub(crate) fn model_turn(&self, response: &CompletionResponse) -> ModelTurn { self.names().model_turn(response) } + + /// The assembler for this turn's provider stream, carrying the names this + /// turn advertised. + /// + /// The streaming counterpart of [`Self::model_turn`], and the way a + /// [`AgentDriver`](super::AgentDriver) caller should build one: it takes + /// the names off the turn the matching + /// [`DriveStep::SendRequest`](super::DriveStep::SendRequest) carried, so + /// there is no pair of same-typed sets for a caller to transpose or widen. + /// Feed [`StreamedTurnAssembler::finish`]'s result to + /// [`AgentDriver::accept_streamed_turn`](super::AgentDriver::accept_streamed_turn). + pub fn streamed_turn_assembler(&self) -> StreamedTurnAssembler { + StreamedTurnAssembler::new(&self.names()) + } } #[cfg(test)] diff --git a/tests/providers/anthropic/cassette/agent_driver.rs b/tests/providers/anthropic/cassette/agent_driver.rs index 3531ce9e82..ab0b0f4f03 100644 --- a/tests/providers/anthropic/cassette/agent_driver.rs +++ b/tests/providers/anthropic/cassette/agent_driver.rs @@ -10,7 +10,7 @@ //! rationale; this module carries the provider-portable core of that suite. use futures::StreamExt; -use rig::agent::run::{OutputMode, StreamedTurnAssembler}; +use rig::agent::run::OutputMode; use rig::agent::{AgentRun, RequestPatch}; use rig::completion::PromptError; use rig::message::{Message, ToolChoice}; @@ -376,10 +376,7 @@ async fn a_streamed_turn_drives_through_the_driver() { let mut driver = agent.drive(ADD_PROMPT); let (request, tools, _) = expect_send(&mut driver).await; - let mut assembler = StreamedTurnAssembler::new( - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ); + let mut assembler = tools.streamed_turn_assembler(); let mut stream = request.stream().await.expect("stream should open"); while let Some(item) = stream.next().await { assembler diff --git a/tests/providers/gemini/cassette/agent_driver.rs b/tests/providers/gemini/cassette/agent_driver.rs index ad4241a9d2..aa296d068d 100644 --- a/tests/providers/gemini/cassette/agent_driver.rs +++ b/tests/providers/gemini/cassette/agent_driver.rs @@ -10,7 +10,7 @@ //! rationale; this module carries the provider-portable core of that suite. use futures::StreamExt; -use rig::agent::run::{OutputMode, StreamedTurnAssembler}; +use rig::agent::run::OutputMode; use rig::agent::{AgentRun, RequestPatch}; use rig::completion::PromptError; use rig::message::{Message, ToolChoice}; @@ -376,10 +376,7 @@ async fn a_streamed_turn_drives_through_the_driver() { let mut driver = agent.drive(ADD_PROMPT); let (request, tools, _) = expect_send(&mut driver).await; - let mut assembler = StreamedTurnAssembler::new( - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ); + let mut assembler = tools.streamed_turn_assembler(); let mut stream = request.stream().await.expect("stream should open"); while let Some(item) = stream.next().await { assembler diff --git a/tests/providers/gemini/cassette/agent_run_streamed.rs b/tests/providers/gemini/cassette/agent_run_streamed.rs index cc66a0a324..41101ca9bd 100644 --- a/tests/providers/gemini/cassette/agent_run_streamed.rs +++ b/tests/providers/gemini/cassette/agent_run_streamed.rs @@ -12,7 +12,7 @@ use rig::agent::run::{ }; use rig::agent::{ AgentHook, InvalidToolCallAction, MultiTurnStreamItem, StreamingError, - ToolCall as ToolCallEvent, ToolCallAction, + ToolCall as ToolCallEvent, ToolCallAction, TurnToolNames, }; use rig::completion::{PromptError, Usage}; use rig::message::{Message, ToolChoice, ToolResult}; @@ -60,7 +60,11 @@ async fn run_streamed_turn( .stream() .await .expect("gemini stream should open"); - let mut assembler = StreamedTurnAssembler::new(executable.clone(), allowed.clone()); + // A raw `AgentRun` commits no prepared-turn metadata, so these carried + // sets are what the machine validates against — the fallback arm of + // `AgentRun::streamed_turn`'s authority rule. + let mut assembler = + StreamedTurnAssembler::new(&TurnToolNames::new(executable.clone(), allowed.clone())); let mut recorded = false; while let Some(item) = stream.next().await { diff --git a/tests/providers/openai/cassette/agent_driver.rs b/tests/providers/openai/cassette/agent_driver.rs index 726366d75a..42cfd60683 100644 --- a/tests/providers/openai/cassette/agent_driver.rs +++ b/tests/providers/openai/cassette/agent_driver.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use futures::StreamExt; -use rig::agent::run::{OutputMode, StreamedTurnAssembler}; +use rig::agent::run::OutputMode; use rig::agent::{AgentRun, RequestPatch, TurnPreparation}; use rig::completion::PromptError; use rig::message::{Message, ToolChoice}; @@ -918,10 +918,7 @@ async fn a_streamed_turn_drives_through_the_driver() { let mut driver = agent.drive(ADD_PROMPT); let (request, tools, _) = expect_send(&mut driver).await; - let mut assembler = StreamedTurnAssembler::new( - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ); + let mut assembler = tools.streamed_turn_assembler(); let mut stream = request.stream().await.expect("stream should open"); while let Some(item) = stream.next().await { let item = item.expect("stream item"); @@ -963,10 +960,7 @@ async fn a_streamed_text_turn_finalizes_the_run() { let mut driver = agent.drive("Say hello."); let (request, tools, _) = expect_send(&mut driver).await; - let mut assembler = StreamedTurnAssembler::new( - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ); + let mut assembler = tools.streamed_turn_assembler(); let mut stream = request.stream().await.expect("stream should open"); while let Some(item) = stream.next().await { assembler @@ -1017,10 +1011,7 @@ async fn a_streamed_turn_honors_the_request_patch() { .await; assert!(!tools.executable_tool_names().contains("subtract")); - let mut assembler = StreamedTurnAssembler::new( - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ); + let mut assembler = tools.streamed_turn_assembler(); let mut stream = request.stream().await.expect("stream should open"); while let Some(item) = stream.next().await { assembler @@ -1072,10 +1063,7 @@ async fn a_truncated_stream_ends_the_turn_without_a_transport_error() { let mut driver = agent.drive(ADD_PROMPT); let (request, tools, _) = expect_send(&mut driver).await; - let mut assembler = StreamedTurnAssembler::new( - tools.executable_tool_names().clone(), - tools.allowed_tool_names().clone(), - ); + let mut assembler = tools.streamed_turn_assembler(); let mut stream = request.stream().await.expect("stream should open"); let mut stream_error = None; while let Some(item) = stream.next().await { From 7a196e2b504459d809d6898a61cc6f1c0b3ed3ac Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 18:16:22 -0700 Subject: [PATCH 20/21] refactor!(agent): resume policy is stated at the resume entry point MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `allow_missing_resumed_tools` was a builder field on `AgentDriver` that `Agent::drive_run` reset to `false`, so a serialized run resumed without reapplying it dispatched differently and nothing on the run said so — the same shape as the sticky `RequestPatch` this branch already deleted, one field narrower. The module's durability paragraph claimed the driver "holds nothing it could lose", which was one field short of true. Replace it with `Agent::resume_run(run, ResumedToolDrift)`, following `RunState.from_json`'s keyword-only resume policy in openai-agents: a process resuming a payload states how it handles drift at the point of resuming, where it cannot be omitted by forgetting a builder call. `drive_run` stays for the common case and means `Reject`. The policy is deliberately not serialized with the run: the process that suspended it cannot know what registry a later one will have, so it has no standing to decide how that process handles its own drift. The durability paragraph now says what the driver actually holds — a rebuildable snapshot cache and this process's own resume policy — rather than claiming it holds nothing. --- crates/rig-agent/CHANGELOG.md | 2 +- crates/rig-agent/src/agent/driver.rs | 126 ++++++++++++++++++++++----- crates/rig-agent/src/agent/mod.rs | 4 +- 3 files changed, 106 insertions(+), 26 deletions(-) diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 1144a8e25e..91b2305aab 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -37,7 +37,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - *(agent)* add `Agent::drive` / `Agent::drive_run`, returning the new public `AgentDriver` (with `DriveStep`, `TurnTools` and `TurnToolNames`): hand-drive the sans-IO `AgentRun` machine with the agent's own configuration while owning every side effect. The driver seeds the run from the agent (`default_max_turns`, `tool_choice`, output schema), yields each turn's fully configured completion request for the caller to send (or hand to a custom transport), pairs every pending tool batch with the exact registry snapshot that advertised it, assembles the model turn internally, and maintains the run's committed structured-output tool across turns exactly as the runner does. Dispatch is gated on the turn's advertised names, so a tool registered after a turn was prepared — or after a run was suspended — cannot be dispatched through a turn that never advertised it. Impossible `tool_choice`/tool-set combinations fail at prepare time with no provider round-trip, and such a failure costs nothing: the run advances only once a request exists, so it keeps its state and its turn budget and the step can be retried in place. The driver performs no provider IO and no dispatch of its own: hooks, memory, retrieval policy, and telemetry still run only under `AgentRunner` -- *(agent)* a driven run is resumable at **every** step boundary, not only while tool calls are pending. The turn's advertised tool names travel with the run — `TurnToolNames`, the serializable half of `TurnTools`, readable via `AgentRun::advertised_tools()` and constructible with `TurnToolNames::new` — so a run serialized after `DriveStep::SendRequest`, the natural suspension point for a queued or long-running provider call, resumes in another process and accepts its reply, validated against the set the request actually carried rather than whatever the resuming registry holds. The names describe the turn in flight and are cleared on every route back to preparing, so a run parked for a fresh call never reports a turn that is over. Resuming rebuilds the dispatch target by resolving those names against the registry directly, without re-running retrieval: a registered *dynamic* tool the turn used is found whether or not a fresh query would rank it, absence means the tool is genuinely gone (surfaced as an error, opt out with `allow_missing_resumed_tools`), and resume does not fail when a vector index is unavailable +- *(agent)* a driven run is resumable at **every** step boundary, not only while tool calls are pending. The turn's advertised tool names travel with the run — `TurnToolNames`, the serializable half of `TurnTools`, readable via `AgentRun::advertised_tools()` and constructible with `TurnToolNames::new` — so a run serialized after `DriveStep::SendRequest`, the natural suspension point for a queued or long-running provider call, resumes in another process and accepts its reply, validated against the set the request actually carried rather than whatever the resuming registry holds. The names describe the turn in flight and are cleared on every route back to preparing, so a run parked for a fresh call never reports a turn that is over. Resuming rebuilds the dispatch target by resolving those names against the registry directly, without re-running retrieval: a registered *dynamic* tool the turn used is found whether or not a fresh query would rank it, absence means the tool is genuinely gone (surfaced as an error, dispatched anyway by resuming with `Agent::resume_run` and `ResumedToolDrift::Dispatch`), and resume does not fail when a vector index is unavailable - *(agent)* [**behavior**] the runner spends a turn only once its request exists. Its loop advanced the run first and then ran completion-call hooks, model selection and request preparation — three steps that can terminate the run — so a hook stop, a selection stop or a preparation failure consumed a turn against a call that never happened. It now advances with `AgentRun::advance`, peeks the inputs, and commits after preparation succeeds, which is the boundary `AgentDriver` already used. Not observable through `AgentRunner`'s own API, since every stop path ends the run either way; it matters because the two coordinators now agree on when a turn is spent, and because anything that later resumes or retries a runner-driven run would have inherited the discrepancy diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index 77538891f1..fafe346d82 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -34,11 +34,15 @@ //! //! # Durability //! -//! The serializable state is *all* of the state: the driver holds nothing it -//! could lose. Serialize [`AgentDriver::run`] at any step boundary — while +//! The serializable state is *all* of the state: what the driver holds besides +//! the run is a tool-registry snapshot cache, rebuilt on demand, and the +//! resume-time [`ResumedToolDrift`] policy the resuming process itself stated. +//! Neither is anything a suspending process recorded, so neither can be lost +//! by serializing. Serialize [`AgentDriver::run`] at any step boundary — while //! tool calls are pending, or while a model call is in flight with a //! long-running or queued provider — and resume in another process with -//! [`Agent::drive_run`]. Every step is a resume point, including +//! [`Agent::drive_run`] or [`Agent::resume_run`]. Every step is a resume point, +//! including //! [`DriveStep::SendRequest`]: the turn's advertised tool names travel with //! the run, so the resuming process validates the model's reply against the //! set the request actually carried rather than against whatever its registry @@ -66,7 +70,7 @@ //! registry snapshot to dispatch pending calls through. If that snapshot no //! longer contains a pending call's tool, the driver surfaces an error instead //! of silently feeding a not-found result to the model (see -//! [`AgentDriver::allow_missing_resumed_tools`]) — the model chose that tool +//! [`ResumedToolDrift`]) — the model chose that tool //! from a registry this process no longer has, and re-prompting cannot fix //! deployment drift. If you suspend runs across deploys, version your agent //! definitions alongside the serialized run; the run's own format is versioned @@ -118,16 +122,52 @@ impl Agent { /// Hand-drive an existing [`AgentRun`] with this agent's configuration — /// the resume path for a run deserialized in a new process, or the entry /// point for a custom-configured run (which is taken as-is, not re-seeded). + /// + /// Resumed pending calls whose tools this process no longer registers are + /// reported as drift. To dispatch them anyway, resume with + /// [`Agent::resume_run`] and [`ResumedToolDrift::Dispatch`]. pub fn drive_run(&self, run: AgentRun) -> AgentDriver { + self.resume_run(run, ResumedToolDrift::Reject) + } + + /// Hand-drive an existing [`AgentRun`], stating how this process handles + /// tool drift in the run it is resuming. + /// + /// The policy is an argument here rather than a setting on the driver + /// because it is the resuming process's decision and nothing else records + /// it: a driver that carried it would answer one way for a run driven + /// straight through and another for the same run resumed from a payload, + /// with nothing on the run to say which. Naming it at the point of + /// resuming is the only place a caller cannot forget it. + pub fn resume_run(&self, run: AgentRun, drift: ResumedToolDrift) -> AgentDriver { AgentDriver { agent: self.clone(), run, snapshot: None, - allow_missing_resumed_tools: false, + drift, } } } +/// How a resumed run treats pending tool calls whose tools this process no +/// longer registers. +/// +/// Resume-time policy, deliberately not serialized with the run: the process +/// that suspended a run cannot know what registry a later one will have, so it +/// has no standing to decide how that process handles its own drift. +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub enum ResumedToolDrift { + /// Surface the drift as an error naming the missing tools. The model chose + /// them from a registry this process no longer has, and re-prompting + /// cannot fix a deployment mismatch. + #[default] + Reject, + /// Dispatch the calls anyway, feeding the resulting not-found results to + /// the model. + Dispatch, +} + /// What the caller must do next to advance an [`AgentDriver`]. /// /// Deliberately exhaustive, like [`AgentRunStep`](super::run::AgentRunStep): a @@ -248,8 +288,9 @@ impl TurnPreparation { } /// Hand-drives one [`AgentRun`] with one [`Agent`]'s configuration. Built by -/// [`Agent::drive`] / [`Agent::drive_run`]; see the [module docs](self) for -/// the driving protocol and the boundary with [`AgentRunner`](super::AgentRunner). +/// [`Agent::drive`] / [`Agent::drive_run`] / [`Agent::resume_run`]; see the +/// [module docs](self) for the driving protocol and the boundary with +/// [`AgentRunner`](super::AgentRunner). pub struct AgentDriver { agent: Agent, run: AgentRun, @@ -261,7 +302,11 @@ pub struct AgentDriver { /// through the exact implementations the provider was shown, and is /// rebuilt on demand when a resumed run reaches its pending tool calls. snapshot: Option>, - allow_missing_resumed_tools: bool, + /// How this process handles drift in a run it resumed — the one thing + /// besides the cache above that is not on [`Self::run`], because it is + /// this process's decision about a payload rather than anything the + /// suspending process recorded. Stated at [`Agent::resume_run`]. + drift: ResumedToolDrift, } impl AgentDriver { @@ -290,15 +335,6 @@ impl AgentDriver { self } - /// Opt out of the resumed-run drift check: dispatch pending calls whose - /// tools are missing from this process's registry anyway, feeding the - /// resulting not-found errors to the model instead of surfacing the drift - /// to the caller. See the [module docs](self) on durability. - pub fn allow_missing_resumed_tools(mut self) -> Self { - self.allow_missing_resumed_tools = true; - self - } - /// The sans-IO run state. Serialize this to suspend the run — for example /// while tool calls are pending approval — and resume it elsewhere with /// [`Agent::drive_run`]. @@ -600,7 +636,7 @@ impl AgentDriver { snapshot.retain_names(&names.executable); let snapshot = Arc::new(snapshot); - if !self.allow_missing_resumed_tools { + if matches!(self.drift, ResumedToolDrift::Reject) { let missing: Vec<&str> = calls .iter() .filter(|call| call.preresolved_result.is_none()) @@ -618,8 +654,8 @@ impl AgentDriver { format!( "resumed run has pending tool calls {missing:?} that are no longer \ registered in this process; register the tools on the agent before \ - resuming, or call `allow_missing_resumed_tools()` to dispatch anyway \ - and feed not-found results to the model" + resuming, or resume with `ResumedToolDrift::Dispatch` to dispatch \ + anyway and feed not-found results to the model" ) .into(), ))); @@ -994,13 +1030,13 @@ mod tests { let message = err.to_string(); assert!(message.contains("add"), "error names the tool: {message}"); assert!( - message.contains("allow_missing_resumed_tools"), + message.contains("ResumedToolDrift::Dispatch"), "error names the opt-out: {message}" ); // The opt-out dispatches anyway and yields a not-found result. let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); - let mut driver = bare_agent.drive_run(run).allow_missing_resumed_tools(); + let mut driver = bare_agent.resume_run(run, ResumedToolDrift::Dispatch); let (calls, tools) = expect_execute_tools!(driver); let mut context = ToolContext::new(); let result = tools @@ -1673,12 +1709,12 @@ mod tests { .expect_err("the missing pending tool must surface"); let message = err.to_string(); assert!( - message.contains("add") && message.contains("allow_missing_resumed_tools"), + message.contains("add") && message.contains("ResumedToolDrift::Dispatch"), "drift must be reported, not the unsatisfiable tool choice: {message}" ); let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); - let mut driver = bare_agent.drive_run(run).allow_missing_resumed_tools(); + let mut driver = bare_agent.resume_run(run, ResumedToolDrift::Dispatch); let (calls, tools) = expect_execute_tools!(driver); let mut context = ToolContext::new(); let result = tools.execute_call(&calls[0], &mut context).await; @@ -1688,6 +1724,48 @@ mod tests { ); } + /// Criterion: `Dispatch` is not merely "no error" — the not-found result + /// has to reach the model, which is the whole reason to choose it. + #[tokio::test] + async fn dispatched_drift_feeds_the_not_found_result_to_the_model() { + let model = MockCompletionModel::new([ + MockTurn::tool_call("call_1", "add", json!({"x": 2, "y": 5})), + MockTurn::text("I could not add those."), + ]); + let agent = AgentBuilder::new(model) + .default_max_turns(2) + .tool(MockAddTool) + .build(); + let mut driver = agent.drive("what is 2 + 5?"); + let (request, _, _) = expect_send!(driver); + let response = request.send().await.expect("scripted turn"); + expect_continue(driver.model_response(&response).expect("turn accepted")); + let _ = expect_execute_tools!(driver); + let serialized = serde_json::to_string(driver.run()).expect("run serializes"); + + let bare_agent = AgentBuilder::new(MockCompletionModel::new([MockTurn::text( + "I could not add those.", + )])) + .default_max_turns(2) + .build(); + let run: AgentRun = serde_json::from_str(&serialized).expect("run deserializes"); + let mut driver = bare_agent.resume_run(run, ResumedToolDrift::Dispatch); + + let (calls, tools) = expect_execute_tools!(driver); + let mut context = ToolContext::new(); + let results = vec![tools.execute_call(&calls[0], &mut context).await]; + driver.tool_results(results).expect("results accepted"); + + // The next request carries the not-found result: the model is told the + // tool is gone rather than left waiting on a call nobody answered. + let (request, _, _) = expect_send!(driver); + let history = format!("{:?}", request.build().chat_history); + assert!( + history.contains("not advertised") || history.contains("not found"), + "the model must see why the call produced nothing: {history}" + ); + } + // ── Streamed ingress authority ────────────────────────────────────── // // `StreamedTurnAssembler` is built by the caller, so the sets it validates diff --git a/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index 8de4457ef0..d310b32750 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -116,7 +116,9 @@ pub(crate) const UNKNOWN_AGENT_NAME: &str = "Unnamed Agent"; pub use builder::{AgentBuilder, NoToolConfig, WithBuilderTools, WithToolServerHandle}; pub use completion::Agent; -pub use driver::{AgentDriver, DriveStep, TurnPreparation, TurnPreparationContext}; +pub use driver::{ + AgentDriver, DriveStep, ResumedToolDrift, TurnPreparation, TurnPreparationContext, +}; pub use hook::CompletionCall as CompletionCallEvent; pub use hook::{ AgentHook, CompletionCallAction, CompletionResponse as CompletionResponseEvent, HookContext, From 9a82b8615ddc60c4066068ff1508ec51b6f1bbdd Mon Sep 17 00:00:00 2001 From: stephen Date: Tue, 11 Aug 2026 18:23:26 -0700 Subject: [PATCH 21/21] docs(agent): say what the parity suite guards and what two coordinators cost MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The suite's module docs described a commit-timing divergence that the very next commit fixed, so the harness misdescribed itself. Replace it with the claim it actually supports: the coordinators agree on the wire, the commit boundary is now shared and pinned by unit tests, and what remains is structural duplication these cassettes *detect* but cannot prevent — which is why the suite is worth extending as either coordinator grows behavior. The driver's module docs said it owns the run/turn pairing "in one place" without noting that `AgentRunner` still implements the same protocol internally. One place for callers is not one place in the crate; say so, name the parity cassettes as what keeps the duplication honest, and state that folding the runner onto the driver is the intended end state. Also record both API breaks from this round in MIGRATING and the changelog. --- MIGRATING.md | 29 +++++++++++++++++++ crates/rig-agent/CHANGELOG.md | 4 +++ crates/rig-agent/src/agent/driver.rs | 9 ++++++ crates/rig-agent/src/agent/run/streamed.rs | 4 +-- crates/rig-agent/src/agent/turn_tools.rs | 5 ++-- .../openai/cassette/coordinator_parity.rs | 23 +++++++++++---- 6 files changed, 64 insertions(+), 10 deletions(-) diff --git a/MIGRATING.md b/MIGRATING.md index 39ddb94a7b..bbda2a5042 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1634,6 +1634,35 @@ resolves the advertised names against the registry directly rather than re-running retrieval, so a registered dynamic tool is found whether or not a fresh query would rank it, and absence really does mean the tool is gone. +Absence is reported, not papered over: `Agent::drive_run` fails with the missing +names rather than silently feeding not-found results to the model. To dispatch +anyway, resume with `Agent::resume_run(run, ResumedToolDrift::Dispatch)`. The +policy is an argument at the resume entry point rather than a setting on the +driver, and it is deliberately not serialized with the run — the process that +suspended a run cannot know what registry a later one will have, so it has no +standing to decide how that process handles its own drift. + +**`StreamedTurnAssembler::new` takes one argument.** It was +`new(executable_tool_names, allowed_tool_names)` — two same-typed +`BTreeSet` values, transposable at the call site — and is now +`new(&TurnToolNames)`. Under `AgentDriver`, do not spell the names at all: call +`tools.streamed_turn_assembler()` on the `TurnTools` the matching +`DriveStep::SendRequest` carried. Elsewhere, wrap what you passed before: + +```rust +// before +StreamedTurnAssembler::new(executable, allowed) +// after +StreamedTurnAssembler::new(&TurnToolNames::new(executable, allowed)) +``` + +Those names now govern *mid-stream* validation only. What a turn is finally +allowed to call is answered by the run, from the metadata committed when its +request was built, so a mis-built assembler can no longer get a tool dispatched +that the request's `tool_choice` forbade — it can only forfeit the early exit +the assembler exists to provide. Runs hand-driven through `AgentRun` commit no +such metadata and keep validating against the names they carry. + An `Agent`'s default model is set at construction. Per-run overrides now go through `runner(...).using_model(...)`, `Agent::set_model`, or a `ModelSelection` hook (see the "runtime model swapping" section for the diff --git a/crates/rig-agent/CHANGELOG.md b/crates/rig-agent/CHANGELOG.md index 91b2305aab..9661989c3e 100644 --- a/crates/rig-agent/CHANGELOG.md +++ b/crates/rig-agent/CHANGELOG.md @@ -9,6 +9,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Changed +- *(agent)* [**breaking**] `StreamedTurnAssembler::new` takes the paired `TurnToolNames` instead of two same-typed `BTreeSet` sets a caller could transpose, and `TurnTools::streamed_turn_assembler()` builds one straight from the turn a `DriveStep::SendRequest` carried, so a driven turn never spells the sets. Those names now govern mid-stream validation only: what a turn is finally allowed to call is answered by the run, from the metadata committed when its request was built. A driven streamed turn therefore cannot be widened past its request's `tool_choice` by a mis-built assembler — previously the streamed ingress validated against the caller's sets while the blocking ingress used the committed ones, so a turn narrowed to `Specific(["add"])` would accept and dispatch a streamed call to any tool the assembler happened to list. Runs hand-driven through `AgentRun` commit no metadata and keep validating against their carried sets + +- *(agent)* [**breaking**] the resumed-run drift policy is an argument at the resume entry point — `Agent::resume_run(run, ResumedToolDrift)` — instead of the `AgentDriver::allow_missing_resumed_tools()` builder, which `drive_run` reset and which therefore made a resumed run dispatch differently with nothing on the run to say so. `drive_run` keeps meaning `Reject`. The policy is not serialized with the run on purpose: the process that suspended it cannot know what registry a later one will have + - *(agent)* [**breaking**] persisted histories and `AgentRun`/`PromptResponse` JSON carry rig-core's tagged assistant content (`{"type": "text", ...}`); the untagged shape does not load — see rig-core's entry and MIGRATING. The flatten `Some({})` round-trip artifact is gone, so `is_empty_assistant_turn`'s classification is identical before and after a persist/restore with no special-casing - *(agent)* [**behavior**] the streamed assembler counts the stream items it excludes from assembly that carry assistant content — replayed tagged assistant blocks (the tagged `AssistantContent` serialization is not a stream-item shape) and text items whose `additional_params` is malformed — and logs a single warning per turn, on every termination path, instead of one per stream item; `StreamedTurnAssembler::excluded_assistant_content` exposes the count, and the full decode-outcome contract is pinned by an enum-driven matrix test. A stream item whose text block carries stray sibling keys decodes as stream *text* — the text is assembled and only the stray keys drop diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs index fafe346d82..14a17f9690 100644 --- a/crates/rig-agent/src/agent/driver.rs +++ b/crates/rig-agent/src/agent/driver.rs @@ -22,6 +22,15 @@ //! advertise/dispatch skew happen; the driver owns it in one place while every //! side effect stays with the caller. //! +//! One place for *callers*, not yet one place in the crate: +//! [`AgentRunner`](super::AgentRunner) still implements the same protocol +//! internally rather than driving this type, so two coordinators exist and can +//! drift apart. They agree today, and the `coordinator_parity` cassettes keep +//! them honest by recording each scenario through both and matching request +//! bodies — but agreement checked by test is weaker than agreement by +//! construction, and folding the runner onto the driver is the intended end +//! state. +//! //! It owns that pairing without *holding* it. Everything durable lives on the //! [`AgentRun`] — including what each committed turn resolved to — and the //! driver's only other field is a cache of the live registry snapshot, which diff --git a/crates/rig-agent/src/agent/run/streamed.rs b/crates/rig-agent/src/agent/run/streamed.rs index a53734c703..f1b10bd385 100644 --- a/crates/rig-agent/src/agent/run/streamed.rs +++ b/crates/rig-agent/src/agent/run/streamed.rs @@ -406,8 +406,8 @@ impl StreamedTurnAssembler { /// /// Takes the paired [`TurnToolNames`] rather than two same-typed sets, so /// the executable and allowed sets cannot be transposed here — the same - /// reason [`TurnToolNames::model_turn`] is the only construction site for - /// driver-facing blocking turns. Under [`AgentDriver`](crate::agent::AgentDriver) + /// reason the blocking path assembles its `ModelTurn` from those names at + /// a single site. Under [`AgentDriver`](crate::agent::AgentDriver) /// prefer [`TurnTools::streamed_turn_assembler`](crate::agent::TurnTools::streamed_turn_assembler), /// which takes the names straight off the turn the matching /// `SendRequest` carried. diff --git a/crates/rig-agent/src/agent/turn_tools.rs b/crates/rig-agent/src/agent/turn_tools.rs index 1c1cbe6a42..183006fc3c 100644 --- a/crates/rig-agent/src/agent/turn_tools.rs +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -334,8 +334,9 @@ impl TurnTools { /// The assembler for this turn's provider stream, carrying the names this /// turn advertised. /// - /// The streaming counterpart of [`Self::model_turn`], and the way a - /// [`AgentDriver`](super::AgentDriver) caller should build one: it takes + /// The streaming counterpart of the blocking path's `ModelTurn` assembly, + /// and the way an [`AgentDriver`](super::AgentDriver) caller should build + /// an assembler: it takes /// the names off the turn the matching /// [`DriveStep::SendRequest`](super::DriveStep::SendRequest) carried, so /// there is no pair of same-typed sets for a caller to transpose or widen. diff --git a/tests/providers/openai/cassette/coordinator_parity.rs b/tests/providers/openai/cassette/coordinator_parity.rs index af2e43478d..850c4215a2 100644 --- a/tests/providers/openai/cassette/coordinator_parity.rs +++ b/tests/providers/openai/cassette/coordinator_parity.rs @@ -24,13 +24,24 @@ //! is renumbered by the scrubber. So as of this suite's first recording the two //! coordinators agree completely on *what* they send. //! -//! Where they still differ is *when* they commit a turn: the runner spends the -//! turn before running its completion-call hooks, its model selection and its +//! # What this suite guards, and what it cannot +//! +//! The commit boundary used to be the interesting divergence: the runner spent +//! its turn before its completion-call hooks, its model selection and its //! request preparation, each of which can terminate the run, while the driver -//! prepares first and commits last. That divergence is invisible here, because -//! it changes no request — which is exactly why it survived several reviews, -//! and why these tests pin the wire while the commit-boundary unit tests pin -//! the rest. +//! prepared first and committed last. That is fixed — both now commit only +//! once a request exists — and the commit-boundary unit tests pin it, because +//! it is invisible on the wire, which is exactly why it survived several +//! reviews. +//! +//! What remains is **structural** duplication: two coordinators implementing +//! one protocol, agreeing today because two code paths currently happen to +//! agree. This suite detects the next disagreement; nothing here prevents one. +//! Only consolidating the coordinators would, and until that lands these +//! cassettes are the guard that makes the delay safe. They are therefore worth +//! extending whenever either coordinator grows a behavior the four scenarios +//! below do not exercise — hook termination, per-turn patches, structured +//! output, recovery resolutions, resumed runs. use rig::agent::AgentRun; use rig::completion::PromptError;