refactor(agent): separate a prepared agent from its runs

CLAUDE.md

@@ -157,10 +157,12 @@ export AWS_REGION="your-region"       # For Knowledge Base
 - **Rig Fork**: `mezmo/rig` branch `mshearer/LOG-23351-openai-reasoning`
 
 ### Key Modules
+- `builder.rs` - `PreparedAgent` (built once from config: provider client, tools, MCP connections) and `Agent` (one run of it, via `PreparedAgent::begin_run`)
+- `forwarded_headers.rs` - `ForwardedHeaders`, the `headers_from_request` values a prepared agent's MCP connections and HITL route were opened with; `begin_run` refuses a request that forwards different values, so a prepared agent never runs under another request's credentials
 - `provider_agent.rs` - Type-erased streaming across providers
 - `stream_events.rs` - Custom aura SSE events
 - `request_cancellation.rs` - The signal that stops a run, as awaiting work sees it
-- `run_context.rs` - The run a task is working on: its event channel, and the FIFO queue for tool_call_id correlation (see critical assumption below)
+- `run_context.rs` - `RunContext`, the run a task is working on: its event channel, the FIFO queue for tool_call_id correlation (see critical assumption below), and the scratchpad budget, turn-nudge counters and skill-invocation recorder an agent keeps for it; `BoundRun` is the slot through which prepare-time tools and wrappers (turn nudge, scratchpad, skills, HITL) reach it, and `RunLease` is what `begin_run` counts to serve one run at a time
 - `orchestration/` - Multi-agent coordinator, workers, DAG execution, orchestration SSE events
 
 ### Critical Assumption: Rig Sequential Tool Execution

DEVELOPMENT.md

@@ -214,6 +214,7 @@ Prompt routing and execution model:
 - Direct Mode (`orchestration.enabled = false`): single `Agent` handles the turn.
 - Orchestration Mode (`orchestration.enabled = true`): `Orchestrator` coordinates worker execution.
 - Both `Agent` and `Orchestrator` implement `StreamingAgent`, so they are interchangeable at the API boundary.
+- An agent is two halves: `PreparedAgent` is built once from config (provider client, discovered tools, MCP connections) and `Agent` is one run of it, begun with `PreparedAgent::begin_run`. Per-run state (the `RunContext`: request id, event channel, tool-call queue, scratchpad budget, turn-nudge counters, skill-invocation recorder) reaches tools through the prepared agent's `BoundRun` slot, and a prepared agent serves one run at a time; see `crates/aura/src/run_context.rs`. A prepared agent also forwards the `headers_from_request` values of the request that prepared it, and `begin_run` refuses a request that forwards different ones (`crates/aura/src/forwarded_headers.rs`). Today every request prepares a fresh agent; the split is what lets a session reuse one.
 
 Ingress paths into the web server (all end in the same `StreamingAgent::stream` call):
 

crates/aura/src/config.rs

@@ -6,6 +6,7 @@
 //! persistence handles, shared chat history, scratchpad runtime state, and the
 //! session id.
 
+use crate::forwarded_headers::ForwardedHeaders;
 use crate::hitl::HitlRuntime;
 use crate::scratchpad::ScratchpadToolsConfig;
 use crate::tool_wrapper::{ToolCallContext, ToolWrapper};
@@ -123,9 +124,6 @@ pub struct AgentRuntimeConfig {
     /// `Some` when scratchpad is wired up for this agent or worker.
     pub scratchpad_tools_config: Option<ScratchpadToolsConfig>,
 
-    /// Shared turn-limit nudge state for this agent's tool calls.
-    pub turn_nudge: Option<Arc<crate::turn_nudge::TurnNudgeState>>,
-
     /// Shared decision state for worker `submit_result` tool.
     /// When set, workers get the `submit_result` tool for structured output.
     pub orchestration_submit_result: Option<crate::orchestration::SubmitResultDecision>,
@@ -135,25 +133,23 @@ pub struct AgentRuntimeConfig {
     /// `None` disables approval gating.
     pub hitl: Option<HitlRuntime>,
 
-    /// Request id (`req_…`) for this build, used to stamp HITL approval requests
-    /// and route their SSE events. Threaded from the web server so the
-    /// single-agent and orchestration paths share one value.
+    /// Request id (`req_…`) of the request this build serves.
     pub request_id: Option<String>,
 
+    /// The request headers this build forwards; see [`ForwardedHeaders`].
+    pub forwarded_headers: ForwardedHeaders,
+
     /// Computed instance UUID for this agent, derived from agent config and
     /// host identity. Threaded into HITL approval requests so webhook
     /// receivers can identify which instance raised each approval.
     pub instance_id: String,
 
-    /// The `request_approval` tool, pre-built with the appropriate
-    /// [`AgentScope`]. Orchestration workers set this in `create_worker` with
-    /// `AgentScope::Worker`; single-agent mode sets it in `Agent::new` with
-    /// `AgentScope::Single`. `None` when `[hitl]` is not configured.
+    /// The `request_approval` tool, pre-built with its [`AgentScope`].
     ///
     /// [`AgentScope`]: crate::hitl::AgentScope
     pub hitl_request_approval_tool: Option<crate::hitl::RequestApprovalTool>,
 
-    /// Recorder for this session's skill-tool invocations.
+    /// Recorder for this request's skill-tool invocations.
     pub skill_recorder: Option<Arc<crate::skill_tool::SkillInvocationRecorder>>,
 }
 
@@ -177,10 +173,10 @@ impl Clone for AgentRuntimeConfig {
             orchestration_persistence: self.orchestration_persistence.clone(),
             session_id: self.session_id.clone(),
             scratchpad_tools_config: self.scratchpad_tools_config.clone(),
-            turn_nudge: self.turn_nudge.clone(),
             orchestration_submit_result: self.orchestration_submit_result.clone(),
             hitl: self.hitl.clone(),
             request_id: self.request_id.clone(),
+            forwarded_headers: self.forwarded_headers.clone(),
             instance_id: self.instance_id.clone(),
             hitl_request_approval_tool: self.hitl_request_approval_tool.clone(),
             skill_recorder: self.skill_recorder.clone(),
@@ -216,7 +212,6 @@ impl std::fmt::Debug for AgentRuntimeConfig {
                     .map(|_| "<persistence>"),
             )
             .field("session_id", &self.session_id)
-            .field("turn_nudge", &self.turn_nudge.as_ref().map(|_| "<state>"))
             .field(
                 "orchestration_submit_result",
                 &self
@@ -226,6 +221,7 @@ impl std::fmt::Debug for AgentRuntimeConfig {
             )
             .field("hitl", &self.hitl.as_ref().map(|_| "<hitl>"))
             .field("request_id", &self.request_id)
+            .field("forwarded_headers", &self.forwarded_headers)
             .field("instance_id", &self.instance_id)
             .field(
                 "hitl_request_approval_tool",

crates/aura/src/forwarded_headers.rs

@@ -0,0 +1,235 @@
+//! The request headers a prepared agent forwards.
+//!
+//! `headers_from_request` mappings copy inbound request headers onto the MCP
+//! servers and the HITL webhook route when an agent is prepared, and the MCP
+//! connections are opened with them. A prepared agent is therefore bound to
+//! the request that prepared it as far as those headers go, and can only serve
+//! another request that forwards the same values.
+
+use std::collections::{BTreeMap, BTreeSet, HashMap};
+
+use aura_config::{Config, DecisionRouteConfig, McpServerConfig};
+
+/// The inbound request headers `headers_from_request` mappings read, with the
+/// values the preparing request carried.
+#[derive(Clone, Default, PartialEq, Eq)]
+pub struct ForwardedHeaders {
+    /// Lowercased header name → the request's value for it.
+    values: BTreeMap<String, Option<String>>,
+}
+
+impl ForwardedHeaders {
+    /// The headers `req_headers` forwards under `config`: every inbound name a
+    /// `headers_from_request` mapping reads, on an MCP server or the HITL
+    /// webhook route, looked up case-insensitively.
+    pub fn resolve(config: &Config, req_headers: Option<&HashMap<String, String>>) -> Self {
+        let mut names = BTreeSet::new();
+        if let Some(mcp) = &config.mcp {
+            for server in mcp.servers.values() {
+                match server {
+                    McpServerConfig::HttpStreamable {
+                        headers_from_request,
+                        ..
+                    }
+                    | McpServerConfig::Sse {
+                        headers_from_request,
+                        ..
+                    } => names.extend(headers_from_request.values().map(|n| n.to_lowercase())),
+                    McpServerConfig::Stdio { .. } => {}
+                }
+            }
+        }
+        if let Some(DecisionRouteConfig::Webhook {
+            headers_from_request,
+            ..
+        }) = config.hitl.as_ref().map(|hitl| &hitl.route)
+        {
+            names.extend(headers_from_request.values().map(|n| n.to_lowercase()));
+        }
+        Self::of(names, req_headers)
+    }
+
+    /// `req_headers` seen through the names this forwards.
+    pub fn project(&self, req_headers: Option<&HashMap<String, String>>) -> Self {
+        Self::of(self.values.keys().cloned(), req_headers)
+    }
+
+    /// `req_headers` seen through `names`, which are lowercase: each name with
+    /// the value the request carries under it in any letter case, or `None`
+    /// when the request does not carry it.
+    pub(crate) fn of(
+        names: impl IntoIterator<Item = String>,
+        req_headers: Option<&HashMap<String, String>>,
+    ) -> Self {
+        let values = names
+            .into_iter()
+            .map(|name| {
+                let value = req_headers.and_then(|headers| {
+                    headers
+                        .iter()
+                        .find(|(k, _)| k.to_lowercase() == name)
+                        .map(|(_, v)| v.clone())
+                });
+                (name, value)
+            })
+            .collect();
+        Self { values }
+    }
+
+    /// The forwarded headers as a request carries them, omitting the absent
+    /// ones.
+    pub fn as_request(&self) -> HashMap<String, String> {
+        self.values
+            .iter()
+            .filter_map(|(name, value)| value.clone().map(|v| (name.clone(), v)))
+            .collect()
+    }
+
+    /// The first header, by name, whose value differs between `self` and
+    /// `req_headers` seen through the same names.
+    pub fn first_difference(&self, req_headers: Option<&HashMap<String, String>>) -> Option<&str> {
+        let theirs = self.project(req_headers);
+        self.values
+            .iter()
+            .find(|(name, value)| theirs.values.get(*name) != Some(value))
+            .map(|(name, _)| name.as_str())
+    }
+}
+
+/// Names only: the values are credentials.
+impl std::fmt::Debug for ForwardedHeaders {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        let mut map = f.debug_map();
+        for (name, value) in &self.values {
+            map.entry(
+                name,
+                &value.as_ref().map(|_| "<present>").unwrap_or("<absent>"),
+            );
+        }
+        map.finish()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn config(mcp_mappings: &[(&str, &str)], hitl_mappings: &[(&str, &str)]) -> Config {
+        let mut toml = String::from(
+            r#"
+[agent]
+name = "t"
+system_prompt = "t"
+
+[agent.llm]
+provider = "openai"
+model = "gpt-4o"
+api_key = "k"
+"#,
+        );
+        if !mcp_mappings.is_empty() {
+            toml.push_str("\n[mcp.servers.s]\ntransport = \"http_streamable\"\nurl = \"https://example.com/mcp\"\n[mcp.servers.s.headers_from_request]\n");
+            for (out, inbound) in mcp_mappings {
+                toml.push_str(&format!("{out} = \"{inbound}\"\n"));
+            }
+        }
+        if !hitl_mappings.is_empty() {
+            toml.push_str(
+                "\n[hitl]\nrequire_approval = [\"x_*\"]\n[hitl.route]\nmode = \"webhook\"\nurl = \"https://example.com/hook\"\n[hitl.route.headers_from_request]\n",
+            );
+            for (out, inbound) in hitl_mappings {
+                toml.push_str(&format!("{out} = \"{inbound}\"\n"));
+            }
+        }
+        aura_config::load_config_from_str(&toml).expect("config parses")
+    }
+
+    fn headers(pairs: &[(&str, &str)]) -> HashMap<String, String> {
+        pairs
+            .iter()
+            .map(|(k, v)| (k.to_string(), v.to_string()))
+            .collect()
+    }
+
+    #[test]
+    fn nothing_is_forwarded_without_mappings() {
+        let forwarded =
+            ForwardedHeaders::resolve(&config(&[], &[]), Some(&headers(&[("x-user", "a")])));
+        assert_eq!(forwarded, ForwardedHeaders::default());
+        assert_eq!(
+            forwarded.first_difference(Some(&headers(&[("x-user", "b")]))),
+            None
+        );
+    }
+
+    #[test]
+    fn every_mapping_source_contributes_its_inbound_names() {
+        let config = config(
+            &[("authorization", "X-User-Token")],
+            &[("x-approver", "x-actor")],
+        );
+        let forwarded =
+            ForwardedHeaders::resolve(&config, Some(&headers(&[("x-user-token", "t1")])));
+
+        assert_eq!(
+            forwarded.as_request(),
+            headers(&[("x-user-token", "t1")]),
+            "the MCP mapping's inbound name is read case-insensitively; the absent HITL one is omitted",
+        );
+        assert_eq!(
+            format!("{forwarded:?}"),
+            r#"{"x-actor": "<absent>", "x-user-token": "<present>"}"#,
+            "debug output names the headers and never their values",
+        );
+    }
+
+    #[test]
+    fn a_request_forwarding_the_same_values_matches() {
+        let config = config(&[("authorization", "x-user-token")], &[]);
+        let forwarded =
+            ForwardedHeaders::resolve(&config, Some(&headers(&[("X-User-Token", "t1")])));
+
+        assert_eq!(
+            forwarded.first_difference(Some(&headers(&[("x-user-token", "t1")]))),
+            None
+        );
+        assert_eq!(
+            forwarded.first_difference(Some(&headers(&[
+                ("x-user-token", "t1"),
+                ("x-other", "ignored")
+            ]))),
+            None,
+            "headers no mapping reads do not matter",
+        );
+    }
+
+    #[test]
+    fn a_request_forwarding_a_different_value_differs() {
+        let config = config(&[("authorization", "x-user-token")], &[]);
+        let forwarded =
+            ForwardedHeaders::resolve(&config, Some(&headers(&[("x-user-token", "t1")])));
+
+        assert_eq!(
+            forwarded.first_difference(Some(&headers(&[("x-user-token", "t2")]))),
+            Some("x-user-token"),
+        );
+        assert_eq!(
+            forwarded.first_difference(None),
+            Some("x-user-token"),
+            "a request without the header the agent forwarded differs too",
+        );
+    }
+
+    #[test]
+    fn an_agent_prepared_without_the_header_differs_from_a_request_carrying_it() {
+        let config = config(&[("authorization", "x-user-token")], &[]);
+        let forwarded = ForwardedHeaders::resolve(&config, None);
+
+        assert_eq!(forwarded.first_difference(None), None);
+        assert_eq!(
+            forwarded.first_difference(Some(&headers(&[("x-user-token", "t1")]))),
+            Some("x-user-token"),
+            "the static fallback the agent connected with is not this request's token",
+        );
+    }
+}