diff --git a/MIGRATING.md b/MIGRATING.md index 5c51ba5540..bbda2a5042 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -1561,6 +1561,108 @@ 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::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`. + +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 +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. + +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. + +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 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. + +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 78f3a0e60f..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 @@ -17,6 +21,10 @@ 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 - *(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 @@ -31,6 +39,25 @@ 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::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, 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 + +- *(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 + +- *(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 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 + - *(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..63593c021e 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"; @@ -216,28 +201,102 @@ 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 — 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>, +} + /// 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 @@ -518,10 +577,15 @@ 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, + // 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), + allowed_tool_names: Arc::new(allowed_tool_names), + output_tool_name, + }, }) } diff --git a/crates/rig-agent/src/agent/driver.rs b/crates/rig-agent/src/agent/driver.rs new file mode 100644 index 0000000000..14a17f9690 --- /dev/null +++ b/crates/rig-agent/src/agent/driver.rs @@ -0,0 +1,1931 @@ +//! 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. +//! +//! 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 +//! 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 +//! +//! 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`] 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 +//! holds now. +//! +//! 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. Deciding to use it takes +//! two answers, not one: +//! [`CompletionError::is_retryable`](crate::completion::CompletionError::is_retryable) +//! 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 +//! 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 +//! [`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 +//! by [`RUN_SCHEMA_VERSION`](super::run::RUN_SCHEMA_VERSION). + +use std::collections::BTreeSet; +use std::sync::Arc; + +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, 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, Usage, +}; +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 + /// 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). + /// + /// 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, + 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 +/// 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 — with the + /// [`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. + request: Box>, + /// 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: 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 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, + }, + /// 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(), + } + } +} + +/// 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`] / [`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, + /// 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>, + /// 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 { + /// 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 + } + + /// 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 + } + + /// 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 + } + + /// Record one provider completion call for a streamed turn. + /// + /// 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`]. + /// + /// 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. + /// + /// 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) + } + + /// 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. + 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). + /// + /// **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 { + 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, 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 mut baseline = TurnBaseline::from_agent(&self.agent); + if let Some(model) = preparation.model.as_ref() { + baseline.model = model; + } + let prepared = build_prepared_completion_request( + baseline, + TurnRequest { + prompt, + chat_history: &history, + committed_output_tool: committed.as_deref(), + patch: Some(&patch), + }, + ) + .await + .map_err(PromptError::CompletionError)?; + + let metadata = prepared.turn_metadata(); + let PreparedCompletionRequest { builder, tools, .. } = prepared; + self.snapshot = Some(tools.snapshot.clone()); + let turn = self.run.commit_model_call(Some(metadata))?; + Ok(DriveStep::SendRequest { + request: Box::new(builder), + tools, + turn, + }) + } + Advance::CallTools(calls) => { + let tools = self.dispatch_tools_for_turn(&calls).await?; + Ok(DriveStep::ExecuteTools { calls, tools }) + } + Advance::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 { + // 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(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. + /// + /// 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_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()), + /// } + /// } + /// ``` + /// + /// 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 + // 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( + &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) + } + + /// 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 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 matches!(self.drift, ResumedToolDrift::Reject) { + 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 no longer \ + registered in this process; register the tools on the agent before \ + resuming, or resume with `ResumedToolDrift::Dispatch` 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, containing + /// exactly the names the turn advertised and still has. + /// + /// **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 { + self.agent + .tool_server_handle + .snapshot_tool_defs_including(None, required) + .await + .map_err(|_| { + PromptError::CompletionError(CompletionError::RequestError( + "Failed to get tool definitions".into(), + )) + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::agent::AgentBuilder; + 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::{ToolCall, ToolCallId, ToolChoice, ToolFunction}; + 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"] + })) + } + + /// 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` 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) => { + 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"); + 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; + 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"); + 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:?}"), + } + } + + /// 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"); + expect_continue(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"); + 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; + + // 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"); + 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); + 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"); + 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:?}"), + } + } + + /// 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"); + 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"); + + // 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("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.resume_run(run, ResumedToolDrift::Dispatch); + 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); + } + + /// 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"); + expect_continue(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); + 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(); + 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"); + 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:?}"), + } + } + + /// 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)); + + 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"), + "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" + )); + } + + /// 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"); + 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:?}"), + } + } + + /// 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"); + expect_continue(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"); + 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"); + + // 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()); + } + + /// 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. + #[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 + /// 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"); + 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"); + + // 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"); + 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 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] + 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"); + 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"); + + // 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("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.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; + assert!( + matches!(result, UserContent::ToolResult(_)), + "the opt-out must dispatch and produce a tool result" + ); + } + + /// 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 + // 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/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/crates/rig-agent/src/agent/mod.rs b/crates/rig-agent/src/agent/mod.rs index b4a7828df5..d310b32750 100644 --- a/crates/rig-agent/src/agent/mod.rs +++ b/crates/rig-agent/src/agent/mod.rs @@ -101,12 +101,14 @@ //! ``` mod builder; mod completion; +pub mod driver; pub mod hook; pub mod model; 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. @@ -114,6 +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, ResumedToolDrift, TurnPreparation, TurnPreparationContext, +}; pub use hook::CompletionCall as CompletionCallEvent; pub use hook::{ AgentHook, CompletionCallAction, CompletionResponse as CompletionResponseEvent, HookContext, @@ -131,5 +136,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::{PreparedTurnMetadata, TurnToolNames, TurnTools}; diff --git a/crates/rig-agent/src/agent/prompt_request/streaming.rs b/crates/rig-agent/src/agent/prompt_request/streaming.rs index dd6b579ac5..1485a174fa 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::{PreparedCompletionRequest, build_prepared_completion_request}, + agent::completion::{TurnBaseline, TurnRequest, build_prepared_completion_request}, agent::hook::{ AgentHook, HookContext, HookStack, InvalidToolCallAction, ModelSelection, ModelSelectionAction, ModelTurnFinished, ReasoningDelta, StepEventKind, @@ -12,14 +12,15 @@ use crate::{ }, agent::prompt_request::{assistant_text_from_choice, is_empty_assistant_turn}, agent::run::{ - AgentRun, AgentRunStep, PendingToolCall, - streamed::{StreamedResolution, StreamedTurnAssembler, StreamedTurnEvent}, + AgentRun, PendingToolCall, + streamed::{StreamedResolution, StreamedTurnEvent}, }, agent::runner::{ AgentRunner, CompletionCallOutcome, ModelTurnDecision, ToolExecution, acquire_agent_span, 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}, }; @@ -30,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}; @@ -489,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); @@ -499,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); @@ -559,23 +578,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 { @@ -586,8 +609,17 @@ where break 'outer; } }; - run.set_output_tool_name(prepared.output_tool_name.clone()); - let turn_tool_snapshot = prepared.tool_snapshot.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(); rig_core::telemetry::record_model_input(&chat_span, &input_messages, true); @@ -629,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( @@ -662,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!( @@ -1081,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.executable_tool_names.clone(), - prepared.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; @@ -1692,6 +1721,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}; diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index 763529be6d..ab9e6b4e6a 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -15,19 +15,37 @@ //! //! 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 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 -//! 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, 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`. +//! 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; @@ -59,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)?; +//! // 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; @@ -80,6 +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::{PreparedTurnMetadata, TurnToolNames}, completion::{Message, PromptError, Usage}, json_utils, }; @@ -108,6 +160,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. @@ -200,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 @@ -285,6 +377,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, @@ -320,14 +415,73 @@ pub struct AgentRun { /// [`AgentRunStep::CallModel`] is emitted. #[serde(default)] streamed_completion_call_recorded: bool, + /// 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, 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; those runs fall + /// back to the baseline, which is all they ever had. + #[serde(default)] + 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 + /// 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 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 whose request preparation can fail does that work +/// *before* anything advances (see [`AgentRun::commit_model_call`]). +/// +/// 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. + 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, @@ -344,6 +498,8 @@ impl AgentRun { invalid_tool_call_retries: 0, rollback_pending: false, streamed_completion_call_recorded: false, + prepared_turn: None, + model_call_rollbacks: 0, state: RunState::PreparingRequest, } } @@ -418,12 +574,24 @@ 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.park_for_new_request(); + } + + /// Park the run for a fresh model call, ending the current turn. + /// + /// 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.prepared_turn = None; self.state = RunState::PreparingRequest; - self.next_step() } /// Set the retry budget for [`InvalidToolCallAction::Retry`] @@ -442,6 +610,81 @@ 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() + } + + /// What the current turn's model call resolved to, when the call was + /// committed with it (see [`Self::commit_model_call`]). + /// + /// 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 + /// 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 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.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()) + } + + /// 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. @@ -558,7 +801,7 @@ impl AgentRun { } } - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(()) } @@ -616,236 +859,453 @@ 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, }) } - /// Advance the machine and return the next action for the driver. + /// 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) + } + + /// 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, 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(), + "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(), - }); - } + self.check_turn_budget()?; - 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, - }) - } - 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 - ) - }) - .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); - } - return self.reprompt_for_output(); - } + Ok(ModelCallInputs { + prompt: prompt.clone(), + history: build_history_for_request(self.chat_history.as_deref(), history_for_turn), + }) + } - // 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(), - }); + /// 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(), + }) + } - 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)); - } + /// 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. + /// + /// `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. + /// + /// # 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::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, + prepared: Option, + ) -> Result { + if !self.is_preparing_request() { + return Err( + self.protocol_violation("commit_model_call called without a peeked model call") + ); + } + self.check_turn_budget()?; + 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; + self.state = RunState::AwaitingModel; + Ok(self.current_turn) + } - // 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(), - }); - } + /// 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. + /// + /// # Retryable is not replay-safe + /// + /// 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 + /// 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. + // 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.park_for_new_request(); + 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 + } - 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, - }) + /// 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 fn advance(&mut self) -> Result { + loop { + if self.is_preparing_request() { + return Ok(Advance::NeedsModelCall); + } + 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, }) - .collect(); - 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 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; + } - let response = - PromptResponse::new(assistant_text_from_choice(&items), self.usage) + // 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(), + }); + + let response = PromptResponse::new(output, 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)) - } - } - 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" + .with_output_tool_calls(output_tool_calls) + .with_content(final_items); + self.state = RunState::Done(Box::new(response.clone())); + return Ok(Advance::Done(response)); } - _ => { - "next_step called while an invalid tool-call resolution is pending; answer it via resolve_invalid_tool_call first" + + // 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(), + }); } - }; - self.state = state; - Err(self.protocol_violation(reason)) + + 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; + } + + 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)); + } + } + 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()); + 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::Failed => Err(self.protocol_violation( - "next_step called after the run already failed or was misdriven", - )), + } + } + + /// 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)?; + Ok(AgentRunStep::CallModel { + prompt, + history, + turn, + }) + } + Advance::CallTools(calls) => Ok(AgentRunStep::CallTools { calls }), + Advance::Done(response) => Ok(AgentRunStep::Done(response)), } } @@ -873,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, @@ -1009,7 +1479,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 } => { @@ -1035,7 +1505,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, @@ -1158,7 +1628,7 @@ impl AgentRun { } self.new_messages.push(Message::User { content: results }); - self.state = RunState::PreparingRequest; + self.park_for_new_request(); Ok(()) } @@ -1271,14 +1741,20 @@ 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(), - tool_choice: self.tool_choice.clone(), + 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())), is_streaming: true, @@ -1306,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 => { @@ -1344,13 +1826,13 @@ 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, }) } 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, @@ -1366,7 +1848,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(), @@ -1397,7 +1879,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)), }) @@ -1434,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 { @@ -1451,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, )); } @@ -1581,6 +2071,332 @@ 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(PreparedTurnMetadata::new( + TurnToolNames::new(["add"], ["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); + } + + /// `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(PreparedTurnMetadata::new(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(PreparedTurnMetadata::new(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" + ); + } + + /// 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 + /// 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) + .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. + #[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) + .expect_err("there is no peeked call to commit"); + } + fn expect_continue(outcome: ModelTurnOutcome) -> bool { match outcome { ModelTurnOutcome::Continue { @@ -2304,19 +3120,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); @@ -2326,6 +3146,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/run/streamed.rs b/crates/rig-agent/src/agent/run/streamed.rs index c8e4d044b2..f1b10bd385 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 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. + /// + /// 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/runner.rs b/crates/rig-agent/src/agent/runner.rs index 23a357a058..9ef2ac45a8 100644 --- a/crates/rig-agent/src/agent/runner.rs +++ b/crates/rig-agent/src/agent/runner.rs @@ -33,7 +33,7 @@ use futures::StreamExt; use tracing::{Instrument, info_span, span::Id}; use super::{ - completion::{Agent, PreparedCompletionRequest}, + completion::Agent, hook::{ AgentHook, CompletionCall, CompletionCallAction, CompletionResponse as CompletionResponseEvent, HookContext, HookStack, @@ -49,9 +49,8 @@ use super::{ }, tool_result_output, }, - run::{ - AgentRun, DEFAULT_OUTPUT_RETRIES, ModelTurn, ModelTurnOutcome, OutputMode, PendingToolCall, - }, + run::{AgentRun, DEFAULT_OUTPUT_RETRIES, ModelTurnOutcome, OutputMode, PendingToolCall}, + turn_tools::PreparedCompletionRequest, }; 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..183006fc3c --- /dev/null +++ b/crates/rig-agent/src/agent/turn_tools.rs @@ -0,0 +1,556 @@ +//! 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. +//! +//! 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::{ToolChoice, UserContent}; + +use super::model::ModelHandle; +use super::run::{ModelTurn, PendingToolCall, StreamedTurnAssembler}; +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, + /// 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 +/// 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 — 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, + /// 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 { + /// 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 + /// 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(), + ) + } +} + +/// 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: +/// +/// - [`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 { + /// 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 + } + + /// 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." + ))); + } + // 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 + } + + /// 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. 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 + /// [`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 { + // 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; + 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. Delegates to [`TurnToolNames::model_turn`], the single + /// construction site for driver-facing turns. + 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 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. + /// 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)] +mod tests { + 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; + + /// 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" + ); + } + + /// 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 + /// 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/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..e26d47f236 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. @@ -474,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; @@ -518,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 7370fda877..34b64a5ae0 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)* `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 - *(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/completion/request.rs b/crates/rig-core/src/completion/request.rs index 25aa2fd6f9..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,15 +143,117 @@ 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()) } } + + /// Whether re-issuing an equivalent request could plausibly succeed. + /// + /// 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. + /// + /// # What it answers + /// + /// - 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. + /// 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. + /// + /// # What it cannot answer + /// + /// **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)?; } + /// // `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()), + /// } + /// ``` + 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() + } + None => match self { + // 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 + // "fix" it by defaulting to true. + _ => false, + }, + } + } } #[derive(Clone, Debug, PartialEq, Deserialize, Serialize)] @@ -1277,6 +1400,204 @@ 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 transient_transport_failures_are_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 + /// 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 + /// 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}; diff --git a/crates/rig-core/src/http_client/mod.rs b/crates/rig-core/src/http_client/mod.rs index 02e1386472..24a54ef9f2 100644 --- a/crates/rig-core/src/http_client/mod.rs +++ b/crates/rig-core/src/http_client/mod.rs @@ -52,6 +52,80 @@ 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 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, + } + } +} + +/// 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; 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 05aa3d2d80..46ae4a7015 100644 --- a/examples/agent_run_stepping/src/main.rs +++ b/examples/agent_run_stepping/src/main.rs @@ -1,11 +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). +//! `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 //! @@ -17,18 +20,15 @@ //! //! Requires `OPENAI_API_KEY`. -use std::collections::BTreeSet; - use anyhow::Result; -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::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,94 +92,61 @@ 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.") + .default_max_turns(2) .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 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`. 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(); - - let mut outcome = run.model_response(ModelTurn::new( - response.message_id.clone(), - response.choice.clone(), - response.usage, - tool_names.clone(), - tool_names, - ))?; + let response = request.send().await?; + 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. - 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"); + // 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 = local_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 19ccfbfbc4..d1541b130d 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::run::{AgentRun, ModelTurnOutcome}; +use rig::agent::{DriveStep, InvalidToolCallAction}; 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,20 @@ 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 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 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.", + ) + .default_max_turns(10) + .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}"); @@ -164,53 +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); + 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 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 mut outcome = run.model_response(ModelTurn::new( - response.message_id.clone(), - response.choice.clone(), - response.usage, - tool_names.clone(), - tool_names, - ))?; + let response = request.send().await?; + 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 ----- - 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 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(); @@ -233,7 +219,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 +238,7 @@ async fn main() -> Result<()> { let execution = tools .execute( &name, - value.to_string(), + &value.to_string(), &mut rig::tool::ToolContext::new(), ) .await; @@ -307,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(()); } 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/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 \ 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..6704875b3c --- /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"},{"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_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/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 new file mode 100644 index 0000000000..543c27cafb --- /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 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: + - 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..6fbb9ee313 --- /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 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/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 new file mode 100644 index 0000000000..28769034d7 --- /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 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/streamed_interrupted.yaml b/tests/cassettes/openai/agent_driver/streamed_interrupted.yaml new file mode 100644 index 0000000000..44284b7e7f --- /dev/null +++ b/tests/cassettes/openai/agent_driver/streamed_interrupted.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":"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/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 new file mode 100644 index 0000000000..b22d6b0ed0 --- /dev/null +++ b/tests/cassettes/openai/agent_driver/streamed_turn.yaml @@ -0,0 +1,59 @@ +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} + + 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":93,"prompt_tokens_details":{"audio_tokens":0,"cached_tokens":0},"total_tokens":110}} + + 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 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":"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 new file mode 100644 index 0000000000..8d492204c6 --- /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 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":"{\"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 + 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":{"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":"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/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/common/driver_support.rs b/tests/common/driver_support.rs new file mode 100644 index 0000000000..4b50c4fe6c --- /dev/null +++ b/tests/common/driver_support.rs @@ -0,0 +1,160 @@ +//! 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, InvalidToolCallAction, ModelTurnOutcome, PendingToolCall, RequestPatch, + TurnPreparation, TurnTools, +}; +use rig::completion::{CompletionResponse, 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 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, +) -> ( + 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"); +} + +/// 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 — +/// 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 { + loop { + match driver.next_step().await? { + DriveStep::SendRequest { request, .. } => { + let response = request.send().await.map_err(PromptError::CompletionError)?; + 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(); + 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..ab0b0f4f03 --- /dev/null +++ b/tests/providers/anthropic/cassette/agent_driver.rs @@ -0,0 +1,464 @@ +//! 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; +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, expect_send_patched, expect_turn_accepted, +}; +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"); + 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"); + 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"); + expect_turn_accepted(&mut driver, &response); + + 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"); + expect_turn_accepted(&mut driver, &response); + + // 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"); + expect_turn_accepted(&mut driver, &response); + 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."); + + 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; + 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); + + 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")); + + let response = request.send().await.expect("should send"); + expect_turn_accepted(&mut driver, &response); + 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); + + 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); + 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"); + expect_turn_accepted(&mut driver, &response); + 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); + expect_turn_accepted(&mut resumed, &response); + + 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"); + expect_turn_accepted(&mut driver, &response); + 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"); + expect_turn_accepted(&mut driver, &response); + 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 = tools.streamed_turn_assembler(); + 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 + .record_stream_usage(stream.usage()) + .expect("usage recorded"); + driver + .accept_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"); + expect_turn_accepted(&mut driver, &response); + + 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..aa296d068d --- /dev/null +++ b/tests/providers/gemini/cassette/agent_driver.rs @@ -0,0 +1,460 @@ +//! 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; +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, expect_send_patched, expect_turn_accepted, +}; +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"); + 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"); + 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"); + expect_turn_accepted(&mut driver, &response); + + 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"); + expect_turn_accepted(&mut driver, &response); + + // 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"); + expect_turn_accepted(&mut driver, &response); + 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."); + + 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; + 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); + + 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")); + + let response = request.send().await.expect("should send"); + expect_turn_accepted(&mut driver, &response); + 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); + + 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); + 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"); + expect_turn_accepted(&mut driver, &response); + 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); + expect_turn_accepted(&mut resumed, &response); + + 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"); + expect_turn_accepted(&mut driver, &response); + 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"); + expect_turn_accepted(&mut driver, &response); + 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 = tools.streamed_turn_assembler(); + 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 + .record_stream_usage(stream.usage()) + .expect("usage recorded"); + driver + .accept_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"); + expect_turn_accepted(&mut driver, &response); + + 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/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/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 new file mode 100644 index 0000000000..42cfd60683 --- /dev/null +++ b/tests/providers/openai/cassette/agent_driver.rs @@ -0,0 +1,1102 @@ +//! 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; 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 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::OutputMode; +use rig::agent::{AgentRun, RequestPatch, TurnPreparation}; +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, expect_send_patched, expect_turn_accepted, +}; +use crate::support::{Adder, Subtract}; + +/// A counting `add` so tests can assert the tool ran exactly once. +#[derive(Clone)] +struct CountingAdder { + calls: Arc, +} + +#[derive(serde::Deserialize)] +struct AddArgs { + x: i32, + y: i32, +} + +impl Tool for CountingAdder { + const NAME: &'static str = "add"; + type Error = std::io::Error; + type Args = AddArgs; + type Output = i32; + + fn description(&self) -> String { + "Add x and y together".to_string() + } + + fn parameters(&self) -> serde_json::Value { + serde_json::json!({ + "type": "object", + "properties": { "x": { "type": "number" }, "y": { "type": "number" } }, + "required": ["x", "y"] + }) + } + + fn call( + &self, + _context: &mut ToolContext, + args: Self::Args, + ) -> impl std::future::Future> + Send { + 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 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(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(CountingAdder { + calls: calls.clone(), + }) + .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"); + 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"); + dispatch_and_feed(&mut driver, &pending, &tools).await; + + let (request, _, turn) = expect_send(&mut driver).await; + assert_eq!(turn, 2); + let response = request.send().await.expect("second turn should send"); + expect_turn_accepted(&mut driver, &response); + + 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"); + 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"); + 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"); + expect_turn_accepted(&mut driver, &response); + + 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. +/// +/// 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(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"); + expect_turn_accepted(&mut driver, &response); + + 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"); + expect_turn_accepted(&mut driver, &response); + 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"); + expect_turn_accepted(&mut driver, &response); + 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."); + + 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; + 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); + + 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" + ); + 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()); + }) + .await; +} + +/// `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_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) + .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_patched(&mut driver, RequestPatch::new().active_tools(["add"])).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"); + expect_turn_accepted(&mut driver, &response); + let (pending, tools) = expect_execute_tools(&mut driver).await; + dispatch_and_feed(&mut driver, &pending, &tools).await; + + // 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; +} + +/// 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."); + + 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; + }) + .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?"); + + 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; + assert!(!response.output.trim().is_empty()); + }) + .await; +} + +/// 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?"); + + 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; + 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"); + expect_turn_accepted(&mut driver, &response); + 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. 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(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(), + "a suspended run carries the turn's advertised names" + ); + let mut resumed = agent.drive_run(restored); + expect_turn_accepted(&mut resumed, &response); + + 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"); + expect_turn_accepted(&mut driver, &response); + 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"); + 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 + .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; +} + +// ── 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"); + 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"); + 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(); + + 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_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(_))); + assert_eq!( + driver.run().turn(), + 0, + "a request that never left the process must not consume a turn" + ); + + // 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")); + let response = request.send().await.expect("the retry should send"); + expect_turn_accepted(&mut driver, &response); + }) + .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"); + expect_turn_accepted(&mut driver, &response); + + 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"); + expect_turn_accepted(&mut driver, &response); + + // 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"); + expect_turn_accepted(&mut driver, &response); + 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 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 { + let calls = Arc::new(AtomicUsize::new(0)); + let agent = client + .agent(openai::GPT_4O) + .preamble(FORCE_TOOLS_PREAMBLE) + .default_max_turns(3) + .tool(CountingAdder { + calls: calls.clone(), + }) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send(&mut driver).await; + + 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"); + assembler.ingest(&item).expect("ingest should succeed"); + } + let final_content = stream.choice.clone(); + let streamed = assembler.finish(stream.message_id.clone(), &final_content); + + driver + .record_stream_usage(stream.usage()) + .expect("usage recorded"); + driver + .accept_streamed_turn(streamed) + .expect("streamed 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()); + assert_eq!(calls.load(Ordering::SeqCst), 1); + }) + .await; +} + +/// 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; + + 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 + .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 + .record_stream_usage(stream.usage()) + .expect("usage recorded"); + driver + .accept_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 streamed turn under a `RequestPatch`: the patch reaches the streaming +/// request exactly as it reaches a blocking one. +#[tokio::test] +async fn a_streamed_turn_honors_the_request_patch() { + with_openai_completions_cassette("agent_driver/streamed_patched", |client| async move { + let agent = client + .agent(openai::GPT_4O) + .preamble("BASELINE — must not appear") + .default_max_turns(2) + .tool(Adder) + .tool(Subtract) + .build(); + + let mut driver = agent.drive(ADD_PROMPT); + let (request, tools, _) = expect_send_patched( + &mut driver, + RequestPatch::new() + .preamble(FORCE_TOOLS_PREAMBLE) + .active_tools(["add"]), + ) + .await; + assert!(!tools.executable_tool_names().contains("subtract")); + + 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 + .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 + .record_stream_usage(stream.usage()) + .expect("usage recorded"); + driver + .accept_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 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 = 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 { + 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 + .record_stream_usage(stream.usage()) + .expect("usage recorded"); + driver + .accept_streamed_turn(streamed) + .expect("the truncated turn is still a turn"); + }) + .await; +} diff --git a/tests/providers/openai/cassette/coordinator_parity.rs b/tests/providers/openai/cassette/coordinator_parity.rs new file mode 100644 index 0000000000..850c4215a2 --- /dev/null +++ b/tests/providers/openai/cassette/coordinator_parity.rs @@ -0,0 +1,266 @@ +//! 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. +//! +//! # 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 +//! 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; +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 57f38aca6a..e23bb1cbbc 100644 --- a/tests/providers/openai/mod.rs +++ b/tests/providers/openai/mod.rs @@ -4,8 +4,10 @@ mod regressions; mod cassette { mod agent; + mod agent_driver; mod chat_history; mod completions_api; + mod coordinator_parity; mod document_ordering; mod extractor; mod extractor_usage;