From 89490d848dba435ceb59921756aabde2e804a7e9 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 18:01:41 +0300 Subject: [PATCH 01/18] test(graph): cover bound continuation across resume and retry Co-authored-by: Medulla --- .../tinyagents-graph/src/compiled/README.md | 6 + crates/tinyagents-graph/src/compiled/test.rs | 238 +++++++++++- .../tinyagents-graph/src/subgraph/README.md | 6 + crates/tinyagents-graph/src/subgraph/test.rs | 348 ++++++++++++++++++ 4 files changed, 597 insertions(+), 1 deletion(-) diff --git a/crates/tinyagents-graph/src/compiled/README.md b/crates/tinyagents-graph/src/compiled/README.md index 5158ec7f..f7cf7653 100644 --- a/crates/tinyagents-graph/src/compiled/README.md +++ b/crates/tinyagents-graph/src/compiled/README.md @@ -99,6 +99,12 @@ Accessors: `graph_id()`, `name()`, `namespace()`. failure boundary. - `retry(thread_id)` — re-run the failed node(s) recorded in the last failure-boundary checkpoint. +- `run_with_agent_binding(..)`, `run_with_thread_agent_binding(..)`, + `resume_with_agent_binding(..)`, `resume_from_with_agent_binding(..)`, and + `retry_with_agent_binding(..)` — execution-scoped variants for graphs that + reach a `SubAgentNode`. The supplied `AgentInvocationBinding` is forwarded + to resumed nodes and nested subgraphs, but is never stored on the reusable + graph or in a checkpoint; an unbound sub-agent continuation fails closed. ### State inspection / time travel diff --git a/crates/tinyagents-graph/src/compiled/test.rs b/crates/tinyagents-graph/src/compiled/test.rs index 87a4e69e..1b5e52a7 100644 --- a/crates/tinyagents-graph/src/compiled/test.rs +++ b/crates/tinyagents-graph/src/compiled/test.rs @@ -8,10 +8,13 @@ use crate::checkpoint::{Checkpointer, InMemoryCheckpointer}; use crate::command::{Command, Interrupt, NodeResult, Send}; use crate::reducer::ClosureStateReducer; use crate::stream::{CollectingSink, GraphEvent}; +use async_trait::async_trait; use serde_json::json; -use std::sync::Arc; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering as AtomicOrdering}; +use std::sync::{Arc, Mutex}; use std::time::Duration; +use tinyagents_harness::cancel::CancellationToken; +use tinyagents_harness::events::{AgentEvent, EventSink, RecordingListener}; use tinyagents_harness::ids::ExecutionStatus; use tinyagents_harness::retry::RetryPolicy; @@ -21,6 +24,42 @@ struct Counter { log: Vec, } +/// A host invoker whose recorded request makes a continuation's live binding +/// observable without putting host capabilities in durable graph state. +#[derive(Clone, Default)] +struct BindingRecordingInvoker(Arc>>); + +#[async_trait] +impl crate::AgentInvoker for BindingRecordingInvoker { + async fn invoke( + &self, + request: crate::AgentInvocation, + ) -> crate::Result { + self.0.lock().unwrap().push(request.clone()); + request.events.emit(AgentEvent::StateUpdate); + Ok(crate::SubAgentOutput { + text: request.input.prompt, + ..Default::default() + }) + } +} + +struct BindingFailingInvoker; + +#[async_trait] +impl crate::AgentInvoker for BindingFailingInvoker { + async fn invoke( + &self, + _request: crate::AgentInvocation, + ) -> crate::Result { + Err(TinyAgentsError::Model("continuation failure".to_string())) + } +} + +fn agent_binding(invoker: Arc) -> crate::AgentInvocationBinding { + crate::AgentInvocationBinding::new(invoker, EventSink::new(), CancellationToken::new()) +} + /// Builds a graph whose nodes return partial `i32` updates merged by a custom /// reducer that adds to `value` and records a log entry. fn adding_graph() -> CompiledGraph { @@ -796,6 +835,203 @@ async fn resume_from_older_checkpoint_replays_forward() { assert!(matches!(err, TinyAgentsError::Resume(_))); } +#[tokio::test] +async fn resume_from_with_agent_binding_keeps_host_capabilities_live_only() { + // A bound run pauses before delegation. Its checkpoint must be enough to + // resume graph state, but must never retain the invoker, event sink, or + // cancellation handle that happened to start the run. + let checkpointer = Arc::new(InMemoryCheckpointer::::new()); + let graph = GraphBuilder::::overwrite() + .add_node("gate", |state: String, ctx: NodeContext| async move { + if state == "unbound" || ctx.resume.is_some() { + Ok(NodeResult::Update(state)) + } else { + Ok(NodeResult::Interrupt(Interrupt::new( + "gate", + json!({ "ask": "continue?" }), + ))) + } + }) + .add_node( + "delegate", + crate::subagent_node(crate::SubAgentNode::from_fns( + "researcher", + |state: &String| crate::SubAgentInput::prompt(state.clone()), + |output: crate::SubAgentOutput| output.text, + )), + ) + .set_entry("gate") + .add_edge("gate", "delegate") + .set_finish("delegate") + .compile() + .unwrap() + .with_checkpointer(checkpointer.clone()); + + let initial_invoker = Arc::new(BindingRecordingInvoker::default()); + let paused = graph + .run_with_thread_agent_binding( + "resume-binding", + "question".to_string(), + agent_binding(initial_invoker), + ) + .await + .unwrap(); + assert!(paused.is_interrupted()); + + let checkpoint = checkpointer + .get("resume-binding", None) + .await + .unwrap() + .expect("interrupt is checkpointed"); + let interrupted_checkpoint_id = checkpoint.checkpoint_id.clone(); + let checkpoint_json = serde_json::to_value(checkpoint).unwrap(); + assert!(checkpoint_json.get("agent_binding").is_none()); + assert!( + !checkpoint_json + .to_string() + .contains("AgentInvocationBinding"), + "a durable checkpoint must contain no live invocation capability" + ); + + // Reloading this bound checkpoint without a fresh binding reaches the + // SubAgentNode but must fail closed. Retrying the resulting failure + // boundary without a binding must do the same; neither continuation may + // recover capabilities from the original bound execution. + let unbound_resume = graph + .resume("resume-binding", Command::resume(json!("approved"))) + .await + .unwrap_err(); + assert!(matches!(unbound_resume, TinyAgentsError::Capability(_))); + assert!( + unbound_resume + .to_string() + .contains("sub-agent `researcher`"), + "unbound resume must fail at SubAgentNode: {unbound_resume}" + ); + + let unbound_retry = graph.retry("resume-binding").await.unwrap_err(); + assert!(matches!(unbound_retry, TinyAgentsError::Capability(_))); + assert!( + unbound_retry.to_string().contains("sub-agent `researcher`"), + "unbound retry must fail at SubAgentNode: {unbound_retry}" + ); + + let latest_after_unbound_failures = checkpointer + .get("resume-binding", None) + .await + .unwrap() + .expect("unbound failures are checkpointed") + .checkpoint_id; + assert_ne!( + latest_after_unbound_failures, interrupted_checkpoint_id, + "the selected interrupt checkpoint must no longer be latest" + ); + + let invoker = Arc::new(BindingRecordingInvoker::default()); + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone()); + let cancellation = CancellationToken::new(); + let resumed = graph + .resume_from_with_agent_binding( + "resume-binding", + ResumeTarget::Checkpoint(interrupted_checkpoint_id), + Command::resume(json!("approved")), + crate::AgentInvocationBinding::new(invoker.clone(), events, cancellation.clone()), + ) + .await + .unwrap(); + + assert_eq!(resumed.state, "question"); + let request = invoker.0.lock().unwrap().pop().expect("delegate ran"); + assert_eq!(request.parent_run_id, resumed.run_id); + assert_eq!(request.root_run_id, resumed.root_run_id); + let request_cancellation = request.cancellation.expect("binding supplies cancellation"); + assert!( + !request_cancellation.is_cancelled(), + "the captured request must initially observe the supplied live token" + ); + cancellation.cancel(); + assert!( + request_cancellation.is_cancelled(), + "cancelling the supplied token after invocation must reach the captured request" + ); + assert_eq!(listener.len(), 1); +} + +#[tokio::test] +async fn retry_with_agent_binding_replaces_failed_run_capabilities() { + let checkpointer = Arc::new(InMemoryCheckpointer::::new()); + let graph = GraphBuilder::::overwrite() + .add_node( + "delegate", + crate::subagent_node(crate::SubAgentNode::from_fns( + "researcher", + |state: &String| crate::SubAgentInput::prompt(state.clone()), + |output: crate::SubAgentOutput| output.text, + )), + ) + .set_entry("delegate") + .set_finish("delegate") + .compile() + .unwrap() + .with_checkpointer(checkpointer); + + let failed = graph + .run_with_thread_agent_binding( + "retry-binding", + "question".to_string(), + agent_binding(Arc::new(BindingFailingInvoker)), + ) + .await + .unwrap_err(); + assert!(matches!(failed, TinyAgentsError::Model(_))); + + // Give every capability a distinct observable behavior: the replacement + // invoker records the retry and its event sink has this listener only. The + // live cancellation handle is checked after the invocation below, proving + // the request received this exact replacement token. + let replacement = Arc::new(BindingRecordingInvoker::default()); + let events = EventSink::new(); + let listener = Arc::new(RecordingListener::new()); + events.subscribe(listener.clone()); + let cancellation = CancellationToken::new(); + let retried = graph + .retry_with_agent_binding( + "retry-binding", + crate::AgentInvocationBinding::new(replacement.clone(), events, cancellation.clone()), + ) + .await + .unwrap(); + + assert_eq!(retried.state, "question"); + let request = replacement + .0 + .lock() + .unwrap() + .pop() + .expect("retry delegated"); + assert_eq!(request.parent_run_id, retried.run_id); + assert_eq!(request.root_run_id, retried.root_run_id); + let request_cancellation = request + .cancellation + .expect("retry binding supplies a cancellation token"); + assert!( + !request_cancellation.is_cancelled(), + "the captured retry request must initially observe a live token" + ); + cancellation.cancel(); + assert!( + request_cancellation.is_cancelled(), + "cancelling the supplied token after invocation must reach the captured retry request" + ); + assert_eq!( + listener.len(), + 1, + "retry must forward the supplied event sink to the replacement invoker" + ); +} + // --- Parallel (fan-out / fan-in) execution --------------------------------- #[derive(Clone, Debug, Default, PartialEq)] diff --git a/crates/tinyagents-graph/src/subgraph/README.md b/crates/tinyagents-graph/src/subgraph/README.md index 8e9aa797..76c5ba71 100644 --- a/crates/tinyagents-graph/src/subgraph/README.md +++ b/crates/tinyagents-graph/src/subgraph/README.md @@ -33,6 +33,12 @@ both adapters check `execution.is_interrupted()` and return `NodeResult::Interrupt(..)` instead of folding the paused child's state through `from_child` (or returning it directly, for the shared-state case). +When a parent is resumed through one of `CompiledGraph`'s binding-aware +continuation APIs, both adapters forward that same live +`AgentInvocationBinding` into every resumed child branch. Bindings remain +execution-only data — they are not checkpointed — so a resumed child that +reaches a `SubAgentNode` without one fails closed. + ## Namespace and recursion bookkeeping Internal helpers (not part of the public surface, but load-bearing for diff --git a/crates/tinyagents-graph/src/subgraph/test.rs b/crates/tinyagents-graph/src/subgraph/test.rs index da4924fd..3ac0a04e 100644 --- a/crates/tinyagents-graph/src/subgraph/test.rs +++ b/crates/tinyagents-graph/src/subgraph/test.rs @@ -35,6 +35,49 @@ impl crate::subagent_node::AgentInvoker for NestedRecordingInvoker { } } +struct NestedFailingInvoker; + +#[async_trait] +impl crate::subagent_node::AgentInvoker for NestedFailingInvoker { + async fn invoke( + &self, + _request: crate::subagent_node::AgentInvocation, + ) -> crate::Result { + Err(crate::TinyAgentsError::Model( + "nested continuation failure".to_string(), + )) + } +} + +fn delegating_child( + checkpointer: Arc>, +) -> CompiledGraph { + GraphBuilder::::overwrite() + .add_node( + "delegate", + crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( + "researcher", + |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), + |output: crate::subagent_node::SubAgentOutput| output.text, + )), + ) + .set_entry("delegate") + .set_finish("delegate") + .compile() + .unwrap() + .with_checkpointer(checkpointer) +} + +fn nested_binding( + invoker: Arc, +) -> crate::subagent_node::AgentInvocationBinding { + crate::subagent_node::AgentInvocationBinding::new( + invoker, + tinyagents_harness::events::EventSink::new(), + tinyagents_harness::cancel::CancellationToken::new(), + ) +} + /// Builds a minimal [`NodeContext`] standing in for the embedding node `id`. fn ctx_for(id: &str) -> NodeContext { NodeContext { @@ -453,6 +496,311 @@ async fn resumed_subgraph_passes_the_supplied_binding_to_its_subagent() { ); } +#[tokio::test] +async fn resumed_adapter_subgraph_passes_the_supplied_binding_to_its_subagent() { + // Exercise the adapter route separately: it has its own child-driving + // closure, so a shared-state test alone cannot prove a bound resume is not + // dropped while mapping parent state into child state and back. + let ckpt = Arc::new(InMemoryCheckpointer::::new()); + let child = GraphBuilder::::overwrite() + .add_node("gate", |state: String, ctx: NodeContext| async move { + if ctx.resume.is_some() { + Ok(NodeResult::Update(state)) + } else { + Ok(NodeResult::Interrupt(crate::command::Interrupt::new( + "gate", + serde_json::json!({ "ask": "continue?" }), + ))) + } + }) + .add_node( + "delegate", + crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( + "researcher", + |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), + |output: crate::subagent_node::SubAgentOutput| output.text, + )), + ) + .set_entry("gate") + .add_edge("gate", "delegate") + .set_finish("delegate") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + let parent = GraphBuilder::::overwrite() + .add_node( + "child", + adapter_subgraph_node( + child, + |state: &String| state.clone(), + |_parent, child| child, + ), + ) + .set_entry("child") + .set_finish("child") + .compile() + .unwrap() + .with_checkpointer(ckpt); + + assert!( + parent + .run_with_thread("adapter-resume", "question".to_string()) + .await + .unwrap() + .is_interrupted() + ); + + let invoker = Arc::new(NestedRecordingInvoker::default()); + let events = tinyagents_harness::events::EventSink::new(); + let listener = Arc::new(tinyagents_harness::events::RecordingListener::new()); + events.subscribe(listener.clone()); + let cancellation = tinyagents_harness::cancel::CancellationToken::new(); + cancellation.cancel(); + let resumed = parent + .resume_with_agent_binding( + "adapter-resume", + crate::command::Command::resume(serde_json::json!("go")), + crate::subagent_node::AgentInvocationBinding::new( + invoker.clone(), + events, + cancellation, + ), + ) + .await + .unwrap(); + + assert_eq!(resumed.state, "question"); + let request = invoker.0.lock().unwrap().pop().expect("child delegated"); + assert_eq!(request.parent_run_id, resumed.child_runs[0].run_id); + assert_eq!(request.root_run_id, resumed.root_run_id); + assert!(request.cancellation.unwrap().is_cancelled()); + assert_eq!(listener.len(), 1); +} + +#[tokio::test] +async fn binding_reaches_a_grandchild_that_itself_resumes() { + // The grandchild interrupts first. Resuming the root must therefore carry + // one binding through *two* resumed drive_child branches, preserving the + // same event sink and cancellation token at the eventual SubAgentNode. + let ckpt = Arc::new(InMemoryCheckpointer::::new()); + let grandchild = GraphBuilder::::overwrite() + .add_node("gate", |state: String, ctx: NodeContext| async move { + if ctx.resume.is_some() { + Ok(NodeResult::Update(state)) + } else { + Ok(NodeResult::Interrupt(crate::command::Interrupt::new( + "gate", + serde_json::json!({ "ask": "continue?" }), + ))) + } + }) + .add_node( + "delegate", + crate::subagent_node::subagent_node(crate::subagent_node::SubAgentNode::from_fns( + "researcher", + |state: &String| crate::subagent_node::SubAgentInput::prompt(state.clone()), + |output: crate::subagent_node::SubAgentOutput| output.text, + )), + ) + .set_entry("gate") + .add_edge("gate", "delegate") + .set_finish("delegate") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + let child = GraphBuilder::::overwrite() + .add_node("grandchild", shared_subgraph_node(grandchild)) + .set_entry("grandchild") + .set_finish("grandchild") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + let parent = GraphBuilder::::overwrite() + .add_node("child", shared_subgraph_node(child)) + .set_entry("child") + .set_finish("child") + .compile() + .unwrap() + .with_checkpointer(ckpt); + + assert!( + parent + .run_with_thread("grandchild-resume", "question".to_string()) + .await + .unwrap() + .is_interrupted() + ); + + let invoker = Arc::new(NestedRecordingInvoker::default()); + let events = tinyagents_harness::events::EventSink::new(); + let listener = Arc::new(tinyagents_harness::events::RecordingListener::new()); + events.subscribe(listener.clone()); + let cancellation = tinyagents_harness::cancel::CancellationToken::new(); + cancellation.cancel(); + let resumed = parent + .resume_with_agent_binding( + "grandchild-resume", + crate::command::Command::resume(serde_json::json!("go")), + crate::subagent_node::AgentInvocationBinding::new( + invoker.clone(), + events, + cancellation, + ), + ) + .await + .unwrap(); + + assert_eq!(resumed.state, "question"); + let request = invoker + .0 + .lock() + .unwrap() + .pop() + .expect("grandchild delegated"); + assert_eq!(request.root_run_id, resumed.root_run_id); + assert!(request.cancellation.unwrap().is_cancelled()); + assert_eq!(listener.len(), 1); +} + +#[tokio::test] +async fn retrying_shared_subgraph_passes_fresh_binding_to_failed_child() { + // A failed child leaves the parent node pending. Retrying that parent with + // a replacement binding drives the shared-state child through the fresh + // threaded binding branch, rather than reviving the failed invoker. + let ckpt = Arc::new(InMemoryCheckpointer::::new()); + let child = delegating_child(ckpt.clone()); + let parent = GraphBuilder::::overwrite() + .add_node("child", shared_subgraph_node(child)) + .set_entry("child") + .set_finish("child") + .compile() + .unwrap() + .with_checkpointer(ckpt); + + let failed = parent + .run_with_thread_agent_binding( + "shared-retry", + "question".to_string(), + nested_binding(Arc::new(NestedFailingInvoker)), + ) + .await + .unwrap_err(); + assert!(matches!(failed, crate::TinyAgentsError::Model(_))); + + let replacement = Arc::new(NestedRecordingInvoker::default()); + let retried = parent + .retry_with_agent_binding("shared-retry", nested_binding(replacement.clone())) + .await + .unwrap(); + + assert_eq!(retried.state, "question"); + let request = replacement + .0 + .lock() + .unwrap() + .pop() + .expect("shared child used the retry binding"); + assert_eq!(request.parent_run_id, retried.child_runs[0].run_id); + assert_eq!(request.root_run_id, retried.root_run_id); +} + +#[tokio::test] +async fn retrying_adapter_subgraph_passes_fresh_binding_to_failed_child() { + // This covers the adapter's distinct child-driving closure. Mapping parent + // state into and out of the child must not drop the retry binding. + let ckpt = Arc::new(InMemoryCheckpointer::::new()); + let child = delegating_child(ckpt.clone()); + let parent = GraphBuilder::::overwrite() + .add_node( + "child", + adapter_subgraph_node( + child, + |state: &String| state.clone(), + |_parent, child| child, + ), + ) + .set_entry("child") + .set_finish("child") + .compile() + .unwrap() + .with_checkpointer(ckpt); + + let failed = parent + .run_with_thread_agent_binding( + "adapter-retry", + "question".to_string(), + nested_binding(Arc::new(NestedFailingInvoker)), + ) + .await + .unwrap_err(); + assert!(matches!(failed, crate::TinyAgentsError::Model(_))); + + let replacement = Arc::new(NestedRecordingInvoker::default()); + let retried = parent + .retry_with_agent_binding("adapter-retry", nested_binding(replacement.clone())) + .await + .unwrap(); + + assert_eq!(retried.state, "question"); + let request = replacement + .0 + .lock() + .unwrap() + .pop() + .expect("adapter child used the retry binding"); + assert_eq!(request.parent_run_id, retried.child_runs[0].run_id); + assert_eq!(request.root_run_id, retried.root_run_id); +} + +#[tokio::test] +async fn retrying_nested_shared_subgraphs_rebinds_the_failed_grandchild() { + // Retrying from the root traverses two independent shared-subgraph + // drive_child calls. The grandchild failure must be repaired solely by the + // fresh root binding, with no capability retained in either checkpoint. + let ckpt = Arc::new(InMemoryCheckpointer::::new()); + let grandchild = delegating_child(ckpt.clone()); + let child = GraphBuilder::::overwrite() + .add_node("grandchild", shared_subgraph_node(grandchild)) + .set_entry("grandchild") + .set_finish("grandchild") + .compile() + .unwrap() + .with_checkpointer(ckpt.clone()); + let parent = GraphBuilder::::overwrite() + .add_node("child", shared_subgraph_node(child)) + .set_entry("child") + .set_finish("child") + .compile() + .unwrap() + .with_checkpointer(ckpt); + + let failed = parent + .run_with_thread_agent_binding( + "grandchild-retry", + "question".to_string(), + nested_binding(Arc::new(NestedFailingInvoker)), + ) + .await + .unwrap_err(); + assert!(matches!(failed, crate::TinyAgentsError::Model(_))); + + let replacement = Arc::new(NestedRecordingInvoker::default()); + let retried = parent + .retry_with_agent_binding("grandchild-retry", nested_binding(replacement.clone())) + .await + .unwrap(); + + assert_eq!(retried.state, "question"); + let request = replacement + .0 + .lock() + .unwrap() + .pop() + .expect("grandchild used the retry binding"); + assert_eq!(request.root_run_id, retried.root_run_id); + assert_eq!(retried.child_runs.len(), 1); +} + #[tokio::test] async fn subgraph_child_run_distinct_and_shares_root() { // A parent embedding one child: the parent run records exactly one child run From 420cab0fc99ce91222aec1ab8b7b99abfbba0325 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 18:13:09 +0300 Subject: [PATCH 02/18] feat(session): own transcript layout migration Co-authored-by: Medulla --- crates/tinyagents-session/src/transcript.rs | 3 + .../src/transcript/migration.rs | 338 ++++++++++++++++++ .../src/transcript/migration_test.rs | 213 +++++++++++ 3 files changed, 554 insertions(+) create mode 100644 crates/tinyagents-session/src/transcript/migration.rs create mode 100644 crates/tinyagents-session/src/transcript/migration_test.rs diff --git a/crates/tinyagents-session/src/transcript.rs b/crates/tinyagents-session/src/transcript.rs index 76b355ad..58649d15 100644 --- a/crates/tinyagents-session/src/transcript.rs +++ b/crates/tinyagents-session/src/transcript.rs @@ -104,11 +104,13 @@ //! | `paths` | `session_raw` / `sessions` path resolution and resume scan. | //! | `markdown` | Human-readable `.md` companion rendering. | //! | `legacy_md` | Legacy HTML-comment `.md` reader. | +//! | `migration` | One-shot legacy date-grouped layout conversion. | mod history; mod jsonl; mod legacy_md; mod markdown; +mod migration; mod paths; mod reader; mod thread_lookup; @@ -120,6 +122,7 @@ pub use history::{ TranscriptRead, TranscriptTurn, }; pub use legacy_md::read_transcript_legacy_md; +pub use migration::{TranscriptLayoutMigration, migrate_layout_if_needed}; pub use paths::{find_latest_transcript, resolve_keyed_transcript_path}; pub use reader::{read_transcript, read_transcript_display}; pub use thread_lookup::{ diff --git a/crates/tinyagents-session/src/transcript/migration.rs b/crates/tinyagents-session/src/transcript/migration.rs new file mode 100644 index 00000000..9cfb9754 --- /dev/null +++ b/crates/tinyagents-session/src/transcript/migration.rs @@ -0,0 +1,338 @@ +//! One-shot migration of legacy date-grouped transcript layouts. +//! +//! Older hosts wrote transcripts beneath `session_raw/DDMMYYYY/` and markdown +//! companions beneath `sessions/DDMMYYYY/`. [`migrate_layout_if_needed`] +//! moves those artifacts into the flat transcript layout without overwriting a +//! file already present at the destination. Its workspace-local marker makes +//! completed migrations idempotent. + +use anyhow::{Context, Result}; +use std::fs; +use std::path::{Path, PathBuf}; + +/// Marker file that signals the v1 session-layout migration has run for a +/// workspace. It lives under `state/migrations/` to keep the workspace root +/// tidy. +const MIGRATION_MARKER: &str = "state/migrations/session_layout_v1.done"; + +/// Counts and non-fatal diagnostics from a session-layout migration. +#[derive(Debug, Default, Clone)] +pub struct TranscriptLayoutMigration { + /// Legacy JSONL transcripts moved into flat `session_raw/`. + pub jsonl_moved: usize, + /// JSONL files retained in the legacy directory because the destination existed. + pub jsonl_skipped: usize, + /// Markdown files moved into ISO-date directories. + pub md_moved: usize, + /// Markdown files retained in the legacy directory because the destination existed. + pub md_skipped: usize, + /// Empty legacy date directories removed after moving their contents. + pub legacy_dirs_pruned: usize, + /// Whether the migration marker was present before this invocation. + pub already_done: bool, + /// Non-fatal filesystem diagnostics encountered while migrating. + pub warnings: Vec, +} + +/// Migrate a workspace's legacy transcript layout if it has not run before. +/// +/// Detects `session_raw/{DDMMYYYY}/...jsonl` and `sessions/{DDMMYYYY}/...md`, +/// moves JSONL files into flat `session_raw/`, and moves markdown files into +/// `sessions/{YYYY_MM_DD}/`. Existing destination files are never +/// overwritten: they are retained in the legacy location and reported in +/// [`TranscriptLayoutMigration::warnings`]. A successful run writes a marker, including +/// when there were no legacy artifacts, so later starts perform no scan. +/// +/// Individual filesystem failures are collected as warnings rather than +/// returned, allowing a host to continue startup and use legacy read fallback. +pub fn migrate_layout_if_needed(workspace_dir: &Path) -> Result { + let marker_path = workspace_dir.join(MIGRATION_MARKER); + if marker_path.exists() { + log::debug!( + "[session-migration] marker present at {} — skipping", + marker_path.display() + ); + return Ok(TranscriptLayoutMigration { + already_done: true, + ..Default::default() + }); + } + + let mut outcome = TranscriptLayoutMigration::default(); + + let raw_root = workspace_dir.join("session_raw"); + if raw_root.is_dir() { + migrate_raw_jsonl(&raw_root, &mut outcome)?; + } + + let sessions_root = workspace_dir.join("sessions"); + if sessions_root.is_dir() { + migrate_md_directories(&sessions_root, &mut outcome)?; + } + + write_marker(&marker_path, &outcome).context("write session-migration marker")?; + + log::info!( + "[session-migration] complete: jsonl moved={} skipped={}, md moved={} skipped={}, legacy dirs pruned={}, warnings={}", + outcome.jsonl_moved, + outcome.jsonl_skipped, + outcome.md_moved, + outcome.md_skipped, + outcome.legacy_dirs_pruned, + outcome.warnings.len(), + ); + + Ok(outcome) +} + +fn migrate_raw_jsonl(raw_root: &Path, outcome: &mut TranscriptLayoutMigration) -> Result<()> { + let entries = match fs::read_dir(raw_root) { + Ok(it) => it, + Err(err) => { + outcome + .warnings + .push(format!("read_dir({}) failed: {err}", raw_root.display())); + return Ok(()); + } + }; + + for entry in entries.flatten() { + let path = entry.path(); + if !path.is_dir() { + continue; + } + let Some(name) = path.file_name().and_then(|s| s.to_str()) else { + continue; + }; + if !is_ddmmyyyy(name) { + continue; + } + move_jsonl_files_up(&path, raw_root, outcome); + prune_if_empty(&path, outcome); + } + + Ok(()) +} + +fn move_jsonl_files_up( + legacy_dir: &Path, + flat_dir: &Path, + outcome: &mut TranscriptLayoutMigration, +) { + let entries = match fs::read_dir(legacy_dir) { + Ok(it) => it, + Err(err) => { + outcome + .warnings + .push(format!("read_dir({}) failed: {err}", legacy_dir.display())); + return; + } + }; + + for entry in entries.flatten() { + let path = entry.path(); + if !path.is_file() || path.extension().and_then(|s| s.to_str()) != Some("jsonl") { + continue; + } + let Some(file_name) = path.file_name() else { + continue; + }; + let dest = flat_dir.join(file_name); + match hard_link_then_remove(&path, &dest) { + Ok(()) => { + outcome.jsonl_moved += 1; + log::debug!( + "[session-migration] moved {} → {}", + path.display(), + dest.display() + ); + } + Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => { + outcome.jsonl_skipped += 1; + outcome.warnings.push(format!( + "skip move: destination already exists at {} (legacy file kept at {})", + dest.display(), + path.display() + )); + } + Err(err) => outcome.warnings.push(format!( + "atomic move({} → {}) failed: {err}", + path.display(), + dest.display() + )), + } + } +} + +fn migrate_md_directories( + sessions_root: &Path, + outcome: &mut TranscriptLayoutMigration, +) -> Result<()> { + let entries = match fs::read_dir(sessions_root) { + Ok(it) => it, + Err(err) => { + outcome.warnings.push(format!( + "read_dir({}) failed: {err}", + sessions_root.display() + )); + return Ok(()); + } + }; + + for entry in entries.flatten() { + let path = entry.path(); + if !path.is_dir() { + continue; + } + let Some(name) = path.file_name().and_then(|s| s.to_str()) else { + continue; + }; + let Some(iso) = ddmmyyyy_to_yyyy_mm_dd(name) else { + continue; + }; + let dest = sessions_root.join(&iso); + match reserve_md_destination(&dest) { + Ok(created) => { + let moved_before = outcome.md_moved; + if created { + log::debug!( + "[session-migration] reserved markdown destination {}", + dest.display() + ); + } + merge_md_dirs(&path, &dest, outcome); + prune_if_empty(&path, outcome); + if created && !path.exists() { + // The legacy implementation moved this entire directory in + // one rename and recorded one move regardless of its file + // count. Keep that marker/report contract stable for a + // complete fresh-directory migration; merging into an + // existing destination remains per-file. + outcome.md_moved = moved_before + 1; + } + } + Err(err) => outcome.warnings.push(format!( + "reserve markdown destination {} failed: {err}", + dest.display() + )), + } + } + + Ok(()) +} + +/// Reserve a markdown destination directory without replacing a concurrent +/// creator's directory. `create_dir` is the atomic no-replace operation: an +/// `AlreadyExists` result means another process (or an earlier migration) owns +/// the directory and its files must be merged rather than replaced. +fn reserve_md_destination(destination: &Path) -> std::io::Result { + if let Some(parent) = destination.parent() { + fs::create_dir_all(parent)?; + } + match fs::create_dir(destination) { + Ok(()) => Ok(true), + Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => Ok(false), + Err(err) => Err(err), + } +} + +fn merge_md_dirs(legacy: &Path, dest: &Path, outcome: &mut TranscriptLayoutMigration) { + let entries = match fs::read_dir(legacy) { + Ok(it) => it, + Err(err) => { + outcome + .warnings + .push(format!("read_dir({}) failed: {err}", legacy.display())); + return; + } + }; + for entry in entries.flatten() { + let src = entry.path(); + if !src.is_file() { + continue; + } + let Some(file_name) = src.file_name() else { + continue; + }; + let target = dest.join(file_name); + match hard_link_then_remove(&src, &target) { + Ok(()) => outcome.md_moved += 1, + Err(err) if err.kind() == std::io::ErrorKind::AlreadyExists => { + outcome.md_skipped += 1; + outcome.warnings.push(format!( + "skip md merge: {} already exists (legacy at {} kept)", + target.display(), + src.display() + )); + } + Err(err) => outcome.warnings.push(format!( + "atomic md move({} → {}) failed: {err}", + src.display(), + target.display() + )), + } + } +} + +/// Atomically create `destination` as a second link to `source`, then remove +/// `source`. `hard_link` refuses an existing destination, so a competing +/// writer can never be replaced between discovery and transfer. +fn hard_link_then_remove(source: &Path, destination: &Path) -> std::io::Result<()> { + fs::hard_link(source, destination)?; + fs::remove_file(source) +} + +fn prune_if_empty(dir: &Path, outcome: &mut TranscriptLayoutMigration) { + match fs::read_dir(dir) { + Ok(mut it) => { + if it.next().is_some() { + return; + } + } + Err(_) => return, + } + if fs::remove_dir(dir).is_ok() { + outcome.legacy_dirs_pruned += 1; + } +} + +fn write_marker(marker_path: &Path, outcome: &TranscriptLayoutMigration) -> Result<()> { + if let Some(parent) = marker_path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("create marker dir {}", parent.display()))?; + } + let body = format!( + "openhuman session_layout migration v1\nrun_at: {}\njsonl_moved: {}\nmd_moved: {}\nlegacy_dirs_pruned: {}\nwarnings: {}\n", + chrono::Utc::now().to_rfc3339(), + outcome.jsonl_moved, + outcome.md_moved, + outcome.legacy_dirs_pruned, + outcome.warnings.len(), + ); + fs::write(marker_path, body) + .with_context(|| format!("write marker {}", marker_path.display()))?; + Ok(()) +} + +fn is_ddmmyyyy(name: &str) -> bool { + name.len() == 8 && name.chars().all(|c| c.is_ascii_digit()) +} + +fn ddmmyyyy_to_yyyy_mm_dd(name: &str) -> Option { + if !is_ddmmyyyy(name) { + return None; + } + let dd = &name[0..2]; + let mm = &name[2..4]; + let yyyy = &name[4..8]; + Some(format!("{yyyy}_{mm}_{dd}")) +} + +/// Return the migration marker path for `workspace_dir`. +fn marker_path_for(workspace_dir: &Path) -> PathBuf { + workspace_dir.join(MIGRATION_MARKER) +} + +#[cfg(test)] +#[path = "migration_test.rs"] +mod tests; diff --git a/crates/tinyagents-session/src/transcript/migration_test.rs b/crates/tinyagents-session/src/transcript/migration_test.rs new file mode 100644 index 00000000..9dd03f55 --- /dev/null +++ b/crates/tinyagents-session/src/transcript/migration_test.rs @@ -0,0 +1,213 @@ +use super::*; +use std::fs; +use tempfile::TempDir; + +fn write_file(path: &std::path::Path, body: &str) { + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(path, body).unwrap(); +} + +#[test] +fn fresh_workspace_writes_marker_with_no_moves() { + let dir = TempDir::new().unwrap(); + let outcome = migrate_layout_if_needed(dir.path()).unwrap(); + assert!(!outcome.already_done); + assert_eq!(outcome.jsonl_moved, 0); + assert_eq!(outcome.md_moved, 0); + assert!(marker_path_for(dir.path()).exists()); +} + +#[test] +fn second_run_is_a_noop() { + let dir = TempDir::new().unwrap(); + let _first = migrate_layout_if_needed(dir.path()).unwrap(); + let second = migrate_layout_if_needed(dir.path()).unwrap(); + assert!(second.already_done); + assert_eq!(second.jsonl_moved, 0); + assert_eq!(second.warnings.len(), 0); +} + +#[test] +fn moves_legacy_jsonl_files_up_to_flat_session_raw() { + let dir = TempDir::new().unwrap(); + let ws = dir.path(); + let legacy_a = ws.join("session_raw").join("01052026"); + let legacy_b = ws.join("session_raw").join("02052026"); + write_file(&legacy_a.join("1714000000_main.jsonl"), "a"); + write_file(&legacy_a.join("1714000001_welcome.jsonl"), "b"); + write_file(&legacy_b.join("1714999999_orchestrator.jsonl"), "c"); + + let outcome = migrate_layout_if_needed(ws).unwrap(); + assert_eq!(outcome.jsonl_moved, 3); + assert_eq!(outcome.legacy_dirs_pruned, 2); + + let raw_root = ws.join("session_raw"); + assert!(raw_root.join("1714000000_main.jsonl").exists()); + assert!(raw_root.join("1714000001_welcome.jsonl").exists()); + assert!(raw_root.join("1714999999_orchestrator.jsonl").exists()); + assert!(!legacy_a.exists(), "legacy date dir should be removed"); + assert!(!legacy_b.exists(), "legacy date dir should be removed"); +} + +#[test] +fn jsonl_destination_collision_is_skipped_with_warning() { + let dir = TempDir::new().unwrap(); + let ws = dir.path(); + let raw_root = ws.join("session_raw"); + write_file(&raw_root.join("1714000000_main.jsonl"), "new"); + write_file( + &raw_root.join("01052026").join("1714000000_main.jsonl"), + "old", + ); + + let outcome = migrate_layout_if_needed(ws).unwrap(); + assert_eq!(outcome.jsonl_moved, 0); + assert_eq!(outcome.jsonl_skipped, 1); + assert!( + outcome + .warnings + .iter() + .any(|w| w.contains("already exists")) + ); + assert_eq!( + fs::read_to_string(raw_root.join("1714000000_main.jsonl")).unwrap(), + "new" + ); + assert_eq!( + fs::read_to_string(raw_root.join("01052026").join("1714000000_main.jsonl")).unwrap(), + "old" + ); +} + +#[test] +fn atomic_transfer_preserves_destination_created_before_transfer() { + let dir = TempDir::new().unwrap(); + let source = dir.path().join("legacy.jsonl"); + let destination = dir.path().join("canonical.jsonl"); + write_file(&source, "legacy contents"); + + // This models another process creating the canonical file after migration + // has discovered `source`, but before its atomic transfer operation. + write_file(&destination, "concurrent contents"); + let error = hard_link_then_remove(&source, &destination).unwrap_err(); + + assert_eq!(error.kind(), std::io::ErrorKind::AlreadyExists); + assert_eq!( + fs::read_to_string(&destination).unwrap(), + "concurrent contents" + ); + assert_eq!(fs::read_to_string(&source).unwrap(), "legacy contents"); +} + +#[test] +fn renames_md_ddmmyyyy_dirs_to_iso() { + let dir = TempDir::new().unwrap(); + let ws = dir.path(); + let legacy_md = ws.join("sessions").join("01052026"); + write_file(&legacy_md.join("main_0.md"), "x"); + write_file(&legacy_md.join("main_1.md"), "y"); + + let outcome = migrate_layout_if_needed(ws).unwrap(); + assert_eq!( + outcome.md_moved, 1, + "a freshly reserved date directory preserves the legacy one-directory count" + ); + let iso = ws.join("sessions").join("2026_05_01"); + assert!(iso.is_dir()); + assert!(iso.join("main_0.md").exists()); + assert!(iso.join("main_1.md").exists()); + assert!( + !legacy_md.exists(), + "DDMMYYYY md dir should be gone after rename" + ); +} + +#[test] +fn merges_md_when_iso_dir_already_exists() { + let dir = TempDir::new().unwrap(); + let ws = dir.path(); + let legacy = ws.join("sessions").join("01052026"); + let iso = ws.join("sessions").join("2026_05_01"); + write_file(&legacy.join("main_0.md"), "legacy"); + write_file(&legacy.join("main_1.md"), "legacy"); + write_file(&iso.join("main_1.md"), "newer"); + + let outcome = migrate_layout_if_needed(ws).unwrap(); + assert_eq!(outcome.md_moved, 1); + assert_eq!(outcome.md_skipped, 1); + assert_eq!(fs::read_to_string(iso.join("main_0.md")).unwrap(), "legacy"); + assert_eq!(fs::read_to_string(iso.join("main_1.md")).unwrap(), "newer"); +} + +#[test] +fn markdown_destination_reserved_by_a_concurrent_creator_is_merged_safely() { + let dir = TempDir::new().unwrap(); + let ws = dir.path(); + let legacy = ws.join("sessions").join("01052026"); + let destination = ws.join("sessions").join("2026_05_01"); + write_file(&legacy.join("new.md"), "legacy file"); + write_file(&legacy.join("shared.md"), "legacy copy"); + + // Model a second process reserving and populating the ISO directory after + // this migration has discovered the legacy directory but before it can + // reserve the destination leaf. + fs::create_dir_all(&destination).unwrap(); + write_file(&destination.join("shared.md"), "concurrent copy"); + + let outcome = migrate_layout_if_needed(ws).unwrap(); + assert_eq!(outcome.md_moved, 1); + assert_eq!(outcome.md_skipped, 1); + assert_eq!( + fs::read_to_string(destination.join("new.md")).unwrap(), + "legacy file" + ); + assert_eq!( + fs::read_to_string(destination.join("shared.md")).unwrap(), + "concurrent copy" + ); + assert_eq!( + fs::read_to_string(legacy.join("shared.md")).unwrap(), + "legacy copy" + ); +} + +#[test] +fn ignores_non_date_subdirectories_in_session_raw() { + let dir = TempDir::new().unwrap(); + let ws = dir.path(); + let weird = ws.join("session_raw").join("my_notes"); + write_file(&weird.join("random.jsonl"), "keep me"); + + let outcome = migrate_layout_if_needed(ws).unwrap(); + assert_eq!(outcome.jsonl_moved, 0); + assert!(weird.is_dir(), "non-date subdir must be left alone"); + assert!(weird.join("random.jsonl").exists()); +} + +#[test] +fn ddmmyyyy_to_iso_handles_boundary_dates() { + assert_eq!( + ddmmyyyy_to_yyyy_mm_dd("01012026").as_deref(), + Some("2026_01_01") + ); + assert_eq!( + ddmmyyyy_to_yyyy_mm_dd("31122099").as_deref(), + Some("2099_12_31") + ); + assert!(ddmmyyyy_to_yyyy_mm_dd("abc12345").is_none()); + assert!(ddmmyyyy_to_yyyy_mm_dd("1234567").is_none(), "7 digits"); + assert!(ddmmyyyy_to_yyyy_mm_dd("123456789").is_none(), "9 digits"); +} + +#[test] +fn marker_persists_run_metadata() { + let dir = TempDir::new().unwrap(); + let ws = dir.path(); + let legacy = ws.join("session_raw").join("01052026"); + write_file(&legacy.join("1714000000_main.jsonl"), "a"); + + migrate_layout_if_needed(ws).unwrap(); + let marker = fs::read_to_string(marker_path_for(ws)).unwrap(); + assert!(marker.contains("jsonl_moved: 1")); + assert!(marker.contains("openhuman session_layout migration v1")); +} From 0a44f3268efe2780bb3b30a7459838574d6db424 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 18:47:30 +0300 Subject: [PATCH 03/18] feat(session): append interrupted partials atomically Co-authored-by: Medulla --- crates/tinyagents-session/src/transcript.rs | 7 +- .../src/transcript/history.rs | 66 +++++++++++++++ .../src/transcript/writer.rs | 81 ++++++++++++++++--- 3 files changed, 139 insertions(+), 15 deletions(-) diff --git a/crates/tinyagents-session/src/transcript.rs b/crates/tinyagents-session/src/transcript.rs index 58649d15..9584924e 100644 --- a/crates/tinyagents-session/src/transcript.rs +++ b/crates/tinyagents-session/src/transcript.rs @@ -119,7 +119,7 @@ mod writer; pub use history::{ FileTranscriptHistory, FileTranscriptLocator, TranscriptHistory, TranscriptLocator, - TranscriptRead, TranscriptTurn, + TranscriptPartial, TranscriptRead, TranscriptTurn, }; pub use legacy_md::read_transcript_legacy_md; pub use migration::{TranscriptLayoutMigration, migrate_layout_if_needed}; @@ -134,7 +134,10 @@ pub use types::{ SessionTranscript, ToolFailure, TranscriptMessage, TranscriptMeta, TranscriptToolCall, TurnUsage, }; -pub use writer::{append_interrupted_partial, append_transcript_turn, write_transcript}; +pub use writer::{ + append_interrupted_partial, append_transcript_turn, append_transcript_turn_with_partial, + write_transcript, +}; // Private helpers the colocated tests exercise directly. #[cfg(test)] diff --git a/crates/tinyagents-session/src/transcript/history.rs b/crates/tinyagents-session/src/transcript/history.rs index e00b0087..6e4d0e41 100644 --- a/crates/tinyagents-session/src/transcript/history.rs +++ b/crates/tinyagents-session/src/transcript/history.rs @@ -54,6 +54,32 @@ pub struct TranscriptTurn<'a> { pub request_id: Option<&'a str>, } +/// Display-only content produced before a turn stopped without a final answer. +/// +/// This deliberately uses transcript-neutral fields. It is never added to a +/// model-context replay: the file writer records it as an interrupted message +/// line for the display projection only. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TranscriptPartial { + /// Visible assistant text accumulated before the interruption. + pub content: String, + /// Optional provider reasoning text associated with the partial. + pub reasoning_content: Option, + /// Optional one-based engine iteration associated with the partial. + pub iteration: Option, +} + +impl TranscriptPartial { + /// Creates a display-only partial with no provider-specific metadata. + pub fn new(content: impl Into) -> Self { + Self { + content: content.into(), + reasoning_content: None, + iteration: None, + } + } +} + /// The seam a host turn path holds as `Arc`. /// /// `append_turn` is deliberately **sync**: `persist_session_transcript` is a @@ -63,6 +89,23 @@ pub trait TranscriptHistory: TranscriptRead { /// Appends one turn, forwarding every argument to the format owner. fn append_turn(&self, turn: TranscriptTurn<'_>) -> anyhow::Result<()>; + /// Appends the logical turn and its optional display-only partial as one + /// history operation. + /// + /// Implementors that cannot make the combined mutation atomic must reject + /// a partial rather than persist either half. Existing implementors which + /// only support logical turns remain source-compatible through this default. + fn append_turn_with_partial( + &self, + turn: TranscriptTurn<'_>, + partial: Option<&TranscriptPartial>, + ) -> anyhow::Result<()> { + if partial.is_some() { + anyhow::bail!("transcript history does not support atomic display partials"); + } + self.append_turn(turn) + } + /// Returns the lossless model-context replay of this transcript. fn messages(&self) -> anyhow::Result>; @@ -426,6 +469,29 @@ impl TranscriptHistory for FileTranscriptHistory { turn.request_id, ) } + + fn append_turn_with_partial( + &self, + turn: TranscriptTurn<'_>, + partial: Option<&TranscriptPartial>, + ) -> anyhow::Result<()> { + log::debug!( + "[transcript-history] append_turn_with_partial prev={} next={} partial={} path={}", + turn.prev.len(), + turn.next.len(), + partial.is_some(), + self.path.display() + ); + crate::transcript::append_transcript_turn_with_partial( + &self.path, + turn.prev, + turn.next, + turn.meta, + turn.turn_usage, + turn.request_id, + partial, + ) + } fn messages(&self) -> anyhow::Result> { self.persisted() } diff --git a/crates/tinyagents-session/src/transcript/writer.rs b/crates/tinyagents-session/src/transcript/writer.rs index 84de361b..76f83820 100644 --- a/crates/tinyagents-session/src/transcript/writer.rs +++ b/crates/tinyagents-session/src/transcript/writer.rs @@ -2,6 +2,7 @@ //! delta (message tail or compaction record), interrupted partials, and the //! derived `.md` companion. +use super::history::TranscriptPartial; use super::jsonl::{ COMPACTION_KIND, CompactionLine, MessageLine, build_message_line, meta_line_json, serialise_message_lines, @@ -79,6 +80,34 @@ pub fn append_transcript_turn( meta: &TranscriptMeta, turn_usage: Option<&TurnUsage>, request_id: Option<&str>, +) -> Result<()> { + append_transcript_turn_with_partial( + jsonl_path, + prev_persisted, + messages, + meta, + turn_usage, + request_id, + None, + ) +} + +/// Appends one logical turn and, when present, its display-only interruption +/// row from one serialized buffer and one file-write operation. +/// +/// The partial is written after the logical delta and refreshed metadata. It +/// has `interrupted: true`, so the model-context reader skips it while the +/// display reader preserves it. Serialization happens before opening the file +/// for append, preventing a serialization failure from leaving a logical turn +/// without its associated display partial. +pub fn append_transcript_turn_with_partial( + jsonl_path: &Path, + prev_persisted: &[TranscriptMessage], + messages: &[TranscriptMessage], + meta: &TranscriptMeta, + turn_usage: Option<&TurnUsage>, + request_id: Option<&str>, + partial: Option<&TranscriptPartial>, ) -> Result<()> { if let Some(parent) = jsonl_path.parent() { fs::create_dir_all(parent) @@ -93,6 +122,7 @@ pub fn append_transcript_turn( buf.push_str(&meta_line_json(meta)?); buf.push('\n'); serialise_message_lines(messages, turn_usage, request_id, &mut buf)?; + serialise_interrupted_partial(partial, request_id, &mut buf)?; fs::write(jsonl_path, buf.as_bytes()) .with_context(|| format!("create transcript {}", jsonl_path.display()))?; log::debug!( @@ -159,12 +189,41 @@ pub fn append_transcript_turn( // last one). Keeps append-only + O(1)-per-turn (no full-file rewrite). buf.push_str(&meta_line_json(meta)?); buf.push('\n'); + serialise_interrupted_partial(partial, request_id, &mut buf)?; append_bytes(jsonl_path, buf.as_bytes())?; render_md_companion(jsonl_path, messages, meta, turn_usage); Ok(()) } +/// Appends the optional display-only row to an already assembled turn buffer. +fn serialise_interrupted_partial( + partial: Option<&TranscriptPartial>, + request_id: Option<&str>, + buf: &mut String, +) -> Result<()> { + let Some(partial) = partial.filter(|partial| !partial.content.is_empty()) else { + return Ok(()); + }; + let mut line = build_message_line( + &TranscriptMessage::assistant(&partial.content), + None, + request_id, + true, + ); + line.iteration = partial.iteration; + line.reasoning_content = partial + .reasoning_content + .as_deref() + .map(str::trim) + .filter(|content| !content.is_empty()) + .map(str::to_owned); + line.ts = Some(chrono::Utc::now().to_rfc3339()); + buf.push_str(&serde_json::to_string(&line).context("serialise interrupted partial line")?); + buf.push('\n'); + Ok(()) +} + /// Append a partial assistant answer, flagged `interrupted: true`, captured /// when a streaming turn was cancelled/interrupted before completion. /// @@ -186,20 +245,16 @@ pub fn append_interrupted_partial( fs::create_dir_all(parent) .with_context(|| format!("create transcript dir {}", parent.display()))?; } - let mut line = build_message_line( - &TranscriptMessage::assistant(partial_content), - None, + let mut buf = String::new(); + serialise_interrupted_partial( + Some(&TranscriptPartial { + content: partial_content.to_owned(), + reasoning_content: reasoning_content.map(str::to_owned), + iteration, + }), request_id, - true, - ); - line.iteration = iteration; - line.reasoning_content = reasoning_content - .map(str::trim) - .filter(|s| !s.is_empty()) - .map(str::to_string); - line.ts = Some(chrono::Utc::now().to_rfc3339()); - let mut buf = serde_json::to_string(&line).context("serialise interrupted partial line")?; - buf.push('\n'); + &mut buf, + )?; append_bytes(jsonl_path, buf.as_bytes())?; log::debug!( "[transcript] appended interrupted partial ({} chars, request_id={:?}) to {}", From 6d611da4a2aaf3daddfd6a68b0d4ab0e1ae13213 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 18:47:41 +0300 Subject: [PATCH 04/18] feat(runtime): add host-neutral stateful sessions Co-authored-by: Medulla --- Cargo.lock | 17 + Cargo.toml | 1 + README.md | 14 + crates/tinyagents-runtime/Cargo.toml | 27 + crates/tinyagents-runtime/README.md | 54 + crates/tinyagents-runtime/src/builder.rs | 111 ++ crates/tinyagents-runtime/src/driver.rs | 138 +++ crates/tinyagents-runtime/src/error.rs | 28 + crates/tinyagents-runtime/src/hooks.rs | 46 + crates/tinyagents-runtime/src/lib.rs | 58 + crates/tinyagents-runtime/src/prefix.rs | 19 + crates/tinyagents-runtime/src/session.rs | 359 ++++++ crates/tinyagents-runtime/src/test.rs | 1345 ++++++++++++++++++++++ crates/tinyagents-runtime/src/tools.rs | 43 + crates/tinyagents-runtime/src/types.rs | 126 ++ docs/modules/runtime/README.md | 41 + docs/spec/README.md | 5 +- 17 files changed, 2431 insertions(+), 1 deletion(-) create mode 100644 crates/tinyagents-runtime/Cargo.toml create mode 100644 crates/tinyagents-runtime/README.md create mode 100644 crates/tinyagents-runtime/src/builder.rs create mode 100644 crates/tinyagents-runtime/src/driver.rs create mode 100644 crates/tinyagents-runtime/src/error.rs create mode 100644 crates/tinyagents-runtime/src/hooks.rs create mode 100644 crates/tinyagents-runtime/src/lib.rs create mode 100644 crates/tinyagents-runtime/src/prefix.rs create mode 100644 crates/tinyagents-runtime/src/session.rs create mode 100644 crates/tinyagents-runtime/src/test.rs create mode 100644 crates/tinyagents-runtime/src/tools.rs create mode 100644 crates/tinyagents-runtime/src/types.rs create mode 100644 docs/modules/runtime/README.md diff --git a/Cargo.lock b/Cargo.lock index 1cb59d85..ad2657e6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1608,6 +1608,23 @@ dependencies = [ "tokio", ] +[[package]] +name = "tinyagents-runtime" +version = "2.1.2" +dependencies = [ + "anyhow", + "async-trait", + "chrono", + "serde_json", + "tempfile", + "thiserror 2.0.20", + "tinyagents-harness", + "tinyagents-session", + "tinyinference-llm", + "tinytools", + "tokio", +] + [[package]] name = "tinyagents-session" version = "2.1.2" diff --git a/Cargo.toml b/Cargo.toml index aa93c250..6c650e91 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,6 +6,7 @@ default-members = [ "crates/tinyagents-graph", "crates/tinyagents-registry", "crates/tinyagents-session", + "crates/tinyagents-runtime", "crates/tinyagents-orchestration", ] exclude = ["vendor", "worktrees"] diff --git a/README.md b/README.md index 18468c94..bddf1170 100644 --- a/README.md +++ b/README.md @@ -42,6 +42,9 @@ TinyAgents is a Cargo workspace, not one crate. Depend on the pieces you need: name, plus an offline model price/capability catalog. - **`tinyagents-session`** — a SQLite-backed store for session history, messages, tool calls, cost, and run lineage. +- **`tinyagents-runtime`** — host-neutral stateful turns over the harness and + append-only transcript seam; hosts retain policy, prompt composition, + authorization, and durable-dialect conversion. - **`tinyagents-tracing`** — the `tracing` macros the other crates gate behind their `tracing` feature. Compiled out by default. - **`tinyagents-integration-tests`** — cross-crate tests and the runnable @@ -146,6 +149,17 @@ wrapped as a tool and handed to another agent (`SubAgent` / `SubAgentSession` / `SubAgentTool`), which is how multi-agent orchestration is composed — plain function composition, not a distinct execution mode. +## Session runtime + +`tinyagents-runtime` owns mutable model history for one host-owned +conversation, a stable prompt prefix, a frozen tool declaration snapshot, and +the sequencing around one append-only transcript commit. A host supplies the +driver, its lossless transcript codec, and lifecycle hooks. On a driver error, +the runtime can commit recoverable logical history with an interrupted, +display-only partial in the same history operation; model-context replay omits +that partial. A post-commit hook observes durable successes but cannot change +their result. See [the runtime module](docs/modules/runtime/README.md). + ## Registry `tinyagents-registry` is a name-addressable catalog of models, tools, agents, diff --git a/crates/tinyagents-runtime/Cargo.toml b/crates/tinyagents-runtime/Cargo.toml new file mode 100644 index 00000000..8fbf83b8 --- /dev/null +++ b/crates/tinyagents-runtime/Cargo.toml @@ -0,0 +1,27 @@ +[package] +publish = false +name = "tinyagents-runtime" +version.workspace = true +edition.workspace = true +license.workspace = true +repository.workspace = true +description = "Host-neutral stateful sessions over TinyAgents harness and transcript storage." + +[dependencies] +async-trait = "0.1" +chrono = "0.4" +thiserror = "2" +tinyagents-harness = { path = "../tinyagents-harness", version = "2.1.2" } +tinyagents-session = { path = "../tinyagents-session", version = "2.1.2" } +tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.3.0" } +tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.2.0" } +tokio = { version = "1", features = ["macros", "rt", "sync"] } + +[dev-dependencies] +anyhow = "1" +serde_json = "1" +tempfile = "3" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "sync", "time"] } + +[lints] +workspace = true diff --git a/crates/tinyagents-runtime/README.md b/crates/tinyagents-runtime/README.md new file mode 100644 index 00000000..08d4b5a2 --- /dev/null +++ b/crates/tinyagents-runtime/README.md @@ -0,0 +1,54 @@ +# tinyagents-runtime + +`tinyagents-runtime` provides the stateful session layer that sits between a +host's turn policy and TinyAgents' provider-neutral harness. A `Session` owns +the mutable model history, a stable prefix, a frozen tool declaration snapshot, +and append-only transcript state. It has no host configuration, credentials, +prompt construction, tool authorization, model selection, or event system. + +## Host responsibilities + +The host supplies three narrow seams: + +- `SessionDriver` executes one prepared history snapshot. `HarnessDriver` + adapts an `AgentHarness` and passes the host's explicit `C` through + unchanged. It fails closed unless the harness registry exactly matches the + frozen request tool snapshot. +- `TranscriptCodec` decodes the host's durable dialect and reconciles the + previous durable rows with a model-history transition. The codec, rather + than the runtime, retains fields inference messages cannot express. `C` is + `Clone` so reconciliation receives the current host context plus request, + thread, stream, and resume options after the live `RunContext` moves into the + driver. +- `SessionHooks` prepares a request and validates the candidate before the + commit point, then observes the durable result through `after_commit` and + one terminal state. It does not make policy decisions. + +```rust,no_run +use std::sync::Arc; +use tinyagents_runtime::{SessionBuilder, SessionDriver, TranscriptCodec}; + +# fn build(driver: Arc>, codec: Arc>) { +let session = SessionBuilder::new(driver) + .codec(codec) + .build(); +# let _ = session; +# } +``` + +To persist, add `SessionBuilder::transcript(locator, stem, meta)`. The runtime +uses `tinyagents-session`'s `TranscriptHistory::append_turn_with_partial`, so a normal +extension appends only the new tail and a reduced context writes one compaction +record. A supplied partial driver outcome is represented through that single +history operation: logical history is replayable and interrupted display text +is not. Histories that cannot provide the combined operation reject a partial +rather than risk a two-step write. A persistence failure leaves the session's +in-memory history and persisted snapshot unchanged. + +Every turn receives explicit `TurnOptions`, including its cancellation token +and `RunContext`; no task-local data crosses the runtime boundary. The +stable prefix is reconciled after resume and driver compaction without +duplication. Cancellation before the commit point leaves no durable mutation; +once it succeeds, the turn remains successful. `after_commit` and terminal +hooks get the committed outcome, but their error or a cooperative cancellation +cannot relabel it. diff --git a/crates/tinyagents-runtime/src/builder.rs b/crates/tinyagents-runtime/src/builder.rs new file mode 100644 index 00000000..6259c884 --- /dev/null +++ b/crates/tinyagents-runtime/src/builder.rs @@ -0,0 +1,111 @@ +use std::sync::Arc; + +use tinyagents_session::transcript::{TranscriptLocator, TranscriptMeta}; + +use crate::{ + NoopSessionHooks, PrefixSnapshot, RuntimeError, Session, SessionDriver, SessionHooks, + ToolSnapshot, TranscriptCodec, +}; + +/// Configures a directly-owned, reusable [`Session`]. +pub struct SessionBuilder { + driver: Arc>, + codec: Option>>, + hooks: Arc, + prefix: PrefixSnapshot, + tools: ToolSnapshot, + transcript: Option, +} + +struct TranscriptConfig { + locator: Arc, + stem: String, + meta: TranscriptMeta, +} + +impl SessionBuilder { + /// Starts a builder over an object-safe execution driver. + pub fn new(driver: Arc>) -> Self { + Self { + driver, + codec: None, + hooks: Arc::new(NoopSessionHooks), + prefix: PrefixSnapshot::default(), + tools: ToolSnapshot::default(), + transcript: None, + } + } + + /// Installs the host-owned lossless transcript conversion. + pub fn codec(mut self, codec: Arc>) -> Self { + self.codec = Some(codec); + self + } + + /// Installs optional host preparation/observation hooks. + pub fn hooks(mut self, hooks: Arc) -> Self { + self.hooks = hooks; + self + } + + /// Freezes the prefix used to initialize this session's history. + pub fn prefix(mut self, prefix: PrefixSnapshot) -> Self { + self.prefix = prefix; + self + } + + /// Freezes the tool declarations exposed to each driver invocation. + pub fn tool_snapshot(mut self, tools: ToolSnapshot) -> Self { + self.tools = tools; + self + } + + /// Enables append-only transcript persistence through a session-owned + /// locator, stem, and neutral metadata seed. + pub fn transcript( + mut self, + locator: Arc, + stem: impl Into, + meta: TranscriptMeta, + ) -> Self { + self.transcript = Some(TranscriptConfig { + locator, + stem: stem.into(), + meta, + }); + self + } + + /// Builds a session. A codec is required only when transcript persistence + /// or transcript resume is configured. + pub fn build(self) -> Result, RuntimeError> { + let (locator, stem, meta, history_handle) = if let Some(config) = self.transcript { + let handle = config + .locator + .open_stem(&config.stem, config.meta.clone()) + .map_err(|error| RuntimeError::Persistence(error.to_string()))?; + ( + Some(config.locator), + Some(config.stem), + Some(config.meta), + Some(handle), + ) + } else { + (None, None, None, None) + }; + if history_handle.is_some() && self.codec.is_none() { + return Err(RuntimeError::MissingDependency("TranscriptCodec")); + } + Ok(Session::::new( + self.driver, + self.codec, + self.hooks, + self.prefix, + self.tools, + locator, + stem, + meta, + history_handle, + )) + } +} diff --git a/crates/tinyagents-runtime/src/driver.rs b/crates/tinyagents-runtime/src/driver.rs new file mode 100644 index 00000000..a85a9a26 --- /dev/null +++ b/crates/tinyagents-runtime/src/driver.rs @@ -0,0 +1,138 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use tinyagents_harness::{context::RunContext, runtime::AgentHarness}; +use tinyinference_llm::message::Message; + +use crate::{RuntimeError, ToolSnapshot, TranscriptPartial}; + +/// Immutable values passed from a [`crate::Session`] to one driver invocation. +pub struct DriverRequest { + /// Full model history, including the stable prefix and deduplicated input. + pub history: Vec, + /// Frozen model-visible tool declarations selected by the host. + pub tools: ToolSnapshot, + /// Explicit live run context for this invocation. + pub run_context: RunContext, + /// Requests the driver's streaming execution path. + pub stream: bool, +} + +/// A driver's complete or partial model result. +#[derive(Clone, Debug, PartialEq)] +pub struct DriverOutcome { + /// Complete history accumulated by the harness/driver. + pub history: Vec, + /// Final visible output, when available. + pub output: Option, + /// Display-only partial text when execution was interrupted. + pub partial: Option, + /// Whether the driver ended at an interruptible point. + pub interrupted: bool, +} + +/// A driver error which may retain model history and display-only partial text. +#[derive(Clone, Debug)] +pub struct DriverFailure { + /// The reason execution failed. + pub error: RuntimeError, + /// Work produced before failure, safe for partial transcript persistence. + pub partial: Option, +} + +/// Object-safe model/tool-loop invocation boundary. +/// +/// The generic `AgentHarness` cannot be stored directly in a +/// reusable `Session` without leaking a host's state type. Implement this seam +/// directly, or use [`HarnessDriver`] for `AgentHarness`. +#[async_trait] +pub trait SessionDriver: Send + Sync { + /// Executes one turn from an immutable runtime snapshot. + async fn execute(&self, request: DriverRequest) -> Result; +} + +/// Adapter for the ordinary, explicit-model TinyAgents harness entry points. +/// +/// Tool selection remains outside this adapter: callers configure the harness +/// from the same already-authorized [`ToolSnapshot`] that they expose here. +pub struct HarnessDriver { + harness: Arc>, + state: Arc, +} + +impl HarnessDriver { + /// Binds reusable harness mechanics and immutable host state to a driver. + pub fn new(harness: Arc>, state: Arc) -> Self { + Self { harness, state } + } +} + +#[async_trait] +impl SessionDriver + for HarnessDriver +{ + async fn execute(&self, request: DriverRequest) -> Result { + // The harness has no per-invocation registry override. Executing when + // its provider-visible declarations differ from the frozen request + // would advertise or dispatch capabilities the session did not grant. + if !same_tools( + &self.harness.tools().declared_specs(), + request.tools.specs(), + ) { + return Err(DriverFailure { + error: RuntimeError::ToolSnapshotMismatch, + partial: None, + }); + } + let partial = if request.stream { + self.harness + .invoke_streaming_in_context_collecting_partial( + self.state.as_ref(), + request.run_context, + request.history, + ) + .await + } else { + self.harness + .invoke_in_context_collecting_partial( + self.state.as_ref(), + request.run_context, + request.history, + ) + .await + }; + let output = + partial.run.messages.iter().rev().find_map(|message| { + matches!(message, Message::Assistant(_)).then(|| message.text()) + }); + let outcome = DriverOutcome { + history: partial.run.messages, + output: output.clone(), + // The pinned harness retains partial model history but exposes no + // separate streamed text delta, reasoning, or iteration here. The + // accumulated final assistant text is the only display-safe value. + partial: partial + .error + .as_ref() + .and(output) + .map(TranscriptPartial::new), + interrupted: partial.run.paused.is_some(), + }; + match partial.error { + Some(error) => Err(DriverFailure { + error: RuntimeError::Driver(error.to_string()), + partial: Some(outcome), + }), + None => Ok(outcome), + } + } +} + +fn same_tools(left: &[tinytools::ToolSpec], right: &[tinytools::ToolSpec]) -> bool { + left.len() == right.len() + && left.iter().zip(right).all(|(left, right)| { + left.name == right.name + && left.description == right.description + && left.parameters == right.parameters + }) +} diff --git a/crates/tinyagents-runtime/src/error.rs b/crates/tinyagents-runtime/src/error.rs new file mode 100644 index 00000000..9831d3c7 --- /dev/null +++ b/crates/tinyagents-runtime/src/error.rs @@ -0,0 +1,28 @@ +/// Errors emitted by the host-neutral session runtime. +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum RuntimeError { + /// A caller cancelled the turn at a runtime await boundary. + #[error("session turn cancelled")] + Cancelled, + /// A driver could not invoke the model/tool loop. + #[error("driver failed: {0}")] + Driver(String), + /// Transcript decoding or encoding failed. + #[error("transcript codec failed: {0}")] + Codec(String), + /// Durable transcript persistence failed. + #[error("transcript persistence failed: {0}")] + Persistence(String), + /// A host lifecycle hook failed. + #[error("session hook failed: {0}")] + Hook(String), + /// The host supplied two distinct tool declarations with one name. + #[error("tool snapshot has conflicting declarations for `{0}`")] + ToolNameCollision(String), + /// A harness cannot bind its registered tools to this turn's frozen snapshot. + #[error("harness tools do not match the session tool snapshot")] + ToolSnapshotMismatch, + /// The builder did not receive a required dependency. + #[error("session builder requires {0}")] + MissingDependency(&'static str), +} diff --git a/crates/tinyagents-runtime/src/hooks.rs b/crates/tinyagents-runtime/src/hooks.rs new file mode 100644 index 00000000..fd182ccb --- /dev/null +++ b/crates/tinyagents-runtime/src/hooks.rs @@ -0,0 +1,46 @@ +use async_trait::async_trait; + +use crate::{RuntimeError, SessionTerminal, SessionTurnOutcome, SessionTurnRequest}; + +/// Host observation/preparation around a session turn. +/// +/// Hooks do not grant tools, select models, compose product prompts, or own +/// transcript state. They can prepare the input and observe committed results. +#[async_trait] +pub trait SessionHooks: Send + Sync { + /// Runs before the driver sees the request. + async fn before_turn(&self, request: &mut SessionTurnRequest) -> Result<(), RuntimeError>; + /// Runs after the driver has produced a candidate and before it commits. + /// Returning an error or observing cancellation therefore leaves no + /// durable session mutation behind. + async fn after_turn(&self, outcome: &SessionTurnOutcome) -> Result<(), RuntimeError>; + /// Runs exactly once after a successful durable transcript commit. + /// + /// Errors and cancellation observed here are deliberately observational: + /// the result has already become durable and remains successful. + async fn after_commit(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + Ok(()) + } + /// Runs exactly once for every terminal turn result. + async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError>; +} + +/// A no-op hook set for hosts that need no lifecycle observation. +#[derive(Default)] +pub struct NoopSessionHooks; + +#[async_trait] +impl SessionHooks for NoopSessionHooks { + async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { + Ok(()) + } + async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + Ok(()) + } + async fn after_commit(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, _: &SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } +} diff --git a/crates/tinyagents-runtime/src/lib.rs b/crates/tinyagents-runtime/src/lib.rs new file mode 100644 index 00000000..7dbcf481 --- /dev/null +++ b/crates/tinyagents-runtime/src/lib.rs @@ -0,0 +1,58 @@ +//! Stateful, host-neutral sessions for TinyAgents. +//! +//! A [`Session`] owns conversation history and its durable transcript view. +//! Hosts supply the model/tool execution through [`SessionDriver`] and own all +//! policy, prompt composition, model selection, and authorization. This +//! crate deliberately has no host configuration or product dependencies. + +mod builder; +mod driver; +mod error; +mod hooks; +mod prefix; +mod session; +mod tools; +mod types; + +pub use builder::SessionBuilder; +pub use driver::{DriverFailure, DriverOutcome, DriverRequest, HarnessDriver, SessionDriver}; +pub use error::RuntimeError; +pub use hooks::{NoopSessionHooks, SessionHooks}; +pub use prefix::PrefixSnapshot; +pub use session::Session; +pub use tinyagents_session::transcript::TranscriptPartial; +pub use tools::ToolSnapshot; +pub use types::{ + ResumeMode, SessionResume, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, + TranscriptTurnOptions, TurnOptions, +}; + +/// Converts between a host's lossless durable transcript dialect and the +/// inference messages driven by a [`SessionDriver`]. +/// +/// It is object-safe so a host can keep its conversion and metadata ownership +/// outside this reusable runtime. The runtime never narrows transcript data +/// itself; a codec must explicitly decide how every field is represented. +pub trait TranscriptCodec: Send + Sync { + /// Decodes a durable transcript for model execution. + fn decode_history( + &self, + transcript: &tinyagents_session::transcript::SessionTranscript, + ) -> Result, RuntimeError>; + + /// Reconciles a model-history transition with the prior lossless durable + /// rows. `prior` must be treated as authoritative for fields absent from + /// inference messages (provider metadata, raw arguments, reasoning, ids, + /// and host extensions). The returned rows are the complete next logical + /// durable set; the history layer writes its delta atomically. + fn reconcile( + &self, + prior: &[tinyagents_session::transcript::TranscriptMessage], + previous: &[tinyinference_llm::message::Message], + next: &[tinyinference_llm::message::Message], + options: &TranscriptTurnOptions, + ) -> Result, RuntimeError>; +} + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-runtime/src/prefix.rs b/crates/tinyagents-runtime/src/prefix.rs new file mode 100644 index 00000000..6114eaa5 --- /dev/null +++ b/crates/tinyagents-runtime/src/prefix.rs @@ -0,0 +1,19 @@ +use tinyinference_llm::message::Message; + +/// Immutable messages which remain at the front of a session's history. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct PrefixSnapshot { + messages: Vec, +} + +impl PrefixSnapshot { + /// Captures the prefix once, before turn history starts growing. + pub fn new(messages: Vec) -> Self { + Self { messages } + } + + /// Returns the captured messages in their original order. + pub fn messages(&self) -> &[Message] { + &self.messages + } +} diff --git a/crates/tinyagents-runtime/src/session.rs b/crates/tinyagents-runtime/src/session.rs new file mode 100644 index 00000000..10df4124 --- /dev/null +++ b/crates/tinyagents-runtime/src/session.rs @@ -0,0 +1,359 @@ +use std::{future::Future, sync::Arc}; + +use tinyagents_harness::CancellationToken; +use tinyagents_session::transcript::{ + TranscriptHistory, TranscriptLocator, TranscriptMeta, TranscriptPartial, TranscriptTurn, +}; +use tinyinference_llm::message::Message; + +use crate::{ + DriverRequest, PrefixSnapshot, ResumeMode, RuntimeError, SessionDriver, SessionHooks, + SessionResume, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, ToolSnapshot, + TranscriptCodec, TurnOptions, +}; + +/// Host-neutral mutable state for one conversation session. +pub struct Session { + driver: Arc>, + codec: Option>>, + hooks: Arc, + prefix: PrefixSnapshot, + tools: ToolSnapshot, + history: Vec, + persisted: Vec, + locator: Option>, + stem: Option, + meta: Option, + transcript: Option>, +} + +impl Session { + #[allow(clippy::too_many_arguments)] + pub(crate) fn new( + driver: Arc>, + codec: Option>>, + hooks: Arc, + prefix: PrefixSnapshot, + tools: ToolSnapshot, + locator: Option>, + stem: Option, + meta: Option, + transcript: Option>, + ) -> Self { + Self { + driver, + codec, + hooks, + history: prefix.messages().to_vec(), + prefix, + tools, + persisted: Vec::new(), + locator, + stem, + meta, + transcript, + } + } + + /// Returns the currently committed model history. + pub fn history(&self) -> &[Message] { + &self.history + } + + /// Returns the unchanging history prefix captured by the builder. + pub fn prefix_snapshot(&self) -> &PrefixSnapshot { + &self.prefix + } + + /// Returns the immutable model-visible tool declaration set. + pub fn tool_snapshot(&self) -> &ToolSnapshot { + &self.tools + } + + /// Loads the selected durable transcript, retaining its lossless raw rows + /// as the base for the next append-only delta. + pub async fn resume( + &mut self, + options: &TurnOptions, + ) -> Result { + if options.cancellation.is_cancelled() { + return Err(RuntimeError::Cancelled); + } + let Some(locator) = self.locator.as_ref() else { + return Ok(SessionResume { + loaded: false, + history: self.history.clone(), + }); + }; + let read = match options.resume { + ResumeMode::Never => None, + ResumeMode::LatestForAgent => self + .stem + .as_deref() + .and_then(|stem| locator.latest_for_agent(stem)), + ResumeMode::Thread => options + .thread_id + .as_deref() + .and_then(|thread| locator.root_for_thread(thread)), + }; + let Some(read) = read else { + return Ok(SessionResume { + loaded: false, + history: self.history.clone(), + }); + }; + let Some(transcript) = read + .read_session() + .map_err(|error| RuntimeError::Persistence(error.to_string()))? + else { + return Ok(SessionResume { + loaded: false, + history: self.history.clone(), + }); + }; + let codec = self + .codec + .as_ref() + .ok_or(RuntimeError::MissingDependency("TranscriptCodec"))?; + let history = self.with_prefix(codec.decode_history(&transcript)?); + self.history = history.clone(); + self.persisted = transcript.messages; + Ok(SessionResume { + loaded: true, + history, + }) + } + + /// Executes and durably commits one state transition. + pub async fn turn( + &mut self, + mut request: SessionTurnRequest, + options: TurnOptions, + ) -> Result { + let mut terminal_guard = TerminalGuard::new(self.hooks.clone()); + let result = self + .turn_inner(&mut request, options, &mut terminal_guard) + .await; + if !terminal_guard.is_committed() { + let terminal = match &result { + Ok(outcome) => SessionTerminal::Completed(outcome.clone()), + Err(RuntimeError::Cancelled) => SessionTerminal::Cancelled, + Err(error) => SessionTerminal::Failed(error.to_string()), + }; + terminal_guard.set(terminal); + } + // Terminal observation cannot revoke a successful durable commit. + // `finish` still schedules it exactly once; hook failures are + // deliberately observational rather than a second terminal result. + let _ = terminal_guard.finish().await; + result + } + + async fn turn_inner( + &mut self, + request: &mut SessionTurnRequest, + options: TurnOptions, + terminal_guard: &mut TerminalGuard, + ) -> Result { + if options.resume != ResumeMode::Never { + self.resume(&options).await?; + } + cancelable(&options.cancellation, self.hooks.before_turn(request)).await?; + let mut input = self.history.clone(); + if input.last() != Some(&request.input) { + input.push(request.input.clone()); + } + // `RunContext` is intentionally consumed exactly once. There is no + // task-local fallback: the host context selected for this turn is what + // reaches model, middleware, and tool execution. + let codec_options = options.transcript_options(); + let TurnOptions { + request_id, + thread_id, + stream, + cancellation, + run_context, + .. + } = options; + let run_context = run_context.with_cancellation(cancellation.clone()); + let driver_result = tokio::select! { + _ = cancellation.cancelled() => return Err(RuntimeError::Cancelled), + result = self.driver.execute(DriverRequest { + history: input, + tools: self.tools.clone(), + run_context, + stream, + }) => result, + }; + let outcome = match driver_result { + Ok(outcome) => outcome, + Err(failure) => { + if let Some(partial) = failure.partial { + let partial_history = self.with_prefix(partial.history); + let raw = self.encode(&self.history, &partial_history, &codec_options)?; + // `append_turn` is the only durable mutation. Do not + // append display partials first: a later append failure + // would leave an unreportable half-commit on disk. + self.persist( + &raw, + request_id.as_deref(), + thread_id.as_deref(), + partial.partial.as_ref(), + )?; + self.history = partial_history; + self.persisted = raw; + } + return Err(failure.error); + } + }; + let candidate = self.with_prefix(outcome.history); + // `after_turn` is a pre-commit hook. It can reject or be cancelled + // without any durable mutation; after `persist` returns success this + // turn is committed and cancellation can no longer change its result. + let committed = SessionTurnOutcome { + history: candidate.clone(), + output: outcome.output, + interrupted: outcome.interrupted, + }; + cancelable(&cancellation, self.hooks.after_turn(&committed)).await?; + if cancellation.is_cancelled() { + return Err(RuntimeError::Cancelled); + } + let raw = self.encode(&self.history, &candidate, &codec_options)?; + self.persist(&raw, request_id.as_deref(), thread_id.as_deref(), None)?; + self.history = committed.history.clone(); + self.persisted = raw; + // Set the truthful durable terminal before invoking an observational + // finalizer. If the caller drops this future while it is running, the + // guard's Drop implementation still reports the completed commit. + terminal_guard.mark_committed(committed.clone()); + // This runs after `append_turn_with_partial` has made the logical + // transition durable. Failure or cooperative cancellation in a host + // finalizer is observational: it cannot relabel that committed turn. + let hooks = self.hooks.clone(); + let finalization = committed.clone(); + let _ = tokio::spawn(async move { hooks.after_commit(&finalization).await }).await; + Ok(committed) + } + + fn encode( + &self, + previous: &[Message], + next: &[Message], + options: &crate::TranscriptTurnOptions, + ) -> Result, RuntimeError> { + match &self.codec { + Some(codec) => codec.reconcile(&self.persisted, previous, next, options), + None => Ok(Vec::new()), + } + } + + fn persist( + &mut self, + raw: &[tinyagents_session::transcript::TranscriptMessage], + request_id: Option<&str>, + thread_id: Option<&str>, + partial: Option<&TranscriptPartial>, + ) -> Result<(), RuntimeError> { + let (Some(transcript), Some(meta)) = (&self.transcript, &self.meta) else { + return Ok(()); + }; + let mut meta = meta.clone(); + meta.turn_count += 1; + meta.updated = chrono::Utc::now().to_rfc3339(); + meta.thread_id = thread_id.map(str::to_owned).or(meta.thread_id); + transcript + .append_turn_with_partial( + TranscriptTurn { + prev: &self.persisted, + next: raw, + meta: &meta, + turn_usage: None, + request_id, + }, + partial, + ) + .map_err(|error| RuntimeError::Persistence(error.to_string()))?; + self.meta = Some(meta); + Ok(()) + } + + fn with_prefix(&self, history: Vec) -> Vec { + let prefix = self.prefix.messages(); + let overlap = (0..=prefix.len().min(history.len())) + .rev() + .find(|&len| prefix[prefix.len() - len..] == history[..len]) + .unwrap_or_default(); + let mut reconciled = prefix[..prefix.len() - overlap].to_vec(); + reconciled.extend(history); + reconciled + } +} + +/// Ensures a terminal hook is scheduled once even if a caller drops a turn +/// future while it is awaiting preparation, driving, persistence, or hooks. +struct TerminalGuard { + hooks: Arc, + terminal: Option, + committed: bool, +} + +impl TerminalGuard { + fn new(hooks: Arc) -> Self { + Self { + hooks, + terminal: Some(SessionTerminal::Failed("session turn dropped".into())), + committed: false, + } + } + + fn set(&mut self, terminal: SessionTerminal) { + self.terminal = Some(terminal); + } + + fn mark_committed(&mut self, outcome: SessionTurnOutcome) { + self.terminal = Some(SessionTerminal::Completed(outcome)); + self.committed = true; + } + + fn is_committed(&self) -> bool { + self.committed + } + + async fn finish(mut self) -> Result<(), RuntimeError> { + let terminal = self.terminal.take().ok_or(RuntimeError::Hook( + "terminal guard already completed".into(), + ))?; + let hooks = self.hooks.clone(); + tokio::spawn(async move { hooks.on_terminal(&terminal).await }) + .await + .map_err(|error| RuntimeError::Hook(format!("terminal hook task failed: {error}")))? + } +} + +impl Drop for TerminalGuard { + fn drop(&mut self) { + let Some(terminal) = self.terminal.take() else { + return; + }; + let hooks = self.hooks.clone(); + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + runtime.spawn(async move { + let _ = hooks.on_terminal(&terminal).await; + }); + } + } +} + +async fn cancelable( + cancellation: &CancellationToken, + future: impl Future>, +) -> Result { + if cancellation.is_cancelled() { + return Err(RuntimeError::Cancelled); + } + tokio::select! { + _ = cancellation.cancelled() => Err(RuntimeError::Cancelled), + result = future => result, + } +} diff --git a/crates/tinyagents-runtime/src/test.rs b/crates/tinyagents-runtime/src/test.rs new file mode 100644 index 00000000..54c3bcca --- /dev/null +++ b/crates/tinyagents-runtime/src/test.rs @@ -0,0 +1,1345 @@ +#![allow(clippy::expect_used, clippy::unwrap_used)] + +use std::{ + collections::VecDeque, + path::{Path, PathBuf}, + sync::{Arc, Mutex}, +}; + +use async_trait::async_trait; +use tinyagents_harness::{ + context::{RunConfig, RunContext}, + runtime::AgentHarness, +}; +use tinyagents_session::transcript::{ + DisplayRecord, FileTranscriptHistory, FileTranscriptLocator, SessionTranscript, + TranscriptHistory, TranscriptLocator, TranscriptMessage, TranscriptMeta, TranscriptRead, + TranscriptTurn, read_transcript, read_transcript_display, +}; +use tinyinference_llm::message::Message; +use tinyinference_llm::providers::MockModel; +use tinytools::{Tool, ToolResult, ToolSpec}; + +use crate::{ + DriverFailure, DriverOutcome, DriverRequest, HarnessDriver, PrefixSnapshot, ResumeMode, + RuntimeError, SessionBuilder, SessionDriver, SessionHooks, SessionTerminal, SessionTurnOutcome, + SessionTurnRequest, ToolSnapshot, TranscriptCodec, TurnOptions, +}; + +struct FakeDriver { + results: Mutex>>, + requests: Mutex>, +} + +impl FakeDriver { + fn new(results: Vec>) -> Self { + Self { + results: Mutex::new(results.into()), + requests: Mutex::new(Vec::new()), + } + } +} + +#[async_trait] +impl SessionDriver for FakeDriver { + async fn execute(&self, request: DriverRequest) -> Result { + self.requests.lock().unwrap().push(request); + self.results + .lock() + .unwrap() + .pop_front() + .expect("planned driver result") + } +} + +struct WaitingDriver; + +#[async_trait] +impl SessionDriver for WaitingDriver { + async fn execute(&self, _: DriverRequest) -> Result { + std::future::pending().await + } +} + +struct DropDriver { + started: Arc, +} + +#[async_trait] +impl SessionDriver for DropDriver { + async fn execute(&self, _: DriverRequest) -> Result { + self.started.notify_waiters(); + std::future::pending().await + } +} + +// Reconciliation retains an explicit clone for the codec after the live +// `RunContext` is consumed by the driver. +#[derive(Clone)] +struct HostContext(String); + +struct ContextDriver; + +#[async_trait] +impl SessionDriver for ContextDriver { + async fn execute( + &self, + request: DriverRequest, + ) -> Result { + Ok(outcome(vec![Message::assistant( + request.run_context.data.0, + )])) + } +} + +struct RegisteredTool; + +#[async_trait] +impl Tool for RegisteredTool { + fn name(&self) -> &str { + "registered" + } + + fn description(&self) -> &str { + "a registered matching tool" + } + + fn parameters_schema(&self) -> serde_json::Value { + serde_json::json!({"type": "object"}) + } + + async fn execute(&self, _: serde_json::Value) -> anyhow::Result { + Ok(ToolResult::success("ok")) + } +} + +#[derive(Default)] +struct BasicCodec { + decoded: Mutex>, +} + +impl TranscriptCodec for BasicCodec { + fn decode_history(&self, transcript: &SessionTranscript) -> Result, RuntimeError> { + self.decoded.lock().unwrap().push(transcript.clone()); + Ok(transcript + .messages + .iter() + .map(|message| match message.role.as_str() { + "assistant" => Message::assistant(&message.content), + "system" => Message::system(&message.content), + _ => Message::user(&message.content), + }) + .collect()) + } + + fn reconcile( + &self, + prior: &[TranscriptMessage], + previous: &[Message], + next: &[Message], + _: &crate::TranscriptTurnOptions, + ) -> Result, RuntimeError> { + // Preserve rows that still correspond to the prior model prefix; + // fresh model suffixes receive only the codec's explicit projection. + let raw_offset = (!prior.is_empty()) + .then(|| { + previous.windows(prior.len()).position(|window| { + window + .iter() + .zip(prior) + .all(|(model, row)| model.text() == row.content) + }) + }) + .flatten() + .unwrap_or(usize::MAX); + Ok(next + .iter() + .enumerate() + .map(|(index, message)| { + let raw_index = index.checked_sub(raw_offset); + if let Some(raw_index) = raw_index.filter(|index| { + *index < prior.len() && previous.get(index + raw_offset) == Some(message) + }) { + return prior[raw_index].clone(); + } + let role = match message { + Message::System(_) => "system", + Message::User(_) => "user", + Message::Assistant(_) => "assistant", + Message::Tool(_) => "tool", + }; + TranscriptMessage::new(role, message.text()) + }) + .collect()) + } +} + +#[derive(Default)] +struct RecordingHooks { + events: Mutex>, + fail_before: bool, + fail_after: bool, + terminal_notified: Option>, +} + +struct WaitingAfterHooks { + started: tokio::sync::Notify, + events: Mutex>, +} + +struct WaitingBeforeHooks { + started: tokio::sync::Notify, +} + +struct WaitingPostCommitHooks { + started: tokio::sync::Notify, + terminal: tokio::sync::Notify, + terminals: Mutex>, +} + +#[async_trait] +impl SessionHooks for WaitingBeforeHooks { + async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { + self.started.notify_waiters(); + std::future::pending().await + } + + async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn on_terminal(&self, _: &SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } +} + +#[async_trait] +impl SessionHooks for WaitingAfterHooks { + async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + self.started.notify_waiters(); + std::future::pending().await + } + + async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { + self.events.lock().unwrap().push(format!("{terminal:?}")); + Ok(()) + } +} + +#[async_trait] +impl SessionHooks for WaitingPostCommitHooks { + async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn after_commit(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + self.started.notify_waiters(); + std::future::pending().await + } + + async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { + self.terminals.lock().unwrap().push(terminal.clone()); + self.terminal.notify_waiters(); + Ok(()) + } +} + +#[async_trait] +impl SessionHooks for RecordingHooks { + async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { + self.events.lock().unwrap().push("before".into()); + if self.fail_before { + Err(RuntimeError::Hook("before".into())) + } else { + Ok(()) + } + } + async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + self.events.lock().unwrap().push("after".into()); + if self.fail_after { + Err(RuntimeError::Hook("after".into())) + } else { + Ok(()) + } + } + async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { + let terminal = match terminal { + SessionTerminal::Completed(_) => "Completed".to_owned(), + other => format!("{other:?}"), + }; + self.events + .lock() + .unwrap() + .push(format!("terminal:{terminal}")); + if let Some(notify) = &self.terminal_notified { + notify.notify_waiters(); + } + Ok(()) + } +} + +#[derive(Default)] +struct FinalizationHooks { + post_commits: Mutex>, + terminals: Mutex>, + fail_post_commit: bool, + cancel_on_post_commit: Option, +} + +#[async_trait] +impl SessionHooks for FinalizationHooks { + async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn after_commit(&self, outcome: &SessionTurnOutcome) -> Result<(), RuntimeError> { + self.post_commits.lock().unwrap().push(outcome.clone()); + if let Some(cancellation) = &self.cancel_on_post_commit { + cancellation.cancel(); + } + if self.fail_post_commit { + Err(RuntimeError::Hook("post-commit".into())) + } else { + Ok(()) + } + } + + async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { + self.terminals.lock().unwrap().push(terminal.clone()); + Ok(()) + } +} + +struct MemoryHistory { + path: PathBuf, + session: Mutex>, + turns: Mutex>>, + fail: bool, + cancel_after_append: Mutex>, +} + +impl TranscriptRead for MemoryHistory { + fn path(&self) -> &Path { + &self.path + } + fn read_session(&self) -> anyhow::Result> { + Ok(self.session.lock().unwrap().clone()) + } +} + +impl TranscriptHistory for MemoryHistory { + fn append_turn(&self, turn: TranscriptTurn<'_>) -> anyhow::Result<()> { + if self.fail { + anyhow::bail!("planned persistence failure"); + } + self.turns.lock().unwrap().push(turn.next.to_vec()); + *self.session.lock().unwrap() = Some(SessionTranscript { + meta: turn.meta.clone(), + messages: turn.next.to_vec(), + }); + if let Some(cancellation) = self.cancel_after_append.lock().unwrap().as_ref() { + cancellation.cancel(); + } + Ok(()) + } + fn messages(&self) -> anyhow::Result> { + Ok(self + .session + .lock() + .unwrap() + .as_ref() + .map(|value| value.messages.clone()) + .unwrap_or_default()) + } + fn append(&self, _: TranscriptMessage) -> anyhow::Result<()> { + Ok(()) + } + fn replace(&self, _: &[TranscriptMessage]) -> anyhow::Result<()> { + Ok(()) + } + fn clear(&self) -> anyhow::Result<()> { + Ok(()) + } +} + +struct MemoryLocator { + history: Arc, +} + +/// A real transcript path whose single atomic turn operation fails before it +/// writes. This catches the old two-write partial path: if runtime wrote a +/// display partial first, the path below would exist after the failure. +struct FailingFileHistory(FileTranscriptHistory); + +impl TranscriptRead for FailingFileHistory { + fn path(&self) -> &Path { + self.0.path() + } + + fn read_session(&self) -> anyhow::Result> { + TranscriptRead::read_session(&self.0) + } +} + +impl TranscriptHistory for FailingFileHistory { + fn append_turn(&self, _: TranscriptTurn<'_>) -> anyhow::Result<()> { + anyhow::bail!("planned atomic append failure") + } + + fn messages(&self) -> anyhow::Result> { + TranscriptHistory::messages(&self.0) + } + + fn append(&self, _: TranscriptMessage) -> anyhow::Result<()> { + anyhow::bail!("planned atomic append failure") + } + + fn replace(&self, _: &[TranscriptMessage]) -> anyhow::Result<()> { + anyhow::bail!("planned atomic append failure") + } + + fn clear(&self) -> anyhow::Result<()> { + anyhow::bail!("planned atomic append failure") + } +} + +struct FailingFileLocator { + history: Arc, +} + +impl TranscriptLocator for FailingFileLocator { + fn latest_for_agent(&self, _: &str) -> Option> { + Some(self.history.clone()) + } + + fn root_for_thread(&self, _: &str) -> Option> { + Some(self.history.clone()) + } + + fn open_stem(&self, _: &str, _: TranscriptMeta) -> anyhow::Result> { + Ok(self.history.clone()) + } +} + +impl TranscriptLocator for MemoryLocator { + fn latest_for_agent(&self, _: &str) -> Option> { + Some(self.history.clone()) + } + fn root_for_thread(&self, _: &str) -> Option> { + Some(self.history.clone()) + } + fn open_stem(&self, _: &str, _: TranscriptMeta) -> anyhow::Result> { + Ok(self.history.clone()) + } +} + +fn meta() -> TranscriptMeta { + TranscriptMeta { + agent_name: "agent".into(), + agent_id: None, + agent_type: None, + dispatcher: "test".into(), + provider: None, + model: None, + created: "now".into(), + updated: "now".into(), + turn_count: 0, + input_tokens: 0, + output_tokens: 0, + cached_input_tokens: 0, + charged_amount_usd: 0.0, + thread_id: None, + task_id: None, + } +} + +fn outcome(history: Vec) -> DriverOutcome { + DriverOutcome { + history, + output: Some("done".into()), + partial: None, + interrupted: false, + } +} + +fn memory_locator( + session: Option, + fail: bool, +) -> (Arc, Arc) { + let dir = tempfile::tempdir().unwrap().keep(); + let history = Arc::new(MemoryHistory { + path: dir.join("session.jsonl"), + session: Mutex::new(session), + turns: Mutex::new(Vec::new()), + fail, + cancel_after_append: Mutex::new(None), + }); + ( + Arc::new(MemoryLocator { + history: history.clone(), + }), + history, + ) +} + +#[tokio::test] +async fn first_turn_commits_history_and_preserves_prefix() { + let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::system("stable"), + Message::user("hi"), + Message::assistant("hello"), + ]))])); + let mut session = SessionBuilder::new(driver.clone()) + .prefix(PrefixSnapshot::new(vec![Message::system("stable")])) + .build() + .unwrap(); + let committed = session + .turn( + SessionTurnRequest::new(Message::user("hi")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!(committed.history, session.history()); + assert_eq!( + session.prefix_snapshot().messages(), + &[Message::system("stable")] + ); + assert_eq!( + driver.requests.lock().unwrap()[0].history, + vec![Message::system("stable"), Message::user("hi")] + ); +} + +#[tokio::test] +async fn generic_session_passes_the_explicit_host_context_to_driver() { + let mut session = SessionBuilder::::new(Arc::new(ContextDriver)) + .build() + .unwrap(); + let options = TurnOptions { + request_id: None, + thread_id: None, + stream: false, + resume: ResumeMode::Never, + cancellation: tinyagents_harness::CancellationToken::new(), + run_context: RunContext::new(RunConfig::new("host"), HostContext("host context".into())), + }; + let result = session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await + .unwrap(); + assert_eq!( + result.history.last().map(Message::text).as_deref(), + Some("host context") + ); +} + +#[tokio::test] +async fn harness_driver_uses_the_pinned_explicit_model_entry_point() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("from harness"))); + harness.register_tool(Arc::new(RegisteredTool)); + let snapshot = ToolSnapshot::new(harness.tools().declared_specs()).unwrap(); + let driver = Arc::new(HarnessDriver::new(Arc::new(harness), Arc::new(()))); + let mut session = SessionBuilder::new(driver) + .tool_snapshot(snapshot) + .build() + .unwrap(); + let result = session + .turn( + SessionTurnRequest::new(Message::user("hello")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!(result.output.as_deref(), Some("from harness")); + assert_eq!( + result.history.last().map(Message::text).as_deref(), + Some("from harness") + ); +} + +#[tokio::test] +async fn trailing_input_is_deduplicated_and_tool_snapshot_is_immutable() { + let driver = Arc::new(FakeDriver::new(vec![ + Ok(outcome(vec![Message::user("same")])), + Ok(outcome(vec![ + Message::user("same"), + Message::assistant("two"), + ])), + ])); + let tools = ToolSnapshot::new(vec![ToolSpec { + name: "echo".into(), + description: "x".into(), + parameters: serde_json::json!({}), + }]) + .unwrap(); + let mut session = SessionBuilder::new(driver.clone()) + .tool_snapshot(tools) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("same")), + TurnOptions::default(), + ) + .await + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("same")), + TurnOptions::default(), + ) + .await + .unwrap(); + let requests = driver.requests.lock().unwrap(); + assert_eq!(requests[1].history, vec![Message::user("same")]); + assert_eq!(requests[0].tools.specs()[0].name, "echo"); + assert_eq!(session.tool_snapshot().specs()[0].name, "echo"); +} + +#[test] +fn tool_snapshots_dedup_identical_names_and_reject_collisions() { + let spec = ToolSpec { + name: "read".into(), + description: "read".into(), + parameters: serde_json::json!({}), + }; + assert_eq!( + ToolSnapshot::new(vec![spec.clone(), spec]) + .unwrap() + .specs() + .len(), + 1 + ); + assert!(matches!( + ToolSnapshot::new(vec![ + ToolSpec { + name: "read".into(), + description: "one".into(), + parameters: serde_json::json!({}) + }, + ToolSpec { + name: "read".into(), + description: "two".into(), + parameters: serde_json::json!({}) + }, + ]), + Err(RuntimeError::ToolNameCollision(name)) if name == "read" + )); +} + +#[tokio::test] +async fn resume_passes_full_durable_transcript_to_codec() { + let mut durable = TranscriptMessage::new("user", "persisted"); + durable.extra_metadata = Some(serde_json::json!({"unmodified": true})); + let transcript = SessionTranscript { + meta: meta(), + messages: vec![durable.clone()], + }; + let (locator, _) = memory_locator(Some(transcript), false); + let codec = Arc::new(BasicCodec::default()); + let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![]))) + .codec(codec.clone()) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + let resumed = session + .resume(&TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }) + .await + .unwrap(); + assert!(resumed.loaded); + assert_eq!(resumed.history, vec![Message::user("persisted")]); + assert_eq!( + codec.decoded.lock().unwrap()[0].messages[0].extra_metadata, + durable.extra_metadata + ); +} + +#[tokio::test] +async fn append_only_delta_and_failure_rollback_are_owned_by_session() { + let (locator, history) = memory_locator(None, false); + let codec = Arc::new(BasicCodec::default()); + let driver = Arc::new(FakeDriver::new(vec![ + Ok(outcome(vec![Message::user("one"), Message::assistant("a")])), + Ok(outcome(vec![Message::assistant("compacted")])), + ])); + let mut session = SessionBuilder::new(driver) + .codec(codec) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("one")), + TurnOptions::default(), + ) + .await + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("two")), + TurnOptions::default(), + ) + .await + .unwrap(); + { + let turns = history.turns.lock().unwrap(); + assert_eq!(turns.len(), 2); + assert_eq!( + turns[1], + vec![TranscriptMessage::new("assistant", "compacted")] + ); + } + + let (bad_locator, _) = memory_locator(None, true); + let bad = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![Message::user( + "will-not-commit", + )]))])); + let mut failing = SessionBuilder::new(bad) + .codec(Arc::new(BasicCodec::default())) + .transcript(bad_locator, "bad", meta()) + .build() + .unwrap(); + assert!(matches!( + failing + .turn( + SessionTurnRequest::new(Message::user("will-not-commit")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::Persistence(_)) + )); + assert!(failing.history().is_empty()); +} + +#[tokio::test] +async fn partial_failure_never_leaves_a_display_partial_on_disk() { + let directory = tempfile::tempdir().unwrap(); + let history = Arc::new(FailingFileHistory( + FileTranscriptHistory::new(directory.path(), "agent", meta()).unwrap(), + )); + let path = history.path().to_path_buf(); + let locator = Arc::new(FailingFileLocator { + history: history.clone(), + }); + let driver = Arc::new(FakeDriver::new(vec![Err(DriverFailure { + error: RuntimeError::Driver("interrupted".into()), + partial: Some(DriverOutcome { + history: vec![Message::assistant("partial model history")], + output: None, + partial: Some(crate::TranscriptPartial::new("display partial")), + interrupted: true, + }), + })])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(BasicCodec::default())) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + assert!(matches!( + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::Persistence(_)) + )); + assert!(!path.exists(), "failed partial turn wrote {path:?}"); + assert!(session.history().is_empty()); +} + +#[tokio::test] +async fn partial_failure_commits_model_history_and_display_partial_together() { + let directory = tempfile::tempdir().unwrap(); + let locator = Arc::new(FileTranscriptLocator::new(directory.path())); + let driver = Arc::new(FakeDriver::new(vec![Err(DriverFailure { + error: RuntimeError::Driver("interrupted".into()), + partial: Some(DriverOutcome { + history: vec![Message::assistant("recoverable model history")], + output: None, + partial: Some(crate::TranscriptPartial { + content: "display partial".into(), + reasoning_content: Some("thinking".into()), + iteration: Some(3), + }), + interrupted: true, + }), + })])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(BasicCodec::default())) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + assert!(matches!( + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::Driver(_)) + )); + + let path = directory.path().join("session_raw/agent.jsonl"); + let model = read_transcript(&path).unwrap(); + assert_eq!( + model + .messages + .iter() + .map(|message| message.content.as_str()) + .collect::>(), + vec!["recoverable model history"] + ); + let display = read_transcript_display(&path).unwrap(); + assert!(display.records.iter().any(|record| matches!( + record, + DisplayRecord::Message(message) + if message.interrupted + && message.message.content == "display partial" + && message.reasoning_content.as_deref() == Some("thinking") + && message.iteration == Some(3) + ))); + assert_eq!( + session.history(), + &[Message::assistant("recoverable model history")] + ); +} + +#[tokio::test] +async fn resume_new_turn_retains_durable_metadata_and_restores_prefix_once() { + let mut durable = TranscriptMessage::new("user", "persisted"); + durable.extra_metadata = Some(serde_json::json!({"provider": {"raw": true}})); + let transcript = SessionTranscript { + meta: meta(), + messages: vec![durable.clone()], + }; + let (locator, history) = memory_locator(Some(transcript), false); + let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + // A compaction/faulty driver omitted the stable prefix. + Message::user("persisted"), + Message::assistant("new"), + ]))])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(BasicCodec::default())) + .prefix(PrefixSnapshot::new(vec![Message::system("stable")])) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + let committed = history.session.lock().unwrap().clone().unwrap(); + assert_eq!(session.history()[0], Message::system("stable")); + assert_eq!( + driver.requests.lock().unwrap()[0].history[0], + Message::system("stable") + ); + assert_eq!(committed.messages[1].extra_metadata, durable.extra_metadata); + assert_eq!( + session + .history() + .iter() + .filter(|m| **m == Message::system("stable")) + .count(), + 1 + ); +} + +#[tokio::test] +async fn prefix_reconciliation_uses_maximal_suffix_prefix_overlap() { + let prefix = PrefixSnapshot::new(vec![ + Message::system("stable-a"), + Message::system("stable-b"), + Message::system("stable-c"), + ]); + for (returned, expected) in [ + ( + vec![Message::system("stable-c"), Message::assistant("partial")], + vec![ + Message::system("stable-a"), + Message::system("stable-b"), + Message::system("stable-c"), + Message::assistant("partial"), + ], + ), + ( + vec![ + Message::system("stable-a"), + Message::system("stable-b"), + Message::system("stable-c"), + Message::assistant("complete"), + ], + vec![ + Message::system("stable-a"), + Message::system("stable-b"), + Message::system("stable-c"), + Message::assistant("complete"), + ], + ), + ( + vec![Message::assistant("missing")], + vec![ + Message::system("stable-a"), + Message::system("stable-b"), + Message::system("stable-c"), + Message::assistant("missing"), + ], + ), + ] { + let mut session = + SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(returned))]))) + .prefix(prefix.clone()) + .build() + .unwrap(); + let committed = session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!(committed.history, expected); + } +} + +#[tokio::test] +async fn hooks_are_ordered_terminal_once_and_driver_errors_are_terminal() { + let hooks = Arc::new(RecordingHooks::default()); + let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("ok"), + ]))])); + let mut session = SessionBuilder::new(driver) + .hooks(hooks.clone()) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!( + *hooks.events.lock().unwrap(), + vec!["before", "after", "terminal:Completed"] + ); + + let hooks = Arc::new(RecordingHooks::default()); + let driver = Arc::new(FakeDriver::new(vec![Err(DriverFailure { + error: RuntimeError::Driver("boom".into()), + partial: None, + })])); + let mut session = SessionBuilder::new(driver) + .hooks(hooks.clone()) + .build() + .unwrap(); + assert!(matches!( + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::Driver(_)) + )); + assert_eq!( + *hooks.events.lock().unwrap(), + vec!["before", "terminal:Failed(\"driver failed: boom\")"] + ); +} + +#[tokio::test] +async fn failed_precommit_hook_and_tool_mismatch_leave_no_commit() { + let hooks = Arc::new(RecordingHooks { + fail_after: true, + ..RecordingHooks::default() + }); + let (locator, history) = memory_locator(None, false); + let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("x"), + ]))]))) + .codec(Arc::new(BasicCodec::default())) + .hooks(hooks) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + assert!(matches!( + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::Hook(_)) + )); + assert!(history.turns.lock().unwrap().is_empty()); + assert!(session.history().is_empty()); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("unreachable"))); + let mut mismatched = SessionBuilder::new(Arc::new(HarnessDriver::new( + Arc::new(harness), + Arc::new(()), + ))) + .tool_snapshot( + ToolSnapshot::new(vec![ToolSpec { + name: "not-registered".into(), + description: "x".into(), + parameters: serde_json::json!({}), + }]) + .unwrap(), + ) + .build() + .unwrap(); + assert!(matches!( + mismatched + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::ToolSnapshotMismatch) + )); +} + +#[tokio::test] +async fn cancellation_at_driver_await_emits_one_terminal_hook() { + let hooks = Arc::new(RecordingHooks::default()); + let mut session = SessionBuilder::new(Arc::new(WaitingDriver)) + .hooks(hooks.clone()) + .build() + .unwrap(); + let options = TurnOptions::default(); + let cancellation = options.cancellation.clone(); + tokio::spawn(async move { + tokio::task::yield_now().await; + cancellation.cancel(); + }); + assert_eq!( + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await, + Err(RuntimeError::Cancelled) + ); + assert_eq!( + *hooks.events.lock().unwrap(), + vec!["before", "terminal:Cancelled"] + ); +} + +#[tokio::test] +async fn dropped_turn_future_observes_one_failed_terminal() { + let started = Arc::new(tokio::sync::Notify::new()); + let terminal = Arc::new(tokio::sync::Notify::new()); + let hooks = Arc::new(RecordingHooks { + terminal_notified: Some(terminal.clone()), + ..RecordingHooks::default() + }); + let mut session = SessionBuilder::new(Arc::new(DropDriver { + started: started.clone(), + })) + .hooks(hooks.clone()) + .build() + .unwrap(); + let entered = started.notified(); + let observed = terminal.notified(); + let task = tokio::spawn(async move { + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + }); + entered.await; + task.abort(); + let _ = task.await; + observed.await; + assert_eq!( + *hooks.events.lock().unwrap(), + vec!["before", "terminal:Failed(\"session turn dropped\")"] + ); +} + +#[tokio::test] +async fn dropped_turn_after_commit_keeps_a_truthful_completed_terminal() { + let (locator, history) = memory_locator(None, false); + let hooks = Arc::new(WaitingPostCommitHooks { + started: tokio::sync::Notify::new(), + terminal: tokio::sync::Notify::new(), + terminals: Mutex::new(Vec::new()), + }); + let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("durable"), + ]))]))) + .codec(Arc::new(BasicCodec::default())) + .hooks(hooks.clone()) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + let entered = hooks.started.notified(); + let observed = hooks.terminal.notified(); + let task = tokio::spawn(async move { + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + }); + entered.await; + assert_eq!(history.turns.lock().unwrap().len(), 1); + task.abort(); + let _ = task.await; + observed.await; + assert!(matches!( + hooks.terminals.lock().unwrap().as_slice(), + [SessionTerminal::Completed(outcome)] if outcome.history.last() == Some(&Message::assistant("durable")) + )); +} + +#[tokio::test] +async fn cancellation_during_before_hook_never_starts_or_commits_a_turn() { + let hooks = Arc::new(WaitingBeforeHooks { + started: tokio::sync::Notify::new(), + }); + let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("nope"), + ]))])); + let mut session = SessionBuilder::new(driver.clone()) + .hooks(hooks.clone()) + .build() + .unwrap(); + let options = TurnOptions::default(); + let cancellation = options.cancellation.clone(); + let started = hooks.started.notified(); + let turn = tokio::spawn(async move { + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await + }); + started.await; + cancellation.cancel(); + assert_eq!(turn.await.unwrap(), Err(RuntimeError::Cancelled)); + assert!(driver.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn cancellation_during_precommit_hook_has_no_durable_commit() { + let hooks = Arc::new(WaitingAfterHooks { + started: tokio::sync::Notify::new(), + events: Mutex::new(Vec::new()), + }); + let (locator, history) = memory_locator(None, false); + let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("candidate"), + ]))]))) + .codec(Arc::new(BasicCodec::default())) + .hooks(hooks.clone()) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + let options = TurnOptions::default(); + let cancellation = options.cancellation.clone(); + let started = hooks.started.notified(); + let turn = tokio::spawn(async move { + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await + }); + started.await; + cancellation.cancel(); + assert_eq!(turn.await.unwrap(), Err(RuntimeError::Cancelled)); + assert!(history.turns.lock().unwrap().is_empty()); + assert_eq!(*hooks.events.lock().unwrap(), vec!["Cancelled"]); +} + +#[tokio::test] +async fn cancellation_signalled_by_successful_commit_cannot_relabel_the_turn() { + let (locator, history) = memory_locator(None, false); + let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("done"), + ]))])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(BasicCodec::default())) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + let options = TurnOptions::default(); + *history.cancel_after_append.lock().unwrap() = Some(options.cancellation.clone()); + assert!( + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await + .is_ok() + ); + assert_eq!(history.turns.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn post_commit_runs_once_after_durability_and_cannot_relabel_success() { + let (locator, history) = memory_locator(None, false); + let cancellation = tinyagents_harness::CancellationToken::new(); + let hooks = Arc::new(FinalizationHooks { + fail_post_commit: true, + cancel_on_post_commit: Some(cancellation.clone()), + ..FinalizationHooks::default() + }); + let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("done"), + ]))]))) + .codec(Arc::new(BasicCodec::default())) + .hooks(hooks.clone()) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + let options = TurnOptions { + cancellation, + ..TurnOptions::default() + }; + let committed = session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await + .expect("post-commit failure/cancellation cannot revoke a durable success"); + assert_eq!(history.turns.lock().unwrap().len(), 1); + assert_eq!(*hooks.post_commits.lock().unwrap(), vec![committed.clone()]); + assert!(matches!( + hooks.terminals.lock().unwrap().as_slice(), + [SessionTerminal::Completed(outcome)] if outcome == &committed + )); +} + +#[tokio::test] +async fn persistence_failure_never_calls_post_commit() { + let (locator, _) = memory_locator(None, true); + let hooks = Arc::new(FinalizationHooks::default()); + let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ + Message::assistant("never durable"), + ]))]))) + .codec(Arc::new(BasicCodec::default())) + .hooks(hooks.clone()) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + assert!(matches!( + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::Persistence(_)) + )); + assert!(hooks.post_commits.lock().unwrap().is_empty()); + assert!(matches!( + hooks.terminals.lock().unwrap().as_slice(), + [SessionTerminal::Failed(_)] + )); +} + +type SeenCodecOptions = (String, Option, Option, bool, ResumeMode); + +struct OptionsCodec { + seen: Mutex>, +} + +impl TranscriptCodec for OptionsCodec { + fn decode_history(&self, _: &SessionTranscript) -> Result, RuntimeError> { + Ok(Vec::new()) + } + + fn reconcile( + &self, + _: &[TranscriptMessage], + _: &[Message], + next: &[Message], + options: &crate::TranscriptTurnOptions, + ) -> Result, RuntimeError> { + self.seen.lock().unwrap().push(( + options.context.0.clone(), + options.request_id.clone(), + options.thread_id.clone(), + options.stream, + options.resume, + )); + Ok(next + .iter() + .map(|message| TranscriptMessage::assistant(message.text())) + .collect()) + } +} + +#[tokio::test] +async fn codec_reconciliation_receives_explicit_host_context_and_turn_options() { + let (locator, _) = memory_locator(None, false); + let codec = Arc::new(OptionsCodec { + seen: Mutex::new(Vec::new()), + }); + let mut session = SessionBuilder::::new(Arc::new(ContextDriver)) + .codec(codec.clone()) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + let cancellation = tinyagents_harness::CancellationToken::new(); + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions { + request_id: Some("request-1".into()), + thread_id: Some("thread-1".into()), + stream: true, + resume: ResumeMode::LatestForAgent, + cancellation: cancellation.clone(), + run_context: RunContext::new( + RunConfig::new("codec-host"), + HostContext("host-owned context".into()), + ) + .with_cancellation(cancellation), + }, + ) + .await + .unwrap(); + assert_eq!( + *codec.seen.lock().unwrap(), + vec![( + "host-owned context".into(), + Some("request-1".into()), + Some("thread-1".into()), + true, + ResumeMode::LatestForAgent, + )] + ); +} + +#[test] +fn dependency_boundary_has_no_product_dependency() { + let manifest = include_str!("../Cargo.toml"); + assert!(manifest.contains("tinyagents-harness")); + assert!(manifest.contains("tinyagents-session")); + assert!(manifest.contains("tinytools")); + assert!(!manifest.to_ascii_lowercase().contains("openhuman")); +} diff --git a/crates/tinyagents-runtime/src/tools.rs b/crates/tinyagents-runtime/src/tools.rs new file mode 100644 index 00000000..cd6ffbe5 --- /dev/null +++ b/crates/tinyagents-runtime/src/tools.rs @@ -0,0 +1,43 @@ +use std::collections::BTreeMap; + +use tinytools::ToolSpec; + +use crate::RuntimeError; + +/// An immutable, model-visible tool declaration set for one session turn. +/// +/// It records declarations only; choosing which tools are permitted and wiring +/// their executors remains a host/driver responsibility. +#[derive(Clone, Debug, Default)] +pub struct ToolSnapshot { + specs: Vec, +} + +impl ToolSnapshot { + /// Validates and freezes a tool declaration set. + /// + /// Identical repeated declarations are deduplicated. A shared name with + /// different contents is rejected rather than silently choosing one. + pub fn new(specs: Vec) -> Result { + let mut names = BTreeMap::::new(); + for spec in specs { + if let Some(existing) = names.get(&spec.name) { + let same = existing.description == spec.description + && existing.parameters == spec.parameters; + if !same { + return Err(RuntimeError::ToolNameCollision(spec.name)); + } + continue; + } + names.insert(spec.name.clone(), spec); + } + Ok(Self { + specs: names.into_values().collect(), + }) + } + + /// Returns the frozen declarations in stable name order. + pub fn specs(&self) -> &[ToolSpec] { + &self.specs + } +} diff --git a/crates/tinyagents-runtime/src/types.rs b/crates/tinyagents-runtime/src/types.rs new file mode 100644 index 00000000..c81285cc --- /dev/null +++ b/crates/tinyagents-runtime/src/types.rs @@ -0,0 +1,126 @@ +use tinyagents_harness::{ + CancellationToken, + context::{RunConfig, RunContext}, +}; +use tinyinference_llm::message::Message; + +/// Selects the durable transcript a turn should load before execution. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ResumeMode { + /// Keep this session's current in-memory history. + #[default] + Never, + /// Load the most recent transcript for the configured agent/stem key. + LatestForAgent, + /// Load the most recent root transcript matching `TurnOptions::thread_id`. + Thread, +} + +/// Explicit runtime controls for one session turn. +pub struct TurnOptions { + /// Opaque correlation identifier persisted with transcript rows. + pub request_id: Option, + /// Optional conversation thread identifier used for resume and metadata. + pub thread_id: Option, + /// Whether the driver should use its streaming invocation path. + pub stream: bool, + /// Transcript resume behavior requested for this turn. + pub resume: ResumeMode, + /// Cooperative cancellation shared with the caller. + pub cancellation: CancellationToken, + /// Explicit live execution context consumed by the driver. + pub run_context: RunContext, +} + +/// The codec-visible, durable subset of one turn's explicit options. +/// +/// `RunContext` itself is live and consumed by the driver. A clone of its host +/// context is captured before that handoff so transcript reconciliation can +/// stamp host-owned data after the driver returns without relying on task-local +/// state or a lossy default context. +#[derive(Clone, Debug)] +pub struct TranscriptTurnOptions { + /// Opaque correlation identifier for the current turn. + pub request_id: Option, + /// Conversation thread selected for this turn. + pub thread_id: Option, + /// Whether this turn used the streaming driver path. + pub stream: bool, + /// Resume mode selected before execution. + pub resume: ResumeMode, + /// Host-owned context cloned from `TurnOptions::run_context.data`. + pub context: C, +} + +impl TurnOptions { + pub(crate) fn transcript_options(&self) -> TranscriptTurnOptions { + TranscriptTurnOptions { + request_id: self.request_id.clone(), + thread_id: self.thread_id.clone(), + stream: self.stream, + resume: self.resume, + context: self.run_context.data.clone(), + } + } +} + +impl Default for TurnOptions<()> { + fn default() -> Self { + let cancellation = CancellationToken::new(); + Self { + request_id: None, + thread_id: None, + stream: false, + resume: ResumeMode::Never, + run_context: RunContext::new(RunConfig::new("session"), ()) + .with_cancellation(cancellation.clone()), + cancellation, + } + } +} + +/// The input a host asks a session to execute. +#[derive(Clone, Debug, PartialEq)] +pub struct SessionTurnRequest { + /// The next user/application message. Hooks may replace it before the + /// runtime performs trailing-input deduplication. + pub input: Message, +} + +impl SessionTurnRequest { + /// Creates a request with one next input message. + pub fn new(input: Message) -> Self { + Self { input } + } +} + +/// A committed turn result. +#[derive(Clone, Debug, PartialEq)] +pub struct SessionTurnOutcome { + /// The full logical history after this turn. + pub history: Vec, + /// The driver's final visible output, when it produced one. + pub output: Option, + /// `true` when the driver intentionally ended at an interruptible point. + pub interrupted: bool, +} + +/// The result of loading a transcript into a session. +#[derive(Clone, Debug, PartialEq)] +pub struct SessionResume { + /// Whether a transcript was found and decoded. + pub loaded: bool, + /// The loaded model history, or the existing history when none was found. + pub history: Vec, +} + +/// The one terminal observation emitted for each call to [`crate::Session::turn`]. +#[derive(Clone, Debug, PartialEq)] +pub enum SessionTerminal { + /// The turn committed. The outcome supplies durable finalization data. + Completed(SessionTurnOutcome), + /// The turn was cooperatively cancelled. + Cancelled, + /// The turn ended with an error after any recoverable partial persistence. + Failed(String), +} diff --git a/docs/modules/runtime/README.md b/docs/modules/runtime/README.md new file mode 100644 index 00000000..82b38647 --- /dev/null +++ b/docs/modules/runtime/README.md @@ -0,0 +1,41 @@ +# Session runtime module + +`tinyagents-runtime` provides one host-neutral, stateful conversation session +above `tinyagents-harness` and `tinyagents-session`. It is not a policy layer: +the host owns prompts, model selection, authorization, tool admission, +credentials, product events, and its durable transcript dialect. + +## Turn boundary + +A `Session` holds model-visible history, an immutable stable prefix, and a +frozen tool declaration snapshot. Every `turn` takes explicit `TurnOptions`: +request and thread identifiers, streaming mode, resume mode, cancellation, and +the live `RunContext`. The driver consumes the live context. Because the +codec must reconcile after that driver call, `C: Clone` and the codec receives +the cloned host context together with the relevant per-turn options. + +`SessionDriver` is the execution seam. `HarnessDriver` adapts the pinned +harness partial-run entry points and retains their accumulated history on an +error. Those harness entry points do not currently expose a separate streamed +text delta, partial reasoning, or iteration; the adapter can only use the last +accumulated assistant text as a display partial when it exists. + +## Durability and projections + +With a `TranscriptHistory`, a successful turn first passes `after_turn`, the +pre-commit validation hook. The runtime then submits its logical transcript +delta, metadata, and any supplied `TranscriptPartial` to one +`append_turn_with_partial` operation. `FileTranscriptHistory` serializes those +records into one buffer before one file write. The interrupted partial remains +in the display projection and is excluded from model-context replay. + +Custom history implementations that cannot make this combined operation reject +a supplied partial, so the runtime does not fall back to two independent +writes. This is an operation-level guarantee; it does not claim crash-safe +filesystem transactions beyond the underlying storage implementation. + +After a successful append, `after_commit` receives the committed +`SessionTurnOutcome`, and `on_terminal` receives a completed terminal outcome. +Both are observational: their errors, or a cooperative cancellation that they +observe, cannot change durable success. Failed persistence never invokes +`after_commit`. diff --git a/docs/spec/README.md b/docs/spec/README.md index dc24114b..8ae5afde 100644 --- a/docs/spec/README.md +++ b/docs/spec/README.md @@ -5,13 +5,14 @@ typed state-graph runtime. It takes its shape from LangChain (models, tools, middleware, structured output, streaming, usage/cost) and LangGraph (`START`/`END`, nodes, conditional edges, channels/reducers, checkpoints, interrupts, subgraphs, time travel) — rebuilt as ordinary, typed Rust. The -system is organized as five public crates: +system is organized as six public crates: 1. the harness 2. the graph 3. the registry 4. the expressive language 5. durable sessions +6. host-neutral session runtime The goal is to make agent systems easy to define, inspect, run, test, and serialize without hiding the Rust types that make production systems reliable. @@ -77,6 +78,7 @@ observability, or test contracts. - [Design](../modules/registry/design.md) - [Model catalog and local snapshots](../modules/registry/model-catalog.md) - [Expressive language module](../modules/expressive-language/README.md) +- [Session runtime module](../modules/runtime/README.md) Docs should follow the module layout. Do not place standalone specification files directly in `docs/` or `docs/modules/`; each high-level topic should have @@ -151,6 +153,7 @@ crates/ tinyagents-graph/ # durable typed state graphs tinyagents-registry/ # named capabilities and model catalog tinyagents-session/ # durable session history and run ledger + tinyagents-runtime/ # host-neutral stateful harness sessions tinyagents-tracing/ # shared opt-in tracing macros tinyagents-integration-tests/ # cross-crate tests and runnable examples ``` From c90dca1112ed5fa98b2cc80e0ea6423134c6a397 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 19:15:38 +0300 Subject: [PATCH 05/18] feat(orchestration): add neutral subagent lifecycle Co-authored-by: Medulla --- Cargo.lock | 2 + crates/tinyagents-orchestration/Cargo.toml | 4 +- crates/tinyagents-orchestration/src/lib.rs | 2 + .../src/subagent/README.md | 52 + .../src/subagent/driver.rs | 297 +++++ .../src/subagent/executor.rs | 14 + .../src/subagent/mod.rs | 31 + .../src/subagent/persistence.rs | 34 + .../src/subagent/planner.rs | 17 + .../src/subagent/test.rs | 1061 +++++++++++++++++ .../src/subagent/types.rs | 227 ++++ 11 files changed, 1740 insertions(+), 1 deletion(-) create mode 100644 crates/tinyagents-orchestration/src/subagent/README.md create mode 100644 crates/tinyagents-orchestration/src/subagent/driver.rs create mode 100644 crates/tinyagents-orchestration/src/subagent/executor.rs create mode 100644 crates/tinyagents-orchestration/src/subagent/mod.rs create mode 100644 crates/tinyagents-orchestration/src/subagent/persistence.rs create mode 100644 crates/tinyagents-orchestration/src/subagent/planner.rs create mode 100644 crates/tinyagents-orchestration/src/subagent/test.rs create mode 100644 crates/tinyagents-orchestration/src/subagent/types.rs diff --git a/Cargo.lock b/Cargo.lock index ad2657e6..e58031fe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1586,7 +1586,9 @@ dependencies = [ "tempfile", "tinyagents-graph", "tinyagents-harness", + "tinyagents-runtime", "tinyagents-session", + "tinyinference-llm", "tokio", "uuid", ] diff --git a/crates/tinyagents-orchestration/Cargo.toml b/crates/tinyagents-orchestration/Cargo.toml index 19328e9c..193f8a5c 100644 --- a/crates/tinyagents-orchestration/Cargo.toml +++ b/crates/tinyagents-orchestration/Cargo.toml @@ -17,6 +17,8 @@ serde_json = "1" tinyagents-graph = { path = "../tinyagents-graph", version = "2.1.2", default-features = false } tinyagents-harness = { path = "../tinyagents-harness", version = "2.1.2", default-features = false } tinyagents-session = { path = "../tinyagents-session", version = "2.1.2", default-features = false } +tinyagents-runtime = { path = "../tinyagents-runtime", version = "2.1.2", default-features = false } +tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.3.0" } tokio = { version = "1", default-features = false, features = ["sync", "rt", "macros"] } uuid = { version = "1", features = ["v4"] } @@ -30,7 +32,7 @@ tracing = [ [dev-dependencies] tempfile = "3" -tokio = { version = "1", features = ["macros", "rt-multi-thread"] } +tokio = { version = "1", features = ["macros", "rt-multi-thread", "time"] } [lints] workspace = true diff --git a/crates/tinyagents-orchestration/src/lib.rs b/crates/tinyagents-orchestration/src/lib.rs index 721dd309..41bfd08f 100644 --- a/crates/tinyagents-orchestration/src/lib.rs +++ b/crates/tinyagents-orchestration/src/lib.rs @@ -9,6 +9,7 @@ //! `orchestration -> {graph, harness, session}`. The lower-level crates never //! depend on this composition layer. +pub mod subagent; pub mod teams; pub mod workflow; @@ -23,6 +24,7 @@ mod boundary_tests { include_str!("../../tinyagents-graph/Cargo.toml"), include_str!("../../tinyagents-harness/Cargo.toml"), include_str!("../../tinyagents-session/Cargo.toml"), + include_str!("../../tinyagents-runtime/Cargo.toml"), ] { assert!( !lower_layer.contains("tinyagents-orchestration"), diff --git a/crates/tinyagents-orchestration/src/subagent/README.md b/crates/tinyagents-orchestration/src/subagent/README.md new file mode 100644 index 00000000..dbbbd298 --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/README.md @@ -0,0 +1,52 @@ +# Subagent lifecycle + +`tinyagents_orchestration::subagent` coordinates a host-resolved subagent +lifecycle without making any product-policy decision. Its direct dependency +direction is: + +```text +orchestration::subagent -> tinyagents-runtime -> {tinyagents-harness, tinyagents-session} +``` + +Hosts provide three object-safe seams. `SubagentPlanner` transforms a live +`SubagentRequest` into a complete `PreparedSubagent`: resolved agent +identity, model messages, immutable `ToolSnapshot`, and explicit child +`RunContext`. `SubagentExecutor` executes precisely that plan. The +planner and executor own prompts, model choice, tool authorization, workspace +policy, artifact resolution, and host context data; this module owns none of +them. `SubagentPersistence` owns a host's durable resume and lifecycle records. + +`SubagentTaskKey` is the durable lifecycle identity. It combines the root run, +immediate parent run, optional thread, and host-local task id. Persistence, +in-flight coalescing, and terminal caching all use this scoped key: two parents +may reuse a task id without sharing state, while repeat calls for the same +durable scope deduplicate and resume correctly. + +`SubagentDriver::run` loads a resume only when the caller did not provide one, +then prepares and executes. An awaiting-input result calls `save_pause`; all +other results call `record_terminal`. These operations are mutually exclusive. +The driver caches only successfully persisted terminal outcomes by scoped +`SubagentTaskKey`, so +repeated calls through the same driver return them without executing or +recording again. A pause is never cached: the next call loads its resume state +and executes the continuation. Concurrent calls for the same task coalesce; +different scoped task keys, including nested child tasks on the same driver, +execute independently. Persistence implementations must also enforce +idempotency by `SubagentTaskKey` across processes and driver instances. + +Cancellation is cooperative. It is observed before planning, after resume +load and planning, and immediately after execution. A cancellation that wins +after the executor produced an outcome preserves that outcome's output, +history, usage, and neutral artifact references while reporting `Cancelled`. +The supplied cancellation token is installed on the prepared `RunContext`, so +the executor and the actual harness context observe one cancellation tree. +Persistence `Ok(())` is the commit boundary: if cancellation wins before it, +the driver abandons that uncommitted operation and records one `Cancelled` +terminal outcome; if persistence commits first, the committed result remains +truthful. Planner and executor task ids must exactly match the request task id +or the driver returns a typed error before persistence or caching. +Each coalesced caller retains its own cancellation token: cancelling a follower +returns a local cancelled outcome promptly, without cancelling the leader or +creating an additional persistence action. +An absent host seam is rejected at driver construction with a typed +`MissingCapability` error; no partial lifecycle runs. diff --git a/crates/tinyagents-orchestration/src/subagent/driver.rs b/crates/tinyagents-orchestration/src/subagent/driver.rs new file mode 100644 index 00000000..811f6795 --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/driver.rs @@ -0,0 +1,297 @@ +use std::{collections::HashMap, future::Future, sync::Arc}; + +use tokio::sync::{Mutex, Notify}; + +use super::{ + PersistedSubagentPause, SubagentError, SubagentExecution, SubagentExecutor, SubagentOutcome, + SubagentPersistence, SubagentPlanner, SubagentRequest, SubagentStatus, SubagentTaskKey, +}; +use tinyagents_harness::CancellationToken; + +/// Optional host seams accepted by [`SubagentDriver::new`]. +/// +/// Hosts that cannot provide every seam must receive a typed construction error +/// rather than accidentally executing a partial lifecycle. +pub struct SubagentCapabilities { + /// Host planner, which resolves policy and explicit execution inputs. + pub planner: Option>>, + /// Host executor, which drives the prepared run. + pub executor: Option>>, + /// Host persistence for resume and one lifecycle record. + pub persistence: Option>, +} + +/// Generic lifecycle driver for one host's subagent runs. +/// +/// Concurrent calls for one scoped task key coalesce, while task ids from +/// distinct parents, roots, or threads remain independent. Hosts still need +/// idempotent persistence for multiple processes/drivers. +pub struct SubagentDriver { + planner: Arc>, + executor: Arc>, + persistence: Arc, + terminal_outcomes: Mutex>, + in_flight: Mutex>>, +} + +/// Result shared by callers that arrived while the same task was executing. +/// +/// The map only protects reservation and removal. Planner, executor, and +/// persistence futures never run while it is locked. +struct InFlight { + result: Mutex>>, + notify: Notify, +} + +impl InFlight { + fn new() -> Self { + Self { + result: Mutex::new(None), + notify: Notify::new(), + } + } + + async fn wait( + &self, + cancellation: &CancellationToken, + task_id: &str, + ) -> Result { + loop { + let notified = self.notify.notified(); + tokio::pin!(notified); + notified.as_mut().enable(); + if let Some(result) = self.result.lock().await.clone() { + return result; + } + tokio::select! { + biased; + _ = cancellation.cancelled() => return Ok(SubagentOutcome::cancelled(task_id)), + _ = &mut notified => {} + } + } + } + + async fn complete(&self, result: Result) { + *self.result.lock().await = Some(result); + self.notify.notify_waiters(); + } +} + +impl SubagentDriver { + /// Validates host capability availability before exposing a runnable driver. + pub fn new(capabilities: SubagentCapabilities) -> Result { + Ok(Self { + planner: capabilities + .planner + .ok_or(SubagentError::MissingCapability("planner"))?, + executor: capabilities + .executor + .ok_or(SubagentError::MissingCapability("executor"))?, + persistence: capabilities + .persistence + .ok_or(SubagentError::MissingCapability("persistence"))?, + terminal_outcomes: Mutex::new(HashMap::new()), + in_flight: Mutex::new(HashMap::new()), + }) + } + + /// Runs `load -> prepare -> execute -> one persistence action`. + /// + /// A caller-supplied resume bypasses loading. Cancellation is checked before + /// preparation, after every awaited lifecycle stage, and after execution; + /// a cancellation observed after an executor result preserves its output, + /// history, usage, and artifact references while changing only the status. + pub async fn run( + &self, + request: SubagentRequest, + cancellation: CancellationToken, + ) -> Result { + let task_key = request.task_key(); + if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { + return Ok(outcome); + } + + let (entry, is_leader) = { + let mut in_flight = self.in_flight.lock().await; + match in_flight.get(&task_key) { + Some(entry) => (entry.clone(), false), + None => { + let entry = Arc::new(InFlight::new()); + in_flight.insert(task_key.clone(), entry.clone()); + (entry, true) + } + } + }; + if !is_leader { + return entry.wait(&cancellation, &task_key.task_id).await; + } + // A preceding caller may have committed a terminal result between the + // initial cache check and this reservation. Do not reopen that task + // after its in-flight entry has been removed. + if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { + entry.complete(Ok(outcome.clone())).await; + let mut in_flight = self.in_flight.lock().await; + in_flight.remove(&task_key); + return Ok(outcome); + } + + let result = self + .run_reserved(request, task_key.clone(), cancellation) + .await; + entry.complete(result.clone()).await; + let mut in_flight = self.in_flight.lock().await; + if in_flight + .get(&task_key) + .is_some_and(|current| Arc::ptr_eq(current, &entry)) + { + in_flight.remove(&task_key); + } + result + } + + async fn run_reserved( + &self, + mut request: SubagentRequest, + task_key: SubagentTaskKey, + cancellation: CancellationToken, + ) -> Result { + let task_id = request.task_id.clone(); + + if cancellation.is_cancelled() { + return self.persist_cancelled(task_key, task_id).await; + } + + if request.resume.is_none() { + request.resume = self.persistence.load(&task_key).await?; + } + if cancellation.is_cancelled() { + return self.persist_cancelled(task_key, task_id).await; + } + + let mut prepared = self.planner.prepare(request).await?; + if prepared.task_id != task_id { + return Err(SubagentError::TaskIdMismatch { + expected: task_id, + actual: prepared.task_id, + }); + } + if cancellation.is_cancelled() { + return self.persist_cancelled(task_key, prepared.task_id).await; + } + prepared.run_context = prepared.run_context.with_cancellation(cancellation.clone()); + + let executed = self + .executor + .execute(SubagentExecution { + prepared, + cancellation: cancellation.clone(), + }) + .await; + + let outcome = match executed { + Ok(outcome) => { + if outcome.task_id != task_id { + return Err(SubagentError::TaskIdMismatch { + expected: task_id, + actual: outcome.task_id, + }); + } + if cancellation.is_cancelled() { + outcome.cancelled_preserving() + } else { + outcome + } + } + Err(SubagentError::Cancelled) => SubagentOutcome::cancelled(task_id), + Err(_) if cancellation.is_cancelled() => SubagentOutcome::cancelled(task_id), + Err(error) => return Err(error), + }; + self.persist(task_key, outcome, &cancellation).await + } + + async fn persist_cancelled( + &self, + task_key: SubagentTaskKey, + task_id: String, + ) -> Result { + let outcome = SubagentOutcome::cancelled(task_id); + // Once cancellation has won, this is the one terminal action. Do not + // race it with the already-latched token or a caller could observe an + // indeterminate terminal write. + self.persistence + .record_terminal(&task_key, &outcome) + .await?; + self.cache_terminal(task_key, &outcome).await; + Ok(outcome) + } + + async fn persist( + &self, + task_key: SubagentTaskKey, + outcome: SubagentOutcome, + cancellation: &CancellationToken, + ) -> Result { + if matches!(&outcome.status, SubagentStatus::Cancelled) { + self.persistence + .record_terminal(&task_key, &outcome) + .await?; + self.cache_terminal(task_key, &outcome).await; + return Ok(outcome); + } + let committed = match &outcome.status { + SubagentStatus::AwaitingInput(pause) => { + self.commit_or_cancel( + self.persistence.save_pause(PersistedSubagentPause { + key: task_key.clone(), + pause: pause.clone(), + }), + cancellation, + ) + .await? + } + SubagentStatus::Completed | SubagentStatus::Incomplete(_) => { + self.commit_or_cancel( + self.persistence.record_terminal(&task_key, &outcome), + cancellation, + ) + .await? + } + SubagentStatus::Cancelled => unreachable!("handled before persistence race"), + }; + if !committed { + return self.persist_cancelled(task_key, outcome.task_id).await; + } + if matches!( + &outcome.status, + SubagentStatus::Completed | SubagentStatus::Incomplete(_) + ) { + self.cache_terminal(task_key, &outcome).await; + } + Ok(outcome) + } + + async fn commit_or_cancel( + &self, + operation: F, + cancellation: &CancellationToken, + ) -> Result + where + F: Future>, + { + tokio::select! { + biased; + _ = cancellation.cancelled() => Ok(false), + result = operation => { + result?; + Ok(true) + } + } + } + + async fn cache_terminal(&self, key: SubagentTaskKey, outcome: &SubagentOutcome) { + self.terminal_outcomes + .lock() + .await + .insert(key, outcome.clone()); + } +} diff --git a/crates/tinyagents-orchestration/src/subagent/executor.rs b/crates/tinyagents-orchestration/src/subagent/executor.rs new file mode 100644 index 00000000..1585037b --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/executor.rs @@ -0,0 +1,14 @@ +use async_trait::async_trait; + +use super::{SubagentError, SubagentExecution, SubagentOutcome}; + +/// Host boundary that executes a prepared subagent through its chosen model, +/// tools and runtime services. +#[async_trait] +pub trait SubagentExecutor: Send + Sync { + /// Executes exactly the prepared plan and returns lossless neutral data. + async fn execute( + &self, + execution: SubagentExecution, + ) -> Result; +} diff --git a/crates/tinyagents-orchestration/src/subagent/mod.rs b/crates/tinyagents-orchestration/src/subagent/mod.rs new file mode 100644 index 00000000..552ec74c --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/mod.rs @@ -0,0 +1,31 @@ +//! Host-neutral subagent lifecycle orchestration. +//! +//! A host resolves its own policy, agent definition, prompt, tool allowlist, +//! and persistence implementation through the three object-safe seams exposed +//! here. This module only orders the lifecycle: resume loading, preparation, +//! execution, and one mutually exclusive persistence action. It never creates +//! a context, selects a model, interprets an artifact path, or applies host +//! policy. +//! +//! Dependency direction remains `orchestration -> runtime -> {harness, +//! session}`. In particular, lower TinyAgents layers and hosts must not depend +//! on this lifecycle module. + +mod driver; +mod executor; +mod persistence; +mod planner; +mod types; + +pub use driver::{SubagentCapabilities, SubagentDriver}; +pub use executor::SubagentExecutor; +pub use persistence::SubagentPersistence; +pub use planner::SubagentPlanner; +pub use types::{ + ArtifactReference, PersistedSubagentPause, PreparedSubagent, SubagentError, SubagentExecution, + SubagentIncomplete, SubagentOutcome, SubagentPause, SubagentRequest, SubagentResume, + SubagentStatus, SubagentTaskKey, +}; + +#[cfg(test)] +mod test; diff --git a/crates/tinyagents-orchestration/src/subagent/persistence.rs b/crates/tinyagents-orchestration/src/subagent/persistence.rs new file mode 100644 index 00000000..8751a76d --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/persistence.rs @@ -0,0 +1,34 @@ +use async_trait::async_trait; + +use super::{ + PersistedSubagentPause, SubagentError, SubagentOutcome, SubagentResume, SubagentTaskKey, +}; + +/// Host boundary for durable pause, resume, and terminal lifecycle state. +/// +/// Implementations must make `record_terminal` idempotent by +/// [`SubagentTaskKey`] across process boundaries. The key is scoped by root +/// run, immediate parent run, and (when supplied) thread, so a bare task id is +/// never a global lifecycle identity. The driver additionally suppresses +/// duplicate records from repeated calls made through the same driver instance. A persistence +/// future's successful return is its commit boundary: implementations must not +/// make a write visible and then await again before returning `Ok(())`. The +/// driver races that boundary with cancellation and, when cancellation wins, +/// records one truthful `Cancelled` terminal outcome instead. +#[async_trait] +pub trait SubagentPersistence: Send + Sync { + /// Loads the most recent resumable state, if a caller did not supply one. + async fn load(&self, key: &SubagentTaskKey) -> Result, SubagentError>; + + /// Saves one resumable pause. The driver never also records a terminal for + /// that same committed outcome. A paused outcome is deliberately not cached + /// by the driver; a later call reloads this state and resumes execution. + async fn save_pause(&self, pause: PersistedSubagentPause) -> Result<(), SubagentError>; + + /// Records a non-pause terminal outcome exactly once per scoped task. + async fn record_terminal( + &self, + key: &SubagentTaskKey, + outcome: &SubagentOutcome, + ) -> Result<(), SubagentError>; +} diff --git a/crates/tinyagents-orchestration/src/subagent/planner.rs b/crates/tinyagents-orchestration/src/subagent/planner.rs new file mode 100644 index 00000000..7e24a4ec --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/planner.rs @@ -0,0 +1,17 @@ +use async_trait::async_trait; + +use super::{PreparedSubagent, SubagentError, SubagentRequest}; + +/// Host boundary that resolves a request into a complete execution plan. +/// +/// It owns agent selection, prompts, limits, workspace metadata, context +/// lineage and tool policy. Returning a plan is intentionally all-or-nothing: +/// the generic driver cannot fill missing host data with defaults. +#[async_trait] +pub trait SubagentPlanner: Send + Sync { + /// Resolves one request before any execution or terminal persistence. + async fn prepare( + &self, + request: SubagentRequest, + ) -> Result, SubagentError>; +} diff --git a/crates/tinyagents-orchestration/src/subagent/test.rs b/crates/tinyagents-orchestration/src/subagent/test.rs new file mode 100644 index 00000000..9313bf26 --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/test.rs @@ -0,0 +1,1061 @@ +use std::sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, +}; +use std::time::Duration; + +use async_trait::async_trait; +use tinyagents_harness::{ + CancellationToken, + context::{RunConfig, RunContext}, +}; +use tinyagents_runtime::ToolSnapshot; +use tinyinference_llm::{message::Message, usage::UsageTotals}; + +use super::*; + +#[derive(Clone, Debug, PartialEq, Eq)] +enum Action { + Load, + Prepare, + Execute, + Pause, + Terminal(SubagentStatusName), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +enum SubagentStatusName { + Completed, + Incomplete, + Cancelled, +} + +#[derive(Clone, Copy)] +enum ExecutorMode { + Completed, + Incomplete, + Pause, + WaitForCancellation, + CancelAfterExecution, + Error, +} + +struct FakePlanner { + calls: Mutex, + saw_resume: Mutex, + seen_resumes: Mutex>, + reject: bool, + actions: Arc>>, +} + +#[async_trait] +impl SubagentPlanner for FakePlanner { + async fn prepare( + &self, + request: SubagentRequest, + ) -> Result, SubagentError> { + *self.calls.lock().unwrap() += 1; + *self.saw_resume.lock().unwrap() = request.resume.is_some(); + self.seen_resumes + .lock() + .unwrap() + .push(request.resume.is_some()); + self.actions.lock().unwrap().push(Action::Prepare); + if self.reject { + return Err(SubagentError::Planning("rejected".into())); + } + Ok(PreparedSubagent { + task_id: request.task_id, + agent_key: "resolved-agent".into(), + input: vec![Message::user(request.input)], + tools: ToolSnapshot::new(vec![]).unwrap(), + run_context: request.parent_run, + }) + } +} + +struct FakeExecutor { + calls: Mutex, + context_ids: Mutex>, + context_cancellations: Mutex>, + mode: ExecutorMode, + started: Mutex>>, + actions: Arc>>, +} + +#[async_trait] +impl SubagentExecutor for FakeExecutor { + async fn execute( + &self, + execution: SubagentExecution, + ) -> Result { + *self.calls.lock().unwrap() += 1; + self.context_ids + .lock() + .unwrap() + .push(execution.prepared.run_context.instance_id()); + self.context_cancellations + .lock() + .unwrap() + .push(execution.prepared.run_context.cancellation.clone()); + self.actions.lock().unwrap().push(Action::Execute); + if matches!(self.mode, ExecutorMode::WaitForCancellation) { + if let Some(sender) = self.started.lock().unwrap().take() { + let _ = sender.send(()); + } + execution.cancellation.cancelled().await; + } + if matches!(self.mode, ExecutorMode::CancelAfterExecution) { + execution.cancellation.cancel(); + } + if matches!(self.mode, ExecutorMode::Error) { + return Err(SubagentError::Execution("executor failed".into())); + } + let status = match self.mode { + ExecutorMode::Completed + | ExecutorMode::WaitForCancellation + | ExecutorMode::CancelAfterExecution => SubagentStatus::Completed, + ExecutorMode::Incomplete => SubagentStatus::Incomplete(SubagentIncomplete { + reason: "budget exhausted".into(), + }), + ExecutorMode::Pause => SubagentStatus::AwaitingInput(SubagentPause { + reason: "need approval".into(), + resume: SubagentResume::default(), + }), + ExecutorMode::Error => unreachable!(), + }; + Ok(SubagentOutcome { + task_id: execution.prepared.task_id, + output: "result".into(), + history: vec![Message::assistant("result")], + status, + usage: UsageTotals { + calls: 7, + ..UsageTotals::default() + }, + artifacts: vec![ArtifactReference { + id: "artifact-1".into(), + ..ArtifactReference::default() + }], + }) + } +} + +enum LoadMode { + Empty, + Resume, + Error, +} + +type Fakes = ( + Arc, + Arc, + Arc, + Arc>>, +); + +struct FakePersistence { + actions: Arc>>, + load_mode: LoadMode, + pause_error: bool, + terminal_error: bool, + outcomes: Mutex>, + saved_pause: Mutex>, + keys: Mutex>, +} + +#[async_trait] +impl SubagentPersistence for FakePersistence { + async fn load(&self, key: &SubagentTaskKey) -> Result, SubagentError> { + self.actions.lock().unwrap().push(Action::Load); + self.keys.lock().unwrap().push(key.clone()); + match self.load_mode { + LoadMode::Empty => Ok(self.saved_pause.lock().unwrap().clone()), + LoadMode::Resume => Ok(Some(SubagentResume { + checkpoint: Some("saved".into()), + ..SubagentResume::default() + })), + LoadMode::Error => Err(SubagentError::Persistence("load failed".into())), + } + } + + async fn save_pause(&self, pause: PersistedSubagentPause) -> Result<(), SubagentError> { + self.actions.lock().unwrap().push(Action::Pause); + self.keys.lock().unwrap().push(pause.key); + if self.pause_error { + Err(SubagentError::Persistence("pause save failed".into())) + } else { + *self.saved_pause.lock().unwrap() = Some(pause.pause.resume); + Ok(()) + } + } + + async fn record_terminal( + &self, + key: &SubagentTaskKey, + outcome: &SubagentOutcome, + ) -> Result<(), SubagentError> { + self.actions + .lock() + .unwrap() + .push(Action::Terminal(match outcome.status { + SubagentStatus::Completed => SubagentStatusName::Completed, + SubagentStatus::Incomplete(_) => SubagentStatusName::Incomplete, + SubagentStatus::Cancelled => SubagentStatusName::Cancelled, + SubagentStatus::AwaitingInput(_) => unreachable!(), + })); + self.keys.lock().unwrap().push(key.clone()); + if self.terminal_error { + Err(SubagentError::Persistence("terminal save failed".into())) + } else { + self.outcomes.lock().unwrap().push(outcome.clone()); + Ok(()) + } + } +} + +fn request(task_id: &str, data: &str) -> SubagentRequest { + request_with_parent( + task_id, + RunContext::new(RunConfig::new(format!("run-{task_id}")), data.into()), + ) +} + +fn request_with_parent(task_id: &str, parent_run: RunContext) -> SubagentRequest { + SubagentRequest { + task_id: task_id.into(), + parent_run, + input: "do work".into(), + thread_id: Some("thread-1".into()), + resume: None, + } +} + +fn driver( + planner: Arc>, + executor: Arc>, + persistence: Arc, +) -> SubagentDriver { + SubagentDriver::new(SubagentCapabilities { + planner: Some(planner), + executor: Some(executor), + persistence: Some(persistence), + }) + .unwrap() +} + +fn fakes(mode: ExecutorMode) -> Fakes { + let actions = Arc::new(Mutex::new(Vec::new())); + ( + Arc::new(FakePlanner { + calls: Mutex::new(0), + saw_resume: Mutex::new(false), + seen_resumes: Mutex::new(Vec::new()), + reject: false, + actions: actions.clone(), + }), + Arc::new(FakeExecutor { + calls: Mutex::new(0), + context_ids: Mutex::new(Vec::new()), + context_cancellations: Mutex::new(Vec::new()), + mode, + started: Mutex::new(None), + actions: actions.clone(), + }), + Arc::new(FakePersistence { + actions: actions.clone(), + load_mode: LoadMode::Empty, + pause_error: false, + terminal_error: false, + outcomes: Mutex::new(Vec::new()), + saved_pause: Mutex::new(None), + keys: Mutex::new(Vec::new()), + }), + actions, + ) +} + +/// Persistence fake whose first selected operation cannot commit until the +/// test releases it. This makes the cancellation/commit boundary observable. +#[derive(Clone, Copy)] +enum BlockingStage { + Pause, + Terminal, +} + +struct BlockingPersistence { + stage: BlockingStage, + started: Arc, + release: Arc, + first: AtomicBool, + outcomes: Mutex>, +} + +#[async_trait] +impl SubagentPersistence for BlockingPersistence { + async fn load(&self, _: &SubagentTaskKey) -> Result, SubagentError> { + Ok(None) + } + + async fn save_pause(&self, _: PersistedSubagentPause) -> Result<(), SubagentError> { + if matches!(self.stage, BlockingStage::Pause) && self.first.swap(false, Ordering::AcqRel) { + self.started.notify_one(); + self.release.notified().await; + } + Ok(()) + } + + async fn record_terminal( + &self, + _: &SubagentTaskKey, + outcome: &SubagentOutcome, + ) -> Result<(), SubagentError> { + if matches!(self.stage, BlockingStage::Terminal) && self.first.swap(false, Ordering::AcqRel) + { + self.started.notify_one(); + self.release.notified().await; + } + self.outcomes.lock().unwrap().push(outcome.clone()); + Ok(()) + } +} + +struct MismatchedPlanner; + +#[async_trait] +impl SubagentPlanner for MismatchedPlanner { + async fn prepare( + &self, + request: SubagentRequest, + ) -> Result, SubagentError> { + Ok(PreparedSubagent { + task_id: "other-task".into(), + agent_key: "resolved-agent".into(), + input: vec![Message::user(request.input)], + tools: ToolSnapshot::new(vec![]).unwrap(), + run_context: request.parent_run, + }) + } +} + +struct MismatchedExecutor; + +#[async_trait] +impl SubagentExecutor for MismatchedExecutor { + async fn execute( + &self, + _: SubagentExecution, + ) -> Result { + Ok(SubagentOutcome { + task_id: "other-task".into(), + output: String::new(), + history: Vec::new(), + status: SubagentStatus::Completed, + usage: UsageTotals::default(), + artifacts: Vec::new(), + }) + } +} + +struct NestedExecutor { + driver: Mutex>>>, + calls: Mutex>, +} + +struct PauseThenCompleteExecutor { + calls: Mutex, + actions: Arc>>, +} + +#[async_trait] +impl SubagentExecutor for PauseThenCompleteExecutor { + async fn execute( + &self, + execution: SubagentExecution, + ) -> Result { + let call = { + let mut calls = self.calls.lock().unwrap(); + *calls += 1; + *calls + }; + self.actions.lock().unwrap().push(Action::Execute); + Ok(SubagentOutcome { + task_id: execution.prepared.task_id, + output: if call == 1 { + "waiting".into() + } else { + "completed".into() + }, + history: Vec::new(), + status: if call == 1 { + SubagentStatus::AwaitingInput(SubagentPause { + reason: "need input".into(), + resume: SubagentResume { + checkpoint: Some("resume-token".into()), + ..SubagentResume::default() + }, + }) + } else { + SubagentStatus::Completed + }, + usage: UsageTotals::default(), + artifacts: Vec::new(), + }) + } +} + +#[async_trait] +impl SubagentExecutor for NestedExecutor { + async fn execute( + &self, + execution: SubagentExecution, + ) -> Result { + let task_id = execution.prepared.task_id.clone(); + self.calls.lock().unwrap().push(task_id.clone()); + if task_id == "parent" { + let child_driver = self + .driver + .lock() + .unwrap() + .as_ref() + .and_then(std::sync::Weak::upgrade) + .expect("nested executor is attached to its driver"); + let child = child_driver + .run( + SubagentRequest { + task_id: "child".into(), + parent_run: execution.prepared.run_context, + input: "nested work".into(), + thread_id: None, + resume: None, + }, + execution.cancellation, + ) + .await?; + return Ok(SubagentOutcome { + task_id, + output: "parent result".into(), + history: Vec::new(), + status: SubagentStatus::Completed, + // This models the host's parent-visible roll-up: the child is + // added once alongside the parent's own model call. + usage: UsageTotals { + calls: child.usage.calls + 1, + ..UsageTotals::default() + }, + artifacts: Vec::new(), + }); + } + Ok(SubagentOutcome { + task_id, + output: "child result".into(), + history: Vec::new(), + status: SubagentStatus::Completed, + usage: UsageTotals { + calls: 7, + ..UsageTotals::default() + }, + artifacts: Vec::new(), + }) + } +} + +#[tokio::test] +async fn planner_rejection_does_not_execute_or_persist_terminal_state() { + let (_, executor, persistence, actions) = fakes(ExecutorMode::Completed); + let rejecting = Arc::new(FakePlanner { + calls: Mutex::new(0), + saw_resume: Mutex::new(false), + seen_resumes: Mutex::new(Vec::new()), + reject: true, + actions: actions.clone(), + }); + let result = driver(rejecting.clone(), executor.clone(), persistence.clone()) + .run(request("task", "ctx"), CancellationToken::new()) + .await; + + assert_eq!(result, Err(SubagentError::Planning("rejected".into()))); + assert_eq!(*executor.calls.lock().unwrap(), 0); + assert!(persistence.outcomes.lock().unwrap().is_empty()); + assert_eq!( + *actions.lock().unwrap(), + vec![Action::Load, Action::Prepare] + ); +} + +#[tokio::test] +async fn prepared_context_identity_reaches_executor() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); + let incoming = request("task", "identity"); + let expected = incoming.parent_run.instance_id(); + let outcome = driver(planner, executor.clone(), persistence) + .run(incoming, CancellationToken::new()) + .await + .unwrap(); + + assert_eq!(outcome.status, SubagentStatus::Completed); + assert_eq!(*executor.context_ids.lock().unwrap(), vec![expected]); +} + +#[tokio::test] +async fn completed_incomplete_and_pause_use_one_mutually_exclusive_persistence_action() { + for (mode, expected) in [ + ( + ExecutorMode::Completed, + Action::Terminal(SubagentStatusName::Completed), + ), + ( + ExecutorMode::Incomplete, + Action::Terminal(SubagentStatusName::Incomplete), + ), + (ExecutorMode::Pause, Action::Pause), + ] { + let (planner, executor, persistence, actions) = fakes(mode); + driver(planner, executor, persistence) + .run(request("task", "ctx"), CancellationToken::new()) + .await + .unwrap(); + let records = actions.lock().unwrap(); + assert_eq!( + records + .iter() + .filter(|action| matches!(action, Action::Pause | Action::Terminal(_))) + .count(), + 1 + ); + assert_eq!(records.last(), Some(&expected)); + } +} + +#[tokio::test] +async fn load_pause_and_terminal_errors_remain_typed() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); + let load_failure = Arc::new(FakePersistence { + actions: persistence.actions.clone(), + load_mode: LoadMode::Error, + pause_error: false, + terminal_error: false, + outcomes: Mutex::new(Vec::new()), + saved_pause: Mutex::new(None), + keys: Mutex::new(Vec::new()), + }); + assert_eq!( + driver(planner.clone(), executor.clone(), load_failure) + .run(request("load", "ctx"), CancellationToken::new()) + .await, + Err(SubagentError::Persistence("load failed".into())) + ); + + let (planner, executor, persistence, _) = fakes(ExecutorMode::Pause); + let pause_failure = Arc::new(FakePersistence { + pause_error: true, + actions: persistence.actions.clone(), + load_mode: LoadMode::Empty, + terminal_error: false, + outcomes: Mutex::new(Vec::new()), + saved_pause: Mutex::new(None), + keys: Mutex::new(Vec::new()), + }); + assert_eq!( + driver(planner, executor, pause_failure) + .run(request("pause", "ctx"), CancellationToken::new()) + .await, + Err(SubagentError::Persistence("pause save failed".into())) + ); + + let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); + let terminal_failure = Arc::new(FakePersistence { + terminal_error: true, + actions: persistence.actions.clone(), + load_mode: LoadMode::Empty, + pause_error: false, + outcomes: Mutex::new(Vec::new()), + saved_pause: Mutex::new(None), + keys: Mutex::new(Vec::new()), + }); + assert_eq!( + driver(planner, executor, terminal_failure) + .run(request("terminal", "ctx"), CancellationToken::new()) + .await, + Err(SubagentError::Persistence("terminal save failed".into())) + ); +} + +#[tokio::test] +async fn loaded_resume_reaches_planner_before_execution_and_execution_errors_do_not_persist() { + let (planner, executor, persistence, actions) = fakes(ExecutorMode::Completed); + let resume_persistence = Arc::new(FakePersistence { + actions: persistence.actions.clone(), + load_mode: LoadMode::Resume, + pause_error: false, + terminal_error: false, + outcomes: Mutex::new(Vec::new()), + saved_pause: Mutex::new(None), + keys: Mutex::new(Vec::new()), + }); + driver(planner.clone(), executor, resume_persistence) + .run(request("resume", "ctx"), CancellationToken::new()) + .await + .unwrap(); + assert!(*planner.saw_resume.lock().unwrap()); + assert_eq!( + actions.lock().unwrap()[..2], + [Action::Load, Action::Prepare] + ); + + let (planner, executor, persistence, actions) = fakes(ExecutorMode::Error); + assert_eq!( + driver(planner, executor, persistence) + .run(request("execution-error", "ctx"), CancellationToken::new()) + .await, + Err(SubagentError::Execution("executor failed".into())) + ); + assert_eq!( + *actions.lock().unwrap(), + vec![Action::Load, Action::Prepare, Action::Execute] + ); +} + +#[tokio::test] +async fn duplicate_task_id_returns_recorded_outcome_without_second_execution_or_record() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); + let driver = driver(planner.clone(), executor.clone(), persistence.clone()); + let first = driver + .run(request("same", "one"), CancellationToken::new()) + .await + .unwrap(); + let second = driver + .run(request("same", "two"), CancellationToken::new()) + .await + .unwrap(); + + assert_eq!(first, second); + assert_eq!(*planner.calls.lock().unwrap(), 1); + assert_eq!(*executor.calls.lock().unwrap(), 1); + assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn concurrent_same_task_calls_coalesce_to_one_lifecycle() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::WaitForCancellation); + let driver = Arc::new(driver( + planner.clone(), + executor.clone(), + persistence.clone(), + )); + let cancellation = CancellationToken::new(); + let (started_at_execution, started) = tokio::sync::oneshot::channel(); + *executor.started.lock().unwrap() = Some(started_at_execution); + + let first = tokio::spawn({ + let driver = driver.clone(); + let cancellation = cancellation.clone(); + async move { + driver + .run(request("same-in-flight", "one"), cancellation) + .await + } + }); + started.await.unwrap(); + let second = tokio::spawn({ + let driver = driver.clone(); + async move { + driver + .run(request("same-in-flight", "two"), CancellationToken::new()) + .await + } + }); + tokio::task::yield_now().await; + cancellation.cancel(); + + assert_eq!( + first.await.unwrap().unwrap().status, + SubagentStatus::Cancelled + ); + assert_eq!( + second.await.unwrap().unwrap().status, + SubagentStatus::Cancelled + ); + assert_eq!(*planner.calls.lock().unwrap(), 1); + assert_eq!(*executor.calls.lock().unwrap(), 1); + assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn cancelled_follower_returns_without_cancelling_the_leader_or_persisting() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::WaitForCancellation); + let driver = Arc::new(driver( + planner.clone(), + executor.clone(), + persistence.clone(), + )); + let leader_cancellation = CancellationToken::new(); + let follower_cancellation = CancellationToken::new(); + let (started_at_execution, started) = tokio::sync::oneshot::channel(); + *executor.started.lock().unwrap() = Some(started_at_execution); + + let leader = tokio::spawn({ + let driver = driver.clone(); + let cancellation = leader_cancellation.clone(); + async move { + driver + .run(request("follower-cancellation", "leader"), cancellation) + .await + } + }); + started.await.unwrap(); + let follower = tokio::spawn({ + let driver = driver.clone(); + let cancellation = follower_cancellation.clone(); + async move { + driver + .run(request("follower-cancellation", "follower"), cancellation) + .await + } + }); + follower_cancellation.cancel(); + + let follower_outcome = tokio::time::timeout(Duration::from_secs(1), follower) + .await + .expect("cancelled follower must not wait for the leader") + .unwrap() + .unwrap(); + assert_eq!(follower_outcome.status, SubagentStatus::Cancelled); + assert!(!leader_cancellation.is_cancelled()); + assert_eq!(*executor.calls.lock().unwrap(), 1); + assert!(persistence.outcomes.lock().unwrap().is_empty()); + + leader_cancellation.cancel(); + assert_eq!( + leader.await.unwrap().unwrap().status, + SubagentStatus::Cancelled + ); + assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn same_task_id_from_distinct_parent_runs_never_shares_lifecycle_state() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); + let driver = driver(planner.clone(), executor.clone(), persistence.clone()); + let root = RunContext::new(RunConfig::new("root"), "root-data".to_owned()); + let first_parent = root + .child(RunConfig::new("parent-one"), "first-parent".into()) + .unwrap(); + let second_parent = root + .child(RunConfig::new("parent-two"), "second-parent".into()) + .unwrap(); + + driver + .run( + request_with_parent("same-task", first_parent), + CancellationToken::new(), + ) + .await + .unwrap(); + driver + .run( + request_with_parent("same-task", second_parent), + CancellationToken::new(), + ) + .await + .unwrap(); + + assert_eq!(*planner.calls.lock().unwrap(), 2); + assert_eq!(*executor.calls.lock().unwrap(), 2); + assert_eq!(persistence.outcomes.lock().unwrap().len(), 2); + let keys = persistence.keys.lock().unwrap(); + let terminal_keys = keys + .iter() + .filter(|key| key.task_id == "same-task") + .collect::>(); + assert_eq!( + terminal_keys.len(), + 4, + "load and terminal use each scoped key" + ); + assert!( + terminal_keys + .iter() + .any(|key| key.parent_run_id == "parent-one") + ); + assert!( + terminal_keys + .iter() + .any(|key| key.parent_run_id == "parent-two") + ); + assert!(terminal_keys.iter().all(|key| key.root_run_id == "root")); +} + +#[tokio::test] +async fn awaiting_input_is_not_cached_and_the_next_call_resumes_to_completion() { + let (planner, _, persistence, actions) = fakes(ExecutorMode::Pause); + let executor = Arc::new(PauseThenCompleteExecutor { + calls: Mutex::new(0), + actions: actions.clone(), + }); + let driver = driver(planner.clone(), executor.clone(), persistence.clone()); + + let first = driver + .run(request("resumable", "first"), CancellationToken::new()) + .await + .unwrap(); + let second = driver + .run(request("resumable", "continued"), CancellationToken::new()) + .await + .unwrap(); + + assert!(matches!(first.status, SubagentStatus::AwaitingInput(_))); + assert_eq!(second.status, SubagentStatus::Completed); + assert_eq!(*planner.calls.lock().unwrap(), 2); + assert_eq!(*executor.calls.lock().unwrap(), 2); + assert_eq!(*planner.seen_resumes.lock().unwrap(), vec![false, true]); + assert_eq!( + actions + .lock() + .unwrap() + .iter() + .filter(|action| matches!(action, Action::Pause | Action::Terminal(_))) + .cloned() + .collect::>(), + vec![ + Action::Pause, + Action::Terminal(SubagentStatusName::Completed) + ] + ); + assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn driver_replaces_prepared_context_cancellation_with_execution_token() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::WaitForCancellation); + let driver = Arc::new(driver(planner, executor.clone(), persistence)); + let cancellation = CancellationToken::new(); + let (started_at_execution, started) = tokio::sync::oneshot::channel(); + *executor.started.lock().unwrap() = Some(started_at_execution); + + let task = tokio::spawn({ + let driver = driver.clone(); + let cancellation = cancellation.clone(); + async move { + driver + .run(request("shared-cancellation", "ctx"), cancellation) + .await + } + }); + started.await.unwrap(); + cancellation.cancel(); + let outcome = task.await.unwrap().unwrap(); + + assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert!(executor.context_cancellations.lock().unwrap()[0].is_cancelled()); +} + +#[tokio::test] +async fn cancellation_before_persistence_commit_records_only_cancelled_terminal() { + for (stage, mode, task_id) in [ + ( + BlockingStage::Pause, + ExecutorMode::Pause, + "cancel-save-pause", + ), + ( + BlockingStage::Terminal, + ExecutorMode::Completed, + "cancel-record-terminal", + ), + ] { + let (planner, executor, _, _) = fakes(mode); + let persistence = Arc::new(BlockingPersistence { + stage, + started: Arc::new(tokio::sync::Notify::new()), + release: Arc::new(tokio::sync::Notify::new()), + first: AtomicBool::new(true), + outcomes: Mutex::new(Vec::new()), + }); + let driver = Arc::new(driver(planner, executor, persistence.clone())); + let cancellation = CancellationToken::new(); + let started = persistence.started.clone(); + let run = tokio::spawn({ + let driver = driver.clone(); + let cancellation = cancellation.clone(); + async move { driver.run(request(task_id, "ctx"), cancellation).await } + }); + + tokio::time::timeout(Duration::from_secs(1), started.notified()) + .await + .expect("first persistence action must be pending"); + cancellation.cancel(); + let outcome = tokio::time::timeout(Duration::from_secs(1), run) + .await + .expect("cancellation must resolve the lifecycle") + .unwrap() + .unwrap(); + + assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); + assert_eq!( + persistence.outcomes.lock().unwrap()[0].status, + SubagentStatus::Cancelled + ); + } +} + +#[tokio::test] +async fn planner_and_executor_task_id_mismatches_do_not_persist_or_cache() { + let (_, executor, persistence, _) = fakes(ExecutorMode::Completed); + let planner_error = driver(Arc::new(MismatchedPlanner), executor, persistence.clone()) + .run(request("expected", "ctx"), CancellationToken::new()) + .await; + assert_eq!( + planner_error, + Err(SubagentError::TaskIdMismatch { + expected: "expected".into(), + actual: "other-task".into(), + }) + ); + assert!(persistence.outcomes.lock().unwrap().is_empty()); + + let (planner, _, persistence, _) = fakes(ExecutorMode::Completed); + let executor_error = driver(planner, Arc::new(MismatchedExecutor), persistence.clone()) + .run(request("expected", "ctx"), CancellationToken::new()) + .await; + assert_eq!( + executor_error, + Err(SubagentError::TaskIdMismatch { + expected: "expected".into(), + actual: "other-task".into(), + }) + ); + assert!(persistence.outcomes.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn nested_same_driver_task_uses_child_reservation_and_rolls_usage_up_once() { + let (planner, _, persistence, _) = fakes(ExecutorMode::Completed); + let executor = Arc::new(NestedExecutor { + driver: Mutex::new(None), + calls: Mutex::new(Vec::new()), + }); + let driver = Arc::new(driver(planner, executor.clone(), persistence.clone())); + *executor.driver.lock().unwrap() = Some(Arc::downgrade(&driver)); + + let parent = tokio::time::timeout( + Duration::from_secs(1), + driver.run(request("parent", "ctx"), CancellationToken::new()), + ) + .await + .expect("a nested child task must not wait on a driver-global lock") + .unwrap(); + + assert_eq!(parent.usage.calls, 8); + assert_eq!(*executor.calls.lock().unwrap(), vec!["parent", "child"]); + let records = persistence.outcomes.lock().unwrap(); + assert_eq!(records.len(), 2); + assert_eq!( + records + .iter() + .find(|outcome| outcome.task_id == "parent") + .expect("parent outcome is persisted") + .usage + .calls, + 8 + ); +} + +#[tokio::test] +async fn cancellation_before_execution_records_one_truthful_terminal() { + let (planner, executor, persistence, actions) = fakes(ExecutorMode::Completed); + let cancellation = CancellationToken::new(); + cancellation.cancel(); + let outcome = driver(planner.clone(), executor.clone(), persistence) + .run(request("cancel-before", "ctx"), cancellation) + .await + .unwrap(); + + assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(*planner.calls.lock().unwrap(), 0); + assert_eq!(*executor.calls.lock().unwrap(), 0); + assert_eq!( + *actions.lock().unwrap(), + vec![Action::Terminal(SubagentStatusName::Cancelled)] + ); +} + +#[tokio::test] +async fn cancellation_during_execution_is_truthful_and_terminal_once() { + let (planner, executor, persistence, actions) = fakes(ExecutorMode::WaitForCancellation); + let driver = Arc::new(driver(planner, executor.clone(), persistence)); + let cancellation = CancellationToken::new(); + let (started_at_execution, started) = tokio::sync::oneshot::channel(); + *executor.started.lock().unwrap() = Some(started_at_execution); + let task = tokio::spawn({ + let driver = driver.clone(); + let cancellation = cancellation.clone(); + async move { + driver + .run(request("cancel-during", "ctx"), cancellation) + .await + } + }); + started.await.unwrap(); + cancellation.cancel(); + let outcome = task.await.unwrap().unwrap(); + + assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(*executor.calls.lock().unwrap(), 1); + assert_eq!( + actions.lock().unwrap().last(), + Some(&Action::Terminal(SubagentStatusName::Cancelled)) + ); +} + +#[tokio::test] +async fn cancellation_after_execution_preserves_lossless_result_data() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::CancelAfterExecution); + let outcome = driver(planner, executor, persistence) + .run(request("cancel-after", "ctx"), CancellationToken::new()) + .await + .unwrap(); + + assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(outcome.output, "result"); + assert_eq!(outcome.history, vec![Message::assistant("result")]); + assert_eq!(outcome.usage.calls, 7); + assert_eq!(outcome.artifacts.len(), 1); +} + +#[tokio::test] +async fn nested_usage_is_persisted_once_and_failure_order_is_load_prepare_execute_then_terminal() { + let (planner, executor, persistence, actions) = fakes(ExecutorMode::Completed); + driver(planner, executor, persistence.clone()) + .run(request("usage", "ctx"), CancellationToken::new()) + .await + .unwrap(); + + let outcomes = persistence.outcomes.lock().unwrap(); + assert_eq!(outcomes.len(), 1); + assert_eq!(outcomes[0].usage.calls, 7); + assert_eq!( + *actions.lock().unwrap(), + vec![ + Action::Load, + Action::Prepare, + Action::Execute, + Action::Terminal(SubagentStatusName::Completed), + ] + ); +} + +#[test] +fn absent_host_capabilities_fail_closed() { + let result = SubagentDriver::::new(SubagentCapabilities { + planner: None, + executor: None, + persistence: None, + }); + let error = match result { + Err(error) => error, + Ok(_) => panic!("missing host capabilities must fail closed"), + }; + assert_eq!(error, SubagentError::MissingCapability("planner")); +} diff --git a/crates/tinyagents-orchestration/src/subagent/types.rs b/crates/tinyagents-orchestration/src/subagent/types.rs new file mode 100644 index 00000000..dee60d25 --- /dev/null +++ b/crates/tinyagents-orchestration/src/subagent/types.rs @@ -0,0 +1,227 @@ +use std::collections::BTreeMap; + +use tinyagents_harness::{CancellationToken, context::RunContext}; +use tinyagents_runtime::ToolSnapshot; +use tinyinference_llm::{message::Message, usage::UsageTotals}; + +/// A host-provided request to run a task through a subagent. +/// +/// `parent_run` is live execution data, deliberately not a serializable host +/// DTO. The planner must derive the child's fully resolved context explicitly +/// instead of consulting task-local state. +pub struct SubagentRequest { + /// Host-local task id. Its durable lifecycle identity is scoped by + /// [`SubagentTaskKey`], derived from this request's parent run. + pub task_id: String, + /// The explicit live parent context from which a child context is derived. + pub parent_run: RunContext, + /// Host-visible task input that the planner converts into model messages. + pub input: String, + /// Optional host-owned conversation thread correlation id. + pub thread_id: Option, + /// A caller-supplied checkpoint. When absent the driver asks persistence. + pub resume: Option, +} + +/// Durable, host-neutral identity for one subagent lifecycle. +/// +/// A task id is only unique inside the recursive run that created it. The +/// parent run and root run therefore scope durable persistence, in-memory +/// coalescing, and terminal cache entries. The optional thread adds the host's +/// durable conversation partition when it is available. +#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] +pub struct SubagentTaskKey { + /// Top-level recursive run that owns this task tree. + pub root_run_id: String, + /// Immediate run that requested this subagent lifecycle. + pub parent_run_id: String, + /// Host conversation partition, if either request or parent supplies one. + pub thread_id: Option, + /// Host-local task id within the scoped recursive run. + pub task_id: String, +} + +impl SubagentRequest { + /// Derives the durable lifecycle key without exposing host context data. + pub fn task_key(&self) -> SubagentTaskKey { + SubagentTaskKey { + root_run_id: self.parent_run.lineage().root_run_id.as_str().to_owned(), + parent_run_id: self.parent_run.run_id().as_str().to_owned(), + thread_id: self + .thread_id + .clone() + .or_else(|| self.parent_run.thread_id().map(|id| id.as_str().to_owned())), + task_id: self.task_id.clone(), + } + } +} + +/// A neutral checkpoint offered to a planner for resumption. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SubagentResume { + /// Lossless model history available to the host planner. + pub history: Vec, + /// Opaque checkpoint token; its interpretation remains host-owned. + pub checkpoint: Option, + /// Small neutral metadata. This deliberately excludes paths and credentials. + pub metadata: BTreeMap, +} + +/// Fully resolved, immutable execution input produced by a host planner. +/// +/// The planner, not this orchestration layer, resolves agent identity, prompt +/// messages, limits, workspace policy and the tool declaration allowlist. +pub struct PreparedSubagent { + /// Host-stable task identity. + pub task_id: String, + /// Host-resolved agent identity. + pub agent_key: String, + /// Complete model input, including any resume history and prompt prefix. + pub input: Vec, + /// Frozen model-visible tool declarations for this one execution. + pub tools: ToolSnapshot, + /// Explicit child run context, including lineage and host context data. + pub run_context: RunContext, +} + +/// The one execution handed to a [`crate::subagent::SubagentExecutor`]. +pub struct SubagentExecution { + /// The planner's complete, host-resolved execution description. + pub prepared: PreparedSubagent, + /// Cooperative cancellation shared with the lifecycle owner. + pub cancellation: CancellationToken, +} + +/// A neutral reference to a host-owned artifact. +/// +/// The reference intentionally contains no filesystem path or URL. Hosts own +/// artifact authorization and resolution. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ArtifactReference { + /// Stable host artifact identifier. + pub id: String, + /// Optional neutral media-type hint. + pub media_type: Option, + /// Opaque, non-location metadata for a host to interpret. + pub metadata: BTreeMap, +} + +/// A neutral suspension point that can later be supplied to a planner. +#[derive(Clone, Debug, PartialEq)] +pub struct SubagentPause { + /// Why execution needs input or an external host action. + pub reason: String, + /// Resume state captured at the suspension point. + pub resume: SubagentResume, +} + +/// A neutral non-successful but terminal completion. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SubagentIncomplete { + /// A host-safe explanation for incomplete work. + pub reason: String, +} + +/// The visible status of one subagent run. +#[derive(Clone, Debug, PartialEq)] +pub enum SubagentStatus { + /// The subagent completed normally. + Completed, + /// The subagent stopped at a resumable input boundary. + AwaitingInput(SubagentPause), + /// The subagent terminated without a complete result. + Incomplete(SubagentIncomplete), + /// Cooperative cancellation won the lifecycle race. + Cancelled, +} + +/// Complete neutral result of one subagent execution. +#[derive(Clone, Debug, PartialEq)] +pub struct SubagentOutcome { + /// Host-stable task identity. + pub task_id: String, + /// Final visible text, retained even when cancellation arrives after work. + pub output: String, + /// Complete model history without lossy transcript conversion. + pub history: Vec, + /// Terminal or pause state. + pub status: SubagentStatus, + /// Model usage reported by the nested execution exactly once. + pub usage: UsageTotals, + /// Host-owned artifacts represented by neutral references. + pub artifacts: Vec, +} + +impl SubagentOutcome { + /// Creates the truthful empty result used when cancellation prevents a + /// planner or executor from starting. + pub fn cancelled(task_id: impl Into) -> Self { + Self { + task_id: task_id.into(), + output: String::new(), + history: Vec::new(), + status: SubagentStatus::Cancelled, + usage: UsageTotals::default(), + artifacts: Vec::new(), + } + } + + pub(crate) fn cancelled_preserving(mut self) -> Self { + self.status = SubagentStatus::Cancelled; + self + } +} + +/// The durable input to [`crate::subagent::SubagentPersistence::save_pause`]. +#[derive(Clone, Debug, PartialEq)] +pub struct PersistedSubagentPause { + /// Durable scoped lifecycle identity. + pub key: SubagentTaskKey, + /// The resumable pause state. + pub pause: SubagentPause, +} + +/// Typed lifecycle failures. Adapters classify their errors at the seam that +/// owns them; the driver never flattens them into an untyped host error. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SubagentError { + /// The planner rejected or could not resolve a request. + Planning(String), + /// The executor could not finish a prepared run. + Execution(String), + /// Resume loading or outcome persistence failed. + Persistence(String), + /// A required planner, executor, or persistence seam was not provided. + MissingCapability(&'static str), + /// A host seam returned an outcome for a task other than the one reserved + /// by this lifecycle. The driver rejects it before any persistence or + /// terminal cache write can corrupt another task's record. + TaskIdMismatch { + /// Task id the caller reserved. + expected: String, + /// Task id returned by the planner or executor. + actual: String, + }, + /// Cooperative cancellation interrupted the lifecycle. + Cancelled, +} + +impl std::fmt::Display for SubagentError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Planning(message) => write!(f, "subagent planning failed: {message}"), + Self::Execution(message) => write!(f, "subagent execution failed: {message}"), + Self::Persistence(message) => write!(f, "subagent persistence failed: {message}"), + Self::MissingCapability(capability) => { + write!(f, "subagent host capability is unavailable: {capability}") + } + Self::TaskIdMismatch { expected, actual } => write!( + f, + "subagent host seam returned task id {actual:?}, expected {expected:?}" + ), + Self::Cancelled => write!(f, "subagent execution was cancelled"), + } + } +} + +impl std::error::Error for SubagentError {} From 2d19206a17af9c2bb3160b89762e5b54c710a187 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 19:29:39 +0300 Subject: [PATCH 06/18] feat(runtime): support host-prepared session turns Co-authored-by: Medulla --- crates/tinyagents-runtime/README.md | 21 +- crates/tinyagents-runtime/src/builder.rs | 30 +- crates/tinyagents-runtime/src/error.rs | 4 + crates/tinyagents-runtime/src/hooks.rs | 52 +- crates/tinyagents-runtime/src/lib.rs | 16 +- crates/tinyagents-runtime/src/session.rs | 331 ++-- crates/tinyagents-runtime/src/test.rs | 2026 ++++++++++++---------- crates/tinyagents-runtime/src/types.rs | 101 ++ docs/modules/runtime/README.md | 26 +- 9 files changed, 1499 insertions(+), 1108 deletions(-) diff --git a/crates/tinyagents-runtime/README.md b/crates/tinyagents-runtime/README.md index 08d4b5a2..519ec9ad 100644 --- a/crates/tinyagents-runtime/README.md +++ b/crates/tinyagents-runtime/README.md @@ -2,8 +2,8 @@ `tinyagents-runtime` provides the stateful session layer that sits between a host's turn policy and TinyAgents' provider-neutral harness. A `Session` owns -the mutable model history, a stable prefix, a frozen tool declaration snapshot, -and append-only transcript state. It has no host configuration, credentials, +the mutable model history, a stable prefix, and append-only transcript state. +Tool declarations are prepared afresh for every driver call. It has no host configuration, credentials, prompt construction, tool authorization, model selection, or event system. ## Host responsibilities @@ -20,9 +20,13 @@ The host supplies three narrow seams: `Clone` so reconciliation receives the current host context plus request, thread, stream, and resume options after the live `RunContext` moves into the driver. -- `SessionHooks` prepares a request and validates the candidate before the - commit point, then observes the durable result through `after_commit` and - one terminal state. It does not make policy decisions. +- `SessionHooks` prepares a request and mutable `TurnOptions` before + handoff. Its `TurnPreparation` can install a first-turn prefix, select the + one immutable `ToolSnapshot` for that request, and lazily choose a + `TranscriptTarget`. `before_commit` validates the candidate; `after_commit` + receives an exactly-once `CommitReceipt` containing the explicit context + snapshot and neutral transcript path/delta receipt. `on_terminal` receives + one truthful terminal state. It does not make policy decisions. ```rust,no_run use std::sync::Arc; @@ -36,7 +40,9 @@ let session = SessionBuilder::new(driver) # } ``` -To persist, add `SessionBuilder::transcript(locator, stem, meta)`. The runtime +To supply a default lazy destination, add `SessionBuilder::transcript(locator, stem, meta)`. +It does not open a transcript while building: a selected target binds only on +resume/first append and cannot be redirected after that. The runtime uses `tinyagents-session`'s `TranscriptHistory::append_turn_with_partial`, so a normal extension appends only the new tail and a reduced context writes one compaction record. A supplied partial driver outcome is represented through that single @@ -48,7 +54,8 @@ in-memory history and persisted snapshot unchanged. Every turn receives explicit `TurnOptions`, including its cancellation token and `RunContext`; no task-local data crosses the runtime boundary. The stable prefix is reconciled after resume and driver compaction without -duplication. Cancellation before the commit point leaves no durable mutation; +duplication. `Session::seed_history(history, raw)` is the explicit, lossless +resume/seed boundary; a host must not keep a second shadow history. Cancellation before the commit point leaves no durable mutation; once it succeeds, the turn remains successful. `after_commit` and terminal hooks get the committed outcome, but their error or a cooperative cancellation cannot relabel it. diff --git a/crates/tinyagents-runtime/src/builder.rs b/crates/tinyagents-runtime/src/builder.rs index 6259c884..dc928dce 100644 --- a/crates/tinyagents-runtime/src/builder.rs +++ b/crates/tinyagents-runtime/src/builder.rs @@ -11,7 +11,7 @@ use crate::{ pub struct SessionBuilder { driver: Arc>, codec: Option>>, - hooks: Arc, + hooks: Arc>, prefix: PrefixSnapshot, tools: ToolSnapshot, transcript: Option, @@ -43,7 +43,7 @@ impl SessionBuilder { } /// Installs optional host preparation/observation hooks. - pub fn hooks(mut self, hooks: Arc) -> Self { + pub fn hooks(mut self, hooks: Arc>) -> Self { self.hooks = hooks; self } @@ -79,21 +79,12 @@ impl SessionBuilder { /// Builds a session. A codec is required only when transcript persistence /// or transcript resume is configured. pub fn build(self) -> Result, RuntimeError> { - let (locator, stem, meta, history_handle) = if let Some(config) = self.transcript { - let handle = config - .locator - .open_stem(&config.stem, config.meta.clone()) - .map_err(|error| RuntimeError::Persistence(error.to_string()))?; - ( - Some(config.locator), - Some(config.stem), - Some(config.meta), - Some(handle), - ) - } else { - (None, None, None, None) - }; - if history_handle.is_some() && self.codec.is_none() { + let target = self.transcript.map(|config| crate::TranscriptTarget { + locator: config.locator, + stem: config.stem, + meta: config.meta, + }); + if target.is_some() && self.codec.is_none() { return Err(RuntimeError::MissingDependency("TranscriptCodec")); } Ok(Session::::new( @@ -102,10 +93,7 @@ impl SessionBuilder { self.hooks, self.prefix, self.tools, - locator, - stem, - meta, - history_handle, + target, )) } } diff --git a/crates/tinyagents-runtime/src/error.rs b/crates/tinyagents-runtime/src/error.rs index 9831d3c7..1ef75eeb 100644 --- a/crates/tinyagents-runtime/src/error.rs +++ b/crates/tinyagents-runtime/src/error.rs @@ -25,4 +25,8 @@ pub enum RuntimeError { /// The builder did not receive a required dependency. #[error("session builder requires {0}")] MissingDependency(&'static str), + /// A host attempted a state transition which would invalidate a committed + /// session invariant. + #[error("invalid session state: {0}")] + InvalidSessionState(String), } diff --git a/crates/tinyagents-runtime/src/hooks.rs b/crates/tinyagents-runtime/src/hooks.rs index fd182ccb..0c93e0b5 100644 --- a/crates/tinyagents-runtime/src/hooks.rs +++ b/crates/tinyagents-runtime/src/hooks.rs @@ -1,28 +1,41 @@ use async_trait::async_trait; -use crate::{RuntimeError, SessionTerminal, SessionTurnOutcome, SessionTurnRequest}; +use crate::{ + CommitReceipt, RuntimeError, SessionStateView, SessionTerminal, SessionTurnOutcome, + SessionTurnRequest, TranscriptTurnOptions, TurnOptions, TurnPreparation, +}; /// Host observation/preparation around a session turn. /// /// Hooks do not grant tools, select models, compose product prompts, or own -/// transcript state. They can prepare the input and observe committed results. +/// transcript state. A session serializes these mutable calls, so a host can +/// keep preparation state without task-local runtime state. #[async_trait] -pub trait SessionHooks: Send + Sync { - /// Runs before the driver sees the request. - async fn before_turn(&self, request: &mut SessionTurnRequest) -> Result<(), RuntimeError>; +pub trait SessionHooks: Send + Sync { + /// Runs before the driver sees the request and consumes the explicit + /// options. Returned values apply only to this driver invocation. + async fn before_turn( + &self, + request: &mut SessionTurnRequest, + options: &mut TurnOptions, + state: SessionStateView<'_>, + ) -> Result; /// Runs after the driver has produced a candidate and before it commits. /// Returning an error or observing cancellation therefore leaves no /// durable session mutation behind. - async fn after_turn(&self, outcome: &SessionTurnOutcome) -> Result<(), RuntimeError>; + async fn before_commit( + &self, + outcome: &SessionTurnOutcome, + options: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError>; /// Runs exactly once after a successful durable transcript commit. - /// /// Errors and cancellation observed here are deliberately observational: /// the result has already become durable and remains successful. - async fn after_commit(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + async fn after_commit(&self, _: CommitReceipt) -> Result<(), RuntimeError> { Ok(()) } /// Runs exactly once for every terminal turn result. - async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError>; + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError>; } /// A no-op hook set for hosts that need no lifecycle observation. @@ -30,17 +43,26 @@ pub trait SessionHooks: Send + Sync { pub struct NoopSessionHooks; #[async_trait] -impl SessionHooks for NoopSessionHooks { - async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { - Ok(()) +impl SessionHooks for NoopSessionHooks { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(TurnPreparation::default()) } - async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { Ok(()) } - async fn after_commit(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { + async fn after_commit(&self, _: CommitReceipt) -> Result<(), RuntimeError> { Ok(()) } - async fn on_terminal(&self, _: &SessionTerminal) -> Result<(), RuntimeError> { + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { Ok(()) } } diff --git a/crates/tinyagents-runtime/src/lib.rs b/crates/tinyagents-runtime/src/lib.rs index 7dbcf481..d2800115 100644 --- a/crates/tinyagents-runtime/src/lib.rs +++ b/crates/tinyagents-runtime/src/lib.rs @@ -23,8 +23,9 @@ pub use session::Session; pub use tinyagents_session::transcript::TranscriptPartial; pub use tools::ToolSnapshot; pub use types::{ - ResumeMode, SessionResume, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, - TranscriptTurnOptions, TurnOptions, + CommitReceipt, ResumeMode, SessionResume, SessionStateView, SessionTerminal, + SessionTurnOutcome, SessionTurnRequest, TranscriptCommitReceipt, TranscriptDelta, + TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; /// Converts between a host's lossless durable transcript dialect and the @@ -41,10 +42,13 @@ pub trait TranscriptCodec: Send + Sync { ) -> Result, RuntimeError>; /// Reconciles a model-history transition with the prior lossless durable - /// rows. `prior` must be treated as authoritative for fields absent from - /// inference messages (provider metadata, raw arguments, reasoning, ids, - /// and host extensions). The returned rows are the complete next logical - /// durable set; the history layer writes its delta atomically. + /// rows. `prior` is authoritative for fields absent from inference + /// messages (provider metadata, raw arguments, reasoning, ids, and host + /// extensions): unchanged model positions must retain their corresponding + /// raw row, including when `next` is a compaction replacement. The returned + /// rows are the complete next logical durable set; the history layer writes + /// its delta atomically. `options` is captured after `before_turn` mutates + /// the explicit turn options and before the driver consumes `RunContext`. fn reconcile( &self, prior: &[tinyagents_session::transcript::TranscriptMessage], diff --git a/crates/tinyagents-runtime/src/session.rs b/crates/tinyagents-runtime/src/session.rs index 10df4124..a9128e32 100644 --- a/crates/tinyagents-runtime/src/session.rs +++ b/crates/tinyagents-runtime/src/session.rs @@ -2,29 +2,29 @@ use std::{future::Future, sync::Arc}; use tinyagents_harness::CancellationToken; use tinyagents_session::transcript::{ - TranscriptHistory, TranscriptLocator, TranscriptMeta, TranscriptPartial, TranscriptTurn, + TranscriptHistory, TranscriptMessage, TranscriptPartial, TranscriptTurn, }; use tinyinference_llm::message::Message; use crate::{ - DriverRequest, PrefixSnapshot, ResumeMode, RuntimeError, SessionDriver, SessionHooks, - SessionResume, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, ToolSnapshot, - TranscriptCodec, TurnOptions, + CommitReceipt, DriverRequest, PrefixSnapshot, ResumeMode, RuntimeError, SessionDriver, + SessionHooks, SessionResume, SessionStateView, SessionTerminal, SessionTurnOutcome, + SessionTurnRequest, ToolSnapshot, TranscriptCodec, TranscriptCommitReceipt, TranscriptDelta, + TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; /// Host-neutral mutable state for one conversation session. pub struct Session { driver: Arc>, codec: Option>>, - hooks: Arc, + hooks: Arc>, prefix: PrefixSnapshot, - tools: ToolSnapshot, + default_tools: ToolSnapshot, history: Vec, - persisted: Vec, - locator: Option>, - stem: Option, - meta: Option, + persisted: Vec, + target: Option, transcript: Option>, + committed_turns: usize, } impl Session { @@ -32,13 +32,10 @@ impl Session { pub(crate) fn new( driver: Arc>, codec: Option>>, - hooks: Arc, + hooks: Arc>, prefix: PrefixSnapshot, - tools: ToolSnapshot, - locator: Option>, - stem: Option, - meta: Option, - transcript: Option>, + default_tools: ToolSnapshot, + target: Option, ) -> Self { Self { driver, @@ -46,12 +43,11 @@ impl Session { hooks, history: prefix.messages().to_vec(), prefix, - tools, + default_tools, persisted: Vec::new(), - locator, - stem, - meta, - transcript, + target, + transcript: None, + committed_turns: 0, } } @@ -60,14 +56,35 @@ impl Session { &self.history } - /// Returns the unchanging history prefix captured by the builder. + /// Returns the stable prefix currently applied to this session. pub fn prefix_snapshot(&self) -> &PrefixSnapshot { &self.prefix } - /// Returns the immutable model-visible tool declaration set. + /// Returns the builder compatibility default used only when a preparation + /// supplies no per-turn tool snapshot. pub fn tool_snapshot(&self) -> &ToolSnapshot { - &self.tools + &self.default_tools + } + + /// Seeds an uncommitted session from an explicit, lossless host snapshot. + /// + /// This replaces neither the host's raw rows nor their metadata. It is the + /// supported alternative to a host keeping a shadow history beside the + /// runtime. Seeding after any durable transition is rejected. + pub fn seed_history( + &mut self, + history: Vec, + raw: Vec, + ) -> Result<(), RuntimeError> { + if self.committed_turns != 0 { + return Err(RuntimeError::InvalidSessionState( + "cannot seed history after a committed turn".into(), + )); + } + self.history = self.with_prefix(history); + self.persisted = raw; + Ok(()) } /// Loads the selected durable transcript, retaining its lossless raw rows @@ -79,7 +96,7 @@ impl Session { if options.cancellation.is_cancelled() { return Err(RuntimeError::Cancelled); } - let Some(locator) = self.locator.as_ref() else { + let Some(target) = self.target.as_ref() else { return Ok(SessionResume { loaded: false, history: self.history.clone(), @@ -87,14 +104,11 @@ impl Session { }; let read = match options.resume { ResumeMode::Never => None, - ResumeMode::LatestForAgent => self - .stem - .as_deref() - .and_then(|stem| locator.latest_for_agent(stem)), + ResumeMode::LatestForAgent => target.locator.latest_for_agent(&target.stem), ResumeMode::Thread => options .thread_id .as_deref() - .and_then(|thread| locator.root_for_thread(thread)), + .and_then(|thread| target.locator.root_for_thread(thread)), }; let Some(read) = read else { return Ok(SessionResume { @@ -118,21 +132,37 @@ impl Session { let history = self.with_prefix(codec.decode_history(&transcript)?); self.history = history.clone(); self.persisted = transcript.messages; + // The discovered metadata, not the builder seed, is authoritative for + // the subsequent append. This keeps resume-only host fields intact. + if let Some(target) = self.target.as_mut() { + target.meta = transcript.meta; + } + // A successful explicit resume also fixes the target's history handle + // for later appends. Builder construction itself remains I/O-free. + if self.transcript.is_none() { + let target = self.target.as_ref().expect("target checked above"); + self.transcript = Some( + target + .locator + .open_stem(&target.stem, target.meta.clone()) + .map_err(|error| RuntimeError::Persistence(error.to_string()))?, + ); + } Ok(SessionResume { loaded: true, history, }) } - /// Executes and durably commits one state transition. + /// Executes and commits one state transition. pub async fn turn( &mut self, mut request: SessionTurnRequest, - options: TurnOptions, + mut options: TurnOptions, ) -> Result { let mut terminal_guard = TerminalGuard::new(self.hooks.clone()); let result = self - .turn_inner(&mut request, options, &mut terminal_guard) + .turn_inner(&mut request, &mut options, &mut terminal_guard) .await; if !terminal_guard.is_committed() { let terminal = match &result { @@ -142,9 +172,7 @@ impl Session { }; terminal_guard.set(terminal); } - // Terminal observation cannot revoke a successful durable commit. - // `finish` still schedules it exactly once; hook failures are - // deliberately observational rather than a second terminal result. + // Terminal observation cannot revoke a durable successful commit. let _ = terminal_guard.finish().await; result } @@ -152,38 +180,55 @@ impl Session { async fn turn_inner( &mut self, request: &mut SessionTurnRequest, - options: TurnOptions, - terminal_guard: &mut TerminalGuard, + options: &mut TurnOptions, + terminal_guard: &mut TerminalGuard, ) -> Result { + let state = SessionStateView { + history: &self.history, + raw_history: &self.persisted, + prefix: &self.prefix, + transcript_target: self.target.as_ref(), + committed_turns: self.committed_turns, + }; + let cancellation = options.cancellation.clone(); + let preparation = cancelable( + &cancellation, + self.hooks.before_turn(request, options, state), + ) + .await?; + let (tools, prepared_prefix) = self.apply_preparation(preparation)?; + // Preparation owns the current turn's target and explicit resume mode, + // so resolve only after it has made its changes visible. This is also + // why a selected target is checked for a codec before driver handoff. if options.resume != ResumeMode::Never { - self.resume(&options).await?; + self.resume(options).await?; + } + if let Some(prefix) = prepared_prefix { + self.apply_prefix(prefix)?; } - cancelable(&options.cancellation, self.hooks.before_turn(request)).await?; + let mut input = self.history.clone(); if input.last() != Some(&request.input) { input.push(request.input.clone()); } - // `RunContext` is intentionally consumed exactly once. There is no - // task-local fallback: the host context selected for this turn is what - // reaches model, middleware, and tool execution. let codec_options = options.transcript_options(); - let TurnOptions { - request_id, - thread_id, - stream, - cancellation, - run_context, - .. - } = options; - let run_context = run_context.with_cancellation(cancellation.clone()); + let request_id = options.request_id.clone(); + let thread_id = options.thread_id.clone(); + let stream = options.stream; + let cancellation = options.cancellation.clone(); + // `RunContext` is consumed exactly once. The host context captured in + // `codec_options` is the one after preparation and before handoff. + let run_context = std::mem::replace( + &mut options.run_context, + tinyagents_harness::context::RunContext::new( + tinyagents_harness::context::RunConfig::new("consumed-session-context"), + codec_options.context.clone(), + ), + ) + .with_cancellation(cancellation.clone()); let driver_result = tokio::select! { _ = cancellation.cancelled() => return Err(RuntimeError::Cancelled), - result = self.driver.execute(DriverRequest { - history: input, - tools: self.tools.clone(), - run_context, - stream, - }) => result, + result = self.driver.execute(DriverRequest { history: input, tools, run_context, stream }) => result, }; let outcome = match driver_result { Ok(outcome) => outcome, @@ -191,10 +236,7 @@ impl Session { if let Some(partial) = failure.partial { let partial_history = self.with_prefix(partial.history); let raw = self.encode(&self.history, &partial_history, &codec_options)?; - // `append_turn` is the only durable mutation. Do not - // append display partials first: a later append failure - // would leave an unreportable half-commit on disk. - self.persist( + let receipt = self.persist( &raw, request_id.as_deref(), thread_id.as_deref(), @@ -202,46 +244,101 @@ impl Session { )?; self.history = partial_history; self.persisted = raw; + if receipt.is_some() { + self.committed_turns += 1; + } } return Err(failure.error); } }; let candidate = self.with_prefix(outcome.history); - // `after_turn` is a pre-commit hook. It can reject or be cancelled - // without any durable mutation; after `persist` returns success this - // turn is committed and cancellation can no longer change its result. let committed = SessionTurnOutcome { history: candidate.clone(), output: outcome.output, interrupted: outcome.interrupted, }; - cancelable(&cancellation, self.hooks.after_turn(&committed)).await?; + cancelable( + &cancellation, + self.hooks.before_commit(&committed, &codec_options), + ) + .await?; if cancellation.is_cancelled() { return Err(RuntimeError::Cancelled); } let raw = self.encode(&self.history, &candidate, &codec_options)?; - self.persist(&raw, request_id.as_deref(), thread_id.as_deref(), None)?; + let transcript = self.persist(&raw, request_id.as_deref(), thread_id.as_deref(), None)?; self.history = committed.history.clone(); self.persisted = raw; - // Set the truthful durable terminal before invoking an observational - // finalizer. If the caller drops this future while it is running, the - // guard's Drop implementation still reports the completed commit. - terminal_guard.mark_committed(committed.clone()); - // This runs after `append_turn_with_partial` has made the logical - // transition durable. Failure or cooperative cancellation in a host - // finalizer is observational: it cannot relabel that committed turn. - let hooks = self.hooks.clone(); - let finalization = committed.clone(); - let _ = tokio::spawn(async move { hooks.after_commit(&finalization).await }).await; + self.committed_turns += 1; + // The receipt is constructed only after append and state replacement. + // Its hook and the completed terminal are owned by one task: errors or + // cancellation cannot relabel the successful durable transition, and + // dropping the caller future cannot drop finalization mid-flight. + let receipt = CommitReceipt { + outcome: committed.clone(), + options: codec_options, + transcript, + }; + let finalization = terminal_guard.finalize_commit(receipt); + // This await deliberately does not observe cancellation. If this turn + // future is dropped, dropping `JoinHandle` detaches rather than aborts + // the owned finalization task. + let _ = finalization.await; Ok(committed) } + fn apply_preparation( + &mut self, + preparation: TurnPreparation, + ) -> Result<(ToolSnapshot, Option), RuntimeError> { + if let Some(target) = preparation.transcript { + if self.transcript.is_some() { + if !self + .target + .as_ref() + .is_some_and(|bound| bound.same_binding(&target)) + { + return Err(RuntimeError::InvalidSessionState( + "cannot change a transcript target after it is bound".into(), + )); + } + } else { + self.target = Some(target); + } + } + if self.target.is_some() && self.codec.is_none() { + return Err(RuntimeError::MissingDependency("TranscriptCodec")); + } + // A returned snapshot never updates `default_tools`: it applies only + // to the `DriverRequest` being built by this call. + Ok(( + preparation + .tools + .unwrap_or_else(|| self.default_tools.clone()), + preparation.prefix, + )) + } + + fn apply_prefix(&mut self, prefix: PrefixSnapshot) -> Result<(), RuntimeError> { + if self.committed_turns != 0 + || !self.persisted.is_empty() + || self.history != self.prefix.messages() + { + return Err(RuntimeError::InvalidSessionState( + "cannot change a session prefix after history is present or committed".into(), + )); + } + self.history = prefix.messages().to_vec(); + self.prefix = prefix; + Ok(()) + } + fn encode( &self, previous: &[Message], next: &[Message], - options: &crate::TranscriptTurnOptions, - ) -> Result, RuntimeError> { + options: &TranscriptTurnOptions, + ) -> Result, RuntimeError> { match &self.codec { Some(codec) => codec.reconcile(&self.persisted, previous, next, options), None => Ok(Vec::new()), @@ -250,15 +347,24 @@ impl Session { fn persist( &mut self, - raw: &[tinyagents_session::transcript::TranscriptMessage], + raw: &[TranscriptMessage], request_id: Option<&str>, thread_id: Option<&str>, partial: Option<&TranscriptPartial>, - ) -> Result<(), RuntimeError> { - let (Some(transcript), Some(meta)) = (&self.transcript, &self.meta) else { - return Ok(()); + ) -> Result, RuntimeError> { + let Some(target) = self.target.as_mut() else { + return Ok(None); }; - let mut meta = meta.clone(); + if self.transcript.is_none() { + self.transcript = Some( + target + .locator + .open_stem(&target.stem, target.meta.clone()) + .map_err(|error| RuntimeError::Persistence(error.to_string()))?, + ); + } + let transcript = self.transcript.as_ref().expect("bound above"); + let mut meta = target.meta.clone(); meta.turn_count += 1; meta.updated = chrono::Utc::now().to_rfc3339(); meta.thread_id = thread_id.map(str::to_owned).or(meta.thread_id); @@ -274,8 +380,26 @@ impl Session { partial, ) .map_err(|error| RuntimeError::Persistence(error.to_string()))?; - self.meta = Some(meta); - Ok(()) + target.meta = meta; + let previous_len = self.persisted.len(); + let next_len = raw.len(); + let common_len = previous_len.min(next_len); + let delta = if next_len >= previous_len && raw[..common_len] == self.persisted[..common_len] + { + TranscriptDelta::Append { + previous_len, + appended: previous_len..next_len, + } + } else { + TranscriptDelta::Replace { + previous_len, + next_len, + } + }; + Ok(Some(TranscriptCommitReceipt { + path: transcript.path().to_path_buf(), + delta, + })) } fn with_prefix(&self, history: Vec) -> Vec { @@ -292,14 +416,14 @@ impl Session { /// Ensures a terminal hook is scheduled once even if a caller drops a turn /// future while it is awaiting preparation, driving, persistence, or hooks. -struct TerminalGuard { - hooks: Arc, +struct TerminalGuard { + hooks: Arc>, terminal: Option, committed: bool, } -impl TerminalGuard { - fn new(hooks: Arc) -> Self { +impl TerminalGuard { + fn new(hooks: Arc>) -> Self { Self { hooks, terminal: Some(SessionTerminal::Failed("session turn dropped".into())), @@ -311,9 +435,17 @@ impl TerminalGuard { self.terminal = Some(terminal); } - fn mark_committed(&mut self, outcome: SessionTurnOutcome) { - self.terminal = Some(SessionTerminal::Completed(outcome)); + fn finalize_commit(&mut self, receipt: CommitReceipt) -> tokio::task::JoinHandle<()> { + let terminal = SessionTerminal::Completed(receipt.outcome.clone()); + // Removing the guard's terminal transfers exactly-once ownership to + // the finalizer. `finish` and `Drop` then become no-ops for this turn. + self.terminal = None; self.committed = true; + let hooks = self.hooks.clone(); + tokio::spawn(async move { + let _ = hooks.after_commit(receipt).await; + let _ = hooks.on_terminal(terminal).await; + }) } fn is_committed(&self) -> bool { @@ -321,17 +453,14 @@ impl TerminalGuard { } async fn finish(mut self) -> Result<(), RuntimeError> { - let terminal = self.terminal.take().ok_or(RuntimeError::Hook( - "terminal guard already completed".into(), - ))?; - let hooks = self.hooks.clone(); - tokio::spawn(async move { hooks.on_terminal(&terminal).await }) - .await - .map_err(|error| RuntimeError::Hook(format!("terminal hook task failed: {error}")))? + let Some(terminal) = self.terminal.take() else { + return Ok(()); + }; + self.hooks.on_terminal(terminal).await } } -impl Drop for TerminalGuard { +impl Drop for TerminalGuard { fn drop(&mut self) { let Some(terminal) = self.terminal.take() else { return; @@ -339,7 +468,7 @@ impl Drop for TerminalGuard { let hooks = self.hooks.clone(); if let Ok(runtime) = tokio::runtime::Handle::try_current() { runtime.spawn(async move { - let _ = hooks.on_terminal(&terminal).await; + let _ = hooks.on_terminal(terminal).await; }); } } diff --git a/crates/tinyagents-runtime/src/test.rs b/crates/tinyagents-runtime/src/test.rs index 54c3bcca..0cfd5d46 100644 --- a/crates/tinyagents-runtime/src/test.rs +++ b/crates/tinyagents-runtime/src/test.rs @@ -12,26 +12,27 @@ use tinyagents_harness::{ runtime::AgentHarness, }; use tinyagents_session::transcript::{ - DisplayRecord, FileTranscriptHistory, FileTranscriptLocator, SessionTranscript, - TranscriptHistory, TranscriptLocator, TranscriptMessage, TranscriptMeta, TranscriptRead, - TranscriptTurn, read_transcript, read_transcript_display, + DisplayRecord, FileTranscriptLocator, SessionTranscript, TranscriptHistory, TranscriptLocator, + TranscriptMessage, TranscriptMeta, TranscriptRead, TranscriptTurn, read_transcript, + read_transcript_display, }; use tinyinference_llm::message::Message; use tinyinference_llm::providers::MockModel; use tinytools::{Tool, ToolResult, ToolSpec}; use crate::{ - DriverFailure, DriverOutcome, DriverRequest, HarnessDriver, PrefixSnapshot, ResumeMode, - RuntimeError, SessionBuilder, SessionDriver, SessionHooks, SessionTerminal, SessionTurnOutcome, - SessionTurnRequest, ToolSnapshot, TranscriptCodec, TurnOptions, + CommitReceipt, DriverFailure, DriverOutcome, DriverRequest, HarnessDriver, PrefixSnapshot, + ResumeMode, RuntimeError, SessionBuilder, SessionDriver, SessionHooks, SessionStateView, + SessionTerminal, SessionTurnOutcome, SessionTurnRequest, ToolSnapshot, TranscriptCodec, + TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; -struct FakeDriver { +struct Driver { results: Mutex>>, requests: Mutex>, } -impl FakeDriver { +impl Driver { fn new(results: Vec>) -> Self { Self { results: Mutex::new(results.into()), @@ -41,57 +42,23 @@ impl FakeDriver { } #[async_trait] -impl SessionDriver for FakeDriver { +impl SessionDriver for Driver { async fn execute(&self, request: DriverRequest) -> Result { self.requests.lock().unwrap().push(request); - self.results - .lock() - .unwrap() - .pop_front() - .expect("planned driver result") + self.results.lock().unwrap().pop_front().unwrap() } } -struct WaitingDriver; +struct WaitingDriver(Arc); #[async_trait] impl SessionDriver for WaitingDriver { async fn execute(&self, _: DriverRequest) -> Result { + self.0.notify_waiters(); std::future::pending().await } } -struct DropDriver { - started: Arc, -} - -#[async_trait] -impl SessionDriver for DropDriver { - async fn execute(&self, _: DriverRequest) -> Result { - self.started.notify_waiters(); - std::future::pending().await - } -} - -// Reconciliation retains an explicit clone for the codec after the live -// `RunContext` is consumed by the driver. -#[derive(Clone)] -struct HostContext(String); - -struct ContextDriver; - -#[async_trait] -impl SessionDriver for ContextDriver { - async fn execute( - &self, - request: DriverRequest, - ) -> Result { - Ok(outcome(vec![Message::assistant( - request.run_context.data.0, - )])) - } -} - struct RegisteredTool; #[async_trait] @@ -99,234 +66,52 @@ impl Tool for RegisteredTool { fn name(&self) -> &str { "registered" } - fn description(&self) -> &str { - "a registered matching tool" + "registered tool" } - fn parameters_schema(&self) -> serde_json::Value { - serde_json::json!({"type": "object"}) + serde_json::json!({}) } - async fn execute(&self, _: serde_json::Value) -> anyhow::Result { Ok(ToolResult::success("ok")) } } -#[derive(Default)] -struct BasicCodec { - decoded: Mutex>, -} - -impl TranscriptCodec for BasicCodec { - fn decode_history(&self, transcript: &SessionTranscript) -> Result, RuntimeError> { - self.decoded.lock().unwrap().push(transcript.clone()); - Ok(transcript - .messages - .iter() - .map(|message| match message.role.as_str() { - "assistant" => Message::assistant(&message.content), - "system" => Message::system(&message.content), - _ => Message::user(&message.content), - }) - .collect()) - } - - fn reconcile( - &self, - prior: &[TranscriptMessage], - previous: &[Message], - next: &[Message], - _: &crate::TranscriptTurnOptions, - ) -> Result, RuntimeError> { - // Preserve rows that still correspond to the prior model prefix; - // fresh model suffixes receive only the codec's explicit projection. - let raw_offset = (!prior.is_empty()) - .then(|| { - previous.windows(prior.len()).position(|window| { - window - .iter() - .zip(prior) - .all(|(model, row)| model.text() == row.content) - }) - }) - .flatten() - .unwrap_or(usize::MAX); - Ok(next - .iter() - .enumerate() - .map(|(index, message)| { - let raw_index = index.checked_sub(raw_offset); - if let Some(raw_index) = raw_index.filter(|index| { - *index < prior.len() && previous.get(index + raw_offset) == Some(message) - }) { - return prior[raw_index].clone(); - } - let role = match message { - Message::System(_) => "system", - Message::User(_) => "user", - Message::Assistant(_) => "assistant", - Message::Tool(_) => "tool", - }; - TranscriptMessage::new(role, message.text()) - }) - .collect()) - } -} - -#[derive(Default)] -struct RecordingHooks { - events: Mutex>, - fail_before: bool, - fail_after: bool, - terminal_notified: Option>, -} - -struct WaitingAfterHooks { - started: tokio::sync::Notify, - events: Mutex>, -} - -struct WaitingBeforeHooks { - started: tokio::sync::Notify, -} - -struct WaitingPostCommitHooks { - started: tokio::sync::Notify, - terminal: tokio::sync::Notify, - terminals: Mutex>, -} - -#[async_trait] -impl SessionHooks for WaitingBeforeHooks { - async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { - self.started.notify_waiters(); - std::future::pending().await - } - - async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { - Ok(()) - } - - async fn on_terminal(&self, _: &SessionTerminal) -> Result<(), RuntimeError> { - Ok(()) - } -} - -#[async_trait] -impl SessionHooks for WaitingAfterHooks { - async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { - Ok(()) - } - - async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { - self.started.notify_waiters(); - std::future::pending().await - } - - async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { - self.events.lock().unwrap().push(format!("{terminal:?}")); - Ok(()) - } -} - -#[async_trait] -impl SessionHooks for WaitingPostCommitHooks { - async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { - Ok(()) - } - - async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { - Ok(()) - } - - async fn after_commit(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { - self.started.notify_waiters(); - std::future::pending().await - } - - async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { - self.terminals.lock().unwrap().push(terminal.clone()); - self.terminal.notify_waiters(); - Ok(()) +fn outcome(history: Vec) -> DriverOutcome { + DriverOutcome { + history, + output: Some("ok".into()), + partial: None, + interrupted: false, } } -#[async_trait] -impl SessionHooks for RecordingHooks { - async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { - self.events.lock().unwrap().push("before".into()); - if self.fail_before { - Err(RuntimeError::Hook("before".into())) - } else { - Ok(()) - } - } - async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { - self.events.lock().unwrap().push("after".into()); - if self.fail_after { - Err(RuntimeError::Hook("after".into())) - } else { - Ok(()) - } - } - async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { - let terminal = match terminal { - SessionTerminal::Completed(_) => "Completed".to_owned(), - other => format!("{other:?}"), - }; - self.events - .lock() - .unwrap() - .push(format!("terminal:{terminal}")); - if let Some(notify) = &self.terminal_notified { - notify.notify_waiters(); - } - Ok(()) +fn meta() -> TranscriptMeta { + TranscriptMeta { + agent_name: "agent".into(), + agent_id: Some("agent-id".into()), + agent_type: None, + dispatcher: "test".into(), + provider: None, + model: None, + created: "then".into(), + updated: "then".into(), + turn_count: 0, + input_tokens: 0, + output_tokens: 0, + cached_input_tokens: 0, + charged_amount_usd: 0.0, + thread_id: None, + task_id: None, } } #[derive(Default)] -struct FinalizationHooks { - post_commits: Mutex>, - terminals: Mutex>, - fail_post_commit: bool, - cancel_on_post_commit: Option, -} - -#[async_trait] -impl SessionHooks for FinalizationHooks { - async fn before_turn(&self, _: &mut SessionTurnRequest) -> Result<(), RuntimeError> { - Ok(()) - } - - async fn after_turn(&self, _: &SessionTurnOutcome) -> Result<(), RuntimeError> { - Ok(()) - } - - async fn after_commit(&self, outcome: &SessionTurnOutcome) -> Result<(), RuntimeError> { - self.post_commits.lock().unwrap().push(outcome.clone()); - if let Some(cancellation) = &self.cancel_on_post_commit { - cancellation.cancel(); - } - if self.fail_post_commit { - Err(RuntimeError::Hook("post-commit".into())) - } else { - Ok(()) - } - } - - async fn on_terminal(&self, terminal: &SessionTerminal) -> Result<(), RuntimeError> { - self.terminals.lock().unwrap().push(terminal.clone()); - Ok(()) - } -} - struct MemoryHistory { path: PathBuf, - session: Mutex>, - turns: Mutex>>, - fail: bool, + state: Mutex>, + opens: Mutex, + fail: Mutex, cancel_after_append: Mutex>, } @@ -335,17 +120,16 @@ impl TranscriptRead for MemoryHistory { &self.path } fn read_session(&self) -> anyhow::Result> { - Ok(self.session.lock().unwrap().clone()) + Ok(self.state.lock().unwrap().clone()) } } impl TranscriptHistory for MemoryHistory { fn append_turn(&self, turn: TranscriptTurn<'_>) -> anyhow::Result<()> { - if self.fail { + if *self.fail.lock().unwrap() { anyhow::bail!("planned persistence failure"); } - self.turns.lock().unwrap().push(turn.next.to_vec()); - *self.session.lock().unwrap() = Some(SessionTranscript { + *self.state.lock().unwrap() = Some(SessionTranscript { meta: turn.meta.clone(), messages: turn.next.to_vec(), }); @@ -356,7 +140,7 @@ impl TranscriptHistory for MemoryHistory { } fn messages(&self) -> anyhow::Result> { Ok(self - .session + .state .lock() .unwrap() .as_ref() @@ -374,384 +158,267 @@ impl TranscriptHistory for MemoryHistory { } } -struct MemoryLocator { - history: Arc, -} - -/// A real transcript path whose single atomic turn operation fails before it -/// writes. This catches the old two-write partial path: if runtime wrote a -/// display partial first, the path below would exist after the failure. -struct FailingFileHistory(FileTranscriptHistory); +struct Locator(Arc); -impl TranscriptRead for FailingFileHistory { - fn path(&self) -> &Path { - self.0.path() +impl TranscriptLocator for Locator { + fn latest_for_agent(&self, _: &str) -> Option> { + Some(self.0.clone()) } - - fn read_session(&self) -> anyhow::Result> { - TranscriptRead::read_session(&self.0) + fn root_for_thread(&self, _: &str) -> Option> { + Some(self.0.clone()) } -} - -impl TranscriptHistory for FailingFileHistory { - fn append_turn(&self, _: TranscriptTurn<'_>) -> anyhow::Result<()> { - anyhow::bail!("planned atomic append failure") + fn open_stem(&self, _: &str, _: TranscriptMeta) -> anyhow::Result> { + *self.0.opens.lock().unwrap() += 1; + Ok(self.0.clone()) } +} - fn messages(&self) -> anyhow::Result> { - TranscriptHistory::messages(&self.0) - } +fn locator(session: Option) -> (Arc, Arc) { + let dir = tempfile::tempdir().unwrap().keep(); + let history = Arc::new(MemoryHistory { + path: dir.join("session.jsonl"), + state: Mutex::new(session), + opens: Mutex::new(0), + fail: Mutex::new(false), + cancel_after_append: Mutex::new(None), + }); + (Arc::new(Locator(history.clone())), history) +} - fn append(&self, _: TranscriptMessage) -> anyhow::Result<()> { - anyhow::bail!("planned atomic append failure") - } +#[derive(Default)] +struct Codec { + seen_prior: Mutex>>, + seen_options: Mutex>, +} - fn replace(&self, _: &[TranscriptMessage]) -> anyhow::Result<()> { - anyhow::bail!("planned atomic append failure") +impl TranscriptCodec for Codec { + fn decode_history(&self, transcript: &SessionTranscript) -> Result, RuntimeError> { + Ok(transcript + .messages + .iter() + .map(|row| Message::user(&row.content)) + .collect()) } - - fn clear(&self) -> anyhow::Result<()> { - anyhow::bail!("planned atomic append failure") + fn reconcile( + &self, + prior: &[TranscriptMessage], + _: &[Message], + next: &[Message], + options: &TranscriptTurnOptions, + ) -> Result, RuntimeError> { + self.seen_prior.lock().unwrap().push(prior.to_vec()); + self.seen_options + .lock() + .unwrap() + .push(format!("{:?}", options.request_id)); + if next.len() < prior.len() { + return Ok(next + .iter() + .map(|message| TranscriptMessage::new("assistant", message.text())) + .collect()); + } + let mut rows = prior.to_vec(); + if rows.len() < next.len() { + rows.extend( + next[rows.len()..] + .iter() + .map(|message| TranscriptMessage::new("assistant", message.text())), + ); + } + Ok(rows) } } -struct FailingFileLocator { - history: Arc, +#[derive(Default)] +struct Events { + terminal: Mutex>, + receipts: Mutex>, + before_commit: Mutex, + order: Mutex>, } -impl TranscriptLocator for FailingFileLocator { - fn latest_for_agent(&self, _: &str) -> Option> { - Some(self.history.clone()) - } - - fn root_for_thread(&self, _: &str) -> Option> { - Some(self.history.clone()) - } - - fn open_stem(&self, _: &str, _: TranscriptMeta) -> anyhow::Result> { - Ok(self.history.clone()) - } +struct Hook { + preparations: Mutex>, + events: Arc, + fail_commit: bool, + fail_post: bool, } -impl TranscriptLocator for MemoryLocator { - fn latest_for_agent(&self, _: &str) -> Option> { - Some(self.history.clone()) - } - fn root_for_thread(&self, _: &str) -> Option> { - Some(self.history.clone()) +#[async_trait] +impl SessionHooks for Hook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(self + .preparations + .lock() + .unwrap() + .pop_front() + .unwrap_or_default()) } - fn open_stem(&self, _: &str, _: TranscriptMeta) -> anyhow::Result> { - Ok(self.history.clone()) + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + *self.events.before_commit.lock().unwrap() += 1; + if self.fail_commit { + Err(RuntimeError::Hook("rejected".into())) + } else { + Ok(()) + } } -} - -fn meta() -> TranscriptMeta { - TranscriptMeta { - agent_name: "agent".into(), - agent_id: None, - agent_type: None, - dispatcher: "test".into(), - provider: None, - model: None, - created: "now".into(), - updated: "now".into(), - turn_count: 0, - input_tokens: 0, - output_tokens: 0, - cached_input_tokens: 0, - charged_amount_usd: 0.0, - thread_id: None, - task_id: None, + async fn after_commit(&self, receipt: CommitReceipt) -> Result<(), RuntimeError> { + self.events.receipts.lock().unwrap().push(receipt); + self.events.order.lock().unwrap().push("after_commit"); + if self.fail_post { + Err(RuntimeError::Hook("post-commit".into())) + } else { + Ok(()) + } } -} - -fn outcome(history: Vec) -> DriverOutcome { - DriverOutcome { - history, - output: Some("done".into()), - partial: None, - interrupted: false, + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.events.terminal.lock().unwrap().push(terminal); + self.events.order.lock().unwrap().push("terminal"); + Ok(()) } } -fn memory_locator( - session: Option, - fail: bool, -) -> (Arc, Arc) { - let dir = tempfile::tempdir().unwrap().keep(); - let history = Arc::new(MemoryHistory { - path: dir.join("session.jsonl"), - session: Mutex::new(session), - turns: Mutex::new(Vec::new()), - fail, - cancel_after_append: Mutex::new(None), - }); +fn hook(preparations: Vec) -> (Arc, Arc) { + let events = Arc::new(Events::default()); ( - Arc::new(MemoryLocator { - history: history.clone(), + Arc::new(Hook { + preparations: Mutex::new(preparations.into()), + events: events.clone(), + fail_commit: false, + fail_post: false, }), - history, + events, ) } +fn tools(name: &str) -> ToolSnapshot { + ToolSnapshot::new(vec![ToolSpec { + name: name.into(), + description: name.into(), + parameters: serde_json::json!({}), + }]) + .unwrap() +} + #[tokio::test] -async fn first_turn_commits_history_and_preserves_prefix() { - let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::system("stable"), - Message::user("hi"), - Message::assistant("hello"), +async fn prepared_first_turn_prefix_replaces_empty_builder_prefix() { + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::system("prepared"), + Message::assistant("ok"), ]))])); + let (hook, _) = hook(vec![TurnPreparation { + prefix: Some(PrefixSnapshot::new(vec![Message::system("prepared")])), + ..Default::default() + }]); let mut session = SessionBuilder::new(driver.clone()) - .prefix(PrefixSnapshot::new(vec![Message::system("stable")])) + .hooks(hook) .build() .unwrap(); - let committed = session + session .turn( SessionTurnRequest::new(Message::user("hi")), TurnOptions::default(), ) .await .unwrap(); - assert_eq!(committed.history, session.history()); assert_eq!( - session.prefix_snapshot().messages(), - &[Message::system("stable")] - ); - assert_eq!( - driver.requests.lock().unwrap()[0].history, - vec![Message::system("stable"), Message::user("hi")] + driver.requests.lock().unwrap()[0].history[0], + Message::system("prepared") ); -} - -#[tokio::test] -async fn generic_session_passes_the_explicit_host_context_to_driver() { - let mut session = SessionBuilder::::new(Arc::new(ContextDriver)) - .build() - .unwrap(); - let options = TurnOptions { - request_id: None, - thread_id: None, - stream: false, - resume: ResumeMode::Never, - cancellation: tinyagents_harness::CancellationToken::new(), - run_context: RunContext::new(RunConfig::new("host"), HostContext("host context".into())), - }; - let result = session - .turn(SessionTurnRequest::new(Message::user("x")), options) - .await - .unwrap(); assert_eq!( - result.history.last().map(Message::text).as_deref(), - Some("host context") + session.prefix_snapshot().messages(), + &[Message::system("prepared")] ); } #[tokio::test] -async fn harness_driver_uses_the_pinned_explicit_model_entry_point() { - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("from harness"))); - harness.register_tool(Arc::new(RegisteredTool)); - let snapshot = ToolSnapshot::new(harness.tools().declared_specs()).unwrap(); - let driver = Arc::new(HarnessDriver::new(Arc::new(harness), Arc::new(()))); - let mut session = SessionBuilder::new(driver) - .tool_snapshot(snapshot) +async fn prepared_tools_are_dynamic_and_never_reused() { + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![Message::assistant("one")])), + Ok(outcome(vec![Message::assistant("two")])), + ])); + let (hook, _) = hook(vec![ + TurnPreparation::with_tools(tools("one")), + TurnPreparation::with_tools(tools("two")), + ]); + let mut session = SessionBuilder::new(driver.clone()) + .hooks(hook) .build() .unwrap(); - let result = session - .turn( - SessionTurnRequest::new(Message::user("hello")), - TurnOptions::default(), - ) - .await - .unwrap(); - assert_eq!(result.output.as_deref(), Some("from harness")); - assert_eq!( - result.history.last().map(Message::text).as_deref(), - Some("from harness") - ); + for input in ["a", "b"] { + session + .turn( + SessionTurnRequest::new(Message::user(input)), + TurnOptions::default(), + ) + .await + .unwrap(); + } + let requests = driver.requests.lock().unwrap(); + assert_eq!(requests[0].tools.specs()[0].name, "one"); + assert_eq!(requests[1].tools.specs()[0].name, "two"); } #[tokio::test] -async fn trailing_input_is_deduplicated_and_tool_snapshot_is_immutable() { - let driver = Arc::new(FakeDriver::new(vec![ +async fn trailing_input_is_deduplicated_and_tool_collisions_fail_closed() { + let driver = Arc::new(Driver::new(vec![ Ok(outcome(vec![Message::user("same")])), Ok(outcome(vec![ Message::user("same"), - Message::assistant("two"), + Message::assistant("ok"), ])), ])); - let tools = ToolSnapshot::new(vec![ToolSpec { - name: "echo".into(), - description: "x".into(), - parameters: serde_json::json!({}), - }]) - .unwrap(); - let mut session = SessionBuilder::new(driver.clone()) - .tool_snapshot(tools) - .build() - .unwrap(); - session - .turn( - SessionTurnRequest::new(Message::user("same")), - TurnOptions::default(), - ) - .await - .unwrap(); - session - .turn( - SessionTurnRequest::new(Message::user("same")), - TurnOptions::default(), - ) - .await - .unwrap(); - let requests = driver.requests.lock().unwrap(); - assert_eq!(requests[1].history, vec![Message::user("same")]); - assert_eq!(requests[0].tools.specs()[0].name, "echo"); - assert_eq!(session.tool_snapshot().specs()[0].name, "echo"); -} - -#[test] -fn tool_snapshots_dedup_identical_names_and_reject_collisions() { - let spec = ToolSpec { - name: "read".into(), - description: "read".into(), - parameters: serde_json::json!({}), - }; + let mut session = SessionBuilder::new(driver.clone()).build().unwrap(); + for _ in 0..2 { + session + .turn( + SessionTurnRequest::new(Message::user("same")), + TurnOptions::default(), + ) + .await + .unwrap(); + } assert_eq!( - ToolSnapshot::new(vec![spec.clone(), spec]) - .unwrap() - .specs() - .len(), - 1 + driver.requests.lock().unwrap()[1].history, + vec![Message::user("same")] ); assert!(matches!( ToolSnapshot::new(vec![ ToolSpec { - name: "read".into(), + name: "same".into(), description: "one".into(), parameters: serde_json::json!({}) }, ToolSpec { - name: "read".into(), + name: "same".into(), description: "two".into(), parameters: serde_json::json!({}) }, ]), - Err(RuntimeError::ToolNameCollision(name)) if name == "read" - )); -} - -#[tokio::test] -async fn resume_passes_full_durable_transcript_to_codec() { - let mut durable = TranscriptMessage::new("user", "persisted"); - durable.extra_metadata = Some(serde_json::json!({"unmodified": true})); - let transcript = SessionTranscript { - meta: meta(), - messages: vec![durable.clone()], - }; - let (locator, _) = memory_locator(Some(transcript), false); - let codec = Arc::new(BasicCodec::default()); - let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![]))) - .codec(codec.clone()) - .transcript(locator, "agent", meta()) - .build() - .unwrap(); - let resumed = session - .resume(&TurnOptions { - resume: ResumeMode::LatestForAgent, - ..TurnOptions::default() - }) - .await - .unwrap(); - assert!(resumed.loaded); - assert_eq!(resumed.history, vec![Message::user("persisted")]); - assert_eq!( - codec.decoded.lock().unwrap()[0].messages[0].extra_metadata, - durable.extra_metadata - ); -} - -#[tokio::test] -async fn append_only_delta_and_failure_rollback_are_owned_by_session() { - let (locator, history) = memory_locator(None, false); - let codec = Arc::new(BasicCodec::default()); - let driver = Arc::new(FakeDriver::new(vec![ - Ok(outcome(vec![Message::user("one"), Message::assistant("a")])), - Ok(outcome(vec![Message::assistant("compacted")])), - ])); - let mut session = SessionBuilder::new(driver) - .codec(codec) - .transcript(locator, "agent", meta()) - .build() - .unwrap(); - session - .turn( - SessionTurnRequest::new(Message::user("one")), - TurnOptions::default(), - ) - .await - .unwrap(); - session - .turn( - SessionTurnRequest::new(Message::user("two")), - TurnOptions::default(), - ) - .await - .unwrap(); - { - let turns = history.turns.lock().unwrap(); - assert_eq!(turns.len(), 2); - assert_eq!( - turns[1], - vec![TranscriptMessage::new("assistant", "compacted")] - ); - } - - let (bad_locator, _) = memory_locator(None, true); - let bad = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![Message::user( - "will-not-commit", - )]))])); - let mut failing = SessionBuilder::new(bad) - .codec(Arc::new(BasicCodec::default())) - .transcript(bad_locator, "bad", meta()) - .build() - .unwrap(); - assert!(matches!( - failing - .turn( - SessionTurnRequest::new(Message::user("will-not-commit")), - TurnOptions::default() - ) - .await, - Err(RuntimeError::Persistence(_)) + Err(RuntimeError::ToolNameCollision(_)) )); - assert!(failing.history().is_empty()); } #[tokio::test] -async fn partial_failure_never_leaves_a_display_partial_on_disk() { - let directory = tempfile::tempdir().unwrap(); - let history = Arc::new(FailingFileHistory( - FileTranscriptHistory::new(directory.path(), "agent", meta()).unwrap(), - )); - let path = history.path().to_path_buf(); - let locator = Arc::new(FailingFileLocator { - history: history.clone(), - }); - let driver = Arc::new(FakeDriver::new(vec![Err(DriverFailure { - error: RuntimeError::Driver("interrupted".into()), - partial: Some(DriverOutcome { - history: vec![Message::assistant("partial model history")], - output: None, - partial: Some(crate::TranscriptPartial::new("display partial")), - interrupted: true, - }), - })])); - let mut session = SessionBuilder::new(driver) - .codec(Arc::new(BasicCodec::default())) - .transcript(locator, "agent", meta()) - .build() - .unwrap(); +async fn persistence_failure_rolls_back_and_a_partial_never_falls_back_to_two_writes() { + let (failing_locator, history) = locator(None); + *history.fail.lock().unwrap() = true; + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("never"), + ]))]))) + .codec(Arc::new(Codec::default())) + .transcript(failing_locator, "agent", meta()) + .build() + .unwrap(); assert!(matches!( session .turn( @@ -761,32 +428,24 @@ async fn partial_failure_never_leaves_a_display_partial_on_disk() { .await, Err(RuntimeError::Persistence(_)) )); - assert!(!path.exists(), "failed partial turn wrote {path:?}"); assert!(session.history().is_empty()); -} - -#[tokio::test] -async fn partial_failure_commits_model_history_and_display_partial_together() { - let directory = tempfile::tempdir().unwrap(); - let locator = Arc::new(FileTranscriptLocator::new(directory.path())); - let driver = Arc::new(FakeDriver::new(vec![Err(DriverFailure { + assert!(history.state.lock().unwrap().is_none()); + + let (locator, history) = locator(None); + let partial = DriverOutcome { + history: vec![Message::assistant("recoverable")], + output: None, + partial: Some(crate::TranscriptPartial::new("display only")), + interrupted: true, + }; + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Err(DriverFailure { error: RuntimeError::Driver("interrupted".into()), - partial: Some(DriverOutcome { - history: vec![Message::assistant("recoverable model history")], - output: None, - partial: Some(crate::TranscriptPartial { - content: "display partial".into(), - reasoning_content: Some("thinking".into()), - iteration: Some(3), - }), - interrupted: true, - }), - })])); - let mut session = SessionBuilder::new(driver) - .codec(Arc::new(BasicCodec::default())) - .transcript(locator, "agent", meta()) - .build() - .unwrap(); + partial: Some(partial), + })]))) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); assert!(matches!( session .turn( @@ -794,280 +453,287 @@ async fn partial_failure_commits_model_history_and_display_partial_together() { TurnOptions::default() ) .await, - Err(RuntimeError::Driver(_)) + Err(RuntimeError::Persistence(_)) )); - - let path = directory.path().join("session_raw/agent.jsonl"); - let model = read_transcript(&path).unwrap(); - assert_eq!( - model - .messages - .iter() - .map(|message| message.content.as_str()) - .collect::>(), - vec!["recoverable model history"] - ); - let display = read_transcript_display(&path).unwrap(); - assert!(display.records.iter().any(|record| matches!( - record, - DisplayRecord::Message(message) - if message.interrupted - && message.message.content == "display partial" - && message.reasoning_content.as_deref() == Some("thinking") - && message.iteration == Some(3) - ))); - assert_eq!( - session.history(), - &[Message::assistant("recoverable model history")] - ); + assert!(history.state.lock().unwrap().is_none()); } #[tokio::test] -async fn resume_new_turn_retains_durable_metadata_and_restores_prefix_once() { - let mut durable = TranscriptMessage::new("user", "persisted"); - durable.extra_metadata = Some(serde_json::json!({"provider": {"raw": true}})); - let transcript = SessionTranscript { - meta: meta(), - messages: vec![durable.clone()], - }; - let (locator, history) = memory_locator(Some(transcript), false); - let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - // A compaction/faulty driver omitted the stable prefix. - Message::user("persisted"), - Message::assistant("new"), +async fn prefix_reconciliation_preserves_maximal_overlap_after_driver_compaction() { + let prefix = PrefixSnapshot::new(vec![Message::system("a"), Message::system("b")]); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::system("b"), + Message::assistant("answer"), ]))])); - let mut session = SessionBuilder::new(driver.clone()) - .codec(Arc::new(BasicCodec::default())) - .prefix(PrefixSnapshot::new(vec![Message::system("stable")])) - .transcript(locator, "agent", meta()) - .build() - .unwrap(); - session + let mut session = SessionBuilder::new(driver).prefix(prefix).build().unwrap(); + let outcome = session .turn( - SessionTurnRequest::new(Message::user("next")), - TurnOptions { - resume: ResumeMode::LatestForAgent, - ..TurnOptions::default() - }, + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), ) .await .unwrap(); - let committed = history.session.lock().unwrap().clone().unwrap(); - assert_eq!(session.history()[0], Message::system("stable")); assert_eq!( - driver.requests.lock().unwrap()[0].history[0], - Message::system("stable") + outcome.history, + vec![ + Message::system("a"), + Message::system("b"), + Message::assistant("answer") + ] ); - assert_eq!(committed.messages[1].extra_metadata, durable.extra_metadata); - assert_eq!( +} + +#[tokio::test] +async fn rejected_before_commit_leaves_no_durable_state() { + let (locator, history) = locator(None); + let events = Arc::new(Events::default()); + let hook = Arc::new(Hook { + preparations: Mutex::new(VecDeque::new()), + events, + fail_commit: true, + fail_post: false, + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("candidate"), + ]))]))) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook) + .build() + .unwrap(); + assert!(matches!( session - .history() - .iter() - .filter(|m| **m == Message::system("stable")) - .count(), - 1 - ); + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::Hook(_)) + )); + assert!(history.state.lock().unwrap().is_none()); + assert!(session.history().is_empty()); } #[tokio::test] -async fn prefix_reconciliation_uses_maximal_suffix_prefix_overlap() { - let prefix = PrefixSnapshot::new(vec![ - Message::system("stable-a"), - Message::system("stable-b"), - Message::system("stable-c"), - ]); - for (returned, expected) in [ - ( - vec![Message::system("stable-c"), Message::assistant("partial")], - vec![ - Message::system("stable-a"), - Message::system("stable-b"), - Message::system("stable-c"), - Message::assistant("partial"), - ], - ), - ( - vec![ - Message::system("stable-a"), - Message::system("stable-b"), - Message::system("stable-c"), - Message::assistant("complete"), - ], - vec![ - Message::system("stable-a"), - Message::system("stable-b"), - Message::system("stable-c"), - Message::assistant("complete"), - ], - ), - ( - vec![Message::assistant("missing")], - vec![ - Message::system("stable-a"), - Message::system("stable-b"), - Message::system("stable-c"), - Message::assistant("missing"), - ], - ), - ] { - let mut session = - SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(returned))]))) - .prefix(prefix.clone()) - .build() - .unwrap(); - let committed = session +async fn post_commit_errors_cannot_relabel_a_successful_turn() { + let (locator, history) = locator(None); + let events = Arc::new(Events::default()); + let hook = Arc::new(Hook { + preparations: Mutex::new(VecDeque::new()), + events: events.clone(), + fail_commit: false, + fail_post: true, + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("done"), + ]))]))) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook) + .build() + .unwrap(); + assert!( + session .turn( SessionTurnRequest::new(Message::user("x")), - TurnOptions::default(), + TurnOptions::default() ) .await - .unwrap(); - assert_eq!(committed.history, expected); - } + .is_ok() + ); + assert!(history.state.lock().unwrap().is_some()); + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Completed(_)] + )); } #[tokio::test] -async fn hooks_are_ordered_terminal_once_and_driver_errors_are_terminal() { - let hooks = Arc::new(RecordingHooks::default()); - let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("ok"), - ]))])); +async fn harness_driver_keeps_the_explicit_model_and_tool_snapshot_boundary() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("from harness"))); + harness.register_tool(Arc::new(RegisteredTool)); + let snapshot = ToolSnapshot::new(harness.tools().declared_specs()).unwrap(); + let driver = Arc::new(HarnessDriver::new(Arc::new(harness), Arc::new(()))); let mut session = SessionBuilder::new(driver) - .hooks(hooks.clone()) + .tool_snapshot(snapshot) .build() .unwrap(); - session + let committed = session .turn( SessionTurnRequest::new(Message::user("x")), TurnOptions::default(), ) .await .unwrap(); - assert_eq!( - *hooks.events.lock().unwrap(), - vec!["before", "after", "terminal:Completed"] - ); + assert_eq!(committed.output.as_deref(), Some("from harness")); +} - let hooks = Arc::new(RecordingHooks::default()); - let driver = Arc::new(FakeDriver::new(vec![Err(DriverFailure { - error: RuntimeError::Driver("boom".into()), - partial: None, - })])); +#[tokio::test] +async fn harness_driver_rejects_a_snapshot_not_registered_by_the_harness() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("unreachable"))); + let driver = Arc::new(HarnessDriver::new(Arc::new(harness), Arc::new(()))); let mut session = SessionBuilder::new(driver) - .hooks(hooks.clone()) + .tool_snapshot(tools("not-registered")) .build() .unwrap(); - assert!(matches!( + assert_eq!( session .turn( SessionTurnRequest::new(Message::user("x")), TurnOptions::default() ) .await, - Err(RuntimeError::Driver(_)) - )); - assert_eq!( - *hooks.events.lock().unwrap(), - vec!["before", "terminal:Failed(\"driver failed: boom\")"] + Err(RuntimeError::ToolSnapshotMismatch) ); } #[tokio::test] -async fn failed_precommit_hook_and_tool_mismatch_leave_no_commit() { - let hooks = Arc::new(RecordingHooks { - fail_after: true, - ..RecordingHooks::default() - }); - let (locator, history) = memory_locator(None, false); - let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("x"), - ]))]))) - .codec(Arc::new(BasicCodec::default())) - .hooks(hooks) +async fn file_history_commits_partial_model_history_and_display_only_partial_together() { + let directory = tempfile::tempdir().unwrap(); + let locator = Arc::new(FileTranscriptLocator::new(directory.path())); + let partial = DriverOutcome { + history: vec![Message::assistant("recoverable")], + output: None, + partial: Some(crate::TranscriptPartial::new("display partial")), + interrupted: true, + }; + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Err(DriverFailure { + error: RuntimeError::Driver("interrupted".into()), + partial: Some(partial), + })]))) + .codec(Arc::new(Codec::default())) .transcript(locator, "agent", meta()) .build() .unwrap(); - assert!(matches!( + assert!( session .turn( SessionTurnRequest::new(Message::user("x")), TurnOptions::default() ) - .await, - Err(RuntimeError::Hook(_)) - )); - assert!(history.turns.lock().unwrap().is_empty()); - assert!(session.history().is_empty()); + .await + .is_err() + ); + let path = directory.path().join("session_raw/agent.jsonl"); + assert_eq!( + read_transcript(&path).unwrap().messages[0].content, + "recoverable" + ); + assert!(read_transcript_display(&path).unwrap().records.iter().any(|record| matches!(record, + DisplayRecord::Message(message) if message.interrupted && message.message.content == "display partial" + ))); +} - let mut harness: AgentHarness<()> = AgentHarness::new(); - harness.register_model("mock", Arc::new(MockModel::constant("unreachable"))); - let mut mismatched = SessionBuilder::new(Arc::new(HarnessDriver::new( - Arc::new(harness), - Arc::new(()), - ))) - .tool_snapshot( - ToolSnapshot::new(vec![ToolSpec { - name: "not-registered".into(), - description: "x".into(), - parameters: serde_json::json!({}), - }]) - .unwrap(), - ) - .build() - .unwrap(); +#[tokio::test] +async fn cancellation_while_the_driver_is_waiting_has_one_cancelled_terminal() { + let started = Arc::new(tokio::sync::Notify::new()); + let (hook, events) = hook(vec![]); + let mut session = SessionBuilder::new(Arc::new(WaitingDriver(started.clone()))) + .hooks(hook) + .build() + .unwrap(); + let options = TurnOptions::default(); + let cancellation = options.cancellation.clone(); + let turn = tokio::spawn(async move { + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await + }); + started.notified().await; + cancellation.cancel(); + assert_eq!(turn.await.unwrap(), Err(RuntimeError::Cancelled)); assert!(matches!( - mismatched - .turn( - SessionTurnRequest::new(Message::user("x")), - TurnOptions::default() - ) - .await, - Err(RuntimeError::ToolSnapshotMismatch) + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Cancelled] )); } #[tokio::test] -async fn cancellation_at_driver_await_emits_one_terminal_hook() { - let hooks = Arc::new(RecordingHooks::default()); - let mut session = SessionBuilder::new(Arc::new(WaitingDriver)) - .hooks(hooks.clone()) +async fn dropped_turn_future_has_one_failed_terminal() { + let started = Arc::new(tokio::sync::Notify::new()); + let (hook, events) = hook(vec![]); + let mut session = SessionBuilder::new(Arc::new(WaitingDriver(started.clone()))) + .hooks(hook) .build() .unwrap(); - let options = TurnOptions::default(); - let cancellation = options.cancellation.clone(); - tokio::spawn(async move { - tokio::task::yield_now().await; - cancellation.cancel(); - }); - assert_eq!( + let turn = tokio::spawn(async move { session - .turn(SessionTurnRequest::new(Message::user("x")), options) - .await, - Err(RuntimeError::Cancelled) - ); - assert_eq!( - *hooks.events.lock().unwrap(), - vec!["before", "terminal:Cancelled"] - ); + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + }); + started.notified().await; + turn.abort(); + let _ = turn.await; + tokio::task::yield_now().await; + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Failed(_)] + )); +} + +struct BlockingAfterCommitHook { + started: Arc, + release: Arc, + terminal: Arc, + events: Arc, +} + +#[async_trait] +impl SessionHooks for BlockingAfterCommitHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn after_commit(&self, _: CommitReceipt) -> Result<(), RuntimeError> { + self.started.notify_waiters(); + self.release.notified().await; + Ok(()) + } + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.events.terminal.lock().unwrap().push(terminal); + self.terminal.notify_waiters(); + Ok(()) + } } #[tokio::test] -async fn dropped_turn_future_observes_one_failed_terminal() { +async fn dropped_turn_after_durable_append_keeps_a_completed_terminal() { + let (locator, history) = locator(None); let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); let terminal = Arc::new(tokio::sync::Notify::new()); - let hooks = Arc::new(RecordingHooks { - terminal_notified: Some(terminal.clone()), - ..RecordingHooks::default() - }); - let mut session = SessionBuilder::new(Arc::new(DropDriver { + let events = Arc::new(Events::default()); + let hook = Arc::new(BlockingAfterCommitHook { started: started.clone(), - })) - .hooks(hooks.clone()) + release: release.clone(), + terminal: terminal.clone(), + events: events.clone(), + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("done"), + ]))]))) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook) .build() .unwrap(); - let entered = started.notified(); - let observed = terminal.notified(); - let task = tokio::spawn(async move { + let turn = tokio::spawn(async move { session .turn( SessionTurnRequest::new(Message::user("x")), @@ -1075,233 +741,704 @@ async fn dropped_turn_future_observes_one_failed_terminal() { ) .await }); - entered.await; - task.abort(); - let _ = task.await; - observed.await; + started.notified().await; + assert!(history.state.lock().unwrap().is_some()); + let observed_terminal = terminal.notified(); + turn.abort(); + let _ = turn.await; + release.notify_one(); + observed_terminal.await; + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Completed(_)] + )); +} + +#[tokio::test] +async fn resumed_history_restores_the_prefix_once_before_the_next_driver_call() { + let raw = TranscriptMessage::new("user", "old"); + let (locator, _) = locator(Some(SessionTranscript { + meta: meta(), + messages: vec![raw], + })); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("old"), + Message::assistant("new"), + ]))])); + let prefix = PrefixSnapshot::new(vec![Message::system("stable")]); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .prefix(prefix) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }, + ) + .await + .unwrap(); assert_eq!( - *hooks.events.lock().unwrap(), - vec!["before", "terminal:Failed(\"session turn dropped\")"] + driver.requests.lock().unwrap()[0].history[0], + Message::system("stable") + ); + assert_eq!( + session + .history() + .iter() + .filter(|message| **message == Message::system("stable")) + .count(), + 1 ); } #[tokio::test] -async fn dropped_turn_after_commit_keeps_a_truthful_completed_terminal() { - let (locator, history) = memory_locator(None, false); - let hooks = Arc::new(WaitingPostCommitHooks { - started: tokio::sync::Notify::new(), - terminal: tokio::sync::Notify::new(), - terminals: Mutex::new(Vec::new()), +async fn cancellation_signalled_by_a_successful_append_cannot_relabel_completion() { + let (locator, history) = locator(None); + let (hook, events) = hook(vec![]); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("done"), + ]))]))) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook) + .build() + .unwrap(); + let options = TurnOptions::default(); + *history.cancel_after_append.lock().unwrap() = Some(options.cancellation.clone()); + assert!( + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await + .is_ok() + ); + assert!(history.state.lock().unwrap().is_some()); + assert_eq!(events.receipts.lock().unwrap().len(), 1); + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Completed(_)] + )); +} + +struct BlockingCommitHook { + started: Arc, + events: Arc, +} +#[async_trait] +impl SessionHooks for BlockingCommitHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + self.started.notify_waiters(); + std::future::pending().await + } + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.events.terminal.lock().unwrap().push(terminal); + Ok(()) + } +} + +#[tokio::test] +async fn cancellation_during_before_commit_leaves_no_append() { + let (locator, history) = locator(None); + let started = Arc::new(tokio::sync::Notify::new()); + let events = Arc::new(Events::default()); + let hook = Arc::new(BlockingCommitHook { + started: started.clone(), + events: events.clone(), }); - let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("durable"), + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("candidate"), ]))]))) - .codec(Arc::new(BasicCodec::default())) - .hooks(hooks.clone()) + .codec(Arc::new(Codec::default())) .transcript(locator, "agent", meta()) + .hooks(hook) .build() .unwrap(); - let entered = hooks.started.notified(); - let observed = hooks.terminal.notified(); - let task = tokio::spawn(async move { + let options = TurnOptions::default(); + let cancellation = options.cancellation.clone(); + let turn = tokio::spawn(async move { session - .turn( - SessionTurnRequest::new(Message::user("x")), - TurnOptions::default(), - ) + .turn(SessionTurnRequest::new(Message::user("x")), options) .await }); - entered.await; - assert_eq!(history.turns.lock().unwrap().len(), 1); - task.abort(); - let _ = task.await; - observed.await; + started.notified().await; + cancellation.cancel(); + assert_eq!(turn.await.unwrap(), Err(RuntimeError::Cancelled)); + assert!(history.state.lock().unwrap().is_none()); assert!(matches!( - hooks.terminals.lock().unwrap().as_slice(), - [SessionTerminal::Completed(outcome)] if outcome.history.last() == Some(&Message::assistant("durable")) + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Cancelled] )); } +struct BlockingBeforeTurnHook { + started: Arc, + terminal: Arc, + terminals: Mutex>, + after_commits: Mutex, +} + +#[async_trait] +impl SessionHooks for BlockingBeforeTurnHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + self.started.notify_waiters(); + std::future::pending().await + } + + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn after_commit(&self, _: CommitReceipt) -> Result<(), RuntimeError> { + *self.after_commits.lock().unwrap() += 1; + Ok(()) + } + + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.terminals.lock().unwrap().push(terminal); + self.terminal.notify_waiters(); + Ok(()) + } +} + #[tokio::test] -async fn cancellation_during_before_hook_never_starts_or_commits_a_turn() { - let hooks = Arc::new(WaitingBeforeHooks { - started: tokio::sync::Notify::new(), +async fn cancellation_during_before_turn_has_no_driver_persist_or_after_commit() { + let started = Arc::new(tokio::sync::Notify::new()); + let terminal = Arc::new(tokio::sync::Notify::new()); + let hook = Arc::new(BlockingBeforeTurnHook { + started: started.clone(), + terminal: terminal.clone(), + terminals: Mutex::new(Vec::new()), + after_commits: Mutex::new(0), }); - let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("nope"), - ]))])); + let (locator, history) = locator(None); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![Message::assistant( + "never", + )]))])); let mut session = SessionBuilder::new(driver.clone()) - .hooks(hooks.clone()) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook.clone()) .build() .unwrap(); let options = TurnOptions::default(); let cancellation = options.cancellation.clone(); - let started = hooks.started.notified(); + let observed_terminal = terminal.notified(); let turn = tokio::spawn(async move { session .turn(SessionTurnRequest::new(Message::user("x")), options) .await }); - started.await; + started.notified().await; cancellation.cancel(); assert_eq!(turn.await.unwrap(), Err(RuntimeError::Cancelled)); + observed_terminal.await; assert!(driver.requests.lock().unwrap().is_empty()); + assert!(history.state.lock().unwrap().is_none()); + assert_eq!(*hook.after_commits.lock().unwrap(), 0); + assert!(matches!( + hook.terminals.lock().unwrap().as_slice(), + [SessionTerminal::Cancelled] + )); +} + +struct BlockingAfterCommitCancellationHook { + started: Arc, + release: Arc, + terminal: Arc, + terminals: Mutex>, +} + +#[async_trait] +impl SessionHooks for BlockingAfterCommitCancellationHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(TurnPreparation::default()) + } + + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + + async fn after_commit(&self, _: CommitReceipt) -> Result<(), RuntimeError> { + self.started.notify_waiters(); + self.release.notified().await; + Ok(()) + } + + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.terminals.lock().unwrap().push(terminal); + self.terminal.notify_waiters(); + Ok(()) + } } #[tokio::test] -async fn cancellation_during_precommit_hook_has_no_durable_commit() { - let hooks = Arc::new(WaitingAfterHooks { - started: tokio::sync::Notify::new(), - events: Mutex::new(Vec::new()), +async fn cancellation_during_after_commit_keeps_the_completed_result_and_terminal() { + let started = Arc::new(tokio::sync::Notify::new()); + let release = Arc::new(tokio::sync::Notify::new()); + let terminal = Arc::new(tokio::sync::Notify::new()); + let hook = Arc::new(BlockingAfterCommitCancellationHook { + started: started.clone(), + release: release.clone(), + terminal: terminal.clone(), + terminals: Mutex::new(Vec::new()), }); - let (locator, history) = memory_locator(None, false); - let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("candidate"), + let (locator, history) = locator(None); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("durable"), ]))]))) - .codec(Arc::new(BasicCodec::default())) - .hooks(hooks.clone()) + .codec(Arc::new(Codec::default())) .transcript(locator, "agent", meta()) + .hooks(hook.clone()) .build() .unwrap(); let options = TurnOptions::default(); let cancellation = options.cancellation.clone(); - let started = hooks.started.notified(); + let observed_terminal = terminal.notified(); let turn = tokio::spawn(async move { session .turn(SessionTurnRequest::new(Message::user("x")), options) .await }); - started.await; + started.notified().await; + assert!(history.state.lock().unwrap().is_some()); cancellation.cancel(); - assert_eq!(turn.await.unwrap(), Err(RuntimeError::Cancelled)); - assert!(history.turns.lock().unwrap().is_empty()); - assert_eq!(*hooks.events.lock().unwrap(), vec!["Cancelled"]); + release.notify_one(); + let committed = turn.await.unwrap().unwrap(); + assert_eq!( + committed.history.last(), + Some(&Message::assistant("durable")) + ); + observed_terminal.await; + assert!( + matches!(hook.terminals.lock().unwrap().as_slice(), [SessionTerminal::Completed(outcome)] if outcome == &committed) + ); } #[tokio::test] -async fn cancellation_signalled_by_successful_commit_cannot_relabel_the_turn() { - let (locator, history) = memory_locator(None, false); - let driver = Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("done"), +async fn prefix_mutation_after_a_commit_is_rejected() { + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![Message::assistant("one")])), + Ok(outcome(vec![Message::assistant("two")])), + ])); + let prep = TurnPreparation { + prefix: Some(PrefixSnapshot::new(vec![Message::system("p")])), + ..Default::default() + }; + let (hook, _) = hook(vec![prep.clone(), prep]); + let mut session = SessionBuilder::new(driver.clone()) + .hooks(hook) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("a")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert!(matches!( + session + .turn( + SessionTurnRequest::new(Message::user("b")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::InvalidSessionState(_)) + )); + assert_eq!(driver.requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn lazy_target_is_opened_only_after_before_turn_selects_it() { + let (locator, history) = locator(None); + let target = TranscriptTarget::new(locator, "late", meta()); + let (hook, _) = hook(vec![TurnPreparation { + transcript: Some(target), + ..Default::default() + }]); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("ok"), + ]))]))) + .codec(Arc::new(Codec::default())) + .hooks(hook) + .build() + .unwrap(); + assert_eq!(*history.opens.lock().unwrap(), 0); + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!(*history.opens.lock().unwrap(), 1); +} + +#[tokio::test] +async fn hook_selected_target_and_resume_mode_apply_before_driver_handoff() { + let (locator, _) = locator(Some(SessionTranscript { + meta: meta(), + messages: vec![TranscriptMessage::new("user", "resumed")], + })); + let target = TranscriptTarget::new(locator, "agent", meta()); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("resumed"), + Message::assistant("ok"), ]))])); - let mut session = SessionBuilder::new(driver) - .codec(Arc::new(BasicCodec::default())) - .transcript(locator, "agent", meta()) + struct ResumeHook(TranscriptTarget); + #[async_trait] + impl SessionHooks for ResumeHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + options: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + options.resume = ResumeMode::LatestForAgent; + Ok(TurnPreparation { + transcript: Some(self.0.clone()), + ..Default::default() + }) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } + } + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .hooks(Arc::new(ResumeHook(target))) .build() .unwrap(); - let options = TurnOptions::default(); - *history.cancel_after_append.lock().unwrap() = Some(options.cancellation.clone()); - assert!( + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!( + driver.requests.lock().unwrap()[0].history[0], + Message::user("resumed") + ); +} + +#[tokio::test] +async fn hook_selected_transcript_without_a_codec_fails_before_driver_execution() { + let (locator, _) = locator(None); + let target = TranscriptTarget::new(locator, "agent", meta()); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![Message::assistant( + "never", + )]))])); + let (hook, _) = hook(vec![TurnPreparation { + transcript: Some(target), + ..Default::default() + }]); + let mut session = SessionBuilder::new(driver.clone()) + .hooks(hook) + .build() + .unwrap(); + assert_eq!( session - .turn(SessionTurnRequest::new(Message::user("x")), options) - .await - .is_ok() + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default() + ) + .await, + Err(RuntimeError::MissingDependency("TranscriptCodec")) ); - assert_eq!(history.turns.lock().unwrap().len(), 1); + assert!(driver.requests.lock().unwrap().is_empty()); } #[tokio::test] -async fn post_commit_runs_once_after_durability_and_cannot_relabel_success() { - let (locator, history) = memory_locator(None, false); - let cancellation = tinyagents_harness::CancellationToken::new(); - let hooks = Arc::new(FinalizationHooks { - fail_post_commit: true, - cancel_on_post_commit: Some(cancellation.clone()), - ..FinalizationHooks::default() - }); - let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("done"), +async fn resumed_raw_rows_and_metadata_survive_the_append() { + let mut raw = TranscriptMessage::new("user", "old"); + raw.extra_metadata = Some(serde_json::json!({"native": true})); + let initial = SessionTranscript { + meta: meta(), + messages: vec![raw.clone()], + }; + let (locator, history) = locator(Some(initial)); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("old"), + Message::assistant("new"), ]))]))) - .codec(Arc::new(BasicCodec::default())) - .hooks(hooks.clone()) + .codec(Arc::new(Codec::default())) .transcript(locator, "agent", meta()) .build() .unwrap(); - let options = TurnOptions { - cancellation, - ..TurnOptions::default() - }; - let committed = session - .turn(SessionTurnRequest::new(Message::user("x")), options) + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }, + ) .await - .expect("post-commit failure/cancellation cannot revoke a durable success"); - assert_eq!(history.turns.lock().unwrap().len(), 1); - assert_eq!(*hooks.post_commits.lock().unwrap(), vec![committed.clone()]); - assert!(matches!( - hooks.terminals.lock().unwrap().as_slice(), - [SessionTerminal::Completed(outcome)] if outcome == &committed - )); + .unwrap(); + let persisted = history.state.lock().unwrap().clone().unwrap(); + assert_eq!(persisted.meta.agent_id.as_deref(), Some("agent-id")); + assert_eq!(persisted.messages[0], raw); } #[tokio::test] -async fn persistence_failure_never_calls_post_commit() { - let (locator, _) = memory_locator(None, true); - let hooks = Arc::new(FinalizationHooks::default()); - let mut session = SessionBuilder::new(Arc::new(FakeDriver::new(vec![Ok(outcome(vec![ - Message::assistant("never durable"), +async fn after_commit_receipt_observes_durable_append_and_terminal_follows_it() { + let (locator, history) = locator(None); + let (hook, events) = hook(vec![]); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("ok"), ]))]))) - .codec(Arc::new(BasicCodec::default())) - .hooks(hooks.clone()) + .codec(Arc::new(Codec::default())) .transcript(locator, "agent", meta()) + .hooks(hook) .build() .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert!(history.state.lock().unwrap().is_some()); + assert_eq!(events.receipts.lock().unwrap().len(), 1); + let receipt = events.receipts.lock().unwrap().pop().unwrap(); + let transcript = receipt.transcript.unwrap(); + assert_eq!(transcript.path, history.path); assert!(matches!( - session + transcript.delta, + crate::TranscriptDelta::Append { .. } + )); + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Completed(_)] + )); + assert_eq!( + *events.order.lock().unwrap(), + vec!["after_commit", "terminal"] + ); +} + +#[tokio::test] +async fn receipt_reports_a_compaction_as_replacement_not_an_append_range() { + let (locator, _) = locator(None); + let (hook, events) = hook(vec![]); + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![ + Message::user("old"), + Message::assistant("one"), + ])), + Ok(outcome(vec![Message::assistant("compacted")])), + ])); + let mut session = SessionBuilder::new(driver) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook) + .build() + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("old")), + TurnOptions::default(), + ) + .await + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("new")), + TurnOptions::default(), + ) + .await + .unwrap(); + let receipts = events.receipts.lock().unwrap(); + assert!(matches!( + receipts[1].transcript.as_ref().unwrap().delta, + crate::TranscriptDelta::Replace { .. } + )); +} + +#[tokio::test] +async fn failure_and_cancellation_do_not_run_after_commit_and_emit_one_terminal() { + let (failure_hook, events) = hook(vec![]); + let mut failed = SessionBuilder::new(Arc::new(Driver::new(vec![Err(DriverFailure { + error: RuntimeError::Driver("no".into()), + partial: None, + })]))) + .hooks(failure_hook) + .build() + .unwrap(); + assert!( + failed .turn( SessionTurnRequest::new(Message::user("x")), TurnOptions::default() ) + .await + .is_err() + ); + assert!(events.receipts.lock().unwrap().is_empty()); + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Failed(_)] + )); + + let (cancel_hook, events) = hook(vec![]); + let mut cancelled = SessionBuilder::new(Arc::new(Driver::new(vec![]))) + .hooks(cancel_hook) + .build() + .unwrap(); + let options = TurnOptions::default(); + options.cancellation.cancel(); + assert_eq!( + cancelled + .turn(SessionTurnRequest::new(Message::user("x")), options) .await, - Err(RuntimeError::Persistence(_)) + Err(RuntimeError::Cancelled) + ); + assert!(events.receipts.lock().unwrap().is_empty()); + assert!(matches!( + events.terminal.lock().unwrap().as_slice(), + [SessionTerminal::Cancelled] )); - assert!(hooks.post_commits.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn seed_history_is_the_only_history_and_rejects_a_later_seed() { + let codec = Arc::new(Codec::default()); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("seed"), + Message::assistant("next"), + ]))]))) + .codec(codec.clone()) + .build() + .unwrap(); + let raw = vec![TranscriptMessage::new("user", "seed")]; + session + .seed_history(vec![Message::user("seed")], raw.clone()) + .unwrap(); + session + .turn( + SessionTurnRequest::new(Message::user("seed")), + TurnOptions::default(), + ) + .await + .unwrap(); + assert_eq!(codec.seen_prior.lock().unwrap().as_slice(), [raw]); assert!(matches!( - hooks.terminals.lock().unwrap().as_slice(), - [SessionTerminal::Failed(_)] + session.seed_history(vec![], vec![]), + Err(RuntimeError::InvalidSessionState(_)) )); } -type SeenCodecOptions = (String, Option, Option, bool, ResumeMode); +#[derive(Clone)] +struct Context(String); -struct OptionsCodec { - seen: Mutex>, +struct ContextDriver(Mutex>); +#[async_trait] +impl SessionDriver for ContextDriver { + async fn execute( + &self, + request: DriverRequest, + ) -> Result { + self.0 + .lock() + .unwrap() + .push(request.run_context.data.0.clone()); + Ok(outcome(vec![Message::assistant("ok")])) + } } -impl TranscriptCodec for OptionsCodec { +struct ContextCodec(Mutex>); +impl TranscriptCodec for ContextCodec { fn decode_history(&self, _: &SessionTranscript) -> Result, RuntimeError> { - Ok(Vec::new()) + Ok(vec![]) } - fn reconcile( &self, _: &[TranscriptMessage], _: &[Message], - next: &[Message], - options: &crate::TranscriptTurnOptions, + _: &[Message], + options: &TranscriptTurnOptions, ) -> Result, RuntimeError> { - self.seen.lock().unwrap().push(( - options.context.0.clone(), - options.request_id.clone(), - options.thread_id.clone(), - options.stream, - options.resume, - )); - Ok(next - .iter() - .map(|message| TranscriptMessage::assistant(message.text())) - .collect()) + self.0.lock().unwrap().push(options.context.0.clone()); + Ok(vec![]) + } +} + +struct ContextHook; +#[async_trait] +impl SessionHooks for ContextHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + options: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + options.run_context.data = Context("mutated".into()); + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) } } #[tokio::test] -async fn codec_reconciliation_receives_explicit_host_context_and_turn_options() { - let (locator, _) = memory_locator(None, false); - let codec = Arc::new(OptionsCodec { - seen: Mutex::new(Vec::new()), - }); - let mut session = SessionBuilder::::new(Arc::new(ContextDriver)) +async fn hook_option_context_mutation_reaches_driver_and_codec() { + let (locator, _) = locator(None); + let driver = Arc::new(ContextDriver(Mutex::new(vec![]))); + let codec = Arc::new(ContextCodec(Mutex::new(vec![]))); + let mut session = SessionBuilder::new(driver.clone()) .codec(codec.clone()) .transcript(locator, "agent", meta()) + .hooks(Arc::new(ContextHook)) .build() .unwrap(); let cancellation = tinyagents_harness::CancellationToken::new(); @@ -1309,37 +1446,26 @@ async fn codec_reconciliation_receives_explicit_host_context_and_turn_options() .turn( SessionTurnRequest::new(Message::user("x")), TurnOptions { - request_id: Some("request-1".into()), - thread_id: Some("thread-1".into()), - stream: true, - resume: ResumeMode::LatestForAgent, + request_id: None, + thread_id: None, + stream: false, + resume: ResumeMode::Never, cancellation: cancellation.clone(), - run_context: RunContext::new( - RunConfig::new("codec-host"), - HostContext("host-owned context".into()), - ) - .with_cancellation(cancellation), + run_context: RunContext::new(RunConfig::new("test"), Context("before".into())) + .with_cancellation(cancellation), }, ) .await .unwrap(); - assert_eq!( - *codec.seen.lock().unwrap(), - vec![( - "host-owned context".into(), - Some("request-1".into()), - Some("thread-1".into()), - true, - ResumeMode::LatestForAgent, - )] - ); + assert_eq!(driver.0.lock().unwrap().as_slice(), ["mutated"]); + assert_eq!(codec.0.lock().unwrap().as_slice(), ["mutated"]); } #[test] -fn dependency_boundary_has_no_product_dependency() { - let manifest = include_str!("../Cargo.toml"); - assert!(manifest.contains("tinyagents-harness")); - assert!(manifest.contains("tinyagents-session")); - assert!(manifest.contains("tinytools")); - assert!(!manifest.to_ascii_lowercase().contains("openhuman")); +fn runtime_stays_host_neutral() { + assert!( + !include_str!("../Cargo.toml") + .to_ascii_lowercase() + .contains("openhuman") + ); } diff --git a/crates/tinyagents-runtime/src/types.rs b/crates/tinyagents-runtime/src/types.rs index c81285cc..2996a03c 100644 --- a/crates/tinyagents-runtime/src/types.rs +++ b/crates/tinyagents-runtime/src/types.rs @@ -1,9 +1,14 @@ +use std::{ops::Range, path::PathBuf, sync::Arc}; + use tinyagents_harness::{ CancellationToken, context::{RunConfig, RunContext}, }; +use tinyagents_session::transcript::{TranscriptLocator, TranscriptMessage, TranscriptMeta}; use tinyinference_llm::message::Message; +use crate::{PrefixSnapshot, ToolSnapshot}; + /// Selects the durable transcript a turn should load before execution. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum ResumeMode { @@ -52,6 +57,102 @@ pub struct TranscriptTurnOptions { pub context: C, } +/// A transcript destination selected lazily by a host for a session. +/// +/// Constructing a target performs no I/O. The runtime opens it only when a +/// requested resume or the first append needs a bound history handle. +#[derive(Clone)] +pub struct TranscriptTarget { + pub locator: Arc, + pub stem: String, + pub meta: TranscriptMeta, +} + +impl TranscriptTarget { + pub fn new( + locator: Arc, + stem: impl Into, + meta: TranscriptMeta, + ) -> Self { + Self { + locator, + stem: stem.into(), + meta, + } + } + + pub(crate) fn same_binding(&self, other: &Self) -> bool { + self.stem == other.stem && Arc::ptr_eq(&self.locator, &other.locator) + } +} + +/// Values prepared by `SessionHooks::before_turn` for exactly one driver call. +#[derive(Clone, Default)] +pub struct TurnPreparation { + /// A replacement prefix allowed only while the session is empty and has + /// never committed a transcript transition. + pub prefix: Option, + /// The immutable tool declarations for this driver request. `None` uses + /// the builder's compatibility default and is never retained from a prior + /// preparation. + pub tools: Option, + /// A lazy transcript destination. It can be selected or replaced before + /// the first bind, but cannot be redirected after binding. + pub transcript: Option, +} + +impl TurnPreparation { + pub fn with_tools(tools: ToolSnapshot) -> Self { + Self { + tools: Some(tools), + ..Self::default() + } + } +} + +/// Read-only session state supplied to `before_turn`. +#[derive(Clone, Copy)] +pub struct SessionStateView<'a> { + pub history: &'a [Message], + pub raw_history: &'a [TranscriptMessage], + pub prefix: &'a PrefixSnapshot, + pub transcript_target: Option<&'a TranscriptTarget>, + pub committed_turns: usize, +} + +/// The shape of a successful logical transcript transition. +/// +/// An append extends the prior logical rows. Any rewrite, including a context +/// compaction with a longer replacement, is reported as `Replace` rather than +/// pretending that a suffix range was appended. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum TranscriptDelta { + Append { + previous_len: usize, + appended: Range, + }, + Replace { + previous_len: usize, + next_len: usize, + }, +} + +/// Durable transcript information supplied after a successful append. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TranscriptCommitReceipt { + pub path: PathBuf, + pub delta: TranscriptDelta, +} + +/// Exactly-once post-durability observation data. +#[derive(Clone, Debug)] +pub struct CommitReceipt { + pub outcome: SessionTurnOutcome, + pub options: TranscriptTurnOptions, + /// `None` when the host selected no durable transcript target. + pub transcript: Option, +} + impl TurnOptions { pub(crate) fn transcript_options(&self) -> TranscriptTurnOptions { TranscriptTurnOptions { diff --git a/docs/modules/runtime/README.md b/docs/modules/runtime/README.md index 82b38647..13884099 100644 --- a/docs/modules/runtime/README.md +++ b/docs/modules/runtime/README.md @@ -7,8 +7,8 @@ credentials, product events, and its durable transcript dialect. ## Turn boundary -A `Session` holds model-visible history, an immutable stable prefix, and a -frozen tool declaration snapshot. Every `turn` takes explicit `TurnOptions`: +A `Session` holds model-visible history and an immutable stable prefix. +Every `turn` takes explicit `TurnOptions`: request and thread identifiers, streaming mode, resume mode, cancellation, and the live `RunContext`. The driver consumes the live context. Because the codec must reconcile after that driver call, `C: Clone` and the codec receives @@ -20,9 +20,18 @@ error. Those harness entry points do not currently expose a separate streamed text delta, partial reasoning, or iteration; the adapter can only use the last accumulated assistant text as a display partial when it exists. +Before driver handoff, `SessionHooks::before_turn` receives the mutable +request and options plus a read-only `SessionStateView`. Its `TurnPreparation` +may set a prefix only for an empty, uncommitted session; it may select one +immutable tool snapshot for this call; and it may select a lazy +`TranscriptTarget`. A returned tool snapshot is not stored for the following +turn. The target is bound only when resume/append needs it and cannot change +after binding. `Session::seed_history` accepts a model history and lossless raw +rows before any commit, replacing host-side shadow history. + ## Durability and projections -With a `TranscriptHistory`, a successful turn first passes `after_turn`, the +With a `TranscriptHistory`, a successful turn first passes `before_commit`, the pre-commit validation hook. The runtime then submits its logical transcript delta, metadata, and any supplied `TranscriptPartial` to one `append_turn_with_partial` operation. `FileTranscriptHistory` serializes those @@ -34,8 +43,9 @@ a supplied partial, so the runtime does not fall back to two independent writes. This is an operation-level guarantee; it does not claim crash-safe filesystem transactions beyond the underlying storage implementation. -After a successful append, `after_commit` receives the committed -`SessionTurnOutcome`, and `on_terminal` receives a completed terminal outcome. -Both are observational: their errors, or a cooperative cancellation that they -observe, cannot change durable success. Failed persistence never invokes -`after_commit`. +After a successful append and in-memory state update, `after_commit` receives +a `CommitReceipt` with the committed outcome, the explicit context/options +snapshot, and neutral transcript path/delta data. `on_terminal` follows with a +completed terminal outcome. Both are observational: their errors, or a +cooperative cancellation that they observe, cannot change durable success. +Failed persistence and cancelled/failed turns never invoke `after_commit`. From 766af06795c46caa734fbe2f6514f3e818711f91 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 19:56:22 +0300 Subject: [PATCH 07/18] feat(runtime): stage transcript resume before turn preparation Co-authored-by: Medulla --- crates/tinyagents-runtime/README.md | 17 +- crates/tinyagents-runtime/src/builder.rs | 1 + crates/tinyagents-runtime/src/hooks.rs | 26 +- crates/tinyagents-runtime/src/lib.rs | 2 +- crates/tinyagents-runtime/src/session.rs | 104 ++-- crates/tinyagents-runtime/src/test.rs | 629 ++++++++++++++++++++++- crates/tinyagents-runtime/src/types.rs | 33 +- 7 files changed, 734 insertions(+), 78 deletions(-) diff --git a/crates/tinyagents-runtime/README.md b/crates/tinyagents-runtime/README.md index 519ec9ad..d80ad8c0 100644 --- a/crates/tinyagents-runtime/README.md +++ b/crates/tinyagents-runtime/README.md @@ -20,10 +20,12 @@ The host supplies three narrow seams: `Clone` so reconciliation receives the current host context plus request, thread, stream, and resume options after the live `RunContext` moves into the driver. -- `SessionHooks` prepares a request and mutable `TurnOptions` before - handoff. Its `TurnPreparation` can install a first-turn prefix, select the - one immutable `ToolSnapshot` for that request, and lazily choose a - `TranscriptTarget`. `before_commit` validates the candidate; `after_commit` +- `SessionHooks` prepares a request and mutable `TurnOptions` in two + stages. `before_resume` lazily chooses a `TranscriptTarget`, then the runtime + binds it and loads any requested transcript. `before_turn` sees that decoded + history and raw rows plus `SessionStateView::resumed`; its `TurnPreparation` + can install or replace the prefix before the first commit and selects the + one immutable `ToolSnapshot` for that request. `before_commit` validates the candidate; `after_commit` receives an exactly-once `CommitReceipt` containing the explicit context snapshot and neutral transcript path/delta receipt. `on_terminal` receives one truthful terminal state. It does not make policy decisions. @@ -42,7 +44,9 @@ let session = SessionBuilder::new(driver) To supply a default lazy destination, add `SessionBuilder::transcript(locator, stem, meta)`. It does not open a transcript while building: a selected target binds only on -resume/first append and cannot be redirected after that. The runtime +resume/first append and cannot be redirected after that. `TranscriptTarget` can +use a distinct `resume_agent` for `LatestForAgent` lookup; writes always use its +`stem`. The runtime uses `tinyagents-session`'s `TranscriptHistory::append_turn_with_partial`, so a normal extension appends only the new tail and a reduced context writes one compaction record. A supplied partial driver outcome is represented through that single @@ -54,7 +58,8 @@ in-memory history and persisted snapshot unchanged. Every turn receives explicit `TurnOptions`, including its cancellation token and `RunContext`; no task-local data crosses the runtime boundary. The stable prefix is reconciled after resume and driver compaction without -duplication. `Session::seed_history(history, raw)` is the explicit, lossless +duplication, including a prefix supplied by `before_turn` after a resumed +history. `Session::seed_history(history, raw)` is the explicit, lossless resume/seed boundary; a host must not keep a second shadow history. Cancellation before the commit point leaves no durable mutation; once it succeeds, the turn remains successful. `after_commit` and terminal hooks get the committed outcome, but their error or a cooperative cancellation diff --git a/crates/tinyagents-runtime/src/builder.rs b/crates/tinyagents-runtime/src/builder.rs index dc928dce..955ee686 100644 --- a/crates/tinyagents-runtime/src/builder.rs +++ b/crates/tinyagents-runtime/src/builder.rs @@ -82,6 +82,7 @@ impl SessionBuilder { let target = self.transcript.map(|config| crate::TranscriptTarget { locator: config.locator, stem: config.stem, + resume_agent: None, meta: config.meta, }); if target.is_some() && self.codec.is_none() { diff --git a/crates/tinyagents-runtime/src/hooks.rs b/crates/tinyagents-runtime/src/hooks.rs index 0c93e0b5..dbd62642 100644 --- a/crates/tinyagents-runtime/src/hooks.rs +++ b/crates/tinyagents-runtime/src/hooks.rs @@ -1,8 +1,8 @@ use async_trait::async_trait; use crate::{ - CommitReceipt, RuntimeError, SessionStateView, SessionTerminal, SessionTurnOutcome, - SessionTurnRequest, TranscriptTurnOptions, TurnOptions, TurnPreparation, + CommitReceipt, ResumePreparation, RuntimeError, SessionStateView, SessionTerminal, + SessionTurnOutcome, SessionTurnRequest, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; /// Host observation/preparation around a session turn. @@ -12,8 +12,20 @@ use crate::{ /// keep preparation state without task-local runtime state. #[async_trait] pub trait SessionHooks: Send + Sync { + /// Runs after the turn's initial cancellation check and before transcript + /// target binding or resume. It may mutate the request and live options. + /// Its target remains lazy until resume or the first append needs it. + async fn before_resume( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(ResumePreparation::default()) + } /// Runs before the driver sees the request and consumes the explicit - /// options. Returned values apply only to this driver invocation. + /// options, after any requested transcript has been loaded. Returned + /// values apply only to this driver invocation. async fn before_turn( &self, request: &mut SessionTurnRequest, @@ -44,6 +56,14 @@ pub struct NoopSessionHooks; #[async_trait] impl SessionHooks for NoopSessionHooks { + async fn before_resume( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(ResumePreparation::default()) + } async fn before_turn( &self, _: &mut SessionTurnRequest, diff --git a/crates/tinyagents-runtime/src/lib.rs b/crates/tinyagents-runtime/src/lib.rs index d2800115..23fd86cc 100644 --- a/crates/tinyagents-runtime/src/lib.rs +++ b/crates/tinyagents-runtime/src/lib.rs @@ -23,7 +23,7 @@ pub use session::Session; pub use tinyagents_session::transcript::TranscriptPartial; pub use tools::ToolSnapshot; pub use types::{ - CommitReceipt, ResumeMode, SessionResume, SessionStateView, SessionTerminal, + CommitReceipt, ResumeMode, ResumePreparation, SessionResume, SessionStateView, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, TranscriptCommitReceipt, TranscriptDelta, TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; diff --git a/crates/tinyagents-runtime/src/session.rs b/crates/tinyagents-runtime/src/session.rs index a9128e32..6aec66e2 100644 --- a/crates/tinyagents-runtime/src/session.rs +++ b/crates/tinyagents-runtime/src/session.rs @@ -7,10 +7,10 @@ use tinyagents_session::transcript::{ use tinyinference_llm::message::Message; use crate::{ - CommitReceipt, DriverRequest, PrefixSnapshot, ResumeMode, RuntimeError, SessionDriver, - SessionHooks, SessionResume, SessionStateView, SessionTerminal, SessionTurnOutcome, - SessionTurnRequest, ToolSnapshot, TranscriptCodec, TranscriptCommitReceipt, TranscriptDelta, - TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, + CommitReceipt, DriverRequest, PrefixSnapshot, ResumeMode, ResumePreparation, RuntimeError, + SessionDriver, SessionHooks, SessionResume, SessionStateView, SessionTerminal, + SessionTurnOutcome, SessionTurnRequest, ToolSnapshot, TranscriptCodec, TranscriptCommitReceipt, + TranscriptDelta, TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; /// Host-neutral mutable state for one conversation session. @@ -104,7 +104,9 @@ impl Session { }; let read = match options.resume { ResumeMode::Never => None, - ResumeMode::LatestForAgent => target.locator.latest_for_agent(&target.stem), + ResumeMode::LatestForAgent => target + .locator + .latest_for_agent(target.resume_agent.as_deref().unwrap_or(&target.stem)), ResumeMode::Thread => options .thread_id .as_deref() @@ -183,26 +185,35 @@ impl Session { options: &mut TurnOptions, terminal_guard: &mut TerminalGuard, ) -> Result { - let state = SessionStateView { - history: &self.history, - raw_history: &self.persisted, - prefix: &self.prefix, - transcript_target: self.target.as_ref(), - committed_turns: self.committed_turns, - }; + if options.cancellation.is_cancelled() { + return Err(RuntimeError::Cancelled); + } let cancellation = options.cancellation.clone(); + let resume_preparation = cancelable( + &cancellation, + self.hooks + .before_resume(request, options, self.state_view(false)), + ) + .await?; + self.apply_resume_preparation(resume_preparation)?; + let resumed = if options.resume == ResumeMode::Never { + false + } else { + self.resume(options).await?.loaded + }; + // `resume` is synchronous after its read, so this explicit boundary + // makes cancellation between loading and before-turn preparation + // observable without handing work to the driver. + if options.cancellation.is_cancelled() { + return Err(RuntimeError::Cancelled); + } let preparation = cancelable( &cancellation, - self.hooks.before_turn(request, options, state), + self.hooks + .before_turn(request, options, self.state_view(resumed)), ) .await?; let (tools, prepared_prefix) = self.apply_preparation(preparation)?; - // Preparation owns the current turn's target and explicit resume mode, - // so resolve only after it has made its changes visible. This is also - // why a selected target is checked for a codec before driver handoff. - if options.resume != ResumeMode::Never { - self.resume(options).await?; - } if let Some(prefix) = prepared_prefix { self.apply_prefix(prefix)?; } @@ -291,15 +302,29 @@ impl Session { &mut self, preparation: TurnPreparation, ) -> Result<(ToolSnapshot, Option), RuntimeError> { + // A returned snapshot never updates `default_tools`: it applies only + // to the `DriverRequest` being built by this call. + Ok(( + preparation + .tools + .unwrap_or_else(|| self.default_tools.clone()), + preparation.prefix, + )) + } + + fn apply_resume_preparation( + &mut self, + preparation: ResumePreparation, + ) -> Result<(), RuntimeError> { if let Some(target) = preparation.transcript { - if self.transcript.is_some() { + if self.transcript.is_some() || self.committed_turns != 0 { if !self .target .as_ref() .is_some_and(|bound| bound.same_binding(&target)) { return Err(RuntimeError::InvalidSessionState( - "cannot change a transcript target after it is bound".into(), + "cannot change a transcript target after it is bound or committed".into(), )); } } else { @@ -309,30 +334,39 @@ impl Session { if self.target.is_some() && self.codec.is_none() { return Err(RuntimeError::MissingDependency("TranscriptCodec")); } - // A returned snapshot never updates `default_tools`: it applies only - // to the `DriverRequest` being built by this call. - Ok(( - preparation - .tools - .unwrap_or_else(|| self.default_tools.clone()), - preparation.prefix, - )) + Ok(()) } fn apply_prefix(&mut self, prefix: PrefixSnapshot) -> Result<(), RuntimeError> { - if self.committed_turns != 0 - || !self.persisted.is_empty() - || self.history != self.prefix.messages() - { + if prefix == self.prefix { + return Ok(()); + } + if self.committed_turns != 0 { return Err(RuntimeError::InvalidSessionState( - "cannot change a session prefix after history is present or committed".into(), + "cannot change a session prefix after a committed turn".into(), )); } - self.history = prefix.messages().to_vec(); + let history = std::mem::take(&mut self.history); + let history = history + .strip_prefix(self.prefix.messages()) + .unwrap_or(&history) + .to_vec(); self.prefix = prefix; + self.history = self.with_prefix(history); Ok(()) } + fn state_view(&self, resumed: bool) -> SessionStateView<'_> { + SessionStateView { + history: &self.history, + raw_history: &self.persisted, + prefix: &self.prefix, + transcript_target: self.target.as_ref(), + committed_turns: self.committed_turns, + resumed, + } + } + fn encode( &self, previous: &[Message], diff --git a/crates/tinyagents-runtime/src/test.rs b/crates/tinyagents-runtime/src/test.rs index 0cfd5d46..9548ce5e 100644 --- a/crates/tinyagents-runtime/src/test.rs +++ b/crates/tinyagents-runtime/src/test.rs @@ -22,9 +22,9 @@ use tinytools::{Tool, ToolResult, ToolSpec}; use crate::{ CommitReceipt, DriverFailure, DriverOutcome, DriverRequest, HarnessDriver, PrefixSnapshot, - ResumeMode, RuntimeError, SessionBuilder, SessionDriver, SessionHooks, SessionStateView, - SessionTerminal, SessionTurnOutcome, SessionTurnRequest, ToolSnapshot, TranscriptCodec, - TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, + ResumeMode, ResumePreparation, RuntimeError, SessionBuilder, SessionDriver, SessionHooks, + SessionStateView, SessionTerminal, SessionTurnOutcome, SessionTurnRequest, ToolSnapshot, + TranscriptCodec, TranscriptTarget, TranscriptTurnOptions, TurnOptions, TurnPreparation, }; struct Driver { @@ -113,6 +113,7 @@ struct MemoryHistory { opens: Mutex, fail: Mutex, cancel_after_append: Mutex>, + cancel_after_read: Mutex>, } impl TranscriptRead for MemoryHistory { @@ -120,7 +121,11 @@ impl TranscriptRead for MemoryHistory { &self.path } fn read_session(&self) -> anyhow::Result> { - Ok(self.state.lock().unwrap().clone()) + let transcript = self.state.lock().unwrap().clone(); + if let Some(cancellation) = self.cancel_after_read.lock().unwrap().as_ref() { + cancellation.cancel(); + } + Ok(transcript) } } @@ -158,18 +163,28 @@ impl TranscriptHistory for MemoryHistory { } } -struct Locator(Arc); +struct Locator { + history: Arc, + latest_agents: Mutex>, + opened_stems: Mutex>, +} impl TranscriptLocator for Locator { - fn latest_for_agent(&self, _: &str) -> Option> { - Some(self.0.clone()) + fn latest_for_agent(&self, agent: &str) -> Option> { + self.latest_agents.lock().unwrap().push(agent.into()); + Some(self.history.clone()) } fn root_for_thread(&self, _: &str) -> Option> { - Some(self.0.clone()) + Some(self.history.clone()) } - fn open_stem(&self, _: &str, _: TranscriptMeta) -> anyhow::Result> { - *self.0.opens.lock().unwrap() += 1; - Ok(self.0.clone()) + fn open_stem( + &self, + stem: &str, + _: TranscriptMeta, + ) -> anyhow::Result> { + self.opened_stems.lock().unwrap().push(stem.into()); + *self.history.opens.lock().unwrap() += 1; + Ok(self.history.clone()) } } @@ -181,8 +196,16 @@ fn locator(session: Option) -> (Arc, Arc>, preparations: Mutex>, events: Arc, fail_commit: bool, @@ -246,6 +270,19 @@ struct Hook { #[async_trait] impl SessionHooks for Hook { + async fn before_resume( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(self + .resume_preparations + .lock() + .unwrap() + .pop_front() + .unwrap_or_default()) + } async fn before_turn( &self, _: &mut SessionTurnRequest, @@ -291,6 +328,7 @@ fn hook(preparations: Vec) -> (Arc, Arc) { let events = Arc::new(Events::default()); ( Arc::new(Hook { + resume_preparations: Mutex::new(VecDeque::new()), preparations: Mutex::new(preparations.into()), events: events.clone(), fail_commit: false, @@ -300,6 +338,15 @@ fn hook(preparations: Vec) -> (Arc, Arc) { ) } +fn hook_with_resume( + resume_preparations: Vec, + preparations: Vec, +) -> (Arc, Arc) { + let (hook, events) = hook(preparations); + *hook.resume_preparations.lock().unwrap() = resume_preparations.into(); + (hook, events) +} + fn tools(name: &str) -> ToolSnapshot { ToolSnapshot::new(vec![ToolSpec { name: name.into(), @@ -488,6 +535,7 @@ async fn rejected_before_commit_leaves_no_durable_state() { let (locator, history) = locator(None); let events = Arc::new(Events::default()); let hook = Arc::new(Hook { + resume_preparations: Mutex::new(VecDeque::new()), preparations: Mutex::new(VecDeque::new()), events, fail_commit: true, @@ -519,6 +567,7 @@ async fn post_commit_errors_cannot_relabel_a_successful_turn() { let (locator, history) = locator(None); let events = Arc::new(Events::default()); let hook = Arc::new(Hook { + resume_preparations: Mutex::new(VecDeque::new()), preparations: Mutex::new(VecDeque::new()), events: events.clone(), fail_commit: false, @@ -1054,11 +1103,15 @@ async fn prefix_mutation_after_a_commit_is_rejected() { Ok(outcome(vec![Message::assistant("one")])), Ok(outcome(vec![Message::assistant("two")])), ])); - let prep = TurnPreparation { + let first = TurnPreparation { prefix: Some(PrefixSnapshot::new(vec![Message::system("p")])), ..Default::default() }; - let (hook, _) = hook(vec![prep.clone(), prep]); + let second = TurnPreparation { + prefix: Some(PrefixSnapshot::new(vec![Message::system("changed")])), + ..Default::default() + }; + let (hook, _) = hook(vec![first, second]); let mut session = SessionBuilder::new(driver.clone()) .hooks(hook) .build() @@ -1083,13 +1136,15 @@ async fn prefix_mutation_after_a_commit_is_rejected() { } #[tokio::test] -async fn lazy_target_is_opened_only_after_before_turn_selects_it() { +async fn lazy_target_is_opened_only_after_before_resume_selects_it() { let (locator, history) = locator(None); let target = TranscriptTarget::new(locator, "late", meta()); - let (hook, _) = hook(vec![TurnPreparation { - transcript: Some(target), - ..Default::default() - }]); + let (hook, _) = hook_with_resume( + vec![ResumePreparation { + transcript: Some(target), + }], + vec![], + ); let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ Message::assistant("ok"), ]))]))) @@ -1108,6 +1163,290 @@ async fn lazy_target_is_opened_only_after_before_turn_selects_it() { assert_eq!(*history.opens.lock().unwrap(), 1); } +#[tokio::test] +async fn latest_resume_agent_is_distinct_from_the_write_stem() { + let (locator, _) = locator(Some(SessionTranscript { + meta: meta(), + messages: vec![TranscriptMessage::new("user", "resumed")], + })); + let target = TranscriptTarget::new(locator.clone(), "write-stem", meta()) + .with_resume_agent("resume-agent"); + let (hook, _) = hook_with_resume( + vec![ResumePreparation { + transcript: Some(target), + }], + vec![], + ); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("resumed"), + Message::assistant("next"), + ]))]))) + .codec(Arc::new(Codec::default())) + .hooks(hook) + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + + assert_eq!( + locator.latest_agents.lock().unwrap().as_slice(), + ["resume-agent"] + ); + assert_eq!( + locator.opened_stems.lock().unwrap().as_slice(), + ["write-stem"] + ); +} + +type ObservedResumedState = (bool, Vec, Vec); + +struct ResumedStateHook { + observed: Mutex>, +} + +#[async_trait] +impl SessionHooks for ResumedStateHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + state: SessionStateView<'_>, + ) -> Result { + *self.observed.lock().unwrap() = Some(( + state.resumed, + state.history.to_vec(), + state.raw_history.to_vec(), + )); + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } +} + +#[tokio::test] +async fn before_turn_receives_resumed_decoded_history_and_raw_rows() { + let mut raw = TranscriptMessage::new("user", "old"); + raw.extra_metadata = Some(serde_json::json!({"preserved": true})); + let (locator, _) = locator(Some(SessionTranscript { + meta: meta(), + messages: vec![raw.clone()], + })); + let hook = Arc::new(ResumedStateHook { + observed: Mutex::new(None), + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::user("old"), + Message::assistant("new"), + ]))]))) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook.clone()) + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + + assert_eq!( + *hook.observed.lock().unwrap(), + Some((true, vec![Message::user("old")], vec![raw])) + ); +} + +struct PrefixAfterResumeHook { + prefix: PrefixSnapshot, + calls: Mutex, +} + +#[async_trait] +impl SessionHooks for PrefixAfterResumeHook { + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + let mut calls = self.calls.lock().unwrap(); + *calls += 1; + Ok(TurnPreparation { + prefix: (*calls == 1).then(|| self.prefix.clone()), + ..TurnPreparation::default() + }) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } +} + +struct SystemCodec; + +impl TranscriptCodec for SystemCodec { + fn decode_history(&self, transcript: &SessionTranscript) -> Result, RuntimeError> { + Ok(transcript + .messages + .iter() + .map(|row| match row.role.as_str() { + "system" => Message::system(&row.content), + _ => Message::user(&row.content), + }) + .collect()) + } + fn reconcile( + &self, + _: &[TranscriptMessage], + _: &[Message], + next: &[Message], + _: &TranscriptTurnOptions, + ) -> Result, RuntimeError> { + Ok(next + .iter() + .map(|message| TranscriptMessage::new("assistant", message.text())) + .collect()) + } +} + +#[tokio::test] +async fn first_turn_prefix_accepts_an_exact_resumed_prefix_and_restores_it_after_compaction() { + let prefix = PrefixSnapshot::new(vec![Message::system("stable")]); + let (locator, _) = locator(Some(SessionTranscript { + meta: meta(), + messages: vec![ + TranscriptMessage::new("system", "stable"), + TranscriptMessage::new("user", "old"), + ], + })); + let driver = Arc::new(Driver::new(vec![ + Ok(outcome(vec![ + Message::system("stable"), + Message::user("old"), + Message::assistant("one"), + ])), + Ok(outcome(vec![Message::assistant("compacted")])), + ])); + let hook = Arc::new(PrefixAfterResumeHook { + prefix: prefix.clone(), + calls: Mutex::new(0), + }); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(SystemCodec)) + .transcript(locator, "agent", meta()) + .hooks(hook) + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("first")), + TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + let compacted = session + .turn( + SessionTurnRequest::new(Message::user("second")), + TurnOptions::default(), + ) + .await + .unwrap(); + + assert_eq!( + driver.requests.lock().unwrap()[0] + .history + .iter() + .filter(|message| **message == Message::system("stable")) + .count(), + 1 + ); + assert_eq!( + compacted.history, + vec![Message::system("stable"), Message::assistant("compacted")] + ); +} + +#[tokio::test] +async fn changed_first_turn_prefix_replaces_a_builder_prefix_after_resume() { + let old_prefix = PrefixSnapshot::new(vec![Message::system("old")]); + let new_prefix = PrefixSnapshot::new(vec![Message::system("new")]); + let (locator, _) = locator(Some(SessionTranscript { + meta: meta(), + messages: vec![ + TranscriptMessage::new("system", "old"), + TranscriptMessage::new("user", "resumed"), + ], + })); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::system("new"), + Message::user("resumed"), + Message::assistant("answer"), + ]))])); + let hook = Arc::new(PrefixAfterResumeHook { + prefix: new_prefix, + calls: Mutex::new(0), + }); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(SystemCodec)) + .prefix(old_prefix) + .transcript(locator, "agent", meta()) + .hooks(hook) + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("next")), + TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }, + ) + .await + .unwrap(); + + assert_eq!( + driver.requests.lock().unwrap()[0].history, + vec![ + Message::system("new"), + Message::user("resumed"), + Message::user("next"), + ] + ); +} + #[tokio::test] async fn hook_selected_target_and_resume_mode_apply_before_driver_handoff() { let (locator, _) = locator(Some(SessionTranscript { @@ -1122,18 +1461,25 @@ async fn hook_selected_target_and_resume_mode_apply_before_driver_handoff() { struct ResumeHook(TranscriptTarget); #[async_trait] impl SessionHooks for ResumeHook { - async fn before_turn( + async fn before_resume( &self, _: &mut SessionTurnRequest, options: &mut TurnOptions, _: SessionStateView<'_>, - ) -> Result { + ) -> Result { options.resume = ResumeMode::LatestForAgent; - Ok(TurnPreparation { + Ok(ResumePreparation { transcript: Some(self.0.clone()), - ..Default::default() }) } + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + Ok(TurnPreparation::default()) + } async fn before_commit( &self, _: &SessionTurnOutcome, @@ -1170,10 +1516,12 @@ async fn hook_selected_transcript_without_a_codec_fails_before_driver_execution( let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![Message::assistant( "never", )]))])); - let (hook, _) = hook(vec![TurnPreparation { - transcript: Some(target), - ..Default::default() - }]); + let (hook, _) = hook_with_resume( + vec![ResumePreparation { + transcript: Some(target), + }], + vec![], + ); let mut session = SessionBuilder::new(driver.clone()) .hooks(hook) .build() @@ -1461,6 +1809,233 @@ async fn hook_option_context_mutation_reaches_driver_and_codec() { assert_eq!(codec.0.lock().unwrap().as_slice(), ["mutated"]); } +struct BeforeResumeContextHook { + target: TranscriptTarget, + history: Arc, + before_turn_saw_lazy_target: Mutex, +} + +#[async_trait] +impl SessionHooks for BeforeResumeContextHook { + async fn before_resume( + &self, + _: &mut SessionTurnRequest, + options: &mut TurnOptions, + state: SessionStateView<'_>, + ) -> Result { + assert!(state.history.is_empty()); + options.request_id = Some("prepared-before-resume".into()); + options.run_context.data = Context("prepared-before-resume".into()); + Ok(ResumePreparation { + transcript: Some(self.target.clone()), + }) + } + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + options: &mut TurnOptions, + state: SessionStateView<'_>, + ) -> Result { + *self.before_turn_saw_lazy_target.lock().unwrap() = state + .transcript_target + .is_some_and(|target| target.stem == "late"); + assert_eq!(*self.history.opens.lock().unwrap(), 0); + assert!(!state.resumed); + assert_eq!( + options.request_id.as_deref(), + Some("prepared-before-resume") + ); + assert_eq!(options.run_context.data.0, "prepared-before-resume"); + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn on_terminal(&self, _: SessionTerminal) -> Result<(), RuntimeError> { + Ok(()) + } +} + +#[tokio::test] +async fn before_resume_mutates_context_and_options_while_target_remains_lazy() { + let (locator, history) = locator(None); + let driver = Arc::new(ContextDriver(Mutex::new(vec![]))); + let codec = Arc::new(ContextCodec(Mutex::new(vec![]))); + let hook = Arc::new(BeforeResumeContextHook { + target: TranscriptTarget::new(locator, "late", meta()), + history: history.clone(), + before_turn_saw_lazy_target: Mutex::new(false), + }); + let mut session = SessionBuilder::new(driver.clone()) + .codec(codec.clone()) + .hooks(hook.clone()) + .build() + .unwrap(); + let cancellation = tinyagents_harness::CancellationToken::new(); + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions { + request_id: None, + thread_id: None, + stream: false, + resume: ResumeMode::Never, + cancellation: cancellation.clone(), + run_context: RunContext::new(RunConfig::new("test"), Context("before".into())) + .with_cancellation(cancellation), + }, + ) + .await + .unwrap(); + + assert!(*hook.before_turn_saw_lazy_target.lock().unwrap()); + assert_eq!( + driver.0.lock().unwrap().as_slice(), + ["prepared-before-resume"] + ); + assert_eq!( + codec.0.lock().unwrap().as_slice(), + ["prepared-before-resume"] + ); + assert_eq!(*history.opens.lock().unwrap(), 1); +} + +struct ResumeCancellationHook { + before_resume_calls: Mutex, + before_turn_calls: Mutex, + after_commit_calls: Mutex, + terminals: Mutex>, +} + +#[async_trait] +impl SessionHooks for ResumeCancellationHook { + async fn before_resume( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + *self.before_resume_calls.lock().unwrap() += 1; + Ok(ResumePreparation::default()) + } + async fn before_turn( + &self, + _: &mut SessionTurnRequest, + _: &mut TurnOptions, + _: SessionStateView<'_>, + ) -> Result { + *self.before_turn_calls.lock().unwrap() += 1; + Ok(TurnPreparation::default()) + } + async fn before_commit( + &self, + _: &SessionTurnOutcome, + _: &TranscriptTurnOptions, + ) -> Result<(), RuntimeError> { + Ok(()) + } + async fn after_commit(&self, _: CommitReceipt) -> Result<(), RuntimeError> { + *self.after_commit_calls.lock().unwrap() += 1; + Ok(()) + } + async fn on_terminal(&self, terminal: SessionTerminal) -> Result<(), RuntimeError> { + self.terminals.lock().unwrap().push(terminal); + Ok(()) + } +} + +#[tokio::test] +async fn cancellation_before_resume_skips_hooks_and_preserves_terminal_behavior() { + let hook = Arc::new(ResumeCancellationHook { + before_resume_calls: Mutex::new(0), + before_turn_calls: Mutex::new(0), + after_commit_calls: Mutex::new(0), + terminals: Mutex::new(vec![]), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![Message::assistant( + "never", + )]))])); + let mut session = SessionBuilder::new(driver.clone()) + .hooks(hook.clone()) + .build() + .unwrap(); + let options = TurnOptions::default(); + options.cancellation.cancel(); + + assert_eq!( + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await, + Err(RuntimeError::Cancelled) + ); + assert_eq!(*hook.before_resume_calls.lock().unwrap(), 0); + assert_eq!(*hook.before_turn_calls.lock().unwrap(), 0); + assert_eq!(*hook.after_commit_calls.lock().unwrap(), 0); + assert!(matches!( + hook.terminals.lock().unwrap().as_slice(), + [SessionTerminal::Cancelled] + )); + assert!(driver.requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn cancellation_after_resume_before_before_turn_skips_driver_and_commit() { + let (locator, history) = locator(Some(SessionTranscript { + meta: meta(), + messages: vec![TranscriptMessage::new("user", "old")], + })); + let hook = Arc::new(ResumeCancellationHook { + before_resume_calls: Mutex::new(0), + before_turn_calls: Mutex::new(0), + after_commit_calls: Mutex::new(0), + terminals: Mutex::new(vec![]), + }); + let driver = Arc::new(Driver::new(vec![Ok(outcome(vec![Message::assistant( + "never", + )]))])); + let mut session = SessionBuilder::new(driver.clone()) + .codec(Arc::new(Codec::default())) + .transcript(locator, "agent", meta()) + .hooks(hook.clone()) + .build() + .unwrap(); + let options = TurnOptions { + resume: ResumeMode::LatestForAgent, + ..TurnOptions::default() + }; + *history.cancel_after_read.lock().unwrap() = Some(options.cancellation.clone()); + + assert_eq!( + session + .turn(SessionTurnRequest::new(Message::user("x")), options) + .await, + Err(RuntimeError::Cancelled) + ); + assert_eq!(*hook.before_resume_calls.lock().unwrap(), 1); + assert_eq!(*hook.before_turn_calls.lock().unwrap(), 0); + assert_eq!(*hook.after_commit_calls.lock().unwrap(), 0); + assert!(matches!( + hook.terminals.lock().unwrap().as_slice(), + [SessionTerminal::Cancelled] + )); + assert!(driver.requests.lock().unwrap().is_empty()); + assert_eq!( + history + .state + .lock() + .unwrap() + .as_ref() + .unwrap() + .meta + .turn_count, + 0 + ); +} + #[test] fn runtime_stays_host_neutral() { assert!( diff --git a/crates/tinyagents-runtime/src/types.rs b/crates/tinyagents-runtime/src/types.rs index 2996a03c..3aa50698 100644 --- a/crates/tinyagents-runtime/src/types.rs +++ b/crates/tinyagents-runtime/src/types.rs @@ -64,7 +64,11 @@ pub struct TranscriptTurnOptions { #[derive(Clone)] pub struct TranscriptTarget { pub locator: Arc, + /// The durable stem used for every append and write. pub stem: String, + /// Optional agent key used only by `ResumeMode::LatestForAgent` lookup. + /// When absent, the write stem is also the resume lookup key. + pub resume_agent: Option, pub meta: TranscriptMeta, } @@ -77,28 +81,42 @@ impl TranscriptTarget { Self { locator, stem: stem.into(), + resume_agent: None, meta, } } + /// Uses a distinct agent key when looking up the latest transcript. + pub fn with_resume_agent(mut self, resume_agent: impl Into) -> Self { + self.resume_agent = Some(resume_agent.into()); + self + } + pub(crate) fn same_binding(&self, other: &Self) -> bool { - self.stem == other.stem && Arc::ptr_eq(&self.locator, &other.locator) + self.stem == other.stem + && self.resume_agent == other.resume_agent + && Arc::ptr_eq(&self.locator, &other.locator) } } +/// Values prepared by `SessionHooks::before_resume` before transcript loading. +#[derive(Clone, Default)] +pub struct ResumePreparation { + /// A lazy transcript destination. It can be selected or replaced before + /// the first history handle is bound, but cannot be redirected afterwards. + pub transcript: Option, +} + /// Values prepared by `SessionHooks::before_turn` for exactly one driver call. #[derive(Clone, Default)] pub struct TurnPreparation { - /// A replacement prefix allowed only while the session is empty and has - /// never committed a transcript transition. + /// A replacement prefix allowed before the first committed turn. It is + /// reconciled against any decoded resumed history without duplication. pub prefix: Option, /// The immutable tool declarations for this driver request. `None` uses /// the builder's compatibility default and is never retained from a prior /// preparation. pub tools: Option, - /// A lazy transcript destination. It can be selected or replaced before - /// the first bind, but cannot be redirected after binding. - pub transcript: Option, } impl TurnPreparation { @@ -118,6 +136,9 @@ pub struct SessionStateView<'a> { pub prefix: &'a PrefixSnapshot, pub transcript_target: Option<&'a TranscriptTarget>, pub committed_turns: usize, + /// `true` only when this call loaded and decoded a durable transcript + /// before `before_turn` ran. + pub resumed: bool, } /// The shape of a successful logical transcript transition. From 612e4e8a5c9dbb0dc19f89462e6687406022eddc Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 20:14:54 +0300 Subject: [PATCH 08/18] feat(orchestration): carry explicit subagent request identity Co-authored-by: Medulla --- .../src/subagent/README.md | 33 ++- .../src/subagent/driver.rs | 24 +- .../src/subagent/planner.rs | 4 +- .../src/subagent/test.rs | 219 ++++++++++++++++-- .../src/subagent/types.rs | 181 +++++++++++++-- 5 files changed, 392 insertions(+), 69 deletions(-) diff --git a/crates/tinyagents-orchestration/src/subagent/README.md b/crates/tinyagents-orchestration/src/subagent/README.md index dbbbd298..bc10b56e 100644 --- a/crates/tinyagents-orchestration/src/subagent/README.md +++ b/crates/tinyagents-orchestration/src/subagent/README.md @@ -8,21 +8,33 @@ direction is: orchestration::subagent -> tinyagents-runtime -> {tinyagents-harness, tinyagents-session} ``` -Hosts provide three object-safe seams. `SubagentPlanner` transforms a live -`SubagentRequest` into a complete `PreparedSubagent`: resolved agent -identity, model messages, immutable `ToolSnapshot`, and explicit child -`RunContext`. `SubagentExecutor` executes precisely that plan. The +Hosts provide three object-safe seams. `SubagentPlanner` transforms a +live `SubagentRequest` into a complete `PreparedSubagent`: resolved +agent identity, model messages, immutable `ToolSnapshot`, and an explicit, +owned `RunContext`. `H` is opaque per-call host options and is forwarded +unchanged only to the planner. `SubagentExecutor` executes precisely that plan. The planner and executor own prompts, model choice, tool authorization, workspace policy, artifact resolution, and host context data; this module owns none of them. `SubagentPersistence` owns a host's durable resume and lifecycle records. -`SubagentTaskKey` is the durable lifecycle identity. It combines the root run, -immediate parent run, optional thread, and host-local task id. Persistence, +`SubagentTaskKey` is the durable lifecycle identity. It combines the original +root run, immediate parent run, optional thread, and host-local task id. +Fresh-task callers use `SubagentRequest::fresh_from_parent`, which derives the +key from an actual parent and validates the owned child context. Continuations +use `continue_with_key`: hosts must authorize and recover the original key, +then retain it even with a fresh owned execution context. Request fields are +private, so callers cannot bypass these identity checks. The key is the sole +durable thread identity; a request carries no duplicate thread field. + +Hosts must assign unique durable `RunConfig` ids to parent and child runs. The +constructors reject an equal parent/child id but cannot prove global uniqueness. +Persistence, in-flight coalescing, and terminal caching all use this scoped key: two parents may reuse a task id without sharing state, while repeat calls for the same durable scope deduplicate and resume correctly. -`SubagentDriver::run` loads a resume only when the caller did not provide one, +`SubagentDriver::run` uses the validated request key before it loads or +prepares anything. It then loads a resume only when the caller did not provide one, then prepares and executes. An awaiting-input result calls `save_pause`; all other results call `record_terminal`. These operations are mutually exclusive. The driver caches only successfully persisted terminal outcomes by scoped @@ -39,11 +51,14 @@ load and planning, and immediately after execution. A cancellation that wins after the executor produced an outcome preserves that outcome's output, history, usage, and neutral artifact references while reporting `Cancelled`. The supplied cancellation token is installed on the prepared `RunContext`, so -the executor and the actual harness context observe one cancellation tree. +the executor and the actual harness context observe one cancellation tree. If +host context `C` itself embeds a separate cancellation token, the executor is +responsible for synchronizing that host-owned token with this supplied token; +the neutral lifecycle cannot inspect opaque host data. Persistence `Ok(())` is the commit boundary: if cancellation wins before it, the driver abandons that uncommitted operation and records one `Cancelled` terminal outcome; if persistence commits first, the committed result remains -truthful. Planner and executor task ids must exactly match the request task id +truthful. Planner and executor task ids must exactly match the durable key task id or the driver returns a typed error before persistence or caching. Each coalesced caller retains its own cancellation token: cancelling a follower returns a local cancelled outcome promptly, without cancelling the leader or diff --git a/crates/tinyagents-orchestration/src/subagent/driver.rs b/crates/tinyagents-orchestration/src/subagent/driver.rs index 811f6795..778f9fd4 100644 --- a/crates/tinyagents-orchestration/src/subagent/driver.rs +++ b/crates/tinyagents-orchestration/src/subagent/driver.rs @@ -12,9 +12,9 @@ use tinyagents_harness::CancellationToken; /// /// Hosts that cannot provide every seam must receive a typed construction error /// rather than accidentally executing a partial lifecycle. -pub struct SubagentCapabilities { +pub struct SubagentCapabilities { /// Host planner, which resolves policy and explicit execution inputs. - pub planner: Option>>, + pub planner: Option>>, /// Host executor, which drives the prepared run. pub executor: Option>>, /// Host persistence for resume and one lifecycle record. @@ -26,8 +26,8 @@ pub struct SubagentCapabilities { /// Concurrent calls for one scoped task key coalesce, while task ids from /// distinct parents, roots, or threads remain independent. Hosts still need /// idempotent persistence for multiple processes/drivers. -pub struct SubagentDriver { - planner: Arc>, +pub struct SubagentDriver { + planner: Arc>, executor: Arc>, persistence: Arc, terminal_outcomes: Mutex>, @@ -77,9 +77,9 @@ impl InFlight { } } -impl SubagentDriver { +impl SubagentDriver { /// Validates host capability availability before exposing a runnable driver. - pub fn new(capabilities: SubagentCapabilities) -> Result { + pub fn new(capabilities: SubagentCapabilities) -> Result { Ok(Self { planner: capabilities .planner @@ -103,10 +103,10 @@ impl SubagentDriver { /// history, usage, and artifact references while changing only the status. pub async fn run( &self, - request: SubagentRequest, + request: SubagentRequest, cancellation: CancellationToken, ) -> Result { - let task_key = request.task_key(); + let task_key = request.task_key().clone(); if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { return Ok(outcome); } @@ -151,18 +151,18 @@ impl SubagentDriver { async fn run_reserved( &self, - mut request: SubagentRequest, + mut request: SubagentRequest, task_key: SubagentTaskKey, cancellation: CancellationToken, ) -> Result { - let task_id = request.task_id.clone(); + let task_id = request.task_id().to_owned(); if cancellation.is_cancelled() { return self.persist_cancelled(task_key, task_id).await; } - if request.resume.is_none() { - request.resume = self.persistence.load(&task_key).await?; + if request.resume().is_none() { + request.set_resume(self.persistence.load(&task_key).await?); } if cancellation.is_cancelled() { return self.persist_cancelled(task_key, task_id).await; diff --git a/crates/tinyagents-orchestration/src/subagent/planner.rs b/crates/tinyagents-orchestration/src/subagent/planner.rs index 7e24a4ec..7635df4b 100644 --- a/crates/tinyagents-orchestration/src/subagent/planner.rs +++ b/crates/tinyagents-orchestration/src/subagent/planner.rs @@ -8,10 +8,10 @@ use super::{PreparedSubagent, SubagentError, SubagentRequest}; /// lineage and tool policy. Returning a plan is intentionally all-or-nothing: /// the generic driver cannot fill missing host data with defaults. #[async_trait] -pub trait SubagentPlanner: Send + Sync { +pub trait SubagentPlanner: Send + Sync { /// Resolves one request before any execution or terminal persistence. async fn prepare( &self, - request: SubagentRequest, + request: SubagentRequest, ) -> Result, SubagentError>; } diff --git a/crates/tinyagents-orchestration/src/subagent/test.rs b/crates/tinyagents-orchestration/src/subagent/test.rs index 9313bf26..25e14a8f 100644 --- a/crates/tinyagents-orchestration/src/subagent/test.rs +++ b/crates/tinyagents-orchestration/src/subagent/test.rs @@ -48,12 +48,40 @@ struct FakePlanner { actions: Arc>>, } +#[derive(Clone, Debug, PartialEq, Eq)] +struct HostRequest { + label: String, +} + +struct PayloadPlanner { + payloads: Mutex>, +} + +#[async_trait] +impl SubagentPlanner for PayloadPlanner { + async fn prepare( + &self, + request: SubagentRequest, + ) -> Result, SubagentError> { + let request = request.into_parts(); + self.payloads.lock().unwrap().push(request.host_request); + Ok(PreparedSubagent { + task_id: request.task_key.task_id, + agent_key: "payload-agent".into(), + input: vec![Message::user(request.input)], + tools: ToolSnapshot::new(vec![]).unwrap(), + run_context: request.run_context, + }) + } +} + #[async_trait] impl SubagentPlanner for FakePlanner { async fn prepare( &self, request: SubagentRequest, ) -> Result, SubagentError> { + let request = request.into_parts(); *self.calls.lock().unwrap() += 1; *self.saw_resume.lock().unwrap() = request.resume.is_some(); self.seen_resumes @@ -65,11 +93,11 @@ impl SubagentPlanner for FakePlanner { return Err(SubagentError::Planning("rejected".into())); } Ok(PreparedSubagent { - task_id: request.task_id, + task_id: request.task_key.task_id, agent_key: "resolved-agent".into(), input: vec![Message::user(request.input)], tools: ToolSnapshot::new(vec![]).unwrap(), - run_context: request.parent_run, + run_context: request.run_context, }) } } @@ -77,6 +105,7 @@ impl SubagentPlanner for FakePlanner { struct FakeExecutor { calls: Mutex, context_ids: Mutex>, + root_run_ids: Mutex>, context_cancellations: Mutex>, mode: ExecutorMode, started: Mutex>>, @@ -94,6 +123,15 @@ impl SubagentExecutor for FakeExecutor { .lock() .unwrap() .push(execution.prepared.run_context.instance_id()); + self.root_run_ids.lock().unwrap().push( + execution + .prepared + .run_context + .lineage() + .root_run_id + .as_str() + .to_owned(), + ); self.context_cancellations .lock() .unwrap() @@ -215,20 +253,23 @@ impl SubagentPersistence for FakePersistence { } fn request(task_id: &str, data: &str) -> SubagentRequest { - request_with_parent( - task_id, - RunContext::new(RunConfig::new(format!("run-{task_id}")), data.into()), - ) + let parent = RunContext::new(RunConfig::new(format!("parent-{task_id}")), data.to_owned()); + let child = parent + .child(RunConfig::new(format!("run-{task_id}")), data.to_owned()) + .unwrap(); + SubagentRequest::fresh_from_parent(&parent, child, task_id, (), "do work", None).unwrap() } -fn request_with_parent(task_id: &str, parent_run: RunContext) -> SubagentRequest { - SubagentRequest { - task_id: task_id.into(), - parent_run, - input: "do work".into(), - thread_id: Some("thread-1".into()), - resume: None, - } +fn request_with_parent(task_id: &str, run_context: RunContext) -> SubagentRequest { + let task_key = SubagentTaskKey::from_context(&run_context, task_id, Some("thread-1".into())); + SubagentRequest::continue_with_key(task_key, run_context, (), "do work", None).unwrap() +} + +fn continuation_request( + task_key: SubagentTaskKey, + run_context: RunContext, +) -> SubagentRequest { + SubagentRequest::continue_with_key(task_key, run_context, (), "do work", None).unwrap() } fn driver( @@ -257,6 +298,7 @@ fn fakes(mode: ExecutorMode) -> Fakes { Arc::new(FakeExecutor { calls: Mutex::new(0), context_ids: Mutex::new(Vec::new()), + root_run_ids: Mutex::new(Vec::new()), context_cancellations: Mutex::new(Vec::new()), mode, started: Mutex::new(None), @@ -328,12 +370,13 @@ impl SubagentPlanner for MismatchedPlanner { &self, request: SubagentRequest, ) -> Result, SubagentError> { + let request = request.into_parts(); Ok(PreparedSubagent { task_id: "other-task".into(), agent_key: "resolved-agent".into(), input: vec![Message::user(request.input)], tools: ToolSnapshot::new(vec![]).unwrap(), - run_context: request.parent_run, + run_context: request.run_context, }) } } @@ -422,13 +465,17 @@ impl SubagentExecutor for NestedExecutor { .expect("nested executor is attached to its driver"); let child = child_driver .run( - SubagentRequest { - task_id: "child".into(), - parent_run: execution.prepared.run_context, - input: "nested work".into(), - thread_id: None, - resume: None, - }, + SubagentRequest::continue_with_key( + SubagentTaskKey::from_context( + &execution.prepared.run_context, + "child", + None, + ), + execution.prepared.run_context, + (), + "nested work", + None, + )?, execution.cancellation, ) .await?; @@ -487,7 +534,7 @@ async fn planner_rejection_does_not_execute_or_persist_terminal_state() { async fn prepared_context_identity_reaches_executor() { let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); let incoming = request("task", "identity"); - let expected = incoming.parent_run.instance_id(); + let expected = incoming.run_context().instance_id(); let outcome = driver(planner, executor.clone(), persistence) .run(incoming, CancellationToken::new()) .await @@ -497,6 +544,58 @@ async fn prepared_context_identity_reaches_executor() { assert_eq!(*executor.context_ids.lock().unwrap(), vec![expected]); } +#[tokio::test] +async fn caller_owned_child_context_keeps_its_lineage_into_execution() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); + let root = RunContext::new(RunConfig::new("root-lineage"), "root".to_owned()); + let child = root + .child(RunConfig::new("owned-child"), "child".to_owned()) + .unwrap(); + driver(planner, executor.clone(), persistence) + .run( + request_with_parent("lineage", child), + CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(*executor.root_run_ids.lock().unwrap(), vec!["root-lineage"]); +} + +#[tokio::test] +async fn opaque_host_request_reaches_the_planner_unchanged() { + let (_, executor, persistence, _) = fakes(ExecutorMode::Completed); + let planner = Arc::new(PayloadPlanner { + payloads: Mutex::new(Vec::new()), + }); + let context = RunContext::new(RunConfig::new("payload-run"), "ctx".into()); + let key = SubagentTaskKey::from_context(&context, "payload", Some("thread-1".into())); + let request = SubagentRequest::continue_with_key( + key, + context, + HostRequest { + label: "per-call options".into(), + }, + "do work", + None, + ) + .unwrap(); + SubagentDriver::new(SubagentCapabilities { + planner: Some(planner.clone()), + executor: Some(executor), + persistence: Some(persistence), + }) + .unwrap() + .run(request, CancellationToken::new()) + .await + .unwrap(); + assert_eq!( + *planner.payloads.lock().unwrap(), + vec![HostRequest { + label: "per-call options".into() + }] + ); +} + #[tokio::test] async fn completed_incomplete_and_pause_use_one_mutually_exclusive_persistence_action() { for (mode, expected) in [ @@ -635,6 +734,43 @@ async fn duplicate_task_id_returns_recorded_outcome_without_second_execution_or_ assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); } +#[tokio::test] +async fn fresh_request_rejects_an_unrelated_or_thread_mismatched_context() { + let parent = RunContext::new( + RunConfig::new("parent").with_thread("thread-a"), + "parent".to_owned(), + ); + let unrelated = RunContext::new( + RunConfig::new("unrelated").with_thread("thread-a"), + "other".to_owned(), + ); + assert!(matches!( + SubagentRequest::fresh_from_parent(&parent, unrelated, "task", (), "input", None), + Err(SubagentError::InvalidRequest(_)) + )); + + let wrong_thread_child = parent + .child( + RunConfig::new("wrong-thread-child").with_thread("thread-b"), + "child".to_owned(), + ) + .unwrap(); + assert!(matches!( + SubagentRequest::fresh_from_parent(&parent, wrong_thread_child, "task", (), "input", None,), + Err(SubagentError::InvalidRequest(_)) + )); + + let key = SubagentTaskKey::from_context(&parent, "task", None); + let fresh_wrong_thread = RunContext::new( + RunConfig::new("continuation").with_thread("thread-b"), + "fresh".to_owned(), + ); + assert!(matches!( + SubagentRequest::continue_with_key(key, fresh_wrong_thread, (), "input", None), + Err(SubagentError::InvalidRequest(_)) + )); +} + #[tokio::test] async fn concurrent_same_task_calls_coalesce_to_one_lifecycle() { let (planner, executor, persistence, _) = fakes(ExecutorMode::WaitForCancellation); @@ -825,6 +961,43 @@ async fn awaiting_input_is_not_cached_and_the_next_call_resumes_to_completion() assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); } +#[tokio::test] +async fn continuation_keeps_original_key_with_a_fresh_owned_context() { + let (planner, _, persistence, actions) = fakes(ExecutorMode::Pause); + let executor = Arc::new(PauseThenCompleteExecutor { + calls: Mutex::new(0), + actions, + }); + let lifecycle = driver(planner.clone(), executor, persistence.clone()); + let original = RunContext::new(RunConfig::new("original-run"), "first".into()); + let key = SubagentTaskKey::from_context(&original, "continued", Some("thread-1".into())); + lifecycle + .run( + continuation_request(key.clone(), original), + CancellationToken::new(), + ) + .await + .unwrap(); + let fresh = RunContext::new(RunConfig::new("fresh-turn-context"), "second".into()); + let completed = lifecycle + .run( + continuation_request(key.clone(), fresh), + CancellationToken::new(), + ) + .await + .unwrap(); + assert_eq!(completed.status, SubagentStatus::Completed); + assert_eq!(*planner.seen_resumes.lock().unwrap(), vec![false, true]); + assert!( + persistence + .keys + .lock() + .unwrap() + .iter() + .all(|seen| seen == &key) + ); +} + #[tokio::test] async fn driver_replaces_prepared_context_cancellation_with_execution_token() { let (planner, executor, persistence, _) = fakes(ExecutorMode::WaitForCancellation); diff --git a/crates/tinyagents-orchestration/src/subagent/types.rs b/crates/tinyagents-orchestration/src/subagent/types.rs index dee60d25..d6103e72 100644 --- a/crates/tinyagents-orchestration/src/subagent/types.rs +++ b/crates/tinyagents-orchestration/src/subagent/types.rs @@ -6,20 +6,34 @@ use tinyinference_llm::{message::Message, usage::UsageTotals}; /// A host-provided request to run a task through a subagent. /// -/// `parent_run` is live execution data, deliberately not a serializable host -/// DTO. The planner must derive the child's fully resolved context explicitly -/// instead of consulting task-local state. -pub struct SubagentRequest { - /// Host-local task id. Its durable lifecycle identity is scoped by - /// [`SubagentTaskKey`], derived from this request's parent run. - pub task_id: String, - /// The explicit live parent context from which a child context is derived. - pub parent_run: RunContext, - /// Host-visible task input that the planner converts into model messages. +/// `run_context` is owned live execution data, deliberately not a serializable +/// host DTO. It may already be a lineage-preserving child created from a +/// borrowed parent. Durable identity is carried separately in `task_key` so a +/// continuation can use a fresh execution context without changing its +/// original persistence key. +pub struct SubagentRequest { + task_key: SubagentTaskKey, + run_context: RunContext, + host_request: H, + input: String, + resume: Option, +} + +/// One-way planner view of a validated request. +/// +/// This is intentionally produced only by [`SubagentRequest::into_parts`]; it +/// cannot be turned back into a lifecycle request without a validated +/// constructor. +pub struct SubagentRequestParts { + /// Durable identity selected by the validated constructor. + pub task_key: SubagentTaskKey, + /// Owned execution context for this invocation. + pub run_context: RunContext, + /// Opaque host options for this invocation. + pub host_request: H, + /// Host-visible task input. pub input: String, - /// Optional host-owned conversation thread correlation id. - pub thread_id: Option, - /// A caller-supplied checkpoint. When absent the driver asks persistence. + /// Loaded or caller-supplied resume state. pub resume: Option, } @@ -41,18 +55,136 @@ pub struct SubagentTaskKey { pub task_id: String, } -impl SubagentRequest { - /// Derives the durable lifecycle key without exposing host context data. - pub fn task_key(&self) -> SubagentTaskKey { +impl SubagentTaskKey { + /// Creates an initial durable identity from a context at task creation. + /// Continuations must retain this returned key rather than derive another + /// one from their fresh execution context. + pub fn from_context( + run_context: &RunContext, + task_id: impl Into, + thread_id: Option, + ) -> Self { + let task_id = task_id.into(); SubagentTaskKey { - root_run_id: self.parent_run.lineage().root_run_id.as_str().to_owned(), - parent_run_id: self.parent_run.run_id().as_str().to_owned(), - thread_id: self - .thread_id - .clone() - .or_else(|| self.parent_run.thread_id().map(|id| id.as_str().to_owned())), - task_id: self.task_id.clone(), + root_run_id: run_context.lineage().root_run_id.as_str().to_owned(), + parent_run_id: run_context.run_id().as_str().to_owned(), + thread_id: thread_id + .or_else(|| run_context.thread_id().map(|id| id.as_str().to_owned())), + task_id, + } + } +} + +impl SubagentRequest { + /// Creates a fresh lifecycle from an actual parent and its owned child. + /// + /// Both `RunConfig` run ids must be durable unique host ids; this API can + /// reject an equal parent/child id but cannot prove global uniqueness. + pub fn fresh_from_parent

( + parent: &RunContext

, + owned_child: RunContext, + task_id: impl Into, + host_request: H, + input: impl Into, + resume: Option, + ) -> Result { + let task_key = SubagentTaskKey::from_context(parent, task_id, None); + Self::validate_key(&task_key)?; + if owned_child.lineage().root_run_id != parent.lineage().root_run_id + || owned_child.lineage().parent_run_id.as_ref() != Some(parent.run_id()) + || owned_child.run_id() == parent.run_id() + { + return Err(SubagentError::InvalidRequest( + "owned execution context is not a distinct direct child of the supplied parent" + .into(), + )); + } + Self::validate_thread(&task_key, &owned_child)?; + Ok(Self { + task_key, + run_context: owned_child, + host_request, + input: input.into(), + resume, + }) + } + + /// Continues an authorized durable lifecycle with a fresh owned context. + /// + /// Hosts must authorize and recover `original_key` from their own durable + /// record; do not derive a replacement key from the fresh context. + pub fn continue_with_key( + original_key: SubagentTaskKey, + fresh_owned_context: RunContext, + host_request: H, + input: impl Into, + resume: Option, + ) -> Result { + Self::validate_key(&original_key)?; + Self::validate_thread(&original_key, &fresh_owned_context)?; + Ok(Self { + task_key: original_key, + run_context: fresh_owned_context, + host_request, + input: input.into(), + resume, + }) + } + + /// Returns the sole durable identity for this lifecycle. + pub fn task_key(&self) -> &SubagentTaskKey { + &self.task_key + } + + /// Returns the durable task id, derived from [`Self::task_key`]. + pub fn task_id(&self) -> &str { + &self.task_key.task_id + } + + /// Borrows the owned execution context before the request is consumed. + pub fn run_context(&self) -> &RunContext { + &self.run_context + } + + pub(crate) fn resume(&self) -> Option<&SubagentResume> { + self.resume.as_ref() + } + + pub(crate) fn set_resume(&mut self, resume: Option) { + self.resume = resume; + } + + /// Consumes this validated request for planner use. + pub fn into_parts(self) -> SubagentRequestParts { + SubagentRequestParts { + task_key: self.task_key, + run_context: self.run_context, + host_request: self.host_request, + input: self.input, + resume: self.resume, + } + } + + fn validate_key(key: &SubagentTaskKey) -> Result<(), SubagentError> { + if key.root_run_id.is_empty() || key.parent_run_id.is_empty() || key.task_id.is_empty() { + return Err(SubagentError::InvalidRequest( + "durable task keys require non-empty root, parent, and task ids".into(), + )); + } + Ok(()) + } + + fn validate_thread( + key: &SubagentTaskKey, + context: &RunContext, + ) -> Result<(), SubagentError> { + let context_thread = context.thread_id().map(|thread| thread.as_str()); + if context_thread.is_some() && context_thread != key.thread_id.as_deref() { + return Err(SubagentError::InvalidRequest( + "execution context thread disagrees with durable task key".into(), + )); } + Ok(()) } } @@ -185,6 +317,8 @@ pub struct PersistedSubagentPause { /// owns them; the driver never flattens them into an untyped host error. #[derive(Clone, Debug, PartialEq, Eq)] pub enum SubagentError { + /// A supplied context or durable key violates neutral lifecycle invariants. + InvalidRequest(String), /// The planner rejected or could not resolve a request. Planning(String), /// The executor could not finish a prepared run. @@ -209,6 +343,7 @@ pub enum SubagentError { impl std::fmt::Display for SubagentError { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { + Self::InvalidRequest(message) => write!(f, "invalid subagent request: {message}"), Self::Planning(message) => write!(f, "subagent planning failed: {message}"), Self::Execution(message) => write!(f, "subagent execution failed: {message}"), Self::Persistence(message) => write!(f, "subagent persistence failed: {message}"), From 7e108c7614ea23c9c2fa2f7f6a4b4b3291ee3c64 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 20:15:00 +0300 Subject: [PATCH 09/18] feat(orchestration): export subagent request parts Co-authored-by: Medulla --- crates/tinyagents-orchestration/src/subagent/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/tinyagents-orchestration/src/subagent/mod.rs b/crates/tinyagents-orchestration/src/subagent/mod.rs index 552ec74c..b7cce138 100644 --- a/crates/tinyagents-orchestration/src/subagent/mod.rs +++ b/crates/tinyagents-orchestration/src/subagent/mod.rs @@ -23,8 +23,8 @@ pub use persistence::SubagentPersistence; pub use planner::SubagentPlanner; pub use types::{ ArtifactReference, PersistedSubagentPause, PreparedSubagent, SubagentError, SubagentExecution, - SubagentIncomplete, SubagentOutcome, SubagentPause, SubagentRequest, SubagentResume, - SubagentStatus, SubagentTaskKey, + SubagentIncomplete, SubagentOutcome, SubagentPause, SubagentRequest, SubagentRequestParts, + SubagentResume, SubagentStatus, SubagentTaskKey, }; #[cfg(test)] From 6930dc5577a6a9af44e64e9fed733f936740ed9b Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 21:14:22 +0300 Subject: [PATCH 10/18] feat(runtime): persist codec-provided turn usage atomically Co-authored-by: Medulla --- README.md | 7 +- crates/tinyagents-runtime/README.md | 13 +- crates/tinyagents-runtime/src/lib.rs | 16 +++ crates/tinyagents-runtime/src/session.rs | 26 +++- crates/tinyagents-runtime/src/test.rs | 150 ++++++++++++++++++++++- 5 files changed, 197 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index bddf1170..015ab54f 100644 --- a/README.md +++ b/README.md @@ -157,8 +157,11 @@ the sequencing around one append-only transcript commit. A host supplies the driver, its lossless transcript codec, and lifecycle hooks. On a driver error, the runtime can commit recoverable logical history with an interrupted, display-only partial in the same history operation; model-context replay omits -that partial. A post-commit hook observes durable successes but cannot change -their result. See [the runtime module](docs/modules/runtime/README.md). +that partial. A codec can also derive `TurnUsage` from its explicit host +context after the driver runs; that usage is attached to the same atomic +append's final assistant row for both success and recoverable partials. A +post-commit hook observes durable successes but cannot change their result. +See [the runtime module](docs/modules/runtime/README.md). ## Registry diff --git a/crates/tinyagents-runtime/README.md b/crates/tinyagents-runtime/README.md index d80ad8c0..b0f393da 100644 --- a/crates/tinyagents-runtime/README.md +++ b/crates/tinyagents-runtime/README.md @@ -19,7 +19,9 @@ The host supplies three narrow seams: than the runtime, retains fields inference messages cannot express. `C` is `Clone` so reconciliation receives the current host context plus request, thread, stream, and resume options after the live `RunContext` moves into the - driver. + driver. Its optional `turn_usage` hook reads host sidecars after driver + execution and returns a `TurnUsage` for the same atomic transcript append; + the default returns `None` for generic codecs. - `SessionHooks` prepares a request and mutable `TurnOptions` in two stages. `before_resume` lazily chooses a `TranscriptTarget`, then the runtime binds it and loads any requested transcript. `before_turn` sees that decoded @@ -51,9 +53,12 @@ uses `tinyagents-session`'s `TranscriptHistory::append_turn_with_partial`, so a extension appends only the new tail and a reduced context writes one compaction record. A supplied partial driver outcome is represented through that single history operation: logical history is replayable and interrupted display text -is not. Histories that cannot provide the combined operation reject a partial -rather than risk a two-step write. A persistence failure leaves the session's -in-memory history and persisted snapshot unchanged. +is not. A codec-provided `TurnUsage` is attached to that append's final +assistant row for both successful and recoverable-partial transitions. +Histories that cannot provide the combined operation reject a partial rather +than risk a two-step write. A reconciliation or usage-hook failure happens +before the append and leaves the session's in-memory history and persisted +snapshot unchanged. Every turn receives explicit `TurnOptions`, including its cancellation token and `RunContext`; no task-local data crosses the runtime boundary. The diff --git a/crates/tinyagents-runtime/src/lib.rs b/crates/tinyagents-runtime/src/lib.rs index 23fd86cc..0548d203 100644 --- a/crates/tinyagents-runtime/src/lib.rs +++ b/crates/tinyagents-runtime/src/lib.rs @@ -56,6 +56,22 @@ pub trait TranscriptCodec: Send + Sync { next: &[tinyinference_llm::message::Message], options: &TranscriptTurnOptions, ) -> Result, RuntimeError>; + + /// Returns host-derived provider usage for this transition, if available. + /// + /// The runtime calls this after the driver has completed (and therefore + /// after a host's explicit context sidecars may have been updated), but + /// before opening or appending a transcript. The returned value travels + /// through the same atomic [`tinyagents_session::transcript::TranscriptTurn`] + /// append as the reconciled rows, where the history implementation attaches + /// it to the turn's final assistant row. Generic codecs need no usage + /// policy, so the default remains `None`. + fn turn_usage( + &self, + _: &TranscriptTurnOptions, + ) -> Result, RuntimeError> { + Ok(None) + } } #[cfg(test)] diff --git a/crates/tinyagents-runtime/src/session.rs b/crates/tinyagents-runtime/src/session.rs index 6aec66e2..6572a77c 100644 --- a/crates/tinyagents-runtime/src/session.rs +++ b/crates/tinyagents-runtime/src/session.rs @@ -2,7 +2,7 @@ use std::{future::Future, sync::Arc}; use tinyagents_harness::CancellationToken; use tinyagents_session::transcript::{ - TranscriptHistory, TranscriptMessage, TranscriptPartial, TranscriptTurn, + TranscriptHistory, TranscriptMessage, TranscriptPartial, TranscriptTurn, TurnUsage, }; use tinyinference_llm::message::Message; @@ -247,11 +247,13 @@ impl Session { if let Some(partial) = failure.partial { let partial_history = self.with_prefix(partial.history); let raw = self.encode(&self.history, &partial_history, &codec_options)?; + let turn_usage = self.turn_usage(&codec_options)?; let receipt = self.persist( &raw, request_id.as_deref(), thread_id.as_deref(), partial.partial.as_ref(), + turn_usage.as_ref(), )?; self.history = partial_history; self.persisted = raw; @@ -277,7 +279,14 @@ impl Session { return Err(RuntimeError::Cancelled); } let raw = self.encode(&self.history, &candidate, &codec_options)?; - let transcript = self.persist(&raw, request_id.as_deref(), thread_id.as_deref(), None)?; + let turn_usage = self.turn_usage(&codec_options)?; + let transcript = self.persist( + &raw, + request_id.as_deref(), + thread_id.as_deref(), + None, + turn_usage.as_ref(), + )?; self.history = committed.history.clone(); self.persisted = raw; self.committed_turns += 1; @@ -379,12 +388,23 @@ impl Session { } } + fn turn_usage( + &self, + options: &TranscriptTurnOptions, + ) -> Result, RuntimeError> { + match &self.codec { + Some(codec) => codec.turn_usage(options), + None => Ok(None), + } + } + fn persist( &mut self, raw: &[TranscriptMessage], request_id: Option<&str>, thread_id: Option<&str>, partial: Option<&TranscriptPartial>, + turn_usage: Option<&TurnUsage>, ) -> Result, RuntimeError> { let Some(target) = self.target.as_mut() else { return Ok(None); @@ -408,7 +428,7 @@ impl Session { prev: &self.persisted, next: raw, meta: &meta, - turn_usage: None, + turn_usage, request_id, }, partial, diff --git a/crates/tinyagents-runtime/src/test.rs b/crates/tinyagents-runtime/src/test.rs index 9548ce5e..c1875f64 100644 --- a/crates/tinyagents-runtime/src/test.rs +++ b/crates/tinyagents-runtime/src/test.rs @@ -13,7 +13,7 @@ use tinyagents_harness::{ }; use tinyagents_session::transcript::{ DisplayRecord, FileTranscriptLocator, SessionTranscript, TranscriptHistory, TranscriptLocator, - TranscriptMessage, TranscriptMeta, TranscriptRead, TranscriptTurn, read_transcript, + TranscriptMessage, TranscriptMeta, TranscriptRead, TranscriptTurn, TurnUsage, read_transcript, read_transcript_display, }; use tinyinference_llm::message::Message; @@ -212,6 +212,8 @@ fn locator(session: Option) -> (Arc, Arc>>, seen_options: Mutex>, + turn_usage: Mutex>, + fail_turn_usage: bool, } impl TranscriptCodec for Codec { @@ -250,6 +252,31 @@ impl TranscriptCodec for Codec { } Ok(rows) } + + fn turn_usage(&self, _: &TranscriptTurnOptions) -> Result, RuntimeError> { + if self.fail_turn_usage { + return Err(RuntimeError::Driver("planned turn-usage failure".into())); + } + Ok(self.turn_usage.lock().unwrap().clone()) + } +} + +fn turn_usage() -> TurnUsage { + TurnUsage { + provider: "provider".into(), + model: "model".into(), + usage: tinyagents_session::transcript::MessageUsage { + input: 11, + output: 7, + cached_input: 3, + context_window: 128, + cost_usd: 0.42, + }, + ts: "now".into(), + reasoning_content: Some("because".into()), + tool_calls: Vec::new(), + iteration: 2, + } } #[derive(Default)] @@ -642,6 +669,10 @@ async fn harness_driver_rejects_a_snapshot_not_registered_by_the_harness() { async fn file_history_commits_partial_model_history_and_display_only_partial_together() { let directory = tempfile::tempdir().unwrap(); let locator = Arc::new(FileTranscriptLocator::new(directory.path())); + let codec = Arc::new(Codec { + turn_usage: Mutex::new(Some(turn_usage())), + ..Default::default() + }); let partial = DriverOutcome { history: vec![Message::assistant("recoverable")], output: None, @@ -652,7 +683,7 @@ async fn file_history_commits_partial_model_history_and_display_only_partial_tog error: RuntimeError::Driver("interrupted".into()), partial: Some(partial), })]))) - .codec(Arc::new(Codec::default())) + .codec(codec) .transcript(locator, "agent", meta()) .build() .unwrap(); @@ -666,15 +697,122 @@ async fn file_history_commits_partial_model_history_and_display_only_partial_tog .is_err() ); let path = directory.path().join("session_raw/agent.jsonl"); - assert_eq!( - read_transcript(&path).unwrap().messages[0].content, - "recoverable" - ); + let persisted = read_transcript(&path).unwrap(); + assert_eq!(persisted.messages[0].content, "recoverable"); + assert_eq!(persisted.messages[0].turn_usage, Some(turn_usage())); assert!(read_transcript_display(&path).unwrap().records.iter().any(|record| matches!(record, DisplayRecord::Message(message) if message.interrupted && message.message.content == "display partial" ))); } +#[tokio::test] +async fn codec_turn_usage_is_written_with_the_successful_turn_append() { + let directory = tempfile::tempdir().unwrap(); + let locator = Arc::new(FileTranscriptLocator::new(directory.path())); + let codec = Arc::new(Codec { + turn_usage: Mutex::new(Some(turn_usage())), + ..Default::default() + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("final"), + ]))]))) + .codec(codec) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await + .unwrap(); + + let transcript = read_transcript(&directory.path().join("session_raw/agent.jsonl")).unwrap(); + assert_eq!(transcript.messages.len(), 1); + assert_eq!(transcript.messages[0].content, "final"); + assert_eq!(transcript.messages[0].turn_usage, Some(turn_usage())); +} + +#[tokio::test] +async fn codec_turn_usage_error_prevents_any_durable_commit() { + let (locator, history) = locator(None); + let codec = Arc::new(Codec { + fail_turn_usage: true, + ..Default::default() + }); + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Ok(outcome(vec![ + Message::assistant("final"), + ]))]))) + .codec(codec) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + + assert_eq!( + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await, + Err(RuntimeError::Driver("planned turn-usage failure".into())) + ); + assert!(history.state.lock().unwrap().is_none()); + assert_eq!(*history.opens.lock().unwrap(), 0); + assert!(session.history().is_empty()); +} + +#[tokio::test] +async fn partial_usage_error_leaves_the_session_and_target_entirely_uncommitted() { + let (locator, history) = locator(None); + let codec = Arc::new(Codec { + fail_turn_usage: true, + ..Default::default() + }); + let partial = DriverOutcome { + history: vec![Message::assistant("recoverable")], + output: None, + partial: Some(crate::TranscriptPartial::new("display partial")), + interrupted: true, + }; + let mut session = SessionBuilder::new(Arc::new(Driver::new(vec![Err(DriverFailure { + error: RuntimeError::Driver("driver interrupted".into()), + partial: Some(partial), + })]))) + .codec(codec) + .transcript(locator, "agent", meta()) + .build() + .unwrap(); + + assert_eq!( + session + .turn( + SessionTurnRequest::new(Message::user("x")), + TurnOptions::default(), + ) + .await, + Err(RuntimeError::Driver("planned turn-usage failure".into())) + ); + // `turn_usage` runs before `persist`, so neither a lazy target bind nor + // its atomic append can happen on this failure path. + assert_eq!(*history.opens.lock().unwrap(), 0); + assert!(history.state.lock().unwrap().is_none()); + assert!(session.history().is_empty()); + // A seed is only rejected after `committed_turns` advances. Its success + // also replaces the runtime's raw persisted snapshot, proving the failed + // partial never became the next in-memory durable baseline. + assert!( + session + .seed_history( + vec![Message::assistant("seed")], + vec![TranscriptMessage::assistant("seed")], + ) + .is_ok() + ); +} + #[tokio::test] async fn cancellation_while_the_driver_is_waiting_has_one_cancelled_terminal() { let started = Arc::new(tokio::sync::Notify::new()); From 7d72ecc5623ca027eabc2460caa490485d368b89 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 23:56:59 +0300 Subject: [PATCH 11/18] feat(orchestration): add durable host-neutral subagent lifecycle Co-authored-by: Medulla --- .../src/subagent/driver.rs | 293 ++++++++++++++---- .../src/subagent/mod.rs | 5 +- .../src/subagent/persistence.rs | 36 ++- .../src/subagent/test.rs | 236 +++++++++++--- .../src/subagent/types.rs | 116 ++++++- 5 files changed, 579 insertions(+), 107 deletions(-) diff --git a/crates/tinyagents-orchestration/src/subagent/driver.rs b/crates/tinyagents-orchestration/src/subagent/driver.rs index 778f9fd4..a07830a8 100644 --- a/crates/tinyagents-orchestration/src/subagent/driver.rs +++ b/crates/tinyagents-orchestration/src/subagent/driver.rs @@ -4,7 +4,9 @@ use tokio::sync::{Mutex, Notify}; use super::{ PersistedSubagentPause, SubagentError, SubagentExecution, SubagentExecutor, SubagentOutcome, - SubagentPersistence, SubagentPlanner, SubagentRequest, SubagentStatus, SubagentTaskKey, + SubagentPausePersistenceDisposition, SubagentPersistence, SubagentPersistenceDisposition, + SubagentPlanner, SubagentRequest, SubagentRunResult, SubagentStatus, SubagentTaskKey, + SubagentTerminalPersistenceDisposition, }; use tinyagents_harness::CancellationToken; @@ -39,7 +41,7 @@ pub struct SubagentDriver { /// The map only protects reservation and removal. Planner, executor, and /// persistence futures never run while it is locked. struct InFlight { - result: Mutex>>, + result: Mutex>>, notify: Notify, } @@ -55,7 +57,7 @@ impl InFlight { &self, cancellation: &CancellationToken, task_id: &str, - ) -> Result { + ) -> Result { loop { let notified = self.notify.notified(); tokio::pin!(notified); @@ -65,13 +67,13 @@ impl InFlight { } tokio::select! { biased; - _ = cancellation.cancelled() => return Ok(SubagentOutcome::cancelled(task_id)), + _ = cancellation.cancelled() => return Ok(SubagentRunResult::new(SubagentOutcome::cancelled(task_id), SubagentPersistenceDisposition::ObserverCancelled)), _ = &mut notified => {} } } } - async fn complete(&self, result: Result) { + async fn complete(&self, result: Result) { *self.result.lock().await = Some(result); self.notify.notify_waiters(); } @@ -105,10 +107,20 @@ impl SubagentDriver { &self, request: SubagentRequest, cancellation: CancellationToken, - ) -> Result { + ) -> Result { let task_key = request.task_key().clone(); if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { - return Ok(outcome); + return Ok(SubagentRunResult::new( + outcome, + SubagentPersistenceDisposition::TerminalExisting, + )); + } + if let Some(outcome) = self.persistence.load_terminal(&task_key).await? { + self.cache_terminal(task_key, &outcome).await; + return Ok(SubagentRunResult::new( + outcome, + SubagentPersistenceDisposition::TerminalExisting, + )); } let (entry, is_leader) = { @@ -123,16 +135,45 @@ impl SubagentDriver { } }; if !is_leader { - return entry.wait(&cancellation, &task_key.task_id).await; + return entry + .wait(&cancellation, &task_key.task_id) + .await + .map(|result| { + SubagentRunResult::new( + result.outcome, + match result.disposition { + SubagentPersistenceDisposition::PauseCommitted + | SubagentPersistenceDisposition::PauseReplaced + | SubagentPersistenceDisposition::PauseExisting => { + SubagentPersistenceDisposition::PauseExisting + } + SubagentPersistenceDisposition::TerminalInserted + | SubagentPersistenceDisposition::TerminalExisting => { + SubagentPersistenceDisposition::TerminalExisting + } + SubagentPersistenceDisposition::ObserverCancelled => { + SubagentPersistenceDisposition::ObserverCancelled + } + }, + ) + }); } // A preceding caller may have committed a terminal result between the // initial cache check and this reservation. Do not reopen that task // after its in-flight entry has been removed. if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { - entry.complete(Ok(outcome.clone())).await; + entry + .complete(Ok(SubagentRunResult::new( + outcome.clone(), + SubagentPersistenceDisposition::TerminalExisting, + ))) + .await; let mut in_flight = self.in_flight.lock().await; in_flight.remove(&task_key); - return Ok(outcome); + return Ok(SubagentRunResult::new( + outcome, + SubagentPersistenceDisposition::TerminalExisting, + )); } let result = self @@ -154,18 +195,31 @@ impl SubagentDriver { mut request: SubagentRequest, task_key: SubagentTaskKey, cancellation: CancellationToken, - ) -> Result { + ) -> Result { let task_id = request.task_id().to_owned(); if cancellation.is_cancelled() { - return self.persist_cancelled(task_key, task_id).await; + return self + .persist_cancelled(task_key, SubagentOutcome::cancelled(task_id), None) + .await; } if request.resume().is_none() { request.set_resume(self.persistence.load(&task_key).await?); } + // Keep the exact pause seen at the durable read boundary. The + // persistence seam uses it as a scoped compare-and-swap expectation + // if execution pauses again, so a resumed pause can advance while a + // duplicate continuation returns the newer durable winner. + let expected_pause = request.resume().cloned(); if cancellation.is_cancelled() { - return self.persist_cancelled(task_key, task_id).await; + return self + .persist_cancelled( + task_key, + SubagentOutcome::cancelled(task_id), + expected_pause, + ) + .await; } let mut prepared = self.planner.prepare(request).await?; @@ -176,7 +230,13 @@ impl SubagentDriver { }); } if cancellation.is_cancelled() { - return self.persist_cancelled(task_key, prepared.task_id).await; + return self + .persist_cancelled( + task_key, + SubagentOutcome::cancelled(prepared.task_id), + expected_pause, + ) + .await; } prepared.run_context = prepared.run_context.with_cancellation(cancellation.clone()); @@ -206,84 +266,207 @@ impl SubagentDriver { Err(_) if cancellation.is_cancelled() => SubagentOutcome::cancelled(task_id), Err(error) => return Err(error), }; - self.persist(task_key, outcome, &cancellation).await + self.persist(task_key, outcome, expected_pause, &cancellation) + .await } async fn persist_cancelled( &self, task_key: SubagentTaskKey, - task_id: String, - ) -> Result { - let outcome = SubagentOutcome::cancelled(task_id); + mut outcome: SubagentOutcome, + expected_pause: Option, + ) -> Result { + outcome.status = SubagentStatus::Cancelled; // Once cancellation has won, this is the one terminal action. Do not // race it with the already-latched token or a caller could observe an // indeterminate terminal write. - self.persistence - .record_terminal(&task_key, &outcome) - .await?; - self.cache_terminal(task_key, &outcome).await; - Ok(outcome) + match self + .persistence + .record_terminal(&task_key, &outcome, expected_pause.as_ref()) + .await? + { + SubagentTerminalPersistenceDisposition::Inserted => { + self.cache_terminal(task_key, &outcome).await; + Ok(SubagentRunResult::new( + outcome, + SubagentPersistenceDisposition::TerminalInserted, + )) + } + SubagentTerminalPersistenceDisposition::Existing => { + self.load_terminal_winner(task_key).await + } + SubagentTerminalPersistenceDisposition::PauseExisting => { + self.load_pause_winner(task_key).await + } + } } async fn persist( &self, task_key: SubagentTaskKey, outcome: SubagentOutcome, + expected_pause: Option, cancellation: &CancellationToken, - ) -> Result { + ) -> Result { if matches!(&outcome.status, SubagentStatus::Cancelled) { - self.persistence - .record_terminal(&task_key, &outcome) - .await?; - self.cache_terminal(task_key, &outcome).await; - return Ok(outcome); + return self + .persist_cancelled(task_key, outcome, expected_pause) + .await; } - let committed = match &outcome.status { - SubagentStatus::AwaitingInput(pause) => { - self.commit_or_cancel( - self.persistence.save_pause(PersistedSubagentPause { - key: task_key.clone(), - pause: pause.clone(), - }), - cancellation, - ) - .await? + let disposition = match &outcome.status { + SubagentStatus::AwaitingInput(_) => { + match self + .commit_or_cancel( + self.persistence.save_pause(PersistedSubagentPause { + key: task_key.clone(), + outcome: outcome.clone(), + replaces: expected_pause.clone(), + }), + cancellation, + ) + .await? + { + Some(SubagentPausePersistenceDisposition::Inserted) => { + SubagentPersistenceDisposition::PauseCommitted + } + Some(SubagentPausePersistenceDisposition::Replaced) => { + SubagentPersistenceDisposition::PauseReplaced + } + Some(SubagentPausePersistenceDisposition::Existing) => { + return self.load_pause_winner(task_key).await; + } + Some(SubagentPausePersistenceDisposition::TerminalExisting) => { + return self.load_terminal_winner(task_key).await; + } + None => { + return self + .persist_cancelled(task_key, outcome, expected_pause) + .await; + } + } } SubagentStatus::Completed | SubagentStatus::Incomplete(_) => { - self.commit_or_cancel( - self.persistence.record_terminal(&task_key, &outcome), - cancellation, - ) - .await? + match self + .commit_or_cancel( + async { + self.persistence + .record_terminal(&task_key, &outcome, expected_pause.as_ref()) + .await + }, + cancellation, + ) + .await? + { + Some(SubagentTerminalPersistenceDisposition::Inserted) => { + SubagentPersistenceDisposition::TerminalInserted + } + Some(SubagentTerminalPersistenceDisposition::Existing) => { + SubagentPersistenceDisposition::TerminalExisting + } + Some(SubagentTerminalPersistenceDisposition::PauseExisting) => { + return self.load_pause_winner(task_key).await; + } + None => { + return self + .persist_cancelled(task_key, outcome, expected_pause) + .await; + } + } } SubagentStatus::Cancelled => unreachable!("handled before persistence race"), }; - if !committed { - return self.persist_cancelled(task_key, outcome.task_id).await; - } if matches!( &outcome.status, SubagentStatus::Completed | SubagentStatus::Incomplete(_) ) { - self.cache_terminal(task_key, &outcome).await; + let terminal = if matches!( + disposition, + SubagentPersistenceDisposition::TerminalExisting + ) { + self.persistence + .load_terminal(&task_key) + .await? + .ok_or_else(|| { + SubagentError::Persistence( + "terminal insert lost without durable outcome".into(), + ) + })? + } else { + outcome.clone() + }; + self.cache_terminal(task_key, &terminal).await; + return Ok(SubagentRunResult::new(terminal, disposition)); + } + Ok(SubagentRunResult::new(outcome, disposition)) + } + + async fn load_pause_winner( + &self, + task_key: SubagentTaskKey, + ) -> Result { + // A terminal can win after the pause CAS reports contention. It is + // authoritative over every prior pause and must never be reopened. + if let Some(terminal) = self.persistence.load_terminal(&task_key).await? { + self.cache_terminal(task_key, &terminal).await; + return Ok(SubagentRunResult::new( + terminal, + SubagentPersistenceDisposition::TerminalExisting, + )); } - Ok(outcome) + let paused = self + .persistence + .load_pause(&task_key) + .await? + .ok_or_else(|| { + SubagentError::Persistence( + "pause compare-and-swap lost without a durable pause or terminal outcome" + .into(), + ) + })?; + if !matches!(paused.status, SubagentStatus::AwaitingInput(_)) { + return Err(SubagentError::Persistence( + "durable pause record did not contain an awaiting-input outcome".into(), + )); + } + Ok(SubagentRunResult::new( + paused, + SubagentPersistenceDisposition::PauseExisting, + )) + } + + async fn load_terminal_winner( + &self, + task_key: SubagentTaskKey, + ) -> Result { + let terminal = self + .persistence + .load_terminal(&task_key) + .await? + .ok_or_else(|| { + SubagentError::Persistence( + "terminal compare-and-swap lost without a durable terminal outcome".into(), + ) + })?; + self.cache_terminal(task_key, &terminal).await; + Ok(SubagentRunResult::new( + terminal, + SubagentPersistenceDisposition::TerminalExisting, + )) } - async fn commit_or_cancel( + async fn commit_or_cancel( &self, operation: F, cancellation: &CancellationToken, - ) -> Result + ) -> Result, SubagentError> where - F: Future>, + F: Future>, { tokio::select! { biased; - _ = cancellation.cancelled() => Ok(false), + _ = cancellation.cancelled() => Ok(None), result = operation => { - result?; - Ok(true) + Ok(Some(result?)) } } } diff --git a/crates/tinyagents-orchestration/src/subagent/mod.rs b/crates/tinyagents-orchestration/src/subagent/mod.rs index b7cce138..44a613d0 100644 --- a/crates/tinyagents-orchestration/src/subagent/mod.rs +++ b/crates/tinyagents-orchestration/src/subagent/mod.rs @@ -23,8 +23,9 @@ pub use persistence::SubagentPersistence; pub use planner::SubagentPlanner; pub use types::{ ArtifactReference, PersistedSubagentPause, PreparedSubagent, SubagentError, SubagentExecution, - SubagentIncomplete, SubagentOutcome, SubagentPause, SubagentRequest, SubagentRequestParts, - SubagentResume, SubagentStatus, SubagentTaskKey, + SubagentIncomplete, SubagentOutcome, SubagentPause, SubagentPausePersistenceDisposition, + SubagentPersistenceDisposition, SubagentRequest, SubagentRequestParts, SubagentResume, + SubagentRunResult, SubagentStatus, SubagentTaskKey, SubagentTerminalPersistenceDisposition, }; #[cfg(test)] diff --git a/crates/tinyagents-orchestration/src/subagent/persistence.rs b/crates/tinyagents-orchestration/src/subagent/persistence.rs index 8751a76d..7957d569 100644 --- a/crates/tinyagents-orchestration/src/subagent/persistence.rs +++ b/crates/tinyagents-orchestration/src/subagent/persistence.rs @@ -1,7 +1,8 @@ use async_trait::async_trait; use super::{ - PersistedSubagentPause, SubagentError, SubagentOutcome, SubagentResume, SubagentTaskKey, + PersistedSubagentPause, SubagentError, SubagentOutcome, SubagentPausePersistenceDisposition, + SubagentResume, SubagentTaskKey, SubagentTerminalPersistenceDisposition, }; /// Host boundary for durable pause, resume, and terminal lifecycle state. @@ -12,23 +13,48 @@ use super::{ /// never a global lifecycle identity. The driver additionally suppresses /// duplicate records from repeated calls made through the same driver instance. A persistence /// future's successful return is its commit boundary: implementations must not -/// make a write visible and then await again before returning `Ok(())`. The +/// make a write visible and then await again before returning `Ok(true)`. The /// driver races that boundary with cancellation and, when cancellation wins, /// records one truthful `Cancelled` terminal outcome instead. #[async_trait] pub trait SubagentPersistence: Send + Sync { + /// Returns a terminal outcome committed by another driver/process, if any. + /// The lifecycle consults this before planning so a durable terminal never + /// reopens merely because this process has an empty in-memory cache. + async fn load_terminal( + &self, + key: &SubagentTaskKey, + ) -> Result, SubagentError>; /// Loads the most recent resumable state, if a caller did not supply one. async fn load(&self, key: &SubagentTaskKey) -> Result, SubagentError>; + /// Loads the complete durable pause outcome for observers that lost a + /// pause compare-and-swap. This is deliberately richer than [`Self::load`] + /// so a caller never returns its own discarded output, usage, or artifacts. + async fn load_pause( + &self, + key: &SubagentTaskKey, + ) -> Result, SubagentError>; + /// Saves one resumable pause. The driver never also records a terminal for /// that same committed outcome. A paused outcome is deliberately not cached /// by the driver; a later call reloads this state and resumes execution. - async fn save_pause(&self, pause: PersistedSubagentPause) -> Result<(), SubagentError>; + /// Atomically creates or advances a pause record. A continuation must + /// replace only the exact durable pause it loaded; competing continuations + /// receive `Existing` and return that winner without emitting effects. + async fn save_pause( + &self, + pause: PersistedSubagentPause, + ) -> Result; - /// Records a non-pause terminal outcome exactly once per scoped task. + /// Atomically records a terminal outcome. A continuation can consume only + /// the exact pause it loaded; a fresh lifecycle can close only an unpaused + /// key. This prevents an independent pause and terminal execution from + /// both claiming host effects. async fn record_terminal( &self, key: &SubagentTaskKey, outcome: &SubagentOutcome, - ) -> Result<(), SubagentError>; + replaces: Option<&SubagentResume>, + ) -> Result; } diff --git a/crates/tinyagents-orchestration/src/subagent/test.rs b/crates/tinyagents-orchestration/src/subagent/test.rs index 25e14a8f..3b26dd5f 100644 --- a/crates/tinyagents-orchestration/src/subagent/test.rs +++ b/crates/tinyagents-orchestration/src/subagent/test.rs @@ -1,8 +1,11 @@ -use std::sync::{ - Arc, Mutex, - atomic::{AtomicBool, Ordering}, -}; use std::time::Duration; +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, +}; use async_trait::async_trait; use tinyagents_harness::{ @@ -198,12 +201,21 @@ struct FakePersistence { pause_error: bool, terminal_error: bool, outcomes: Mutex>, + terminals: Mutex>, + pauses: Mutex>, saved_pause: Mutex>, keys: Mutex>, } #[async_trait] impl SubagentPersistence for FakePersistence { + async fn load_terminal( + &self, + key: &SubagentTaskKey, + ) -> Result, SubagentError> { + self.keys.lock().unwrap().push(key.clone()); + Ok(self.terminals.lock().unwrap().get(key).cloned()) + } async fn load(&self, key: &SubagentTaskKey) -> Result, SubagentError> { self.actions.lock().unwrap().push(Action::Load); self.keys.lock().unwrap().push(key.clone()); @@ -217,14 +229,52 @@ impl SubagentPersistence for FakePersistence { } } - async fn save_pause(&self, pause: PersistedSubagentPause) -> Result<(), SubagentError> { + async fn load_pause( + &self, + key: &SubagentTaskKey, + ) -> Result, SubagentError> { + Ok(self.pauses.lock().unwrap().get(key).cloned()) + } + + async fn save_pause( + &self, + pause: PersistedSubagentPause, + ) -> Result { self.actions.lock().unwrap().push(Action::Pause); - self.keys.lock().unwrap().push(pause.key); + self.keys.lock().unwrap().push(pause.key.clone()); if self.pause_error { Err(SubagentError::Persistence("pause save failed".into())) } else { - *self.saved_pause.lock().unwrap() = Some(pause.pause.resume); - Ok(()) + let mut pauses = self.pauses.lock().unwrap(); + match pauses.get(&pause.key) { + None if pause.replaces.is_none() => { + let resume = match &pause.outcome.status { + SubagentStatus::AwaitingInput(pause) => pause.resume.clone(), + _ => unreachable!("fake only persists awaiting outcomes"), + }; + *self.saved_pause.lock().unwrap() = Some(resume); + pauses.insert(pause.key, pause.outcome); + Ok(SubagentPausePersistenceDisposition::Inserted) + } + Some(current) + if pause.replaces.as_ref().is_some_and(|expected| { + matches!( + ¤t.status, + SubagentStatus::AwaitingInput(current_pause) + if current_pause.resume == *expected + ) + }) => + { + let resume = match &pause.outcome.status { + SubagentStatus::AwaitingInput(pause) => pause.resume.clone(), + _ => unreachable!("fake only persists awaiting outcomes"), + }; + *self.saved_pause.lock().unwrap() = Some(resume); + pauses.insert(pause.key, pause.outcome); + Ok(SubagentPausePersistenceDisposition::Replaced) + } + _ => Ok(SubagentPausePersistenceDisposition::Existing), + } } } @@ -232,7 +282,8 @@ impl SubagentPersistence for FakePersistence { &self, key: &SubagentTaskKey, outcome: &SubagentOutcome, - ) -> Result<(), SubagentError> { + replaces: Option<&SubagentResume>, + ) -> Result { self.actions .lock() .unwrap() @@ -246,8 +297,24 @@ impl SubagentPersistence for FakePersistence { if self.terminal_error { Err(SubagentError::Persistence("terminal save failed".into())) } else { - self.outcomes.lock().unwrap().push(outcome.clone()); - Ok(()) + let mut terminals = self.terminals.lock().unwrap(); + if terminals.contains_key(key) { + Ok(SubagentTerminalPersistenceDisposition::Existing) + } else if self.pauses.lock().unwrap().get(key).is_some_and(|paused| { + !replaces.is_some_and(|expected| { + matches!( + &paused.status, + SubagentStatus::AwaitingInput(pause) if pause.resume == *expected + ) + }) + }) { + Ok(SubagentTerminalPersistenceDisposition::PauseExisting) + } else { + terminals.insert(key.clone(), outcome.clone()); + self.pauses.lock().unwrap().remove(key); + self.outcomes.lock().unwrap().push(outcome.clone()); + Ok(SubagentTerminalPersistenceDisposition::Inserted) + } } } } @@ -310,6 +377,8 @@ fn fakes(mode: ExecutorMode) -> Fakes { pause_error: false, terminal_error: false, outcomes: Mutex::new(Vec::new()), + terminals: Mutex::new(HashMap::new()), + pauses: Mutex::new(HashMap::new()), saved_pause: Mutex::new(None), keys: Mutex::new(Vec::new()), }), @@ -335,30 +404,47 @@ struct BlockingPersistence { #[async_trait] impl SubagentPersistence for BlockingPersistence { + async fn load_terminal( + &self, + _key: &SubagentTaskKey, + ) -> Result, SubagentError> { + Ok(None) + } async fn load(&self, _: &SubagentTaskKey) -> Result, SubagentError> { Ok(None) } - async fn save_pause(&self, _: PersistedSubagentPause) -> Result<(), SubagentError> { + async fn load_pause( + &self, + _: &SubagentTaskKey, + ) -> Result, SubagentError> { + Ok(None) + } + + async fn save_pause( + &self, + _: PersistedSubagentPause, + ) -> Result { if matches!(self.stage, BlockingStage::Pause) && self.first.swap(false, Ordering::AcqRel) { self.started.notify_one(); self.release.notified().await; } - Ok(()) + Ok(SubagentPausePersistenceDisposition::Inserted) } async fn record_terminal( &self, _: &SubagentTaskKey, outcome: &SubagentOutcome, - ) -> Result<(), SubagentError> { + _: Option<&SubagentResume>, + ) -> Result { if matches!(self.stage, BlockingStage::Terminal) && self.first.swap(false, Ordering::AcqRel) { self.started.notify_one(); self.release.notified().await; } self.outcomes.lock().unwrap().push(outcome.clone()); - Ok(()) + Ok(SubagentTerminalPersistenceDisposition::Inserted) } } @@ -487,7 +573,7 @@ impl SubagentExecutor for NestedExecutor { // This models the host's parent-visible roll-up: the child is // added once alongside the parent's own model call. usage: UsageTotals { - calls: child.usage.calls + 1, + calls: child.outcome.usage.calls + 1, ..UsageTotals::default() }, artifacts: Vec::new(), @@ -540,7 +626,7 @@ async fn prepared_context_identity_reaches_executor() { .await .unwrap(); - assert_eq!(outcome.status, SubagentStatus::Completed); + assert_eq!(outcome.outcome.status, SubagentStatus::Completed); assert_eq!(*executor.context_ids.lock().unwrap(), vec![expected]); } @@ -635,6 +721,8 @@ async fn load_pause_and_terminal_errors_remain_typed() { pause_error: false, terminal_error: false, outcomes: Mutex::new(Vec::new()), + terminals: Mutex::new(HashMap::new()), + pauses: Mutex::new(HashMap::new()), saved_pause: Mutex::new(None), keys: Mutex::new(Vec::new()), }); @@ -652,6 +740,8 @@ async fn load_pause_and_terminal_errors_remain_typed() { load_mode: LoadMode::Empty, terminal_error: false, outcomes: Mutex::new(Vec::new()), + terminals: Mutex::new(HashMap::new()), + pauses: Mutex::new(HashMap::new()), saved_pause: Mutex::new(None), keys: Mutex::new(Vec::new()), }); @@ -669,6 +759,8 @@ async fn load_pause_and_terminal_errors_remain_typed() { load_mode: LoadMode::Empty, pause_error: false, outcomes: Mutex::new(Vec::new()), + terminals: Mutex::new(HashMap::new()), + pauses: Mutex::new(HashMap::new()), saved_pause: Mutex::new(None), keys: Mutex::new(Vec::new()), }); @@ -689,6 +781,8 @@ async fn loaded_resume_reaches_planner_before_execution_and_execution_errors_do_ pause_error: false, terminal_error: false, outcomes: Mutex::new(Vec::new()), + terminals: Mutex::new(HashMap::new()), + pauses: Mutex::new(HashMap::new()), saved_pause: Mutex::new(None), keys: Mutex::new(Vec::new()), }); @@ -728,12 +822,68 @@ async fn duplicate_task_id_returns_recorded_outcome_without_second_execution_or_ .await .unwrap(); - assert_eq!(first, second); + assert_eq!(first.outcome, second.outcome); + assert_eq!( + first.disposition, + SubagentPersistenceDisposition::TerminalInserted + ); + assert_eq!( + second.disposition, + SubagentPersistenceDisposition::TerminalExisting + ); assert_eq!(*planner.calls.lock().unwrap(), 1); assert_eq!(*executor.calls.lock().unwrap(), 1); assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); } +#[tokio::test] +async fn separate_drivers_share_terminal_winner_and_disposition() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); + let first = driver(planner.clone(), executor.clone(), persistence.clone()) + .run(request("shared-terminal", "one"), CancellationToken::new()) + .await + .unwrap(); + let second = driver(planner, executor, persistence.clone()) + .run(request("shared-terminal", "two"), CancellationToken::new()) + .await + .unwrap(); + + assert_eq!( + first.disposition, + SubagentPersistenceDisposition::TerminalInserted + ); + assert_eq!( + second.disposition, + SubagentPersistenceDisposition::TerminalExisting + ); + assert_eq!(first.outcome, second.outcome); + assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn a_later_continuation_replaces_its_consumed_pause_and_owns_new_effects() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::Pause); + let first = driver(planner.clone(), executor.clone(), persistence.clone()) + .run(request("shared-pause", "one"), CancellationToken::new()) + .await + .unwrap(); + let second = driver(planner, executor, persistence) + .run(request("shared-pause", "two"), CancellationToken::new()) + .await + .unwrap(); + + assert_eq!( + first.disposition, + SubagentPersistenceDisposition::PauseCommitted + ); + assert_eq!( + second.disposition, + SubagentPersistenceDisposition::PauseReplaced + ); + assert!(first.should_emit_host_effects()); + assert!(second.should_emit_host_effects()); +} + #[tokio::test] async fn fresh_request_rejects_an_unrelated_or_thread_mismatched_context() { let parent = RunContext::new( @@ -805,11 +955,11 @@ async fn concurrent_same_task_calls_coalesce_to_one_lifecycle() { cancellation.cancel(); assert_eq!( - first.await.unwrap().unwrap().status, + first.await.unwrap().unwrap().outcome.status, SubagentStatus::Cancelled ); assert_eq!( - second.await.unwrap().unwrap().status, + second.await.unwrap().unwrap().outcome.status, SubagentStatus::Cancelled ); assert_eq!(*planner.calls.lock().unwrap(), 1); @@ -856,14 +1006,19 @@ async fn cancelled_follower_returns_without_cancelling_the_leader_or_persisting( .expect("cancelled follower must not wait for the leader") .unwrap() .unwrap(); - assert_eq!(follower_outcome.status, SubagentStatus::Cancelled); + assert_eq!(follower_outcome.outcome.status, SubagentStatus::Cancelled); + assert_eq!( + follower_outcome.disposition, + SubagentPersistenceDisposition::ObserverCancelled + ); + assert!(!follower_outcome.should_emit_host_effects()); assert!(!leader_cancellation.is_cancelled()); assert_eq!(*executor.calls.lock().unwrap(), 1); assert!(persistence.outcomes.lock().unwrap().is_empty()); leader_cancellation.cancel(); assert_eq!( - leader.await.unwrap().unwrap().status, + leader.await.unwrap().unwrap().outcome.status, SubagentStatus::Cancelled ); assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); @@ -906,8 +1061,8 @@ async fn same_task_id_from_distinct_parent_runs_never_shares_lifecycle_state() { .collect::>(); assert_eq!( terminal_keys.len(), - 4, - "load and terminal use each scoped key" + 6, + "initial terminal lookup, pause lookup, and terminal persistence use each scoped key" ); assert!( terminal_keys @@ -940,8 +1095,11 @@ async fn awaiting_input_is_not_cached_and_the_next_call_resumes_to_completion() .await .unwrap(); - assert!(matches!(first.status, SubagentStatus::AwaitingInput(_))); - assert_eq!(second.status, SubagentStatus::Completed); + assert!(matches!( + first.outcome.status, + SubagentStatus::AwaitingInput(_) + )); + assert_eq!(second.outcome.status, SubagentStatus::Completed); assert_eq!(*planner.calls.lock().unwrap(), 2); assert_eq!(*executor.calls.lock().unwrap(), 2); assert_eq!(*planner.seen_resumes.lock().unwrap(), vec![false, true]); @@ -986,7 +1144,7 @@ async fn continuation_keeps_original_key_with_a_fresh_owned_context() { ) .await .unwrap(); - assert_eq!(completed.status, SubagentStatus::Completed); + assert_eq!(completed.outcome.status, SubagentStatus::Completed); assert_eq!(*planner.seen_resumes.lock().unwrap(), vec![false, true]); assert!( persistence @@ -1019,7 +1177,7 @@ async fn driver_replaces_prepared_context_cancellation_with_execution_token() { cancellation.cancel(); let outcome = task.await.unwrap().unwrap(); - assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(outcome.outcome.status, SubagentStatus::Cancelled); assert!(executor.context_cancellations.lock().unwrap()[0].is_cancelled()); } @@ -1064,7 +1222,11 @@ async fn cancellation_before_persistence_commit_records_only_cancelled_terminal( .unwrap() .unwrap(); - assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(outcome.outcome.status, SubagentStatus::Cancelled); + assert_eq!( + outcome.disposition, + SubagentPersistenceDisposition::TerminalInserted + ); assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); assert_eq!( persistence.outcomes.lock().unwrap()[0].status, @@ -1120,7 +1282,7 @@ async fn nested_same_driver_task_uses_child_reservation_and_rolls_usage_up_once( .expect("a nested child task must not wait on a driver-global lock") .unwrap(); - assert_eq!(parent.usage.calls, 8); + assert_eq!(parent.outcome.usage.calls, 8); assert_eq!(*executor.calls.lock().unwrap(), vec!["parent", "child"]); let records = persistence.outcomes.lock().unwrap(); assert_eq!(records.len(), 2); @@ -1145,7 +1307,7 @@ async fn cancellation_before_execution_records_one_truthful_terminal() { .await .unwrap(); - assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(outcome.outcome.status, SubagentStatus::Cancelled); assert_eq!(*planner.calls.lock().unwrap(), 0); assert_eq!(*executor.calls.lock().unwrap(), 0); assert_eq!( @@ -1174,7 +1336,7 @@ async fn cancellation_during_execution_is_truthful_and_terminal_once() { cancellation.cancel(); let outcome = task.await.unwrap().unwrap(); - assert_eq!(outcome.status, SubagentStatus::Cancelled); + assert_eq!(outcome.outcome.status, SubagentStatus::Cancelled); assert_eq!(*executor.calls.lock().unwrap(), 1); assert_eq!( actions.lock().unwrap().last(), @@ -1190,11 +1352,11 @@ async fn cancellation_after_execution_preserves_lossless_result_data() { .await .unwrap(); - assert_eq!(outcome.status, SubagentStatus::Cancelled); - assert_eq!(outcome.output, "result"); - assert_eq!(outcome.history, vec![Message::assistant("result")]); - assert_eq!(outcome.usage.calls, 7); - assert_eq!(outcome.artifacts.len(), 1); + assert_eq!(outcome.outcome.status, SubagentStatus::Cancelled); + assert_eq!(outcome.outcome.output, "result"); + assert_eq!(outcome.outcome.history, vec![Message::assistant("result")]); + assert_eq!(outcome.outcome.usage.calls, 7); + assert_eq!(outcome.outcome.artifacts.len(), 1); } #[tokio::test] diff --git a/crates/tinyagents-orchestration/src/subagent/types.rs b/crates/tinyagents-orchestration/src/subagent/types.rs index d6103e72..c5df0f65 100644 --- a/crates/tinyagents-orchestration/src/subagent/types.rs +++ b/crates/tinyagents-orchestration/src/subagent/types.rs @@ -189,7 +189,7 @@ impl SubagentRequest { } /// A neutral checkpoint offered to a planner for resumption. -#[derive(Clone, Debug, Default, PartialEq)] +#[derive(Clone, Debug, Default, PartialEq, serde::Serialize, serde::Deserialize)] pub struct SubagentResume { /// Lossless model history available to the host planner. pub history: Vec, @@ -228,7 +228,7 @@ pub struct SubagentExecution { /// /// The reference intentionally contains no filesystem path or URL. Hosts own /// artifact authorization and resolution. -#[derive(Clone, Debug, Default, PartialEq, Eq)] +#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct ArtifactReference { /// Stable host artifact identifier. pub id: String, @@ -239,7 +239,7 @@ pub struct ArtifactReference { } /// A neutral suspension point that can later be supplied to a planner. -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] pub struct SubagentPause { /// Why execution needs input or an external host action. pub reason: String, @@ -248,14 +248,14 @@ pub struct SubagentPause { } /// A neutral non-successful but terminal completion. -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct SubagentIncomplete { /// A host-safe explanation for incomplete work. pub reason: String, } /// The visible status of one subagent run. -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] pub enum SubagentStatus { /// The subagent completed normally. Completed, @@ -267,8 +267,75 @@ pub enum SubagentStatus { Cancelled, } -/// Complete neutral result of one subagent execution. +/// The durable lifecycle action that made a run result observable. +/// +/// Hosts use this to attach side effects such as event-bus notifications and +/// progress records to the same successful commit boundary as the outcome. +/// In particular, a terminal result loaded from durable storage must be +/// returned to a caller without publishing a second terminal notification. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SubagentPersistenceDisposition { + /// This invocation committed a resumable pause record. + PauseCommitted, + /// This invocation atomically advanced an existing durable pause after a + /// continuation. It owns the new pause's host effects just as the first + /// pause writer does; observers that lost the compare-and-swap receive + /// [`Self::PauseExisting`] instead. + PauseReplaced, + /// This invocation atomically inserted the terminal record. + TerminalInserted, + /// A pause record was already made visible by another invocation. + PauseExisting, + /// A terminal record was already made visible by another invocation. + TerminalExisting, + /// This caller stopped observing an in-flight lifecycle before it reached + /// a durable boundary. It must not publish terminal or pause effects. + ObserverCancelled, +} + +impl SubagentPersistenceDisposition { + /// Whether this invocation owns host terminal/pause side effects. + pub const fn should_emit_host_effects(self) -> bool { + matches!( + self, + Self::PauseCommitted | Self::PauseReplaced | Self::TerminalInserted + ) + } +} + +/// Result of a neutral subagent lifecycle. +/// +/// The outcome is lossless host-neutral execution data. The disposition says +/// whether this caller won durable persistence, preventing a second driver or +/// a coalesced caller from repeating host-visible lifecycle effects. #[derive(Clone, Debug, PartialEq)] +pub struct SubagentRunResult { + /// The authoritative execution outcome. + pub outcome: SubagentOutcome, + /// The successful persistence boundary for this invocation. + pub disposition: SubagentPersistenceDisposition, +} + +impl SubagentRunResult { + /// Creates a result at a known persistence disposition. + pub const fn new( + outcome: SubagentOutcome, + disposition: SubagentPersistenceDisposition, + ) -> Self { + Self { + outcome, + disposition, + } + } + + /// Whether the caller must publish host effects for this result. + pub const fn should_emit_host_effects(&self) -> bool { + self.disposition.should_emit_host_effects() + } +} + +/// Complete neutral result of one subagent execution. +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] pub struct SubagentOutcome { /// Host-stable task identity. pub task_id: String, @@ -309,8 +376,41 @@ impl SubagentOutcome { pub struct PersistedSubagentPause { /// Durable scoped lifecycle identity. pub key: SubagentTaskKey, - /// The resumable pause state. - pub pause: SubagentPause, + /// The complete paused outcome. Keeping the visible output, history, + /// usage, and artifacts with the suspension means a duplicate driver can + /// return the durable winner rather than its own uncommitted result. + pub outcome: SubagentOutcome, + /// The exact resumable state this invocation consumed, when it is a + /// continuation. Persistence implementations compare this under their + /// scoped-key transaction before replacing a pause. `None` means this + /// attempt creates the first pause for a fresh lifecycle. + pub replaces: Option, +} + +/// Result of atomically persisting a pause transition. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SubagentPausePersistenceDisposition { + /// A fresh lifecycle installed its first pause. + Inserted, + /// A continuation consumed and replaced the exact prior pause. + Replaced, + /// Another lifecycle already owns the current pause. + Existing, + /// A terminal outcome already closed the lifecycle. + TerminalExisting, +} + +/// Result of atomically persisting a terminal transition. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SubagentTerminalPersistenceDisposition { + /// This invocation atomically closed an unpaused lifecycle or consumed the + /// exact pause it had loaded. + Inserted, + /// Another invocation already committed a terminal outcome. + Existing, + /// Another invocation owns a newer or unconsumed pause. The caller must + /// return that pause outcome rather than manufacture a terminal result. + PauseExisting, } /// Typed lifecycle failures. Adapters classify their errors at the seam that From e1b2035e2ddbea471f9bda6b27aa58b11205ac79 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 23:57:00 +0300 Subject: [PATCH 12/18] refactor(todos): remove static agent assignment surface Co-authored-by: Medulla --- Cargo.lock | 2 +- crates/tinyagents-graph/src/todos/README.md | 4 +- .../src/todos/dispatch/README.md | 4 +- .../src/todos/dispatch/select.rs | 12 +----- .../src/todos/dispatch/test.rs | 40 +++-------------- crates/tinyagents-graph/src/todos/store.rs | 4 -- crates/tinyagents-graph/src/todos/tool.rs | 6 +-- crates/tinyagents-graph/src/todos/types.rs | 14 +----- .../tests/e2e_graph_task_dispatch.rs | 43 +++++++------------ docs/modules/graph/todos.md | 4 +- 10 files changed, 32 insertions(+), 101 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 65859c41..cb9ce574 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1594,7 +1594,7 @@ dependencies = [ "chrono", "serde_json", "tempfile", - "thiserror 2.0.20", + "thiserror", "tinyagents-harness", "tinyagents-session", "tinyinference-llm", diff --git a/crates/tinyagents-graph/src/todos/README.md b/crates/tinyagents-graph/src/todos/README.md index 01f5c2df..3fcbe451 100644 --- a/crates/tinyagents-graph/src/todos/README.md +++ b/crates/tinyagents-graph/src/todos/README.md @@ -13,8 +13,8 @@ the concrete work items. - `TaskCardStatus { Todo, AwaitingApproval, Ready, InProgress, Blocked, Done, Rejected }` and `TaskApprovalMode { Required, NotRequired }` (each `as_str`). -- `TaskBoardCard { id, title, status, objective, plan, assigned_agent, - allowed_tools, approval_mode, acceptance_criteria, evidence, notes, blocker, +- `TaskBoardCard { id, title, status, objective, plan, allowed_tools, + approval_mode, acceptance_criteria, evidence, notes, blocker, session_thread_id, source_metadata, order, updated_at }` (serde `camelCase`). - `TaskBoard { thread_id, cards, updated_at }`. - `CardPatch` — optional `add`/`edit` fields; `approval_mode` is doubly-optional diff --git a/crates/tinyagents-graph/src/todos/dispatch/README.md b/crates/tinyagents-graph/src/todos/dispatch/README.md index de420b4f..06f1aa9c 100644 --- a/crates/tinyagents-graph/src/todos/dispatch/README.md +++ b/crates/tinyagents-graph/src/todos/dispatch/README.md @@ -13,12 +13,10 @@ loading, personality profiles, and event bus). Pure functions over a board snapshot — no store, so they are trivially testable and can be applied to cards a host already holds. -- `pick_next_card(cards, agent_assigned_only)` — the highest-urgency +- `pick_next_card(cards)` — the highest-urgency dispatchable card (`Todo` or approved `Ready`). Urgency comes from `source_metadata.urgency` (`card_urgency`, default `0.0`); ties break toward the lower board `order`, so equal-priority work runs in planned order. - `agent_assigned_only` restricts the pick to cards with an `assigned_agent`, - which is how a host keeps an autonomous sweep off a person's own todos. - `has_card_in_progress(cards)` — the board already has a card being worked, so there is nothing to claim this tick. - `requires_plan_approval(global_required, approval_mode)` — the card's own diff --git a/crates/tinyagents-graph/src/todos/dispatch/select.rs b/crates/tinyagents-graph/src/todos/dispatch/select.rs index 0c8342ca..a47cb6b0 100644 --- a/crates/tinyagents-graph/src/todos/dispatch/select.rs +++ b/crates/tinyagents-graph/src/todos/dispatch/select.rs @@ -36,14 +36,10 @@ pub fn has_card_in_progress(cards: &[TaskBoardCard]) -> bool { /// urgency break toward the lower board `order`, so equal-priority work runs in /// the order it was planned. /// -/// `agent_assigned_only` restricts the pick to cards with an `assigned_agent`. -/// A host uses it for boards that mix human-authored and agent-authored cards, -/// so an autonomous sweep never picks up a card a person wrote for themselves. -pub fn pick_next_card(cards: &[TaskBoardCard], agent_assigned_only: bool) -> Option { +pub fn pick_next_card(cards: &[TaskBoardCard]) -> Option { cards .iter() .filter(|card| matches!(card.status, TaskCardStatus::Todo | TaskCardStatus::Ready)) - .filter(|card| !agent_assigned_only || is_agent_assigned(card)) .max_by(|a, b| { card_urgency(a) .partial_cmp(&card_urgency(b)) @@ -54,12 +50,6 @@ pub fn pick_next_card(cards: &[TaskBoardCard], agent_assigned_only: bool) -> Opt .cloned() } -fn is_agent_assigned(card: &TaskBoardCard) -> bool { - card.assigned_agent - .as_deref() - .is_some_and(|agent| !agent.trim().is_empty()) -} - /// Whether a card must be parked at /// [`AwaitingApproval`](TaskCardStatus::AwaitingApproval) before it runs. /// diff --git a/crates/tinyagents-graph/src/todos/dispatch/test.rs b/crates/tinyagents-graph/src/todos/dispatch/test.rs index f80c6d20..ff95248a 100644 --- a/crates/tinyagents-graph/src/todos/dispatch/test.rs +++ b/crates/tinyagents-graph/src/todos/dispatch/test.rs @@ -26,11 +26,6 @@ fn with_urgency(mut card: TaskBoardCard, urgency: f64) -> TaskBoardCard { card } -fn assigned(mut card: TaskBoardCard, agent: &str) -> TaskBoardCard { - card.assigned_agent = Some(agent.to_string()); - card -} - // ── Selection ─────────────────────────────────────────────────────────────── #[test] @@ -42,11 +37,11 @@ fn only_todo_and_ready_cards_are_dispatchable() { card("rejected", TaskCardStatus::Rejected, 3), card("running", TaskCardStatus::InProgress, 4), ]; - assert!(pick_next_card(&cards, false).is_none()); + assert!(pick_next_card(&cards).is_none()); let mut cards = cards; cards.push(card("ready", TaskCardStatus::Ready, 5)); - assert_eq!(pick_next_card(&cards, false).unwrap().id, "ready"); + assert_eq!(pick_next_card(&cards).unwrap().id, "ready"); } #[test] @@ -56,7 +51,7 @@ fn the_most_urgent_card_wins() { with_urgency(card("high", TaskCardStatus::Todo, 1), 0.9), card("none", TaskCardStatus::Todo, 2), ]; - assert_eq!(pick_next_card(&cards, false).unwrap().id, "high"); + assert_eq!(pick_next_card(&cards).unwrap().id, "high"); } #[test] @@ -65,42 +60,19 @@ fn equal_urgency_runs_in_board_order() { with_urgency(card("second", TaskCardStatus::Todo, 5), 0.5), with_urgency(card("first", TaskCardStatus::Todo, 1), 0.5), ]; - assert_eq!(pick_next_card(&cards, false).unwrap().id, "first"); + assert_eq!(pick_next_card(&cards).unwrap().id, "first"); // Unscored cards tie at 0.0 and follow the same rule. let cards = vec![ card("later", TaskCardStatus::Todo, 9), card("earlier", TaskCardStatus::Todo, 2), ]; - assert_eq!(pick_next_card(&cards, false).unwrap().id, "earlier"); -} - -#[test] -fn agent_assigned_only_skips_human_authored_cards() { - let cards = vec![ - with_urgency(card("mine", TaskCardStatus::Todo, 0), 0.9), - with_urgency( - assigned(card("agents", TaskCardStatus::Todo, 1), "researcher"), - 0.1, - ), - ]; - - // Unfiltered, urgency wins; filtered, the unassigned card is invisible even - // though it is the more urgent one. - assert_eq!(pick_next_card(&cards, false).unwrap().id, "mine"); - assert_eq!(pick_next_card(&cards, true).unwrap().id, "agents"); -} - -#[test] -fn a_blank_assignee_does_not_count_as_assigned() { - let cards = vec![assigned(card("blank", TaskCardStatus::Todo, 0), " ")]; - assert!(pick_next_card(&cards, true).is_none()); - assert_eq!(pick_next_card(&cards, false).unwrap().id, "blank"); + assert_eq!(pick_next_card(&cards).unwrap().id, "earlier"); } #[test] fn an_empty_board_has_nothing_to_dispatch() { - assert!(pick_next_card(&[], false).is_none()); + assert!(pick_next_card(&[]).is_none()); assert!(!has_card_in_progress(&[])); } diff --git a/crates/tinyagents-graph/src/todos/store.rs b/crates/tinyagents-graph/src/todos/store.rs index 4860548b..1dc24874 100644 --- a/crates/tinyagents-graph/src/todos/store.rs +++ b/crates/tinyagents-graph/src/todos/store.rs @@ -182,7 +182,6 @@ pub async fn add( status: patch.status.unwrap_or(TaskCardStatus::Todo), objective: patch.objective.and_then(non_empty), plan: patch.plan.unwrap_or_default(), - assigned_agent: patch.assigned_agent.and_then(non_empty), allowed_tools: patch.allowed_tools.unwrap_or_default(), approval_mode: patch.approval_mode.flatten(), acceptance_criteria: patch.acceptance_criteria.unwrap_or_default(), @@ -232,9 +231,6 @@ pub async fn edit( if let Some(plan) = patch.plan { card.plan = plan; } - if let Some(assigned_agent) = patch.assigned_agent { - card.assigned_agent = non_empty(assigned_agent); - } if let Some(allowed_tools) = patch.allowed_tools { card.allowed_tools = allowed_tools; } diff --git a/crates/tinyagents-graph/src/todos/tool.rs b/crates/tinyagents-graph/src/todos/tool.rs index 48e21095..796437f4 100644 --- a/crates/tinyagents-graph/src/todos/tool.rs +++ b/crates/tinyagents-graph/src/todos/tool.rs @@ -27,8 +27,8 @@ const TODO_DESCRIPTION: &str = "Maintain a visible plan for THIS thread: an orde the start, `add` one card per step; keep exactly ONE card `in_progress` at a time; mark a card \ `done` the moment it finishes; if a step is blocked, set it `blocked` with a `blocker`. `list` \ to re-read the plan. The board is bound automatically to the current thread — do not pass a \ - thread id. Dispatch via `op`: `add` (content, status?, objective?, plan?, assignedAgent?, \ - allowedTools?, approvalMode?, acceptanceCriteria?, evidence?, notes?, blocker?), `edit` (id, \ + thread id. Dispatch via `op`: `add` (content, status?, objective?, plan?, allowedTools?, \ + approvalMode?, acceptanceCriteria?, evidence?, notes?, blocker?), `edit` (id, \ same optional fields), `update_status` (id, status), `decide_plan` (id, approve), \ `revise_plan`, `remove` (id), `replace` (cards), `clear`, or `list`. Returns the updated cards \ plus a markdown rendering."; @@ -219,7 +219,6 @@ fn patch_from_args(args: &Value) -> std::result::Result { status, objective: optional_string(args, "objective"), plan: optional_string_array(args, "plan")?, - assigned_agent: optional_string(args, "assignedAgent"), allowed_tools: optional_string_array(args, "allowedTools")?, approval_mode, acceptance_criteria: optional_string_array(args, "acceptanceCriteria")?, @@ -250,7 +249,6 @@ fn parameters_schema() -> Value { "blocker": { "type": "string" }, "objective": { "type": "string", "description": "Desired outcome for this task." }, "plan": { "type": "array", "items": { "type": "string" }, "description": "Ordered execution steps." }, - "assignedAgent": { "type": "string" }, "allowedTools": { "type": "array", "items": { "type": "string" } }, "approvalMode": { "type": ["string", "null"], "enum": ["required", "not_required", null] }, "acceptanceCriteria": { "type": "array", "items": { "type": "string" } }, diff --git a/crates/tinyagents-graph/src/todos/types.rs b/crates/tinyagents-graph/src/todos/types.rs index 05dac1a4..c3cd3282 100644 --- a/crates/tinyagents-graph/src/todos/types.rs +++ b/crates/tinyagents-graph/src/todos/types.rs @@ -82,10 +82,7 @@ pub struct TaskBoardCard { /// Ordered plan steps. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub plan: Vec, - /// The agent assigned to run this card, if any. - #[serde(default, skip_serializing_if = "Option::is_none")] - pub assigned_agent: Option, - /// Tools the assigned agent is allowed to use. + /// Tools an agent working this card is allowed to use. #[serde(default, skip_serializing_if = "Vec::is_empty")] pub allowed_tools: Vec, /// Plan-approval mode, if the card is gated. @@ -127,7 +124,6 @@ impl TaskBoardCard { status: TaskCardStatus::Todo, objective: None, plan: Vec::new(), - assigned_agent: None, allowed_tools: Vec::new(), approval_mode: None, acceptance_criteria: Vec::new(), @@ -191,8 +187,6 @@ pub struct CardPatch { pub objective: Option, /// New plan steps. pub plan: Option>, - /// New assigned agent (empty clears). - pub assigned_agent: Option, /// New allowed-tools list. pub allowed_tools: Option>, /// New approval mode (`Some(None)` clears). @@ -254,11 +248,6 @@ pub fn render_markdown(cards: &[TaskBoardCard]) -> String { out.push_str(objective); out.push('\n'); } - if let Some(agent) = card.assigned_agent.as_deref() { - out.push_str(" - agent: "); - out.push_str(agent); - out.push('\n'); - } if !card.allowed_tools.is_empty() { out.push_str(" - tools: "); out.push_str(&card.allowed_tools.join(", ")); @@ -326,7 +315,6 @@ pub fn normalise_board(board: &mut TaskBoard) { } card.notes = trim_opt(card.notes.take()); card.objective = trim_opt(card.objective.take()); - card.assigned_agent = trim_opt(card.assigned_agent.take()); trim_string_vec(&mut card.plan); trim_string_vec(&mut card.allowed_tools); trim_string_vec(&mut card.acceptance_criteria); diff --git a/crates/tinyagents-integration-tests/tests/e2e_graph_task_dispatch.rs b/crates/tinyagents-integration-tests/tests/e2e_graph_task_dispatch.rs index 8f3cf67a..8eda5202 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_graph_task_dispatch.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_graph_task_dispatch.rs @@ -83,9 +83,8 @@ async fn add_card(store: &Arc, title: &str, patch: CardPatch) -> Stri .clone() } -fn agent_card(agent: &str, urgency: f64) -> CardPatch { +fn task_card(urgency: f64) -> CardPatch { CardPatch { - assigned_agent: Some(agent.to_string()), source_metadata: Some(json!({ "urgency": urgency })), ..CardPatch::default() } @@ -105,7 +104,7 @@ enum Tick { /// One sweep of the board, wired from the crate's dispatch policy: reclaim /// what has gone stale, refuse to double-book a busy board, pick the most -/// urgent agent-assigned card, and either park it for approval or claim it. +/// urgent card, and either park it for approval or claim it. async fn tick(store: &Arc, approval_required: bool) -> Tick { task_run_store::reclaim_stale(store, THREAD, &RunLimits::default()) .await @@ -115,7 +114,7 @@ async fn tick(store: &Arc, approval_required: bool) -> Tick { if has_card_in_progress(&board.cards) { return Tick::Idle; } - let Some(card) = pick_next_card(&board.cards, true) else { + let Some(card) = pick_next_card(&board.cards) else { return Tick::Idle; }; @@ -205,24 +204,15 @@ async fn run_card(store: &Arc, card: &TaskBoardCard, run_id: &str) -> } #[tokio::test] -async fn the_dispatcher_runs_the_most_urgent_agent_card_and_leaves_human_work_alone() { +async fn the_dispatcher_runs_the_most_urgent_card() { let store: Arc = Arc::new(InMemoryStore::default()); - // A human's own todo (unassigned, most urgent), and two agent-assigned - // cards. Only the agent cards are the dispatcher's to run. - let mine = add_card( - &store, - "call the dentist", - CardPatch { - source_metadata: Some(json!({ "urgency": 0.99 })), - ..CardPatch::default() - }, - ) - .await; - let low = add_card(&store, "tidy the changelog", agent_card("scribe", 0.1)).await; - let high = add_card(&store, "apply the migration", agent_card("dba", 0.8)).await; + // Every card is eligible for an autonomous sweep; urgency determines the + // first card selected. + let low = add_card(&store, "tidy the changelog", task_card(0.1)).await; + let high = add_card(&store, "apply the migration", task_card(0.8)).await; - // Tick 1: the urgent agent card is claimed; the human's card is not touched. + // Tick 1: the highest-urgency card is claimed. let Tick::Dispatched { card_id, run_id } = tick(&store, false).await else { panic!("expected a dispatch"); }; @@ -246,9 +236,9 @@ async fn the_dispatcher_runs_the_most_urgent_agent_card_and_leaves_human_work_al "the run is told which card it owns" ); - // Tick 2: with the first card done, the remaining agent card goes next. + // Tick 2: with the first card done, the remaining card goes next. let Tick::Dispatched { card_id, .. } = tick(&store, false).await else { - panic!("expected the second agent card"); + panic!("expected the second card"); }; assert_eq!(card_id, low); @@ -258,7 +248,6 @@ async fn the_dispatcher_runs_the_most_urgent_agent_card_and_leaves_human_work_al .iter() .map(|card| (card.id.clone(), card.status)) .collect(); - assert!(statuses.contains(&(mine.clone(), TaskCardStatus::Todo))); assert!(statuses.contains(&(high, TaskCardStatus::Done))); assert!(statuses.contains(&(low, TaskCardStatus::InProgress))); @@ -274,7 +263,7 @@ async fn the_dispatcher_runs_the_most_urgent_agent_card_and_leaves_human_work_al #[tokio::test] async fn a_plan_awaiting_approval_never_runs_until_it_is_approved() { let store: Arc = Arc::new(InMemoryStore::default()); - let card_id = add_card(&store, "delete the old bucket", agent_card("ops", 0.5)).await; + let card_id = add_card(&store, "delete the old bucket", task_card(0.5)).await; // With approval on, the tick parks the card instead of claiming it, and // keeps parking nothing afterwards: an awaiting card is not dispatchable. @@ -309,7 +298,7 @@ async fn a_card_stamped_required_is_parked_even_with_the_global_gate_off() { "email the customer", CardPatch { approval_mode: Some(Some(TaskApprovalMode::Required)), - ..agent_card("support", 0.5) + ..task_card(0.5) }, ) .await; @@ -322,7 +311,7 @@ async fn a_card_stamped_required_is_parked_even_with_the_global_gate_off() { #[tokio::test] async fn a_cancelled_run_leaves_its_card_blocked_rather_than_stranded() { let store: Arc = Arc::new(InMemoryStore::default()); - let card_id = add_card(&store, "long crawl", agent_card("crawler", 0.5)).await; + let card_id = add_card(&store, "long crawl", task_card(0.5)).await; let registry: ActiveRunRegistry = ActiveRunRegistry::new(); let Tick::Dispatched { run_id, .. } = tick(&store, false).await else { @@ -398,7 +387,7 @@ async fn a_cancelled_run_leaves_its_card_blocked_rather_than_stranded() { #[tokio::test] async fn an_abandoned_run_is_reclaimed_by_the_next_tick() { let store: Arc = Arc::new(InMemoryStore::default()); - let card_id = add_card(&store, "flaky job", agent_card("runner", 0.5)).await; + let card_id = add_card(&store, "flaky job", task_card(0.5)).await; let Tick::Dispatched { run_id, .. } = tick(&store, false).await else { panic!("expected a dispatch"); @@ -466,7 +455,7 @@ async fn an_idle_board_backs_the_sweep_off_and_fresh_work_resets_it() { ); // Work arrives, the tick dispatches, and the cadence snaps back. - add_card(&store, "new work", agent_card("worker", 0.5)).await; + add_card(&store, "new work", task_card(0.5)).await; assert!(matches!(tick(&store, false).await, Tick::Dispatched { .. })); idle_ticks = 0; assert_eq!(cadence.next_delay(idle_ticks), cadence.base); diff --git a/docs/modules/graph/todos.md b/docs/modules/graph/todos.md index 3a5c3a66..e662e390 100644 --- a/docs/modules/graph/todos.md +++ b/docs/modules/graph/todos.md @@ -10,8 +10,8 @@ surface; this spec captures the design contract. ## Model -- `TaskBoardCard { id, title, status, objective, plan, assigned_agent, - allowed_tools, approval_mode, acceptance_criteria, evidence, notes, blocker, +- `TaskBoardCard { id, title, status, objective, plan, allowed_tools, + approval_mode, acceptance_criteria, evidence, notes, blocker, session_thread_id, source_metadata, order, updated_at }`. - `TaskCardStatus`: `Todo`, `AwaitingApproval`, `Ready`, `InProgress`, `Blocked`, `Done`, `Rejected`. `TaskApprovalMode`: `Required`, `NotRequired`. From 9e5a1ed77372c255dfa0dc936aacada8f8e21b3e Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 23:57:00 +0300 Subject: [PATCH 13/18] fix(session): compile transcript marker helper only in tests Co-authored-by: Medulla --- crates/tinyagents-session/src/transcript/migration.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/crates/tinyagents-session/src/transcript/migration.rs b/crates/tinyagents-session/src/transcript/migration.rs index 9cfb9754..f6f405dd 100644 --- a/crates/tinyagents-session/src/transcript/migration.rs +++ b/crates/tinyagents-session/src/transcript/migration.rs @@ -8,7 +8,7 @@ use anyhow::{Context, Result}; use std::fs; -use std::path::{Path, PathBuf}; +use std::path::Path; /// Marker file that signals the v1 session-layout migration has run for a /// workspace. It lives under `state/migrations/` to keep the workspace root @@ -329,7 +329,8 @@ fn ddmmyyyy_to_yyyy_mm_dd(name: &str) -> Option { } /// Return the migration marker path for `workspace_dir`. -fn marker_path_for(workspace_dir: &Path) -> PathBuf { +#[cfg(test)] +fn marker_path_for(workspace_dir: &Path) -> std::path::PathBuf { workspace_dir.join(MIGRATION_MARKER) } From b1d680476c3f053a2f5dbe7bbf202d2383c1c00e Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sat, 19 Sep 2026 23:58:29 +0300 Subject: [PATCH 14/18] fix(runtime): align tinytools dependency with vendored release Co-authored-by: Medulla --- crates/tinyagents-runtime/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/tinyagents-runtime/Cargo.toml b/crates/tinyagents-runtime/Cargo.toml index 8fbf83b8..d4cbf889 100644 --- a/crates/tinyagents-runtime/Cargo.toml +++ b/crates/tinyagents-runtime/Cargo.toml @@ -14,7 +14,7 @@ thiserror = "2" tinyagents-harness = { path = "../tinyagents-harness", version = "2.1.2" } tinyagents-session = { path = "../tinyagents-session", version = "2.1.2" } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.3.0" } -tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.2.0" } +tinytools = { path = "../../vendor/tinytools/crates/tinytools", version = "0.3.0" } tokio = { version = "1", features = ["macros", "rt", "sync"] } [dev-dependencies] From fc935e60cd37055e8f703603aa806bdc32f303f2 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 20 Sep 2026 00:18:46 +0300 Subject: [PATCH 15/18] fix: close subagent lifecycle review races Co-authored-by: Medulla --- .../src/subagent/driver.rs | 100 +++++++++++++----- .../src/subagent/test.rs | 36 +++++++ crates/tinyagents-runtime/src/session.rs | 9 +- crates/tinyagents-runtime/src/test.rs | 53 ++++++++++ 4 files changed, 170 insertions(+), 28 deletions(-) diff --git a/crates/tinyagents-orchestration/src/subagent/driver.rs b/crates/tinyagents-orchestration/src/subagent/driver.rs index a07830a8..d6740dd5 100644 --- a/crates/tinyagents-orchestration/src/subagent/driver.rs +++ b/crates/tinyagents-orchestration/src/subagent/driver.rs @@ -33,7 +33,7 @@ pub struct SubagentDriver { executor: Arc>, persistence: Arc, terminal_outcomes: Mutex>, - in_flight: Mutex>>, + in_flight: Arc>>>, } /// Result shared by callers that arrived while the same task was executing. @@ -79,6 +79,55 @@ impl InFlight { } } +/// Owns the leader's in-flight reservation until its result is published. +/// +/// Dropping a caller future must also release its reservation: otherwise a +/// later caller becomes a follower of work that will never publish a result. +struct LeaderReservation { + in_flight: Arc>>>, + task_key: SubagentTaskKey, + entry: Arc, + finished: bool, +} + +impl LeaderReservation { + async fn finish(mut self, result: Result) { + self.entry.complete(result).await; + let mut in_flight = self.in_flight.lock().await; + if in_flight + .get(&self.task_key) + .is_some_and(|current| Arc::ptr_eq(current, &self.entry)) + { + in_flight.remove(&self.task_key); + } + self.finished = true; + } +} + +impl Drop for LeaderReservation { + fn drop(&mut self) { + if self.finished { + return; + } + let in_flight = self.in_flight.clone(); + let task_key = self.task_key.clone(); + let entry = self.entry.clone(); + // `run` is async, so a current Tokio runtime is always available while + // this guard can be dropped. Publish a typed result before removal so + // existing followers cannot wait indefinitely. + tokio::spawn(async move { + entry.complete(Err(SubagentError::Cancelled)).await; + let mut in_flight = in_flight.lock().await; + if in_flight + .get(&task_key) + .is_some_and(|current| Arc::ptr_eq(current, &entry)) + { + in_flight.remove(&task_key); + } + }); + } +} + impl SubagentDriver { /// Validates host capability availability before exposing a runnable driver. pub fn new(capabilities: SubagentCapabilities) -> Result { @@ -93,7 +142,7 @@ impl SubagentDriver { .persistence .ok_or(SubagentError::MissingCapability("persistence"))?, terminal_outcomes: Mutex::new(HashMap::new()), - in_flight: Mutex::new(HashMap::new()), + in_flight: Arc::new(Mutex::new(HashMap::new())), }) } @@ -158,18 +207,22 @@ impl SubagentDriver { ) }); } + let reservation = LeaderReservation { + in_flight: self.in_flight.clone(), + task_key: task_key.clone(), + entry: entry.clone(), + finished: false, + }; // A preceding caller may have committed a terminal result between the // initial cache check and this reservation. Do not reopen that task // after its in-flight entry has been removed. if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { - entry - .complete(Ok(SubagentRunResult::new( + reservation + .finish(Ok(SubagentRunResult::new( outcome.clone(), SubagentPersistenceDisposition::TerminalExisting, ))) .await; - let mut in_flight = self.in_flight.lock().await; - in_flight.remove(&task_key); return Ok(SubagentRunResult::new( outcome, SubagentPersistenceDisposition::TerminalExisting, @@ -179,14 +232,7 @@ impl SubagentDriver { let result = self .run_reserved(request, task_key.clone(), cancellation) .await; - entry.complete(result.clone()).await; - let mut in_flight = self.in_flight.lock().await; - if in_flight - .get(&task_key) - .is_some_and(|current| Arc::ptr_eq(current, &entry)) - { - in_flight.remove(&task_key); - } + reservation.finish(result.clone()).await; result } @@ -413,16 +459,22 @@ impl SubagentDriver { SubagentPersistenceDisposition::TerminalExisting, )); } - let paused = self - .persistence - .load_pause(&task_key) - .await? - .ok_or_else(|| { - SubagentError::Persistence( - "pause compare-and-swap lost without a durable pause or terminal outcome" - .into(), - ) - })?; + let Some(paused) = self.persistence.load_pause(&task_key).await? else { + // A different driver can consume the pause into a terminal between + // the first terminal read and this pause read. Recheck the + // authoritative terminal before classifying that interleaving as a + // persistence failure. + if let Some(terminal) = self.persistence.load_terminal(&task_key).await? { + self.cache_terminal(task_key, &terminal).await; + return Ok(SubagentRunResult::new( + terminal, + SubagentPersistenceDisposition::TerminalExisting, + )); + } + return Err(SubagentError::Persistence( + "pause compare-and-swap lost without a durable pause or terminal outcome".into(), + )); + }; if !matches!(paused.status, SubagentStatus::AwaitingInput(_)) { return Err(SubagentError::Persistence( "durable pause record did not contain an awaiting-input outcome".into(), diff --git a/crates/tinyagents-orchestration/src/subagent/test.rs b/crates/tinyagents-orchestration/src/subagent/test.rs index 3b26dd5f..95ce362a 100644 --- a/crates/tinyagents-orchestration/src/subagent/test.rs +++ b/crates/tinyagents-orchestration/src/subagent/test.rs @@ -1024,6 +1024,42 @@ async fn cancelled_follower_returns_without_cancelling_the_leader_or_persisting( assert_eq!(persistence.outcomes.lock().unwrap().len(), 1); } +#[tokio::test] +async fn dropped_leader_releases_its_in_flight_reservation() { + let (planner, executor, persistence, _) = fakes(ExecutorMode::WaitForCancellation); + let driver = Arc::new(driver(planner.clone(), executor.clone(), persistence)); + let (started_at_execution, started) = tokio::sync::oneshot::channel(); + *executor.started.lock().unwrap() = Some(started_at_execution); + + let leader = tokio::spawn({ + let driver = driver.clone(); + async move { + driver + .run( + request("dropped-leader", "leader"), + CancellationToken::new(), + ) + .await + } + }); + started.await.unwrap(); + leader.abort(); + let _ = leader.await; + + let cancellation = CancellationToken::new(); + cancellation.cancel(); + let result = tokio::time::timeout( + Duration::from_secs(1), + driver.run(request("dropped-leader", "replacement"), cancellation), + ) + .await + .expect("replacement leader must not wait on an abandoned reservation") + .unwrap(); + + assert_eq!(result.outcome.status, SubagentStatus::Cancelled); + assert_eq!(*planner.calls.lock().unwrap(), 1); +} + #[tokio::test] async fn same_task_id_from_distinct_parent_runs_never_shares_lifecycle_state() { let (planner, executor, persistence, _) = fakes(ExecutorMode::Completed); diff --git a/crates/tinyagents-runtime/src/session.rs b/crates/tinyagents-runtime/src/session.rs index 6572a77c..73e5c131 100644 --- a/crates/tinyagents-runtime/src/session.rs +++ b/crates/tinyagents-runtime/src/session.rs @@ -107,10 +107,11 @@ impl Session { ResumeMode::LatestForAgent => target .locator .latest_for_agent(target.resume_agent.as_deref().unwrap_or(&target.stem)), - ResumeMode::Thread => options - .thread_id - .as_deref() - .and_then(|thread| target.locator.root_for_thread(thread)), + ResumeMode::Thread => options.thread_id.as_deref().and_then(|thread| { + target + .locator + .root_for_thread_scoped(thread, target.meta.agent_id.as_deref()) + }), }; let Some(read) = read else { return Ok(SessionResume { diff --git a/crates/tinyagents-runtime/src/test.rs b/crates/tinyagents-runtime/src/test.rs index c1875f64..fe793a81 100644 --- a/crates/tinyagents-runtime/src/test.rs +++ b/crates/tinyagents-runtime/src/test.rs @@ -166,6 +166,7 @@ impl TranscriptHistory for MemoryHistory { struct Locator { history: Arc, latest_agents: Mutex>, + scoped_threads: Mutex)>>, opened_stems: Mutex>, } @@ -177,6 +178,17 @@ impl TranscriptLocator for Locator { fn root_for_thread(&self, _: &str) -> Option> { Some(self.history.clone()) } + fn root_for_thread_scoped( + &self, + thread: &str, + agent_id: Option<&str>, + ) -> Option> { + self.scoped_threads + .lock() + .unwrap() + .push((thread.into(), agent_id.map(str::to_owned))); + Some(self.history.clone()) + } fn open_stem( &self, stem: &str, @@ -202,6 +214,7 @@ fn locator(session: Option) -> (Arc, Arc, Vec); struct ResumedStateHook { From 1125ee2df285a268baaed8ef2750a1e9c5318389 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 20 Sep 2026 00:19:42 +0300 Subject: [PATCH 16/18] fix: make dropped reservation cleanup runtime-independent Co-authored-by: Medulla --- .../src/subagent/driver.rs | 66 ++++++++++--------- 1 file changed, 35 insertions(+), 31 deletions(-) diff --git a/crates/tinyagents-orchestration/src/subagent/driver.rs b/crates/tinyagents-orchestration/src/subagent/driver.rs index d6740dd5..69c16ead 100644 --- a/crates/tinyagents-orchestration/src/subagent/driver.rs +++ b/crates/tinyagents-orchestration/src/subagent/driver.rs @@ -1,6 +1,10 @@ -use std::{collections::HashMap, future::Future, sync::Arc}; +use std::{ + collections::HashMap, + future::Future, + sync::{Arc, Mutex}, +}; -use tokio::sync::{Mutex, Notify}; +use tokio::sync::{Mutex as AsyncMutex, Notify}; use super::{ PersistedSubagentPause, SubagentError, SubagentExecution, SubagentExecutor, SubagentOutcome, @@ -32,7 +36,7 @@ pub struct SubagentDriver { planner: Arc>, executor: Arc>, persistence: Arc, - terminal_outcomes: Mutex>, + terminal_outcomes: AsyncMutex>, in_flight: Arc>>>, } @@ -62,7 +66,12 @@ impl InFlight { let notified = self.notify.notified(); tokio::pin!(notified); notified.as_mut().enable(); - if let Some(result) = self.result.lock().await.clone() { + if let Some(result) = self + .result + .lock() + .expect("in-flight result poisoned") + .clone() + { return result; } tokio::select! { @@ -73,8 +82,8 @@ impl InFlight { } } - async fn complete(&self, result: Result) { - *self.result.lock().await = Some(result); + fn complete(&self, result: Result) { + *self.result.lock().expect("in-flight result poisoned") = Some(result); self.notify.notify_waiters(); } } @@ -91,9 +100,9 @@ struct LeaderReservation { } impl LeaderReservation { - async fn finish(mut self, result: Result) { - self.entry.complete(result).await; - let mut in_flight = self.in_flight.lock().await; + fn finish(mut self, result: Result) { + self.entry.complete(result); + let mut in_flight = self.in_flight.lock().expect("in-flight map poisoned"); if in_flight .get(&self.task_key) .is_some_and(|current| Arc::ptr_eq(current, &self.entry)) @@ -112,19 +121,16 @@ impl Drop for LeaderReservation { let in_flight = self.in_flight.clone(); let task_key = self.task_key.clone(); let entry = self.entry.clone(); - // `run` is async, so a current Tokio runtime is always available while - // this guard can be dropped. Publish a typed result before removal so - // existing followers cannot wait indefinitely. - tokio::spawn(async move { - entry.complete(Err(SubagentError::Cancelled)).await; - let mut in_flight = in_flight.lock().await; - if in_flight - .get(&task_key) - .is_some_and(|current| Arc::ptr_eq(current, &entry)) - { - in_flight.remove(&task_key); - } - }); + // Publish a typed result before removal so existing followers cannot + // wait indefinitely, even if the future is dropped outside a runtime. + entry.complete(Err(SubagentError::Cancelled)); + let mut in_flight = in_flight.lock().expect("in-flight map poisoned"); + if in_flight + .get(&task_key) + .is_some_and(|current| Arc::ptr_eq(current, &entry)) + { + in_flight.remove(&task_key); + } } } @@ -141,7 +147,7 @@ impl SubagentDriver { persistence: capabilities .persistence .ok_or(SubagentError::MissingCapability("persistence"))?, - terminal_outcomes: Mutex::new(HashMap::new()), + terminal_outcomes: AsyncMutex::new(HashMap::new()), in_flight: Arc::new(Mutex::new(HashMap::new())), }) } @@ -173,7 +179,7 @@ impl SubagentDriver { } let (entry, is_leader) = { - let mut in_flight = self.in_flight.lock().await; + let mut in_flight = self.in_flight.lock().expect("in-flight map poisoned"); match in_flight.get(&task_key) { Some(entry) => (entry.clone(), false), None => { @@ -217,12 +223,10 @@ impl SubagentDriver { // initial cache check and this reservation. Do not reopen that task // after its in-flight entry has been removed. if let Some(outcome) = self.terminal_outcomes.lock().await.get(&task_key).cloned() { - reservation - .finish(Ok(SubagentRunResult::new( - outcome.clone(), - SubagentPersistenceDisposition::TerminalExisting, - ))) - .await; + reservation.finish(Ok(SubagentRunResult::new( + outcome.clone(), + SubagentPersistenceDisposition::TerminalExisting, + ))); return Ok(SubagentRunResult::new( outcome, SubagentPersistenceDisposition::TerminalExisting, @@ -232,7 +236,7 @@ impl SubagentDriver { let result = self .run_reserved(request, task_key.clone(), cancellation) .await; - reservation.finish(result.clone()).await; + reservation.finish(result.clone()); result } From 746abc66c3606a61c99be3c6cb4c28f624ca25c6 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 20 Sep 2026 00:32:35 +0300 Subject: [PATCH 17/18] fix: address runtime lifecycle review findings Co-authored-by: Medulla --- .../src/subagent/persistence.rs | 6 ++--- crates/tinyagents-runtime/src/driver.rs | 13 +++++++---- crates/tinyagents-runtime/src/session.rs | 23 ++++++++++--------- crates/tinyagents-runtime/src/test.rs | 2 +- 4 files changed, 25 insertions(+), 19 deletions(-) diff --git a/crates/tinyagents-orchestration/src/subagent/persistence.rs b/crates/tinyagents-orchestration/src/subagent/persistence.rs index 7957d569..29172464 100644 --- a/crates/tinyagents-orchestration/src/subagent/persistence.rs +++ b/crates/tinyagents-orchestration/src/subagent/persistence.rs @@ -13,9 +13,9 @@ use super::{ /// never a global lifecycle identity. The driver additionally suppresses /// duplicate records from repeated calls made through the same driver instance. A persistence /// future's successful return is its commit boundary: implementations must not -/// make a write visible and then await again before returning `Ok(true)`. The -/// driver races that boundary with cancellation and, when cancellation wins, -/// records one truthful `Cancelled` terminal outcome instead. +/// make a write visible and then await again before returning their +/// disposition. The driver races that boundary with cancellation and, when +/// cancellation wins, records one truthful `Cancelled` terminal outcome instead. #[async_trait] pub trait SubagentPersistence: Send + Sync { /// Returns a terminal outcome committed by another driver/process, if any. diff --git a/crates/tinyagents-runtime/src/driver.rs b/crates/tinyagents-runtime/src/driver.rs index a85a9a26..8cc75b2e 100644 --- a/crates/tinyagents-runtime/src/driver.rs +++ b/crates/tinyagents-runtime/src/driver.rs @@ -84,6 +84,7 @@ impl SessionDriver partial: None, }); } + let input_len = request.history.len(); let partial = if request.stream { self.harness .invoke_streaming_in_context_collecting_partial( @@ -101,10 +102,14 @@ impl SessionDriver ) .await }; - let output = - partial.run.messages.iter().rev().find_map(|message| { - matches!(message, Message::Assistant(_)).then(|| message.text()) - }); + let output = partial + .run + .messages + .get(input_len..) + .unwrap_or_default() + .iter() + .rev() + .find_map(|message| matches!(message, Message::Assistant(_)).then(|| message.text())); let outcome = DriverOutcome { history: partial.run.messages, output: output.clone(), diff --git a/crates/tinyagents-runtime/src/session.rs b/crates/tinyagents-runtime/src/session.rs index 73e5c131..3500639c 100644 --- a/crates/tinyagents-runtime/src/session.rs +++ b/crates/tinyagents-runtime/src/session.rs @@ -140,17 +140,15 @@ impl Session { if let Some(target) = self.target.as_mut() { target.meta = transcript.meta; } - // A successful explicit resume also fixes the target's history handle - // for later appends. Builder construction itself remains I/O-free. - if self.transcript.is_none() { - let target = self.target.as_ref().expect("target checked above"); - self.transcript = Some( - target - .locator - .open_stem(&target.stem, target.meta.clone()) - .map_err(|error| RuntimeError::Persistence(error.to_string()))?, - ); - } + // A successful explicit resume always rebinds the write handle to the + // selected transcript. Builder construction itself remains I/O-free. + let target = self.target.as_ref().expect("target checked above"); + self.transcript = Some( + target + .locator + .open_stem(&target.stem, target.meta.clone()) + .map_err(|error| RuntimeError::Persistence(error.to_string()))?, + ); Ok(SessionResume { loaded: true, history, @@ -246,6 +244,9 @@ impl Session { Ok(outcome) => outcome, Err(failure) => { if let Some(partial) = failure.partial { + if cancellation.is_cancelled() { + return Err(RuntimeError::Cancelled); + } let partial_history = self.with_prefix(partial.history); let raw = self.encode(&self.history, &partial_history, &codec_options)?; let turn_usage = self.turn_usage(&codec_options)?; diff --git a/crates/tinyagents-runtime/src/test.rs b/crates/tinyagents-runtime/src/test.rs index fe793a81..1bea3ee6 100644 --- a/crates/tinyagents-runtime/src/test.rs +++ b/crates/tinyagents-runtime/src/test.rs @@ -54,7 +54,7 @@ struct WaitingDriver(Arc); #[async_trait] impl SessionDriver for WaitingDriver { async fn execute(&self, _: DriverRequest) -> Result { - self.0.notify_waiters(); + self.0.notify_one(); std::future::pending().await } } From f3f2c64606b8866995ff15df21fdfa2ec5951f33 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Sun, 20 Sep 2026 00:33:42 +0300 Subject: [PATCH 18/18] docs: correct task dispatch selection contract Co-authored-by: Medulla --- docs/modules/graph/todos.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/modules/graph/todos.md b/docs/modules/graph/todos.md index e662e390..eda7baaa 100644 --- a/docs/modules/graph/todos.md +++ b/docs/modules/graph/todos.md @@ -69,7 +69,7 @@ staleness policy itself is the pure, clock-injected `staleness_reason`. See [`crates/tinyagents-graph/src/todos/runs/README.md`](../../../crates/tinyagents-graph/src/todos/runs/README.md). `graph::todos::dispatch` is the scheduling policy: `pick_next_card` (urgency, -then board order, optionally agent-assigned only), `requires_plan_approval` +then board order), `requires_plan_approval` (the card's own mode outranks the global gate), `PollCadence` (idle backoff), `build_task_prompt` / `build_progress_instruction`, and `ActiveRunRegistry` (in-flight runs with race-free removal, so a terminal write-back happens once).