From 025471c554d0da4799749acfdc21865e66d9fefb Mon Sep 17 00:00:00 2001 From: Naz Quadri Date: Mon, 3 Aug 2026 14:45:17 -0400 Subject: [PATCH] feat(agent): resume a suspended AgentRun through the runner surface AgentRun serializes between steps, but the high-level driver could not accept a restored run: resuming a persisted run meant hand-writing the drive loop and losing the runner's hook, tool-server, memory, and telemetry wiring. Add Agent::resume(run) / AgentRunner::resume(agent, run), which hand a deserialized run back to the single engine shared by run() and stream(). The run supplies the loop state and seeds the runner's budgets; the agent supplies the environment. A run restored with tool calls pending binds them to the current tool registry (pinned per-turn handles do not survive serialization); memory is appended at Done but never loaded. --- crates/rig-agent/src/agent/completion.rs | 41 ++- .../src/agent/prompt_request/streaming.rs | 75 +++-- crates/rig-agent/src/agent/run/mod.rs | 38 ++- crates/rig-agent/src/agent/runner.rs | 287 +++++++++++++++++- 4 files changed, 400 insertions(+), 41 deletions(-) diff --git a/crates/rig-agent/src/agent/completion.rs b/crates/rig-agent/src/agent/completion.rs index a6fb079311..4cc03ad39e 100644 --- a/crates/rig-agent/src/agent/completion.rs +++ b/crates/rig-agent/src/agent/completion.rs @@ -1,6 +1,6 @@ use super::hook::{HookStack, RequestPatch}; use super::prompt_request::{self, PromptRequest}; -use super::run::OutputMode; +use super::run::{AgentRun, OutputMode}; use super::runner::AgentRunner; use crate::{ agent::prompt_request::streaming::StreamingPromptRequest, @@ -621,6 +621,45 @@ where AgentRunner::from_agent(self, prompt) } + /// Build an [`AgentRunner`] that continues a previously suspended + /// [`AgentRun`] — one serialized between steps, persisted, and restored + /// (possibly in another process) — with this agent's hooks, tool server, + /// memory, and telemetry. [`AgentRunner::run`] and [`AgentRunner::stream`] + /// share their drive loop with a fresh run, so a resumed run behaves + /// identically to one that was never suspended. + /// + /// The restored run supplies the loop state (prompt, history, pending tool + /// calls, budgets, aggregated usage); this agent supplies the environment + /// it resumes into and should be configured like the one that started the + /// run — pinned per-turn tool handles do not survive serialization, so + /// pending tool calls execute against this agent's current registry. + /// Conversation memory is never loaded for a resumed run (the run carries + /// its own history), but a configured backend still receives the completed + /// run's messages at the end. See the [`run`](crate::agent::run) module + /// docs for serialization caveats. + /// + /// # Example + /// ```no_run + /// use rig_agent::{agent::run::AgentRun, prelude::*}; + /// use rig_core::{client::ProviderClient, providers::openai}; + /// + /// # async fn run() -> Result<(), Box> { + /// let openai = openai::Client::from_env()?; + /// let agent = openai.agent(openai::GPT_5_2).build(); + /// + /// // A run suspended earlier (possibly by another process). + /// let json = std::fs::read_to_string("suspended-run.json")?; + /// let run: AgentRun = serde_json::from_str(&json)?; + /// + /// let response = agent.resume(run).run().await?; + /// println!("{}", response.output); + /// # Ok(()) + /// # } + /// ``` + pub fn resume(&self, run: AgentRun) -> AgentRunner { + AgentRunner::resume(self, run) + } + /// Resolve the provider-facing tool definitions available for a prompt. /// /// This read-only view does not expose tool dispatch. Agent execution and diff --git a/crates/rig-agent/src/agent/prompt_request/streaming.rs b/crates/rig-agent/src/agent/prompt_request/streaming.rs index 6d6be91148..02d79022f2 100644 --- a/crates/rig-agent/src/agent/prompt_request/streaming.rs +++ b/crates/rig-agent/src/agent/prompt_request/streaming.rs @@ -505,6 +505,9 @@ where // immediately following CallTools step. This keeps the sans-IO run state // serializable while pinning execution to the definitions sent that turn. let mut pending_tool_snapshot: Option> = None; + // Only a restored run may open on CallTools without a pinned snapshot; + // on any later step a missing snapshot is a driver bug. + let mut first_step = true; 'outer: loop { let step = match run.next_step() { @@ -612,13 +615,32 @@ where pending_tool_snapshot = Some(turn_tool_snapshot); } AgentRunStep::CallTools { calls } => { - let Some(tool_snapshot) = pending_tool_snapshot.take() else { - store_error_usage(&runner, &run); - yield Err(StreamingError::Completion(CompletionError::ResponseError( - "agent requested tool execution without a prepared registry snapshot" - .to_string(), - ))); - break 'outer; + // A restored run's pinned handles did not survive serialization: + // its opening batch binds to the current registry instead. + let tool_snapshot = match pending_tool_snapshot.take() { + Some(tool_snapshot) => tool_snapshot, + None if first_step => { + match runner.tool_server_handle.snapshot_tool_defs(None).await { + Ok(tool_snapshot) => Arc::new(tool_snapshot), + Err(_) => { + store_error_usage(&runner, &run); + yield Err(StreamingError::Completion( + CompletionError::RequestError( + "Failed to get tool definitions".into(), + ), + )); + break 'outer; + } + } + } + None => { + store_error_usage(&runner, &run); + yield Err(StreamingError::Completion(CompletionError::ResponseError( + "agent requested tool execution without a prepared registry snapshot" + .to_string(), + ))); + break 'outer; + } }; let mut tool_stream = source.run_tool_calls( &runner, @@ -668,6 +690,7 @@ where break 'outer; } } + first_step = false; } } } @@ -1504,7 +1527,7 @@ where /// hook handling with the blocking [`run`](AgentRunner::run) via /// `drive_agent`, so the two behave identically apart from the streamed /// delta events. - pub async fn stream(self) -> StreamingResult { + pub async fn stream(mut self) -> StreamingResult { let (agent_span, created_agent_span) = acquire_agent_span( self.agent_name_or_default(), self.preamble.as_deref(), @@ -1520,25 +1543,29 @@ where // When the caller passes explicit history, memory is fully bypassed for // this request (no load AND no save). Otherwise, if a memory backend and // conversation id are both configured, load prior history. - let (history_override, memory_handle) = match &self.chat_history { - Some(_) => (None, None), - None => match (&self.memory, &self.conversation_id) { - (Some(memory), Some(id)) => match memory.load(id).await { - Ok(loaded) => (Some(loaded), Some((memory.clone(), id.clone()))), - Err(err) => { - let stream = async_stream::stream! { - yield Err(StreamingError::from(err)); - }; - // Instrument under the agent span like the success path so - // a load failure stays tied to invoke_agent. - return Box::pin(stream.instrument(agent_span)); - } + let (history_override, memory_handle) = if self.restored_run.is_some() { + (None, self.resume_memory_handle()) + } else { + match &self.chat_history { + Some(_) => (None, None), + None => match (&self.memory, &self.conversation_id) { + (Some(memory), Some(id)) => match memory.load(id).await { + Ok(loaded) => (Some(loaded), Some((memory.clone(), id.clone()))), + Err(err) => { + let stream = async_stream::stream! { + yield Err(StreamingError::from(err)); + }; + // Instrument under the agent span like the success path so + // a load failure stays tied to invoke_agent. + return Box::pin(stream.instrument(agent_span)); + } + }, + _ => (None, None), }, - _ => (None, None), - }, + } }; - let run = self.build_run(history_override); + let run = self.take_run(history_override); let source = StreamingTurnSource::new( &self.hooks, self.agent_name_or_default().to_string(), diff --git a/crates/rig-agent/src/agent/run/mod.rs b/crates/rig-agent/src/agent/run/mod.rs index 65ebf39738..4111f4b4e3 100644 --- a/crates/rig-agent/src/agent/run/mod.rs +++ b/crates/rig-agent/src/agent/run/mod.rs @@ -16,11 +16,13 @@ //! 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. +//! resume it later in another process — by continuing to drive it by hand, or +//! by handing it back to the full agent driver with +//! [`Agent::resume`](crate::agent::Agent::resume). 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. //! //! `AgentRun` deliberately contains no model, tool registry, memory backend, or //! hook stack. Hand-driving it is a low-level provider integration: the caller @@ -486,6 +488,32 @@ impl AgentRun { &self.new_messages } + /// The run's total model-call budget. + pub(crate) fn max_turns_limit(&self) -> usize { + self.max_turns + } + + /// The run's invalid tool-call retry budget. + pub(crate) fn invalid_tool_call_retry_limit(&self) -> usize { + self.max_invalid_tool_call_retries + } + + /// The tool-choice policy stored on the run. + pub(crate) fn tool_choice(&self) -> Option<&ToolChoice> { + self.tool_choice.as_ref() + } + + /// Replace the tool-choice policy, used when a resuming runner re-applies + /// its loop overrides to a restored run. + pub(crate) fn set_tool_choice(&mut self, tool_choice: Option) { + self.tool_choice = tool_choice; + } + + /// The prompt message that started the run. + pub(crate) fn initial_prompt(&self) -> Option<&Message> { + self.new_messages.first() + } + /// Canonical content for the accepted model turn awaiting advancement. pub(crate) fn accepted_turn_choice(&self) -> Option> { let RunState::AwaitingAdvance(turn) = &self.state else { diff --git a/crates/rig-agent/src/agent/runner.rs b/crates/rig-agent/src/agent/runner.rs index 55dbb1e5e3..3cf0747777 100644 --- a/crates/rig-agent/src/agent/runner.rs +++ b/crates/rig-agent/src/agent/runner.rs @@ -226,6 +226,9 @@ where pub(crate) conversation_id: Option, pub(crate) hooks: HookStack, pub(crate) error_usage: Option>>, + /// A restored run to continue instead of building a fresh one. Set by + /// [`resume`](Self::resume). + pub(crate) restored_run: Option, } impl AgentRunner @@ -262,9 +265,29 @@ where conversation_id: agent.default_conversation_id.clone(), hooks: agent.hooks.clone(), error_usage: None, + restored_run: None, } } + /// Build a runner that continues a previously suspended [`AgentRun`], + /// seeding it with the agent's default hook stack and the run's own + /// turn/retry budgets and tool-choice policy, so builder overrides such as + /// [`max_turns`](Self::max_turns) default to the suspended values. Prefer + /// [`Agent::resume`](crate::agent::Agent::resume), which documents the + /// resume semantics. + pub fn resume(agent: &Agent, run: AgentRun) -> Self { + let prompt = run + .initial_prompt() + .cloned() + .unwrap_or_else(|| Message::user(String::new())); + let mut runner = Self::from_agent(agent, prompt); + runner.max_turns = run.max_turns_limit(); + runner.max_invalid_tool_call_retries = run.invalid_tool_call_retry_limit(); + runner.tool_choice = run.tool_choice().cloned(); + runner.restored_run = Some(run); + runner + } + /// Append a hook to the stack (on top of any the agent already carries). /// Hooks run in registration order; how their results compose is /// event-dependent (`CompletionCall` request patches accumulate and merge, @@ -508,6 +531,31 @@ where None => run, } } + + /// The run to drive: the restored run when resuming (with this runner's + /// loop overrides re-applied), otherwise a freshly built one. + pub(crate) fn take_run(&mut self, history_override: Option>) -> AgentRun { + match self.restored_run.take() { + Some(run) => { + let mut run = run + .max_turns(self.max_turns) + .max_invalid_tool_call_retries(self.max_invalid_tool_call_retries); + run.set_tool_choice(self.tool_choice.clone()); + run + } + None => self.build_run(history_override), + } + } + + /// Memory participation for a resumed run: never load (the run's history + /// is authoritative); a configured backend still gets the completed run's + /// messages at `Done`, which the suspending process never appended. + pub(crate) fn resume_memory_handle(&self) -> Option<(Arc, String)> { + match (&self.memory, &self.conversation_id) { + (Some(memory), Some(id)) => Some((memory.clone(), id.clone())), + _ => None, + } + } } /// Construct an [`AgentRun`] from explicit run configuration. The single place a @@ -1115,7 +1163,7 @@ where /// Drive the agent loop to completion, returning the aggregated /// [`PromptResponse`]. Hooks fire at every observable point; the first hook /// to terminate cancels the run. - pub async fn run(self) -> Result { + pub async fn run(mut self) -> Result { let (agent_span, created_agent_span) = acquire_agent_span( self.agent_name_or_default(), self.preamble.as_deref(), @@ -1131,18 +1179,22 @@ where // When the caller passes explicit history, memory is fully bypassed for // this run (no load AND no save). Otherwise, if a memory backend and // conversation id are both configured, load prior history. - let (history_override, memory_handle) = match &self.chat_history { - Some(_) => (None, None), - None => match (&self.memory, &self.conversation_id) { - (Some(memory), Some(id)) => { - let loaded = memory.load(id).await?; - (Some(loaded), Some((memory.clone(), id.clone()))) - } - _ => (None, None), - }, + let (history_override, memory_handle) = if self.restored_run.is_some() { + (None, self.resume_memory_handle()) + } else { + match &self.chat_history { + Some(_) => (None, None), + None => match (&self.memory, &self.conversation_id) { + (Some(memory), Some(id)) => { + let loaded = memory.load(id).await?; + (Some(loaded), Some((memory.clone(), id.clone()))) + } + _ => (None, None), + }, + } }; - let run = self.build_run(history_override); + let run = self.take_run(history_override); // Fold the shared engine to its final response. The blocking surface // uses a unary model transport and ignores the intermediate items the @@ -10188,4 +10240,217 @@ mod migrated_tests { assert_eq!(stop_calls.load(SeqCst), 1); assert_eq!(after_stop_calls.load(SeqCst), 0); } + + mod resume { + use super::*; + use crate::agent::run::{AgentRun, AgentRunStep, ModelTurn, ModelTurnOutcome}; + use crate::test_utils::CountingMemory; + use rig_core::memory::ConversationMemory; + use std::collections::BTreeSet; + + /// Hand-drive a run to its tool boundary (one pending `add` call) and + /// serialize it to JSON, as a persisting host does before suspending. + fn suspended_mid_tool_batch(max_turns: usize) -> String { + let mut run = AgentRun::new("add 2 and 3").max_turns(max_turns); + let step = run.next_step().expect("first step"); + assert!(matches!(step, AgentRunStep::CallModel { turn: 1, .. })); + + let mut usage = Usage::new(); + usage.input_tokens = 7; + usage.output_tokens = 3; + usage.total_tokens = 10; + let tool_names: BTreeSet = ["add".to_string()].into(); + let outcome = run + .model_response(ModelTurn::new( + None, + OneOrMany::one(AssistantContent::ToolCall(MessageToolCall::new( + "tc1".to_string(), + ToolFunction::new("add".to_string(), json!({"x": 2, "y": 3})), + ))), + usage, + tool_names.clone(), + tool_names, + )) + .expect("model turn should be accepted"); + assert!(matches!(outcome, ModelTurnOutcome::Continue { .. })); + + let step = run.next_step().expect("tool step"); + assert!(matches!(step, AgentRunStep::CallTools { .. })); + serde_json::to_string(&run).expect("run should serialize to JSON") + } + + /// A run serialized at its tool boundary, restored from JSON, completes + /// through the public runner surface: the pending tool executes through + /// the live tool server, hooks fire on the resumed leg, and usage + /// aggregates across the suspend/resume boundary. + #[tokio::test] + async fn resumed_mid_tool_batch_run_completes_with_hooks() { + let snapshot = suspended_mid_tool_batch(3); + + let run: AgentRun = serde_json::from_str(&snapshot).expect("run should deserialize"); + let mut resumed_usage = Usage::new(); + resumed_usage.input_tokens = 5; + resumed_usage.output_tokens = 2; + resumed_usage.total_tokens = 7; + let model = MockCompletionModel::from_turns([ + MockTurn::text("the answer is 5").with_usage(resumed_usage) + ]); + let hook = RecordingHook::default(); + let agent = AgentBuilder::new(model.clone()).tool(MockAddTool).build(); + + let response = agent + .resume(run) + .add_hook(hook.clone()) + .run() + .await + .expect("resumed run should complete"); + + assert_eq!(response.output, "the answer is 5"); + assert_eq!(hook.tool_results(), vec!["5".to_string()]); + assert_eq!( + hook.shared_events(), + vec![ + StepEventKind::ToolCall, + StepEventKind::ToolResult, + StepEventKind::CompletionCall, + ] + ); + // Only the post-resume model call reached the live model. + assert_eq!(model.request_count(), 1); + // Usage spans the suspending and resuming processes. + assert_eq!(response.usage.input_tokens, 12); + assert_eq!(response.usage.output_tokens, 5); + // Full history: prompt, tool call, tool result, final text. + let messages = response.messages.as_deref().unwrap_or_default(); + assert_eq!(messages.len(), 4); + } + + /// Resuming seeds the runner with the suspended run's own turn budget; + /// the runner's builder methods still override it before driving. + #[tokio::test] + async fn resume_seeds_and_overrides_the_suspended_turn_budget() { + // Suspended with a budget of one, already spent on the recorded call. + let snapshot = suspended_mid_tool_batch(1); + + let run: AgentRun = serde_json::from_str(&snapshot).expect("run should deserialize"); + let agent = AgentBuilder::new(MockCompletionModel::default()) + .tool(MockAddTool) + .build(); + let err = agent + .resume(run) + .run() + .await + .expect_err("the exhausted budget should be preserved"); + assert!(matches!( + err, + PromptError::MaxTurnsError { max_turns: 1, .. } + )); + + // The same snapshot resumes to completion once the budget is raised. + let run: AgentRun = serde_json::from_str(&snapshot).expect("run should deserialize"); + let agent = AgentBuilder::new(MockCompletionModel::text("done")) + .tool(MockAddTool) + .build(); + let response = agent + .resume(run) + .max_turns(2) + .run() + .await + .expect("the raised budget should let the run finish"); + assert_eq!(response.output, "done"); + } + + /// A resumed run never loads conversation memory — the restored run + /// carries its own history — but a configured backend still receives + /// the completed run's messages at the end. + #[tokio::test] + async fn resumed_run_appends_to_memory_without_loading() { + let snapshot = suspended_mid_tool_batch(3); + let run: AgentRun = serde_json::from_str(&snapshot).expect("run should deserialize"); + + let memory = CountingMemory::default(); + let agent = AgentBuilder::new(MockCompletionModel::text("the answer is 5")) + .tool(MockAddTool) + .memory(memory.clone()) + .conversation("resumed") + .build(); + + let response = agent.resume(run).run().await.expect("resumed run"); + + assert_eq!(memory.load_count(), 0, "resume must not load memory"); + assert_eq!(memory.append_count(), 1, "completion must append once"); + let stored = memory + .inner() + .load("resumed") + .await + .expect("stored history"); + assert_eq!(Some(stored), response.messages); + } + + /// Resuming a run serialized before its first model call behaves like a + /// run that was never suspended. + #[tokio::test] + async fn resume_at_the_model_boundary_runs_like_a_new_run() { + let snapshot = serde_json::to_string(&AgentRun::new("hi").max_turns(2)) + .expect("run should serialize to JSON"); + + let run: AgentRun = serde_json::from_str(&snapshot).expect("run should deserialize"); + let model = MockCompletionModel::text("hello"); + let agent = AgentBuilder::new(model.clone()).build(); + let response = agent.resume(run).run().await.expect("resumed run"); + + assert_eq!(response.output, "hello"); + assert_eq!(model.request_count(), 1); + } + + /// A restored run resumes through `stream()` as well: the pending + /// tool's result surfaces as a stream item, hooks fire on the resumed + /// leg, and the run ends with a final response. + #[tokio::test] + async fn resumed_run_streams_tool_activity_and_final_response() { + let snapshot = suspended_mid_tool_batch(3); + let run: AgentRun = serde_json::from_str(&snapshot).expect("run should deserialize"); + + let model = MockCompletionModel::from_stream_turns([vec![ + MockStreamEvent::text("the answer is 5"), + MockStreamEvent::final_response_with_total_tokens(0), + ]]); + let hook = RecordingHook::default(); + let agent = AgentBuilder::new(model).tool(MockAddTool).build(); + + let mut stream = agent.resume(run).add_hook(hook.clone()).stream().await; + + let mut saw_tool_result = false; + let mut final_response = None; + while let Some(item) = stream.next().await { + match item.expect("stream item") { + MultiTurnStreamItem::StreamUserItem(StreamedUserContent::ToolResult { + .. + }) => { + saw_tool_result = true; + } + MultiTurnStreamItem::FinalResponse(response) => { + final_response = Some(response); + } + _ => {} + } + } + + assert!( + saw_tool_result, + "the resumed tool batch should surface its result" + ); + let final_response = final_response.expect("final response"); + assert_eq!(final_response.output(), "the answer is 5"); + assert_eq!(hook.tool_results(), vec!["5".to_string()]); + assert_eq!( + hook.shared_events(), + vec![ + StepEventKind::ToolCall, + StepEventKind::ToolResult, + StepEventKind::CompletionCall, + ] + ); + } + } }