diff --git a/crates/tinyagents-graph/src/stream/project.rs b/crates/tinyagents-graph/src/stream/project.rs index 4988795d2..c585ce303 100644 --- a/crates/tinyagents-graph/src/stream/project.rs +++ b/crates/tinyagents-graph/src/stream/project.rs @@ -185,7 +185,9 @@ impl StreamProjection { }, }); } - AgentEvent::ToolStarted { call_id, tool_name } => { + AgentEvent::ToolStarted { + call_id, tool_name, .. + } => { self.push_tool_call(call_id.clone(), tool_name.clone(), ToolCallPhase::Started); } AgentEvent::ToolCompleted { diff --git a/crates/tinyagents-graph/src/stream/project/test.rs b/crates/tinyagents-graph/src/stream/project/test.rs index 4832e5a61..5c4674040 100644 --- a/crates/tinyagents-graph/src/stream/project/test.rs +++ b/crates/tinyagents-graph/src/stream/project/test.rs @@ -74,6 +74,7 @@ fn stream_projection_folds_tool_lifecycle_as_two_entries() { projection.fold_agent_event(&AgentEvent::ToolStarted { call_id: CallId::from("call-1".to_string()), tool_name: "search".into(), + input: None, }); projection.fold_agent_event(&AgentEvent::ToolCompleted { call_id: CallId::from("call-1".to_string()), @@ -148,6 +149,7 @@ fn stream_projection_since_replays_only_items_after_the_given_cursor() { projection.fold_agent_event(&AgentEvent::ToolStarted { call_id: CallId::from("call-1".to_string()), tool_name: "search".into(), + input: None, }); let cursor_after_first = projection.cursor(); projection.fold_agent_event(&AgentEvent::ModelDelta { diff --git a/crates/tinyagents-harness/src/agent_loop/test.rs b/crates/tinyagents-harness/src/agent_loop/test.rs index 4afd22e70..5ad0fd3c1 100644 --- a/crates/tinyagents-harness/src/agent_loop/test.rs +++ b/crates/tinyagents-harness/src/agent_loop/test.rs @@ -5389,6 +5389,88 @@ async fn tool_completed_event_carries_outcome() { assert_eq!(output_bytes, Some(6), "\"kaboom\".len() == 6"); } +#[tokio::test] +async fn tool_started_event_carries_input_when_capture_enabled() { + use crate::events::RecordingListener; + use crate::runtime::PayloadCapture; + + // Hosts render a tool call's arguments as soon as it starts, not only on + // completion, so `ToolStarted` must carry the same captured input the + // policy already puts on `ToolCompleted`. + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "echo", json!({ "q": "weather" })), + text_response("done", 1, 1), + ])), + ); + harness.register_tool(Arc::new(FakeTool::new("echo", "ok"))); + harness.with_policy(RunPolicy { + capture: PayloadCapture::all(), + ..RunPolicy::default() + }); + + let recorder = Arc::new(RecordingListener::new()); + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-ts"), ()); + ctx.events.subscribe(recorder.clone()); + harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let input = recorder + .events() + .into_iter() + .find_map(|record| match record.event { + AgentEvent::ToolStarted { + tool_name, input, .. + } if tool_name == "echo" => Some(input), + _ => None, + }) + .expect("a ToolStarted event for `echo`"); + + assert_eq!(input, Some(json!({ "q": "weather" }))); +} + +#[tokio::test] +async fn tool_started_event_has_no_input_when_capture_disabled() { + use crate::events::RecordingListener; + + // Default policy is payload-free: `ToolStarted.input` stays `None` so no + // tool argument is captured unless the host opts in. + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_call_response("c1", "echo", json!({ "q": "weather" })), + text_response("done", 1, 1), + ])), + ); + harness.register_tool(Arc::new(FakeTool::new("echo", "ok"))); + + let recorder = Arc::new(RecordingListener::new()); + let ctx: RunContext<()> = RunContext::new(RunConfig::new("run-ts-off"), ()); + ctx.events.subscribe(recorder.clone()); + harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let input = recorder + .events() + .into_iter() + .find_map(|record| match record.event { + AgentEvent::ToolStarted { + tool_name, input, .. + } if tool_name == "echo" => Some(input), + _ => None, + }) + .expect("a ToolStarted event for `echo`"); + + assert_eq!(input, None); +} + // ── `ModelResponse::continue_turn` ─────────────────────────────────────────── /// A tool-less response that keeps the floor by carrying `nudge`. diff --git a/crates/tinyagents-harness/src/agent_loop/tools.rs b/crates/tinyagents-harness/src/agent_loop/tools.rs index c78be5c69..6bd5f901c 100644 --- a/crates/tinyagents-harness/src/agent_loop/tools.rs +++ b/crates/tinyagents-harness/src/agent_loop/tools.rs @@ -902,9 +902,15 @@ impl AgentHarness { // Captured here (where the call actually starts) so the completed // event carries a real start time for duration-aware exporters. let started_at_ms = crate::ids::now_ms(); + // Snapshot the arguments for observability before `call` is moved + // into execution, gated by the capture policy. Shared between the + // `ToolStarted` event (so a host sees the arguments as soon as the + // call starts) and the fold-phase `ToolCompleted` event. + let captured_input = self.policy.capture.tool_io.then(|| call.arguments.clone()); let record = ctx.emit(AgentEvent::ToolStarted { call_id: call_id.clone(), tool_name: tool_name.clone(), + input: captured_input.clone(), }); crate::runtime::emit_host_progress::( ctx, @@ -915,9 +921,6 @@ impl AgentHarness { }, ); status.set_last_event(record.id); - // Snapshot the arguments for observability before `call` is moved - // into execution, gated by the capture policy. - let captured_input = self.policy.capture.tool_io.then(|| call.arguments.clone()); PreparedToolCall { call_id, tool_name, diff --git a/crates/tinyagents-harness/src/events/types.rs b/crates/tinyagents-harness/src/events/types.rs index d3182bb31..6d21f16a4 100644 --- a/crates/tinyagents-harness/src/events/types.rs +++ b/crates/tinyagents-harness/src/events/types.rs @@ -228,6 +228,14 @@ pub enum AgentEvent { call_id: CallId, /// Name of the tool being invoked. tool_name: String, + /// The arguments the tool is being invoked with, captured only when + /// [`PayloadCapture::tool_io`][crate::runtime::PayloadCapture::tool_io] + /// is enabled. `None` in the default payload-free mode, and for + /// events serialized before this field existed. Populated so a host + /// or UI can render the call's arguments as soon as it starts, + /// instead of waiting for [`AgentEvent::ToolCompleted`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + input: Option, }, /// A tool invocation returned. diff --git a/crates/tinyagents-harness/src/middleware/library/test.rs b/crates/tinyagents-harness/src/middleware/library/test.rs index e8c47937c..5d3b45cec 100644 --- a/crates/tinyagents-harness/src/middleware/library/test.rs +++ b/crates/tinyagents-harness/src/middleware/library/test.rs @@ -1676,3 +1676,157 @@ async fn tracing_records_are_bounded_by_max_records() { let counts = tracing.counts(); assert_eq!(counts.agent, 50); } + +// ── PlanModeMiddleware ─────────────────────────────────────────────────────── + +fn plan_mode_policies() -> std::collections::HashMap { + let mut policies = std::collections::HashMap::new(); + policies.insert("read_file".to_string(), ToolPolicy::read_only()); + policies.insert( + "write_file".to_string(), + ToolPolicy::classified().with_side_effects(ToolSideEffects { + writes_files: true, + ..ToolSideEffects::default() + }), + ); + policies +} + +#[tokio::test] +async fn build_mode_leaves_every_tool_exposed_and_executable() { + let (mut ctx, _recorder) = ctx_with_recorder(); + let mode = RunModeHandle::new(RunMode::Build); + let mw = plan_mode_middleware(mode, plan_mode_policies()); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + + let mut request = ModelRequest::new(Vec::new()) + .with_tools(vec![schema_named("read_file"), schema_named("write_file")]); + stack + .run_before_model(&mut ctx, &(), &mut request) + .await + .unwrap(); + assert_eq!(request.tools.len(), 2); + + let mut call = tool_call("write_file"); + stack + .run_before_tool(&mut ctx, &(), &mut call) + .await + .expect("build mode admits every tool"); +} + +#[tokio::test] +async fn plan_mode_hides_side_effecting_tools_and_keeps_read_only_ones() { + let (mut ctx, recorder) = ctx_with_recorder(); + let mode = RunModeHandle::new(RunMode::Plan); + let mw = plan_mode_middleware(mode, plan_mode_policies()); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + + let mut request = ModelRequest::new(Vec::new()) + .with_tools(vec![schema_named("read_file"), schema_named("write_file")]); + stack + .run_before_model(&mut ctx, &(), &mut request) + .await + .unwrap(); + assert_eq!(request.tools.len(), 1); + assert_eq!(request.tools[0].name, "read_file"); + + let filtered = recorder.events().into_iter().find_map(|r| match r.event { + AgentEvent::ToolsFiltered { excluded, .. } => Some(excluded), + _ => None, + }); + assert_eq!(filtered, Some(vec!["write_file".to_string()])); +} + +#[tokio::test] +async fn plan_mode_denies_side_effecting_tool_at_execution() { + let (mut ctx, _recorder) = ctx_with_recorder(); + let mode = RunModeHandle::new(RunMode::Plan); + let mw = plan_mode_middleware(mode, plan_mode_policies()); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + + let mut call = tool_call("write_file"); + let err = stack + .run_before_tool(&mut ctx, &(), &mut call) + .await + .expect_err("side-effecting tool denied in plan mode"); + assert!(matches!(err, TinyAgentsError::Validation(_))); + + let mut call = tool_call("read_file"); + stack + .run_before_tool(&mut ctx, &(), &mut call) + .await + .expect("read-only tool stays executable in plan mode"); +} + +#[tokio::test] +async fn plan_mode_allowlist_keeps_specific_tools_available() { + let (mut ctx, _recorder) = ctx_with_recorder(); + let mode = RunModeHandle::new(RunMode::Plan); + let mw = plan_mode_middleware(mode, plan_mode_policies()).allow(["plan_exit"]); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + + let mut request = ModelRequest::new(Vec::new()).with_tools(vec![ + schema_named("read_file"), + schema_named("write_file"), + schema_named("plan_exit"), + ]); + stack + .run_before_model(&mut ctx, &(), &mut request) + .await + .unwrap(); + let names: Vec<_> = request.tools.iter().map(|s| s.name.as_str()).collect(); + assert_eq!(names, vec!["read_file", "plan_exit"]); + + let mut call = tool_call("plan_exit"); + stack + .run_before_tool(&mut ctx, &(), &mut call) + .await + .expect("allowlisted tool stays executable in plan mode"); +} + +#[tokio::test] +async fn plan_mode_treats_unclassified_tools_as_side_effecting() { + // Fail-closed: a tool with no policy entry at all is hidden/denied in + // plan mode unless explicitly allowlisted. + let (mut ctx, _recorder) = ctx_with_recorder(); + let mode = RunModeHandle::new(RunMode::Plan); + let mw = plan_mode_middleware(mode, std::collections::HashMap::new()); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + + let mut call = tool_call("mystery"); + let err = stack + .run_before_tool(&mut ctx, &(), &mut call) + .await + .expect_err("unclassified tool denied by default in plan mode"); + assert!(matches!(err, TinyAgentsError::Validation(_))); +} + +#[tokio::test] +async fn run_mode_handle_set_takes_effect_immediately_on_shared_clones() { + let mode = RunModeHandle::new(RunMode::Build); + let mw_mode = mode.clone(); + let mw = plan_mode_middleware(mw_mode, plan_mode_policies()); + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push(Arc::new(mw)); + let (mut ctx, _recorder) = ctx_with_recorder(); + + let mut call = tool_call("write_file"); + stack + .run_before_tool(&mut ctx, &(), &mut call) + .await + .expect("build mode admits the call"); + + // Flip the shared handle; the middleware's own clone observes it. + mode.set(RunMode::Plan); + + let mut call = tool_call("write_file"); + stack + .run_before_tool(&mut ctx, &(), &mut call) + .await + .expect_err("plan mode now denies the same call"); +} diff --git a/crates/tinyagents-harness/src/middleware/library/tool_policy.rs b/crates/tinyagents-harness/src/middleware/library/tool_policy.rs index df80d0ab1..e53e3d1a5 100644 --- a/crates/tinyagents-harness/src/middleware/library/tool_policy.rs +++ b/crates/tinyagents-harness/src/middleware/library/tool_policy.rs @@ -617,3 +617,111 @@ impl Middleware for HumanAppro self.decide(ctx, call) } } + +// ── PlanModeMiddleware ───────────────────────────────────────────────────────── + +impl PlanModeMiddleware { + /// Creates a plan-mode middleware driven by `mode`, classifying tools + /// from `policies`. The allowlist starts empty; widen it with + /// [`Self::allow`]. + pub fn new( + mode: RunModeHandle, + policies: std::collections::HashMap, + ) -> Self { + Self { + label: "plan_mode", + mode, + policies, + allow: std::collections::HashSet::new(), + } + } + + /// Adds tools that stay exposed and executable in [`RunMode::Plan`] + /// regardless of their declared side effects — read-only tools and + /// plan-mode-specific tools such as a plan-exit or review-request tool. + pub fn allow(mut self, names: impl IntoIterator>) -> Self { + self.allow.extend(names.into_iter().map(Into::into)); + self + } + + /// Returns `true` when `name` declares a side effect (or has no declared + /// policy at all — see the fail-closed note on [`PlanModeMiddleware`]). + fn is_side_effecting(&self, name: &str) -> bool { + let Some(policy) = self.policies.get(name) else { + return true; + }; + let s = &policy.side_effects; + s.writes_files + || s.network + || s.installs_dependencies + || s.destructive + || s.external_service + || s.payment + } + + /// Returns `true` when `name` may be exposed/executed under + /// [`RunMode::Plan`]: allowlisted, or classified with no side effects. + fn allowed_in_plan_mode(&self, name: &str) -> bool { + self.allow.contains(name) || !self.is_side_effecting(name) + } +} + +#[async_trait] +impl Middleware for PlanModeMiddleware { + fn name(&self) -> &str { + self.label + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &State, + request: &mut ModelRequest, + ) -> Result<()> { + if self.mode.get() != RunMode::Plan { + return Ok(()); + } + let mut excluded = Vec::new(); + request.tools.retain(|schema| { + let keep = self.allowed_in_plan_mode(&schema.name); + if !keep { + excluded.push(schema.name.clone()); + } + keep + }); + // Auditable, mirroring `ContextualToolSelectionMiddleware`: a UI or + // log can see exactly which tools plan mode withheld and why. + if !excluded.is_empty() { + ctx.emit(AgentEvent::ToolsFiltered { + by: self.label.to_string(), + explanations: excluded + .iter() + .map(|name| { + ( + name.clone(), + crate::tool::ToolExposureExplanation::FilteredOut, + ) + }) + .collect(), + excluded, + remaining: request.tools.len(), + }); + } + Ok(()) + } + + async fn before_tool( + &self, + _ctx: &mut RunContext, + _state: &State, + call: &mut ToolCall, + ) -> Result<()> { + if self.mode.get() != RunMode::Plan || self.allowed_in_plan_mode(&call.name) { + return Ok(()); + } + Err(TinyAgentsError::Validation(format!( + "tool `{}` is side-effecting and unavailable in plan mode", + call.name + ))) + } +} diff --git a/crates/tinyagents-harness/src/middleware/library/types.rs b/crates/tinyagents-harness/src/middleware/library/types.rs index b61cc7786..935c51220 100644 --- a/crates/tinyagents-harness/src/middleware/library/types.rs +++ b/crates/tinyagents-harness/src/middleware/library/types.rs @@ -19,7 +19,7 @@ //! `mod.rs`; tests live in `test.rs`. Every public item is re-exported through //! `crate::middleware` so callers import from one place. -use std::collections::VecDeque; +use std::collections::{HashSet, VecDeque}; use std::marker::PhantomData; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; @@ -618,3 +618,107 @@ pub struct TraceCounts { /// Number of `on_error` invocations. pub error: usize, } + +// ── PlanModeMiddleware ───────────────────────────────────────────────────────── + +/// A per-run mode gating side-effecting tools, toggled by the host without +/// restarting the run. +/// +/// `Build` (the default) leaves tool exposure and execution unrestricted. +/// `Plan` hides every side-effecting tool from the model +/// ([`PlanModeMiddleware`]'s `before_model`, the same auditable +/// `AgentEvent::ToolsFiltered` mechanism [`ContextualToolSelectionMiddleware`] +/// uses) and denies it at execution (`before_tool`, the same +/// [`ToolPolicy`]/[`ToolSideEffects`] classification +/// [`ToolPolicyMiddleware::deny_side_effects`] enforces) — except for tools on +/// the middleware's allowlist, which stay available in either mode so a host +/// can keep read-only tools and a handful of plan-mode-specific tools (e.g. +/// `plan_exit`, `request_plan_review`, `todo`) reachable while planning. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum RunMode { + /// Every registered tool may be exposed and executed. The default. + #[default] + Build, + /// Only allowlisted and side-effect-free tools may be exposed or executed. + Plan, +} + +/// A shared, host-settable handle to the current [`RunMode`] for a run. +/// +/// Cloning shares the same underlying state: a host can hold one clone to +/// flip modes (e.g. from a UI toggle or a `plan_exit` tool call) while +/// [`PlanModeMiddleware`] holds another to read it. [`Self::set`] takes +/// effect on the very next tool exposure or execution check — no run restart +/// is required, so a host can switch modes mid-run, between turns. +#[derive(Debug, Clone)] +pub struct RunModeHandle(pub(crate) Arc); + +impl RunModeHandle { + /// Creates a handle starting in `mode`. + pub fn new(mode: RunMode) -> Self { + Self(Arc::new(std::sync::atomic::AtomicU8::new(Self::encode( + mode, + )))) + } + + /// Returns the current mode. + pub fn get(&self) -> RunMode { + Self::decode(self.0.load(std::sync::atomic::Ordering::SeqCst)) + } + + /// Sets the current mode. + pub fn set(&self, mode: RunMode) { + self.0 + .store(Self::encode(mode), std::sync::atomic::Ordering::SeqCst); + } + + fn encode(mode: RunMode) -> u8 { + match mode { + RunMode::Build => 0, + RunMode::Plan => 1, + } + } + + fn decode(value: u8) -> RunMode { + match value { + 1 => RunMode::Plan, + _ => RunMode::Build, + } + } +} + +impl Default for RunModeHandle { + fn default() -> Self { + Self::new(RunMode::default()) + } +} + +/// Lifecycle middleware that enforces [`RunMode::Plan`] by hiding and denying +/// side-effecting tools, driven by a live [`RunModeHandle`]. +/// +/// Build with [`plan_mode_middleware`] or [`PlanModeMiddleware::new`], then +/// widen the plan-mode allowlist with [`PlanModeMiddleware::allow`]. A tool is +/// treated as side-effecting when its [`ToolPolicy::side_effects`] declares +/// any of `writes_files`, `network`, `installs_dependencies`, `destructive`, +/// `external_service`, or `payment` — or when `policies` has no entry for it +/// at all (fail-closed: an unclassified tool is assumed capable of side +/// effects until proven otherwise). +pub struct PlanModeMiddleware { + pub(crate) label: &'static str, + pub(crate) mode: RunModeHandle, + pub(crate) policies: std::collections::HashMap, + pub(crate) allow: HashSet, +} + +/// Creates a [`PlanModeMiddleware`] driven by `mode`, classifying tools from +/// `policies` (typically [`ToolRegistry::policies`][crate::tool::ToolRegistry::policies]). +/// +/// Equivalent to [`PlanModeMiddleware::new`]; provided as a free function so a +/// host can wire plan mode into a [`MiddlewareStack`][crate::middleware::MiddlewareStack] +/// with one call. +pub fn plan_mode_middleware( + mode: RunModeHandle, + policies: std::collections::HashMap, +) -> PlanModeMiddleware { + PlanModeMiddleware::new(mode, policies) +} diff --git a/crates/tinyagents-harness/src/observability/mod.rs b/crates/tinyagents-harness/src/observability/mod.rs index adc30b385..caa3add8b 100644 --- a/crates/tinyagents-harness/src/observability/mod.rs +++ b/crates/tinyagents-harness/src/observability/mod.rs @@ -98,7 +98,9 @@ impl AgentLatencyMetrics { }); } } - AgentEvent::ToolStarted { call_id, tool_name } => { + AgentEvent::ToolStarted { + call_id, tool_name, .. + } => { tool_starts.insert(call_id.clone(), (tool_name.clone(), obs.ts_ms)); } AgentEvent::ToolCompleted { call_id, .. } => { diff --git a/crates/tinyagents-harness/src/observability/test.rs b/crates/tinyagents-harness/src/observability/test.rs index 96dff4b44..ed59bfa80 100644 --- a/crates/tinyagents-harness/src/observability/test.rs +++ b/crates/tinyagents-harness/src/observability/test.rs @@ -132,6 +132,7 @@ fn agent_latency_metrics_include_model_tool_and_run_elapsed() { AgentEvent::ToolStarted { call_id: tool_id.clone(), tool_name: "lookup".to_string(), + input: None, }, ), obs( diff --git a/crates/tinyagents-harness/src/run_queue/mod.rs b/crates/tinyagents-harness/src/run_queue/mod.rs index 1c1e300d7..57ab7d98c 100644 --- a/crates/tinyagents-harness/src/run_queue/mod.rs +++ b/crates/tinyagents-harness/src/run_queue/mod.rs @@ -134,6 +134,60 @@ impl RunQueue { inner.collects.clear(); total } + + /// Returns a snapshot of every queued item, tagged with its lane, in + /// lane/delivery order: [`QueueLane::Steer`], then + /// [`QueueLane::Followup`], then [`QueueLane::Collect`], each lane + /// preserving its own FIFO order. Non-destructive — the queue is + /// unchanged, so a host can inspect pending work (e.g. to render it, or + /// to decide whether [`Self::remove_where`] applies) without racing the + /// agent loop's own drain. + pub async fn snapshot(&self) -> Vec<(QueueLane, T)> + where + T: Clone, + { + let inner = self.inner.lock().await; + let mut items = + Vec::with_capacity(inner.steers.len() + inner.followups.len() + inner.collects.len()); + items.extend( + inner + .steers + .iter() + .cloned() + .map(|item| (QueueLane::Steer, item)), + ); + items.extend( + inner + .followups + .iter() + .cloned() + .map(|item| (QueueLane::Followup, item)), + ); + items.extend( + inner + .collects + .iter() + .cloned() + .map(|item| (QueueLane::Collect, item)), + ); + items + } + + /// Removes every queued item across all lanes for which `pred` returns + /// `true`, preserving the relative order of the items that remain in + /// each lane. Returns the number of items removed. + /// + /// Use to retract a specific queued item (e.g. one the host decided not + /// to apply after all) without clearing the rest of the queue. + pub async fn remove_where(&self, pred: impl Fn(&T) -> bool) -> usize { + let mut inner = self.inner.lock().await; + let before = inner.steers.len() + inner.followups.len() + inner.collects.len(); + inner.steers.retain(|item| !pred(item)); + inner.followups.retain(|item| !pred(item)); + inner.collects.retain(|item| !pred(item)); + let after = inner.steers.len() + inner.followups.len() + inner.collects.len(); + before - after + } } impl Default for RunQueue { diff --git a/crates/tinyagents-harness/src/run_queue/test.rs b/crates/tinyagents-harness/src/run_queue/test.rs index 80b2d4c75..d27735230 100644 --- a/crates/tinyagents-harness/src/run_queue/test.rs +++ b/crates/tinyagents-harness/src/run_queue/test.rs @@ -93,3 +93,74 @@ async fn take_all_drains_the_whole_lane() { ); assert_eq!(queue.status().await.followups, 0); } + +#[tokio::test] +async fn snapshot_returns_every_item_in_lane_and_fifo_order_without_draining() { + let queue = RunQueue::new(); + queue.push(QueueLane::Steer, "steer-1").await; + queue.push(QueueLane::Followup, "followup-1").await; + queue.push(QueueLane::Steer, "steer-2").await; + queue.push(QueueLane::Collect, "collect-1").await; + queue.push(QueueLane::Followup, "followup-2").await; + + assert_eq!( + queue.snapshot().await, + vec![ + (QueueLane::Steer, "steer-1"), + (QueueLane::Steer, "steer-2"), + (QueueLane::Followup, "followup-1"), + (QueueLane::Followup, "followup-2"), + (QueueLane::Collect, "collect-1"), + ] + ); + // Non-destructive: the queue is unchanged. + assert_eq!(queue.status().await.total, 5); +} + +#[tokio::test] +async fn snapshot_of_an_empty_queue_is_empty() { + let queue = RunQueue::::new(); + assert_eq!(queue.snapshot().await, Vec::new()); +} + +#[tokio::test] +async fn remove_where_removes_matching_items_across_every_lane() { + let queue = RunQueue::new(); + queue.push(QueueLane::Steer, "keep").await; + queue.push(QueueLane::Steer, "drop").await; + queue.push(QueueLane::Followup, "drop").await; + queue.push(QueueLane::Collect, "keep").await; + + let removed = queue.remove_where(|item| *item == "drop").await; + + assert_eq!(removed, 2); + assert_eq!( + queue.snapshot().await, + vec![(QueueLane::Steer, "keep"), (QueueLane::Collect, "keep"),] + ); +} + +#[tokio::test] +async fn remove_where_preserves_relative_order_of_survivors() { + let queue = RunQueue::new(); + queue.push(QueueLane::Steer, 1).await; + queue.push(QueueLane::Steer, 2).await; + queue.push(QueueLane::Steer, 3).await; + + let removed = queue.remove_where(|item| *item == 2).await; + + assert_eq!(removed, 1); + assert_eq!(queue.drain(QueueLane::Steer).await, vec![1, 3]); +} + +#[tokio::test] +async fn remove_where_with_no_match_removes_nothing() { + let queue = RunQueue::new(); + queue.push(QueueLane::Steer, "a").await; + queue.push(QueueLane::Followup, "b").await; + + let removed = queue.remove_where(|item| *item == "nonexistent").await; + + assert_eq!(removed, 0); + assert_eq!(queue.status().await.total, 2); +} diff --git a/crates/tinyagents-harness/src/stream/test.rs b/crates/tinyagents-harness/src/stream/test.rs index ffb37ad0a..1111b9f53 100644 --- a/crates/tinyagents-harness/src/stream/test.rs +++ b/crates/tinyagents-harness/src/stream/test.rs @@ -236,6 +236,7 @@ mod project { let event = AgentEvent::ToolStarted { call_id: CallId::new("c9"), tool_name: "search".into(), + input: None, }; let chunk = project_event(&event).expect("must project"); assert_eq!(chunk.mode(), StreamMode::Debug); diff --git a/crates/tinyagents-harness/src/testkit/mod.rs b/crates/tinyagents-harness/src/testkit/mod.rs index eecbbdee9..5f6f04ff4 100644 --- a/crates/tinyagents-harness/src/testkit/mod.rs +++ b/crates/tinyagents-harness/src/testkit/mod.rs @@ -642,6 +642,7 @@ impl Trajectory { /// AgentEvent::ToolStarted { /// call_id: CallId::new("c1"), /// tool_name: "search".into(), + /// input: None, /// }, /// ]; /// Trajectory::from_events(events).assert_tool_called("search"); diff --git a/crates/tinyagents-harness/src/testkit/test.rs b/crates/tinyagents-harness/src/testkit/test.rs index 3e6fad27b..b4f40ab01 100644 --- a/crates/tinyagents-harness/src/testkit/test.rs +++ b/crates/tinyagents-harness/src/testkit/test.rs @@ -324,6 +324,7 @@ fn make_trajectory() -> Vec { AgentEvent::ToolStarted { call_id: CallId::new("t1"), tool_name: "search".into(), + input: None, }, AgentEvent::ToolCompleted { call_id: CallId::new("t1"), @@ -380,6 +381,7 @@ fn trajectory_tool_call_count() { events.push(AgentEvent::ToolStarted { call_id: CallId::new("t2"), tool_name: "search".into(), + input: None, }); let traj = Trajectory::from_events(events); assert_eq!(traj.tool_call_count("search"), 2); diff --git a/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs b/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs index b04c03f02..4556c764e 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_registry_observability_contracts.rs @@ -157,6 +157,7 @@ fn component_metadata_and_event_kinds_are_stable_serializable_contracts() { AgentEvent::ToolStarted { call_id: CallId::new("tool-1"), tool_name: "lookup".into(), + input: None, }, AgentEvent::ToolCompleted { call_id: CallId::new("tool-1"), @@ -285,6 +286,7 @@ async fn event_sinks_journals_and_status_stores_preserve_run_lineage() { journal.append(AgentEvent::ToolStarted { call_id: CallId::new("tool-1"), tool_name: "lookup".into(), + input: None, }); journal.append(AgentEvent::ToolCompleted { call_id: CallId::new("tool-1"), diff --git a/crates/tinyagents-integration-tests/tests/feature_infra_observability.rs b/crates/tinyagents-integration-tests/tests/feature_infra_observability.rs index 677b344c5..7180bfb2a 100644 --- a/crates/tinyagents-integration-tests/tests/feature_infra_observability.rs +++ b/crates/tinyagents-integration-tests/tests/feature_infra_observability.rs @@ -77,6 +77,7 @@ fn latency_metrics_correlate_started_and_completed_by_call_id() { AgentEvent::ToolStarted { call_id: CallId::new("t1"), tool_name: "search".into(), + input: None, }, ), obs( @@ -186,6 +187,7 @@ async fn journal_read_filtered_selects_by_event_kind() { AgentEvent::ToolStarted { call_id: CallId::new("t1"), tool_name: "x".into(), + input: None, }, )) .await diff --git a/crates/tinyagents-orchestration/tests/e2e_orchestrator_subagents.rs b/crates/tinyagents-orchestration/tests/e2e_orchestrator_subagents.rs index 287fd3c95..38ce5ce36 100644 --- a/crates/tinyagents-orchestration/tests/e2e_orchestrator_subagents.rs +++ b/crates/tinyagents-orchestration/tests/e2e_orchestrator_subagents.rs @@ -182,6 +182,7 @@ async fn orchestrator_resolves_and_runs_only_the_chosen_subagents() -> Result<() sink.emit(AgentEvent::ToolStarted { call_id: CallId::new(call_id.clone()), tool_name: name.clone(), + input: None, }); let parent = RunContext::new(RunConfig::new(format!("dispatch-{i}")), ()); let result = dispatch diff --git a/crates/tinyagents-session/src/transcript.rs b/crates/tinyagents-session/src/transcript.rs index 72a5ab4ce..20fbf5ca4 100644 --- a/crates/tinyagents-session/src/transcript.rs +++ b/crates/tinyagents-session/src/transcript.rs @@ -124,7 +124,7 @@ mod writer; pub use adoption::{SessionAdoption, adopt_legacy_session_transcripts}; pub use history::{ FileTranscriptHistory, FileTranscriptLocator, TranscriptHistory, TranscriptLocator, - TranscriptPartial, TranscriptRead, TranscriptTurn, + TranscriptPartial, TranscriptRead, TranscriptTurn, TruncateCut, }; pub use legacy_md::read_transcript_legacy_md; pub use migration::{TranscriptLayoutMigration, migrate_layout_if_needed}; diff --git a/crates/tinyagents-session/src/transcript/history.rs b/crates/tinyagents-session/src/transcript/history.rs index 8d3379057..d888a7ad4 100644 --- a/crates/tinyagents-session/src/transcript/history.rs +++ b/crates/tinyagents-session/src/transcript/history.rs @@ -362,6 +362,113 @@ pub trait TranscriptLocator: Send + Sync { let handle = self.open_session(&successor, meta)?; Ok((successor, handle)) } + + /// Forks `session`'s head generation for edit or regenerate, without + /// erasing history. + /// + /// This is [`Self::begin_generation`]'s compaction move, aimed at a + /// different caller: instead of a summarizer replacing old turns with a + /// digest, a host wants to edit a past message or regenerate the last + /// answer. Both need the exact same guarantee compaction already + /// provides — the current generation is sealed **untouched** on disk + /// (nothing is written to it; that is the whole point of never rewriting + /// a sealed file) and the next generation records it as parent — so this + /// is built on the same primitive rather than a second, parallel one. + /// + /// Reads the current [`Self::head_generation`]'s messages, resolves + /// `cut` against them, seals that head and opens its successor via + /// [`Self::begin_generation`], and writes the retained prefix into the + /// successor with [`TranscriptHistory::replace`] — the identical call a + /// compaction makes to persist its own replacement set. Because the + /// successor's parent is the sealed head exactly as `begin_generation` + /// records it, [`Self::session_chain`] walks both generations, so the + /// full pre-truncation history stays recoverable even though the model + /// now reads only the truncated head. + /// + /// Returns the new generation's [`SessionRef`], its bound handle (already + /// carrying the truncated messages), and the truncated messages + /// themselves for the caller's own use (e.g. re-driving the model on the + /// retained context). + /// + /// Fails if `session` has no transcript yet, or if `cut` is + /// [`TruncateCut::BeforeMessageId`] naming an id absent from the head + /// generation — silently falling back to some other cut point would risk + /// truncating the wrong turn. + fn truncate_into_next_generation( + &self, + session: &SessionRef, + cut: TruncateCut, + seed: TranscriptMeta, + ) -> anyhow::Result<( + SessionRef, + Arc, + Vec, + )> { + let head = self.head_generation(session); + let head_read = self.read_session_transcript(&head).ok_or_else(|| { + anyhow::anyhow!( + "session {} has no transcript to truncate", + head.session_id() + ) + })?; + let transcript = head_read.read_session()?.ok_or_else(|| { + anyhow::anyhow!( + "session {} has no transcript to truncate", + head.session_id() + ) + })?; + let keep = cut.resolve(&transcript.messages)?; + let truncated = transcript.messages[..keep].to_vec(); + + let (successor, handle) = self.begin_generation(&head, seed)?; + // Same call a compaction makes to persist its own replacement set — + // see `a_compaction_seals_a_generation_and_leaves_it_untouched`. + handle.replace(&truncated)?; + Ok((successor, handle, truncated)) + } +} + +/// Where to cut a session's head-generation messages when forking it with +/// [`TranscriptLocator::truncate_into_next_generation`]. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TruncateCut { + /// Keep messages `[0, index)`; drop the message at `index` and everything + /// after it. Clamped to the message count, so an out-of-range index keeps + /// every message. + BeforeIndex(usize), + /// Keep everything before the message carrying this id. The id must name + /// a message in the head generation — [`TranscriptMessage::id`] is only + /// ever set by a host that assigns stable ids, so this is the most + /// robust key to cut on when the caller has one; unlike an index, it + /// cannot point at the wrong turn after an earlier truncation shifted + /// everything else. + BeforeMessageId(String), + /// Drop the trailing assistant turn: everything strictly after the last + /// `role == "user"` message, matching the `role == "assistant"` cutpoint + /// convention already used across this crate's writer (e.g. + /// `writer::append_transcript_turn`'s `last_assistant_idx`). Used for + /// "regenerate the last answer." When there is no user message at all, + /// every message is dropped. + LastAssistantTurn, +} + +impl TruncateCut { + /// Resolves this cut to a keep-count (`messages[..keep]` survives) + /// against the head generation's `messages`. + fn resolve(&self, messages: &[TranscriptMessage]) -> anyhow::Result { + match self { + TruncateCut::BeforeIndex(index) => Ok((*index).min(messages.len())), + TruncateCut::BeforeMessageId(id) => messages + .iter() + .position(|message| message.id.as_deref() == Some(id.as_str())) + .ok_or_else(|| anyhow::anyhow!("no message with id `{id}` in the head generation")), + TruncateCut::LastAssistantTurn => Ok(messages + .iter() + .rposition(|message| message.role == "user") + .map(|index| index + 1) + .unwrap_or(0)), + } + } } /// The default [`TranscriptLocator`]: real files under diff --git a/crates/tinyagents-session/src/transcript/test.rs b/crates/tinyagents-session/src/transcript/test.rs index 06e331ddd..4797db5dd 100644 --- a/crates/tinyagents-session/src/transcript/test.rs +++ b/crates/tinyagents-session/src/transcript/test.rs @@ -389,6 +389,160 @@ fn a_compaction_seals_a_generation_and_leaves_it_untouched() { ); } +/// A fork for edit/regenerate must give the same untouched-sealed-file +/// guarantee a compaction gives, and the chain must walk both generations. +#[test] +fn truncate_into_next_generation_seals_the_head_byte_identical_and_chain_walks_both() { + let dir = tempdir().unwrap(); + let locator = FileTranscriptLocator::new(dir.path()); + let session = SessionRef::scoped("thread-1", "orchestrator"); + + let first = locator.open_session(&session, meta()).unwrap(); + for turn in ["one", "two", "three", "four"] { + first.append(TranscriptMessage::new("user", turn)).unwrap(); + } + let sealed_path = first.path().to_path_buf(); + let sealed_bytes = std::fs::read(&sealed_path).unwrap(); + + let (successor, handle, truncated) = locator + .truncate_into_next_generation(&session, TruncateCut::BeforeIndex(2), meta()) + .unwrap(); + + assert_eq!(successor.generation, 1); + assert_eq!( + std::fs::read(&sealed_path).unwrap(), + sealed_bytes, + "the sealed head must be byte-identical afterwards" + ); + assert_ne!(handle.path(), sealed_path); + + // The new head carries only the retained prefix. + let contents: Vec<_> = truncated.iter().map(|m| m.content.as_str()).collect(); + assert_eq!(contents, vec!["one", "two"]); + assert_eq!(handle.messages().unwrap().len(), 2); + + // Parent recorded exactly as `begin_generation` records it. + let successor_meta = handle.read_session().unwrap().unwrap().meta; + assert_eq!( + successor_meta.parent_session_id.as_deref(), + Some(session_stem(&session).as_str()) + ); + + // Both generations are reachable by walking the chain — nothing is lost. + let chain = locator.session_chain(&session); + assert_eq!(chain.len(), 2); + assert_eq!(chain[0], session); + assert_eq!(chain[1], successor); + assert_eq!(locator.head_generation(&session), successor); +} + +#[test] +fn truncate_into_next_generation_by_message_id_cuts_at_that_message() { + let dir = tempdir().unwrap(); + let locator = FileTranscriptLocator::new(dir.path()); + let session = SessionRef::scoped("thread-1", "orchestrator"); + + let first = locator.open_session(&session, meta()).unwrap(); + let mut keep_a = TranscriptMessage::new("user", "keep-a"); + keep_a.id = Some("m1".into()); + let mut cut_here = TranscriptMessage::new("assistant", "cut-here"); + cut_here.id = Some("m2".into()); + let mut dropped = TranscriptMessage::new("user", "dropped"); + dropped.id = Some("m3".into()); + for message in [keep_a, cut_here, dropped] { + first.append(message).unwrap(); + } + + let (_, handle, truncated) = locator + .truncate_into_next_generation(&session, TruncateCut::BeforeMessageId("m2".into()), meta()) + .unwrap(); + + assert_eq!(truncated.len(), 1); + assert_eq!(truncated[0].id.as_deref(), Some("m1")); + assert_eq!(handle.messages().unwrap().len(), 1); +} + +#[test] +fn truncate_into_next_generation_by_unknown_message_id_fails() { + let dir = tempdir().unwrap(); + let locator = FileTranscriptLocator::new(dir.path()); + let session = SessionRef::scoped("thread-1", "orchestrator"); + locator + .open_session(&session, meta()) + .unwrap() + .append(TranscriptMessage::new("user", "one")) + .unwrap(); + + let err = locator.truncate_into_next_generation( + &session, + TruncateCut::BeforeMessageId("nonexistent".into()), + meta(), + ); + assert!( + err.is_err(), + "an unresolvable id must fail rather than silently keep everything" + ); +} + +#[test] +fn truncate_into_next_generation_last_assistant_turn_drops_only_the_trailing_answer() { + let dir = tempdir().unwrap(); + let locator = FileTranscriptLocator::new(dir.path()); + let session = SessionRef::scoped("thread-1", "orchestrator"); + + let first = locator.open_session(&session, meta()).unwrap(); + for (role, content) in [ + ("user", "first question"), + ("assistant", "first answer"), + ("user", "second question"), + ("assistant", "second answer, to regenerate"), + ] { + first.append(TranscriptMessage::new(role, content)).unwrap(); + } + + let (_, handle, truncated) = locator + .truncate_into_next_generation(&session, TruncateCut::LastAssistantTurn, meta()) + .unwrap(); + + let contents: Vec<_> = truncated.iter().map(|m| m.content.as_str()).collect(); + assert_eq!( + contents, + vec!["first question", "first answer", "second question"] + ); + assert_eq!(handle.messages().unwrap().len(), 3); +} + +#[test] +fn truncate_into_next_generation_operates_on_the_current_head_not_the_root() { + // A fork after an earlier compaction must truncate the head generation's + // messages, not the sealed root's. + let dir = tempdir().unwrap(); + let locator = FileTranscriptLocator::new(dir.path()); + let session = SessionRef::scoped("thread-1", "orchestrator"); + + locator + .open_session(&session, meta()) + .unwrap() + .append(TranscriptMessage::new("user", "root-only")) + .unwrap(); + let (compacted, compacted_handle) = locator.begin_generation(&session, meta()).unwrap(); + for turn in ["alpha", "beta", "gamma"] { + compacted_handle + .append(TranscriptMessage::new("user", turn)) + .unwrap(); + } + + let (successor, handle, truncated) = locator + .truncate_into_next_generation(&session, TruncateCut::BeforeIndex(1), meta()) + .unwrap(); + + assert_eq!(successor.generation, 2); + assert_eq!(successor.parent_session_id(), Some(compacted.session_id())); + let contents: Vec<_> = truncated.iter().map(|m| m.content.as_str()).collect(); + assert_eq!(contents, vec!["alpha"]); + assert_eq!(handle.messages().unwrap().len(), 1); +} + /// After a compaction, a resume must land on the newest generation — the one /// the model is actually continuing — not on the sealed original. #[test] diff --git a/docs/modules/harness/middleware.md b/docs/modules/harness/middleware.md index 0b2028eca..e09ebb308 100644 --- a/docs/modules/harness/middleware.md +++ b/docs/modules/harness/middleware.md @@ -294,6 +294,57 @@ Exposure only changes what the model *sees*; pair it with [tool policy enforcement](#tool-policy-enforcement) or `ToolAllowlistMiddleware` so a model that calls a hidden tool is still stopped at execution. +## Plan mode + +`PlanModeMiddleware` (built with `plan_mode_middleware`) is a ready-made +combination of the two mechanisms above, gated by a live, host-settable +`RunMode`: `Build` (the default) leaves every tool exposed and executable; +`Plan` hides every side-effecting tool from the model at `before_model` (the +same `AgentEvent::ToolsFiltered` auditing `ContextualToolSelectionMiddleware` +emits) and denies it at `before_tool` (the same `ToolPolicy`/`ToolSideEffects` +classification `ToolPolicyMiddleware::deny_side_effects` enforces) — except for +tools on the middleware's own allowlist, which a host uses to keep read-only +tools and plan-mode-specific tools (e.g. `plan_exit`, `request_plan_review`, +`todo`) available while planning. + +A tool counts as side-effecting when its policy declares any of +`writes_files`, `network`, `installs_dependencies`, `destructive`, +`external_service`, or `payment` — or when `policies` has no entry for it at +all: an unclassified tool is assumed capable of side effects until it is +either classified read-only or added to the allowlist. + +`RunModeHandle` is `Clone` and cheap to share: a host keeps one clone to flip +modes (a UI toggle, a `plan_exit` tool call, a slash command) while the +middleware holds another to read it. `set` takes effect on the very next tool +exposure or execution check, so a host can switch modes mid-run, between +turns, without restarting it. + +```rust +use std::collections::HashMap; +use std::sync::Arc; +use tinyagents_harness::middleware::{plan_mode_middleware, MiddlewareStack, RunMode, RunModeHandle}; +use tinytools::{ToolPolicy, ToolSideEffects}; + +let mut policies = HashMap::new(); +policies.insert("read_file".to_string(), ToolPolicy::read_only()); +policies.insert( + "write_file".to_string(), + ToolPolicy::classified().with_side_effects(ToolSideEffects { + writes_files: true, + ..ToolSideEffects::default() + }), +); + +let mode = RunModeHandle::new(RunMode::Plan); +let mw = plan_mode_middleware(mode.clone(), policies).allow(["plan_exit"]); + +let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); +stack.push(Arc::new(mw)); +// While `mode.get() == RunMode::Plan`: `write_file` is hidden and denied, +// `read_file` and `plan_exit` stay available. `mode.set(RunMode::Build)` +// lifts the restriction on the very next check. +``` + ## Middleware control (A1) Any middleware (or step) can steer the loop out-of-band via