feat(agent): resume a suspended AgentRun through the runner surface

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<dyn std::error::Error>> {
+    /// 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<M> {
+        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

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<Arc<ToolRegistrySnapshot>> = 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<M::StreamingResponse> {
+    pub async fn stream(mut self) -> StreamingResult<M::StreamingResponse> {
         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(),

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<ToolChoice>) {
+        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<OneOrMany<AssistantContent>> {
         let RunState::AwaitingAdvance(turn) = &self.state else {