diff --git a/docs/modules/harness/cache.md b/docs/modules/harness/cache.md index 6b0d18a..06956d2 100644 --- a/docs/modules/harness/cache.md +++ b/docs/modules/harness/cache.md @@ -107,25 +107,65 @@ stable prefix fingerprint even if the full request changes. ## Cache Policy +As implemented today (`harness::cache::CachePolicy`): + ```rust pub struct CachePolicy { - pub enabled: bool, - pub ttl: Option, - pub scope: CacheScope, - pub include_tools: bool, - pub include_model_responses: bool, - pub preserve_provider_prefix: bool, - pub stable_prefix_min_tokens: Option, + pub response_cache_enabled: bool, + pub protect_prompt_prefix: bool, + pub ttl_ms: Option, + pub namespace: Option, } ``` -Cache keys must include every behavior-affecting input: model, messages, tools, -tool schemas, response format, provider options, and relevant metadata. Unsafe -or side-effecting tool calls should not be cached by default. +`ttl_ms` and `namespace` cover the `ttl` / `scope` this spec asks for. The +remaining aspirational fields (`include_tools`, `include_model_responses`, +`stable_prefix_min_tokens`) are **not implemented**; tools are always part of +the key and there is no minimum-prefix threshold. + +Unsafe or side-effecting tool calls should not be cached by default. + +### Key composition + +The key is a two-part composition, never the prompt alone: -The local response cache key is a SHA-256 digest of canonical request JSON. -Prompt text is not embedded directly in the key, but every serialized -behavior-affecting request field participates in the digest. +```text +scoped_cache_key(cache_key(request), model.cache_identity(), streaming, namespace) +``` + +- `cache_key(request)` is a SHA-256 digest over per-message and per-tool frames + plus an **explicit allowlist projection** of the behaviour-affecting + parameters. The projection destructures `ModelRequest` exhaustively, so adding + a request field is a compile error until someone decides whether it belongs in + the key. Fields that cannot change the answer — `tags`, `timeout_ms`, + `metadata`, `cache_policy`, `prompt_fingerprint`, `cache_segments` — are + deliberately excluded; folding them in gave a caller who put a run id in + `metadata` a permanent 0% hit rate. +- `cache_identity()` names the provider family, model id, API base URL, optional + scope, and a **fingerprint** of the credential. It is computed *after* model + resolution, because the real model is chosen by `ModelRegistry::resolve_request` + and the endpoint and credential live inside the `Arc`, never in + the request. Without it one shared cache serves a hosted harness's answer to a + local one. +- `streaming` is a parameter of the call rather than a request field, so it is + folded explicitly; a warm streaming run is not served an entry written by a + unary run. + +Raw credentials never reach a key: `credential_fingerprint` hashes them first. + +### Write rules + +- Only the **primary** model's answer is written under its own key. When the + fallback chain answers, the write is skipped — otherwise the primary's key is + poisoned (permanently, absent a TTL) with a different model's response. +- A cache read or write failure is logged and ignored. The provider call has + already succeeded and been paid for; discarding its answer because the cache + was unavailable is strictly worse than not caching. +- A cache hit is stamped `ModelResponse::served_from_cache` so token/cost + accounting can tell a replay from a real call and not re-bill it. +- A cache hit on a **streaming** run is replayed as synthetic `ModelDelta` + events (text, then one per tool call) so warm and cold runs are + observationally identical. ## Cache Key Inputs @@ -167,5 +207,25 @@ Every lookup should produce a decision: - write skipped - write completed +Implemented today: `AgentEvent::CacheHit` / `AgentEvent::CacheMiss` are emitted +as events, and the "no lookup happened at all" half is reported as a +`CacheSkipReason` (`no_cache_attached`, `policy_disabled`, +`multi_turn_transcript`) on a `[cache]`-prefixed debug log. A cache also exposes +`ResponseCache::stats() -> CacheStats` (hits, misses, writes, evictions, +expirations, entries, bytes). The remaining decisions are not yet distinct +events. + The usage feature should record provider prompt-cache hits separately from local response-cache hits. + +## Backends + +- `InMemoryResponseCache` — bounded on **both** an entry count and an + approximate byte budget (an entry count alone does not bound memory when + responses are long-context). Recency is an ordered map, not a linear scan. +- `SqliteResponseCache` (feature `sqlite`) — durable, WAL, `(ns, key)` primary + key with an `expiry` column and a lazy purge on read. Namespaces do not + cross-serve and clear independently. +- `SingleFlight` — collapses concurrent identical misses into one provider call + so N simultaneous callers do not all pay for the same answer. Errors are not + shared: a follower whose leader failed runs its own call. diff --git a/examples/agent_loop_tools.rs b/examples/agent_loop_tools.rs index 5ebb57d..3a07841 100644 --- a/examples/agent_loop_tools.rs +++ b/examples/agent_loop_tools.rs @@ -81,6 +81,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -98,6 +99,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/src/error.rs b/src/error.rs index 842b297..823eae7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -96,6 +96,35 @@ pub enum TinyAgentsError { #[error("model error: {0}")] Provider(Box), + /// The request did not fit in the model's context window. + /// + /// Distinguished from the generic [`TinyAgentsError::Provider`] because the + /// remedy is specific and mechanical — compact or drop transcript history + /// and retry — where a generic provider failure has none. A caller that can + /// summarise its own transcript (see + /// [`crate::harness::summarization`]) can match on this variant instead of + /// string-matching a provider message that differs per vendor and changes + /// without notice. Port of LangChain's `ContextOverflowError`. + /// + /// # Detection is best-effort, and asymmetric + /// + /// Hosted providers raise an explicit 400 for this, which + /// [`crate::harness::providers::openai::CONTEXT_OVERFLOW_CODE`] classifies. + /// **Local servers usually truncate the front of the prompt silently + /// instead**, so the absence of this error is not evidence that the prompt + /// fitted — pair it with a probed real context window + /// ([`crate::harness::providers::openai::LocalProbe`]) rather than relying + /// on it alone. + #[error("context overflow: {message}")] + ContextOverflow { + /// Provider family identifier, for example `openai` or `ollama`. + provider: String, + /// Provider model id, when known. + model: Option, + /// The provider's own message, preserved verbatim. + message: String, + }, + /// A tool invocation returned an error. The payload describes the failure. #[error("tool error: {0}")] Tool(String), @@ -235,6 +264,56 @@ pub enum TinyAgentsError { Storage(String), } +impl TinyAgentsError { + /// Builds the right error for a structured provider failure, promoting a + /// recognised context overflow to [`TinyAgentsError::ContextOverflow`]. + /// + /// Provider adapters classify the overflow and stamp + /// [`CONTEXT_OVERFLOW_CODE`][code] on + /// [`ProviderError::code`][pc]; this is where that code becomes a type. Use + /// it in place of `TinyAgentsError::Provider(Box::new(error))` at every + /// site that has a `ProviderError` in hand — the generic variant is still + /// correct for everything else, and is what this returns when the code is + /// absent or unrecognised. + /// + /// [code]: crate::harness::providers::openai::CONTEXT_OVERFLOW_CODE + /// [pc]: crate::harness::model::ProviderError::code + pub fn from_provider_error(error: crate::harness::model::ProviderError) -> Self { + if error.code.as_deref() == Some(crate::harness::providers::openai::CONTEXT_OVERFLOW_CODE) { + tracing::debug!( + "[error] promoting provider `{}` context-overflow code to a typed error", + error.provider + ); + return Self::ContextOverflow { + provider: error.provider, + model: error.model, + message: error.message, + }; + } + Self::Provider(Box::new(error)) + } + + /// Whether this error means the request did not fit the model's context + /// window. + /// + /// Recognises **both** the typed [`TinyAgentsError::ContextOverflow`] and a + /// [`TinyAgentsError::Provider`] still carrying the classification code, so + /// a caller's compact-and-retry logic behaves identically no matter which + /// construction site produced the error. Call sites are migrating to + /// [`Self::from_provider_error`]; until every one has, the two shapes must + /// classify the same or the same failure would be handled two ways. + pub fn is_context_overflow(&self) -> bool { + match self { + Self::ContextOverflow { .. } => true, + Self::Provider(error) => { + error.code.as_deref() + == Some(crate::harness::providers::openai::CONTEXT_OVERFLOW_CODE) + } + _ => false, + } + } +} + /// Converts a raw `rusqlite` failure into [`TinyAgentsError::Storage`] so the /// session store and run ledger can use `?` on driver calls directly. Call /// sites that have useful context to add should still map explicitly rather diff --git a/src/graph/checkpoint/file.rs b/src/graph/checkpoint/file.rs index 2a0de37..76967ba 100644 --- a/src/graph/checkpoint/file.rs +++ b/src/graph/checkpoint/file.rs @@ -29,13 +29,38 @@ struct CheckpointIdHeader { checkpoint_id: String, } -use super::{Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointTuple, Checkpointer}; +use super::{ + Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointTuple, Checkpointer, PendingWrite, + merge_writes, +}; use crate::harness::ids::CheckpointId; use crate::{Result, TinyAgentsError}; /// File extension for per-thread checkpoint logs. const THREAD_EXT: &str = "jsonl"; +/// Filename suffix for a thread's **pending-writes** sidecar. +/// +/// Writes are recorded after their checkpoint is already durable, so they +/// cannot live in the append-only checkpoint log without turning it into a +/// mixed-record format that every reader would have to discriminate. A sibling +/// file keeps the checkpoint log exactly as it was. +const WRITES_SUFFIX: &str = ".writes.jsonl"; + +/// Process-wide counter making temp-file names unique so concurrent atomic +/// rewrites of the same thread never collide on their scratch file. +static TMP_COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); + +/// One line of a thread's pending-writes sidecar: the write plus the +/// `(namespace, checkpoint_id)` it is filed under. +#[derive(Clone, serde::Serialize, serde::Deserialize)] +struct WriteRecord { + #[serde(default)] + namespace: Vec, + checkpoint_id: String, + write: PendingWrite, +} + /// A [`Checkpointer`] that persists checkpoints as JSONL files under a base /// directory, one file per thread. /// @@ -67,9 +92,62 @@ impl FileCheckpointer { /// The thread id is percent-escaped so it is a safe, injective single path /// component (no separators, no collisions between distinct ids). fn thread_path(&self, thread_id: &str) -> PathBuf { + let canonical = self.canonical_thread_path(thread_id); + if canonical.exists() { + canonical + } else { + let legacy = self.legacy_thread_path(thread_id); + if legacy.exists() { legacy } else { canonical } + } + } + + /// Resolves the pending-writes sidecar path for `thread_id`. + fn writes_path(&self, thread_id: &str) -> PathBuf { + let canonical = self.canonical_writes_path(thread_id); + if canonical.exists() { + canonical + } else { + let legacy = self.legacy_writes_path(thread_id); + if legacy.exists() { legacy } else { canonical } + } + } + + fn canonical_thread_path(&self, thread_id: &str) -> PathBuf { self.base_dir .join(format!("{}.{THREAD_EXT}", escape_thread_id(thread_id))) } + + fn canonical_writes_path(&self, thread_id: &str) -> PathBuf { + self.base_dir + .join(format!("{}{WRITES_SUFFIX}", escape_thread_id(thread_id))) + } + + fn legacy_thread_path(&self, thread_id: &str) -> PathBuf { + self.base_dir.join(format!( + "{}.{THREAD_EXT}", + legacy_escape_thread_id(thread_id) + )) + } + + fn legacy_writes_path(&self, thread_id: &str) -> PathBuf { + self.base_dir.join(format!( + "{}{WRITES_SUFFIX}", + legacy_escape_thread_id(thread_id) + )) + } + + /// Reads a thread's write sidecar, tolerating a torn trailing line exactly + /// as [`FileCheckpointer::read_records`] does. + fn read_write_records(path: &Path, thread_id: &str) -> Result> { + let text = match fs::read_to_string(path) { + Ok(t) => t, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), + Err(e) => return Err(io_err("open writes file", e)), + }; + decode_lines(&text, &format!("writes for thread `{thread_id}`"), |line| { + serde_json::from_str::(line) + }) + } } impl Clone for FileCheckpointer { @@ -81,9 +159,42 @@ impl Clone for FileCheckpointer { } } -/// Percent-escapes any byte outside `[A-Za-z0-9._-]` so a thread id maps to a -/// single, collision-free filename component. +/// Percent-escapes any byte outside `[a-z0-9._-]` so a thread id maps to a +/// single filename component that is injective **even on a case-insensitive +/// filesystem**. +/// +/// # Why uppercase is escaped +/// +/// The obvious safe set is `[A-Za-z0-9._-]`, and that is what this used to use. +/// It is injective on a case-*sensitive* filesystem and silently is not on +/// APFS, HFS+ or NTFS: threads `"Alice"` and `"alice"` map to `Alice.jsonl` and +/// `alice.jsonl`, which are the *same file*. Two unrelated runs then append into +/// one lineage, and reads hand each of them the other's checkpoints. +/// +/// Escaping `A-Z` fixes it while staying case-*preserving* (the id is still +/// recoverable byte-for-byte from the name). The only uppercase characters left +/// in the output are the hex digits `A-F` of an escape, and escapes are always +/// emitted as `%` + exactly two uppercase hex digits, so no two outputs can +/// differ only by letter case: lowercasing the whole name is injective on the +/// image, which is exactly what case-insensitive collision-freedom means. +/// fn escape_thread_id(thread_id: &str) -> String { + let mut out = String::with_capacity(thread_id.len()); + for &b in thread_id.as_bytes() { + if b.is_ascii_lowercase() || b.is_ascii_digit() || matches!(b, b'.' | b'_' | b'-') { + out.push(b as char); + } else { + out.push('%'); + out.push_str(&format!("{b:02X}")); + } + } + out +} + +/// The filename escaping used before uppercase letters were made explicit. +/// Kept only as a read/write fallback so persisted threads remain reachable +/// after an upgrade; new thread files always use [`escape_thread_id`]. +fn legacy_escape_thread_id(thread_id: &str) -> String { let mut out = String::with_capacity(thread_id.len()); for &b in thread_id.as_bytes() { if b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-') { @@ -135,24 +246,53 @@ where /// Returns an empty vec when the thread file does not exist. fn read_records(&self, thread_id: &str) -> Result>> { let path = self.thread_path(thread_id); - let file = match File::open(&path) { - Ok(f) => f, + let text = match fs::read_to_string(&path) { + Ok(t) => t, Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()), Err(e) => return Err(io_err("open thread file", e)), }; - let reader = BufReader::new(file); - let mut records = Vec::new(); - for line in reader.lines() { - let line = line.map_err(|e| io_err("read line", e))?; - if line.trim().is_empty() { - continue; + decode_lines(&text, &format!("thread `{thread_id}`"), |line| { + serde_json::from_str::>(line) + }) + } +} + +/// Decodes one JSON object per line, tolerating a **torn trailing line**. +/// +/// A crash between `write_all` and the OS flushing the tail of the buffer +/// leaves a partial final line. It can only ever be the last one — the file is +/// append-only — so that is the only line whose decode failure is forgiven, and +/// only when the file does not end in a newline (a complete record always +/// does). Anything else is real corruption and still errors. +/// +/// This matters more than "one lost record": the previous behaviour failed the +/// whole read, so a single torn byte made a thread permanently unreadable, with +/// no way to get at the hundreds of intact checkpoints in front of it. +fn decode_lines(text: &str, what: &str, mut decode: F) -> Result> +where + F: FnMut(&str) -> std::result::Result, +{ + let complete = text.is_empty() || text.ends_with('\n'); + let lines: Vec<&str> = text.lines().collect(); + let last_index = lines.len().saturating_sub(1); + let mut out = Vec::with_capacity(lines.len()); + for (i, line) in lines.iter().enumerate() { + if line.trim().is_empty() { + continue; + } + match decode(line) { + Ok(record) => out.push(record), + Err(e) if !complete && i == last_index => { + tracing::warn!( + "[checkpoint:file] {what}: discarding torn trailing line \ + ({} bytes, no terminating newline): {e}", + line.len() + ); } - let record: Checkpoint = - serde_json::from_str(&line).map_err(|e| io_err("decode record", e))?; - records.push(record); + Err(e) => return Err(io_err("decode record", e)), } - Ok(records) } + Ok(out) } impl FileCheckpointer @@ -179,11 +319,53 @@ where buf.push_str(&line); buf.push('\n'); } - fs::write(&path, buf).map_err(|e| io_err("write thread file", e)) + write_atomic(&path, buf.as_bytes()) } } } +/// Writes `bytes` to `path` atomically: a uniquely named temp file in the same +/// directory, fsynced, then renamed over the destination. +/// +/// The prune/delete path used to rewrite the thread file **in place** with +/// `fs::write`, which truncates first: a crash anywhere in the following write +/// leaves a truncated or empty file, and the whole history is gone — not the +/// pruned tail, all of it. Rename is atomic for same-directory paths on POSIX +/// and Windows, so a reader sees either the old file or the new one, and a +/// crash leaves the old one intact. This is the same shape `FileStore::put` +/// already uses. +fn write_atomic(path: &Path, bytes: &[u8]) -> Result<()> { + let dir = path.parent().ok_or_else(|| { + TinyAgentsError::Checkpoint(format!( + "file checkpointer: path has no parent directory: {}", + path.display() + )) + })?; + let file_name = path + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("thread"); + let tmp = dir.join(format!( + "{file_name}.tmp.{}.{}", + std::process::id(), + TMP_COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + )); + let write_and_sync = || -> std::io::Result<()> { + let mut file = File::create(&tmp)?; + file.write_all(bytes)?; + file.sync_all() + }; + if let Err(e) = write_and_sync() { + let _ = fs::remove_file(&tmp); + return Err(io_err("write temp thread file", e)); + } + if let Err(e) = fs::rename(&tmp, path) { + let _ = fs::remove_file(&tmp); + return Err(io_err("rename temp thread file", e)); + } + Ok(()) +} + #[async_trait] impl Checkpointer for FileCheckpointer where @@ -207,7 +389,13 @@ where .open(&path) .map_err(|e| io_err("open thread file for append", e))?; file.write_all(line.as_bytes()) - .map_err(|e| io_err("append record", e)) + .map_err(|e| io_err("append record", e))?; + // Without an explicit flush to stable storage a "persisted" + // checkpoint is only in the page cache: a host crash loses + // boundaries the executor has already reported as durable, and can + // leave a torn trailing line behind (which `read_records` now + // tolerates, but should not have to see). + file.sync_all().map_err(|e| io_err("fsync record", e)) }) .await .map_err(|e| io_err("join blocking put task", e))??; @@ -329,7 +517,16 @@ where for entry in entries { let entry = entry.map_err(|e| io_err("read dir entry", e))?; let path = entry.path(); - if path.extension().and_then(|s| s.to_str()) != Some(THREAD_EXT) { + // Match on the filename suffix rather than `Path::extension()`. + // The empty thread id escapes to the empty string, so its file is + // literally `.jsonl` — a dotfile whose `extension()` is `None`, + // which made that thread invisible to listing (and to everything + // built on listing) while `get`/`put` addressed it perfectly well. + let name = entry.file_name(); + let Some(name) = name.to_str() else { + continue; + }; + if !name.ends_with(&format!(".{THREAD_EXT}")) || name.ends_with(WRITES_SUFFIX) { continue; } // Recover the canonical thread id from the first record rather than @@ -349,9 +546,17 @@ where if first.trim().is_empty() { continue; } - let record: Checkpoint = - serde_json::from_str(&first).map_err(|e| io_err("decode header", e))?; - threads.push(record.thread_id); + // One unreadable file must not take down the whole listing. + // `list_threads` decodes the first line of *every* file, so an + // error here made a single poisoned thread break listing — + // and therefore every operation built on it — globally. + match serde_json::from_str::>(&first) { + Ok(record) => threads.push(record.thread_id), + Err(e) => tracing::warn!( + "[checkpoint:file] list_threads: skipping unreadable thread file {}: {e}", + path.display() + ), + } break; } } @@ -359,12 +564,16 @@ where } async fn delete_thread(&self, thread_id: &str) -> Result<()> { - let path = self.thread_path(thread_id); - match fs::remove_file(&path) { - Ok(()) => Ok(()), - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), - Err(e) => Err(io_err("delete thread file", e)), + // The write sidecar goes with the thread: leaving it behind would let a + // later thread of the same id inherit a dead ledger. + for path in [self.thread_path(thread_id), self.writes_path(thread_id)] { + match fs::remove_file(&path) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(io_err("delete thread file", e)), + } } + Ok(()) } async fn delete_checkpoints(&self, thread_id: &str, ids: &[String]) -> Result { @@ -378,21 +587,85 @@ where let removed = before - records.len(); if removed > 0 { self.write_records(thread_id, &records)?; + // Drop the deleted checkpoints' write ledgers with them. + let writes_path = self.writes_path(thread_id); + let write_records = Self::read_write_records(&writes_path, thread_id)?; + let kept: Vec<&WriteRecord> = write_records + .iter() + .filter(|r| !drop.contains(r.checkpoint_id.as_str())) + .collect(); + if kept.len() != write_records.len() { + let mut buf = String::new(); + for record in kept { + let line = + serde_json::to_string(record).map_err(|e| io_err("encode write", e))?; + buf.push_str(&line); + buf.push('\n'); + } + if buf.is_empty() { + match fs::remove_file(&writes_path) { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => return Err(io_err("remove empty writes file", e)), + } + } else { + write_atomic(&writes_path, buf.as_bytes())?; + } + } } Ok(removed) } - async fn copy_thread(&self, source_thread: &str, target_thread: &str) -> Result<()> { - let mut records = self.read_records(source_thread)?; - if records.is_empty() { + async fn put_writes(&self, config: &CheckpointConfig, writes: &[PendingWrite]) -> Result<()> { + let checkpoint_id = super::require_checkpoint_id(config)?; + if writes.is_empty() { return Ok(()); } - for record in &mut records { - record.thread_id = target_thread.to_string(); + let path = self.writes_path(&config.thread_id); + let mut records = Self::read_write_records(&path, &config.thread_id)?; + + // Split out this checkpoint's ledger, merge, then rebuild the file. + let (mut mine, others): (Vec, Vec) = records + .drain(..) + .partition(|r| r.checkpoint_id == checkpoint_id && r.namespace == config.namespace); + let mut existing: Vec = mine.drain(..).map(|r| r.write).collect(); + let changed = merge_writes(&mut existing, writes); + + let mut buf = String::new(); + for record in others.iter() { + let line = serde_json::to_string(record).map_err(|e| io_err("encode write", e))?; + buf.push_str(&line); + buf.push('\n'); + } + for write in existing { + let record = WriteRecord { + namespace: config.namespace.clone(), + checkpoint_id: checkpoint_id.clone(), + write, + }; + let line = serde_json::to_string(&record).map_err(|e| io_err("encode write", e))?; + buf.push_str(&line); + buf.push('\n'); } - // Append onto any existing target file to match `put` semantics. - let mut existing = self.read_records(target_thread)?; - existing.extend(records); - self.write_records(target_thread, &existing) + fs::create_dir_all(&self.base_dir).map_err(|e| io_err("create base dir", e))?; + write_atomic(&path, buf.as_bytes())?; + tracing::debug!( + "[checkpoint:file] put_writes thread={} checkpoint={checkpoint_id} offered={} stored={changed}", + config.thread_id, + writes.len() + ); + Ok(()) + } + + async fn get_writes(&self, config: &CheckpointConfig) -> Result> { + let Some(checkpoint_id) = self.resolve_write_target(config).await? else { + return Ok(Vec::new()); + }; + let path = self.writes_path(&config.thread_id); + Ok(Self::read_write_records(&path, &config.thread_id)? + .into_iter() + .filter(|r| r.checkpoint_id == checkpoint_id && r.namespace == config.namespace) + .map(|r| r.write) + .collect()) } } diff --git a/src/graph/checkpoint/mod.rs b/src/graph/checkpoint/mod.rs index 822beb1..6941a14 100644 --- a/src/graph/checkpoint/mod.rs +++ b/src/graph/checkpoint/mod.rs @@ -23,7 +23,8 @@ pub use file::FileCheckpointer; pub use sqlite::SqliteCheckpointer; pub use types::{ BarrierArrivals, Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, - CheckpointTuple, DurabilityMode, PendingActivation, PendingWrite, + CheckpointTuple, DurabilityMode, PendingActivation, PendingWrite, WRITES_IDX_ERROR, + WRITES_IDX_INTERRUPT, WRITES_IDX_RESUME, merge_writes, }; use std::collections::{HashMap, HashSet}; @@ -61,7 +62,9 @@ where /// returned (last-write-wins, consistent with [`Checkpointer::get`]). /// /// Composed from [`Checkpointer::list`] + [`Checkpointer::get`] so every - /// backend inherits it; override only for a cheaper scoped query. + /// backend inherits it; override for a cheaper scoped query — both durable + /// backends do, because the default costs a full thread scan per call and + /// [`Checkpointer::state_history`] issues one per lineage hop. async fn get_scoped( &self, thread_id: &str, @@ -90,6 +93,63 @@ where /// Lists checkpoint metadata for a thread in insertion order. async fn list(&self, thread_id: &str) -> Result>; + // ---- Pending writes ---------------------------------------------------- + // + // The partial-failure protocol. A superstep can fail after some of its + // tasks have already run; without a per-task record of what they wrote, a + // resume cannot tell "already ran" from "not yet run" and re-executes their + // side effects. `put_writes` records that work against the checkpoint it + // belongs to; `get_writes` reads it back, and `get_tuple` surfaces it as + // `CheckpointTuple::pending_writes` so resume can skip completed tasks. + // + // Both carry default no-op bodies so an out-of-tree `Checkpointer` keeps + // compiling: such a backend simply never persists writes, which is exactly + // the behaviour every backend had before the protocol existed. + + /// Records `writes` against the checkpoint addressed by `config`. + /// + /// `config.checkpoint_id` must name a specific checkpoint — writes belong + /// to the boundary they were produced at, so a `None` id has no meaning and + /// backends reject it. + /// + /// Idempotency follows [`PendingWrite`]'s replace-vs-ignore rule: a data + /// write (`idx >= 0`) whose `(task_id, idx)` is already stored is ignored, + /// while a control-plane write (`idx < 0`) replaces the stored value. Both + /// are implemented through [`merge_writes`], so every backend agrees. + /// + /// The default body is a no-op returning `Ok(())`. + async fn put_writes(&self, _config: &CheckpointConfig, _writes: &[PendingWrite]) -> Result<()> { + Ok(()) + } + + /// Reads back the writes recorded against the checkpoint addressed by + /// `config`, in insertion order. + /// + /// Returns an empty vec for an unknown checkpoint or one that has no + /// writes. When `config.checkpoint_id` is `None` the latest checkpoint in + /// `config.namespace` is resolved first. + /// + /// The default body returns an empty vec. + async fn get_writes(&self, _config: &CheckpointConfig) -> Result> { + Ok(Vec::new()) + } + + /// Resolves the checkpoint id a **read** of writes addresses. + /// + /// Unlike [`Checkpointer::put_writes`] (where an unaddressed id is a caller + /// bug), a read may legitimately mean "the latest checkpoint in this + /// namespace" — the same relaxation [`Checkpointer::get`] makes. Returns + /// `None` when the thread/namespace has no checkpoint at all. + async fn resolve_write_target(&self, config: &CheckpointConfig) -> Result> { + match &config.checkpoint_id { + Some(id) => Ok(Some(id.clone())), + None => Ok(self + .get_scoped(&config.thread_id, None, &config.namespace) + .await? + .map(|c| c.checkpoint_id)), + } + } + /// Loads every checkpoint stored under `thread_id`, in listing order. /// /// This is the bulk-read companion to [`Checkpointer::list`]: it returns @@ -145,7 +205,7 @@ where checkpoint_id: Some(parent.clone()), namespace: checkpoint.namespace.clone(), }); - let pending_writes = checkpoint.pending_writes.clone(); + let pending_writes = self.resolved_writes(&resolved, &checkpoint).await?; Ok(Some(CheckpointTuple { config: resolved, checkpoint, @@ -154,16 +214,50 @@ where })) } + /// The writes to surface on a tuple for `checkpoint`. + /// + /// Prefers the separately persisted [`Checkpointer::get_writes`] records — + /// the authoritative partial-failure ledger — and falls back to the + /// checkpoint record's own inline `pending_writes` for backends that do not + /// implement the write protocol (whose `get_writes` default returns empty) + /// and for records written before it existed. + async fn resolved_writes( + &self, + config: &CheckpointConfig, + checkpoint: &Checkpoint, + ) -> Result> { + let stored = self.get_writes(config).await?; + if stored.is_empty() { + Ok(checkpoint.pending_writes.clone()) + } else { + Ok(stored) + } + } + /// Returns a thread's checkpoint lineage newest-first, following each /// checkpoint's `parent_checkpoint_id` from the latest checkpoint in /// `namespace`. `limit` caps the number of tuples returned (the most recent /// ones). /// /// The default walks [`Checkpointer::get_tuple`] once per hop, so a backend - /// that re-reads the whole thread per lookup (the file/JSONL backend) is - /// O(H²) over the lineage. Such backends override this to read the thread - /// once and walk the lineage in memory (O(H)). The observable result is - /// identical to iterating `get_tuple` by parent pointer. + /// whose scoped lookup re-reads the whole thread is O(H²) over the lineage. + /// Every bundled backend overrides it to read the thread (or the + /// namespace's rows) once and walk the lineage in memory, so none of the + /// three is in that class: the JSONL backend parses its file once, and the + /// SQLite backend issues one indexed range query — the + /// `(thread_id, namespace, seq)` index is what makes the namespace scope + /// expressible in SQL at all, and without it that backend silently fell + /// back to this default. The observable result is identical to iterating + /// `get_tuple` by parent pointer. + /// + /// The walk carries a **visited set**. `parent_checkpoint_id` is caller-set + /// data, not a structurally enforced acyclic pointer: a hand-written + /// checkpoint, a bad fork, or a `copy_thread` that reused ids can point a + /// record at itself or at one of its descendants. With `limit == None` an + /// unguarded walk then never terminates — it does not merely return a wrong + /// answer, it hangs the caller. Revisiting an id ends the walk. + /// [`FileCheckpointer`] has always had this guard (its in-memory `remove` + /// doubles as one); this makes it uniform. async fn state_history( &self, thread_id: &str, @@ -172,6 +266,7 @@ where ) -> Result>> { let mut out = Vec::new(); let mut cursor: Option = None; + let mut visited: HashSet = HashSet::new(); loop { if let Some(limit) = limit && out.len() >= limit @@ -186,6 +281,14 @@ where let Some(tuple) = self.get_tuple(config).await? else { break; }; + if !visited.insert(tuple.checkpoint.checkpoint_id.clone()) { + tracing::warn!( + "[checkpoint] state_history: lineage cycle at checkpoint `{}` \ + (thread `{thread_id}`); truncating the walk", + tuple.checkpoint.checkpoint_id + ); + break; + } let parent = tuple.checkpoint.parent_checkpoint_id.clone(); out.push(tuple); match parent { @@ -252,10 +355,47 @@ where /// Composed from [`Checkpointer::get_thread`] + [`Checkpointer::put`], so /// the source thread is read once (a bulk read, not one /// [`Checkpointer::get`] per checkpoint). + /// + /// # The target must be empty + /// + /// Copying preserves each record's `checkpoint_id` — that is what keeps the + /// lineage spine intact — so appending a lineage onto a thread that already + /// has one produces a file/table containing two disjoint lineages *with + /// reused ids*. Every subsequent `get(Some(id))` then resolves to whichever + /// copy was written last and the parent walk crosses between lineages: the + /// thread is silently corrupt, with no error at the point of damage. + /// + /// So a non-empty target is **rejected** rather than merged into. Callers + /// that genuinely want to overwrite call [`Checkpointer::delete_thread`] + /// first, which makes the destructive intent explicit. Copying an empty or + /// unknown source thread is a no-op (still `Ok`). async fn copy_thread(&self, source_thread: &str, target_thread: &str) -> Result<()> { + let existing = self.list(target_thread).await?; + if !existing.is_empty() { + return Err(TinyAgentsError::Checkpoint(format!( + "copy_thread: target thread `{target_thread}` already has {} checkpoint(s); \ + copying would interleave two lineages with reused checkpoint ids. \ + Delete the target first if replacing it is intended.", + existing.len() + ))); + } for mut checkpoint in self.get_thread(source_thread).await? { + let source_config = CheckpointConfig { + thread_id: source_thread.to_string(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let writes = self.get_writes(&source_config).await?; checkpoint.thread_id = target_thread.to_string(); + let target_config = CheckpointConfig { + thread_id: target_thread.to_string(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; self.put(checkpoint).await?; + if !writes.is_empty() { + self.put_writes(&target_config, &writes).await?; + } } Ok(()) } @@ -337,13 +477,20 @@ where /// Cheap to clone; clones share the same underlying store. pub struct InMemoryCheckpointer { inner: Arc>>>>, + /// Pending writes keyed by `(thread_id, namespace, checkpoint_id)` — the + /// same identity the SQL backends use as a primary key prefix. + writes: Arc>>>, } +/// The address a batch of pending writes is filed under. +type WriteKey = (String, Vec, String); + impl InMemoryCheckpointer { /// Creates an empty checkpointer. pub fn new() -> Self { Self { inner: Arc::new(Mutex::new(HashMap::new())), + writes: Arc::new(Mutex::new(HashMap::new())), } } @@ -366,6 +513,7 @@ impl Clone for InMemoryCheckpointer { fn clone(&self) -> Self { Self { inner: self.inner.clone(), + writes: self.writes.clone(), } } } @@ -439,6 +587,11 @@ where async fn delete_thread(&self, thread_id: &str) -> Result<()> { let mut map = self.inner.lock().map_err(|_| lock_err())?; map.remove(thread_id); + // Writes are keyed by thread too, and deleting a thread must not leave + // its write ledger behind for a later thread of the same name to + // inherit. The conformance suite asserts exactly this. + let mut writes = self.writes.lock().map_err(|_| lock_err())?; + writes.retain(|(thread, _, _), _| thread != thread_id); Ok(()) } @@ -453,9 +606,75 @@ where }; let before = list.len(); list.retain(|c| !drop.contains(c.checkpoint_id.as_str())); - Ok(before - list.len()) + let removed = before - list.len(); + drop_writes_for(&self.writes, thread_id, &drop)?; + Ok(removed) + } + + async fn put_writes(&self, config: &CheckpointConfig, writes: &[PendingWrite]) -> Result<()> { + let checkpoint_id = require_checkpoint_id(config)?; + if writes.is_empty() { + return Ok(()); + } + let key: WriteKey = ( + config.thread_id.clone(), + config.namespace.clone(), + checkpoint_id, + ); + let mut map = self.writes.lock().map_err(|_| lock_err())?; + let slot = map.entry(key).or_default(); + let changed = merge_writes(slot, writes); + tracing::debug!( + "[checkpoint:memory] put_writes thread={} checkpoint={:?} offered={} stored={}", + config.thread_id, + config.checkpoint_id, + writes.len(), + changed + ); + Ok(()) + } + + async fn get_writes(&self, config: &CheckpointConfig) -> Result> { + let Some(checkpoint_id) = self.resolve_write_target(config).await? else { + return Ok(Vec::new()); + }; + let key: WriteKey = ( + config.thread_id.clone(), + config.namespace.clone(), + checkpoint_id, + ); + let map = self.writes.lock().map_err(|_| lock_err())?; + Ok(map.get(&key).cloned().unwrap_or_default()) } } +/// Removes the write ledgers of `ids` within `thread_id`. +fn drop_writes_for( + writes: &Mutex>>, + thread_id: &str, + ids: &HashSet<&str>, +) -> Result<()> { + let mut map = writes.lock().map_err(|_| lock_err())?; + map.retain(|(thread, _, checkpoint), _| { + thread != thread_id || !ids.contains(checkpoint.as_str()) + }); + Ok(()) +} + +/// Extracts the checkpoint id a `put_writes` call addresses. +/// +/// A write belongs to the boundary that produced it, so an unaddressed +/// (`None`) id is a caller bug rather than "the latest": silently filing the +/// writes against whatever checkpoint happens to be newest is precisely the +/// corruption the protocol exists to prevent. +pub(crate) fn require_checkpoint_id(config: &CheckpointConfig) -> Result { + config.checkpoint_id.clone().ok_or_else(|| { + TinyAgentsError::Checkpoint(format!( + "put_writes requires an explicit checkpoint_id (thread `{}`)", + config.thread_id + )) + }) +} + #[cfg(test)] mod test; diff --git a/src/graph/checkpoint/sqlite.rs b/src/graph/checkpoint/sqlite.rs index f0f1f18..05df92d 100644 --- a/src/graph/checkpoint/sqlite.rs +++ b/src/graph/checkpoint/sqlite.rs @@ -32,7 +32,10 @@ use rusqlite::{Connection, OptionalExtension, params}; use serde::Serialize; use serde::de::DeserializeOwned; -use super::{Checkpoint, CheckpointMetadata, CheckpointSource, Checkpointer}; +use super::{ + Checkpoint, CheckpointConfig, CheckpointMetadata, CheckpointSource, CheckpointTuple, + Checkpointer, PendingWrite, merge_writes, +}; use crate::harness::ids::{CheckpointId, NodeId}; use crate::{Result, TinyAgentsError}; @@ -118,7 +121,26 @@ impl SqliteCheckpointer { } /// Table + indexes. `seq` preserves insertion order; the indexes serve thread -/// listing and `(thread_id, checkpoint_id)` parent-chain lookups. +/// listing, `(thread_id, checkpoint_id)` parent-chain lookups, and — since the +/// namespace-scoped overrides landed — `(thread_id, namespace, …)` scoped +/// lookups. +/// +/// # `namespace` is a first-class, indexed column +/// +/// It holds the canonical JSON encoding of the namespace vector, which +/// `serde_json` emits deterministically, so equality on the column is exactly +/// equality on the namespace. It was already stored this way; what was missing +/// were the indexes, and therefore the ability to *push the scope down into +/// SQL* at all. Without them `get_scoped`, `get_tuple` and `state_history` all +/// fell back to the trait defaults, which scan the whole thread once per +/// lineage hop — O(H²) per namespaced read on the one backend that had no +/// business being in that class. Both are `CREATE INDEX IF NOT EXISTS`, so an +/// existing database picks them up on the next open with no migration step. +/// +/// `checkpoint_writes` is the partial-failure ledger. Its primary key +/// `(thread_id, namespace, checkpoint_id, task_id, idx)` mirrors LangGraph's +/// writes table and is what makes `put_writes` idempotent in SQL rather than in +/// application code. const SCHEMA: &str = "\ CREATE TABLE IF NOT EXISTS checkpoints ( seq INTEGER PRIMARY KEY AUTOINCREMENT, @@ -135,6 +157,23 @@ CREATE TABLE IF NOT EXISTS checkpoints ( ); CREATE INDEX IF NOT EXISTS idx_checkpoints_thread ON checkpoints (thread_id, seq); CREATE INDEX IF NOT EXISTS idx_checkpoints_lookup ON checkpoints (thread_id, checkpoint_id); +CREATE INDEX IF NOT EXISTS idx_checkpoints_scoped ON checkpoints (thread_id, namespace, seq); +CREATE INDEX IF NOT EXISTS idx_checkpoints_scoped_lookup + ON checkpoints (thread_id, namespace, checkpoint_id, seq); + +CREATE TABLE IF NOT EXISTS checkpoint_writes ( + thread_id TEXT NOT NULL, + namespace TEXT NOT NULL, + checkpoint_id TEXT NOT NULL, + task_id TEXT NOT NULL, + idx INTEGER NOT NULL, + node TEXT NOT NULL, + channel TEXT NOT NULL, + payload TEXT NOT NULL, + PRIMARY KEY (thread_id, namespace, checkpoint_id, task_id, idx) +); +CREATE INDEX IF NOT EXISTS idx_checkpoint_writes_thread + ON checkpoint_writes (thread_id, checkpoint_id); "; /// The projected listing columns read from one `checkpoints` row. @@ -259,6 +298,134 @@ where } } + async fn get_scoped( + &self, + thread_id: &str, + checkpoint_id: Option<&str>, + namespace: &[String], + ) -> Result>> { + // Pushed down to one indexed query. The trait default lists the whole + // thread and then re-`get`s the winner, which costs a full thread scan + // per call — and `state_history` calls it once per lineage hop. + let namespace_json = + serde_json::to_string(namespace).map_err(|e| sqlite_err("encode namespace", e))?; + let conn = self.lock()?; + let record: Option = match checkpoint_id { + Some(id) => conn + .query_row( + "SELECT record FROM checkpoints + WHERE thread_id = ?1 AND namespace = ?2 AND checkpoint_id = ?3 + ORDER BY seq DESC LIMIT 1", + params![thread_id, namespace_json, id], + |row| row.get(0), + ) + .optional() + .map_err(|e| sqlite_err("query scoped checkpoint", e))?, + None => conn + .query_row( + "SELECT record FROM checkpoints + WHERE thread_id = ?1 AND namespace = ?2 + ORDER BY seq DESC LIMIT 1", + params![thread_id, namespace_json], + |row| row.get(0), + ) + .optional() + .map_err(|e| sqlite_err("query latest scoped checkpoint", e))?, + }; + match record { + Some(json) => Ok(Some( + serde_json::from_str(&json).map_err(|e| sqlite_err("decode record", e))?, + )), + None => Ok(None), + } + } + + async fn state_history( + &self, + thread_id: &str, + namespace: &[String], + limit: Option, + ) -> Result>> { + // One indexed range read of the namespace's rows, then the lineage walk + // in memory — instead of the default's `get_tuple` (and therefore + // `get_scoped`) per hop. + let namespace_json = + serde_json::to_string(namespace).map_err(|e| sqlite_err("encode namespace", e))?; + let (records, writes) = { + let conn = self.lock()?; + let mut stmt = conn + .prepare( + "SELECT record FROM checkpoints + WHERE thread_id = ?1 AND namespace = ?2 ORDER BY seq ASC", + ) + .map_err(|e| sqlite_err("prepare state_history", e))?; + let rows = stmt + .query_map(params![thread_id, namespace_json], |row| { + row.get::<_, String>(0) + }) + .map_err(|e| sqlite_err("query state_history", e))?; + let mut records: Vec> = Vec::new(); + for row in rows { + let json = row.map_err(|e| sqlite_err("read record row", e))?; + records + .push(serde_json::from_str(&json).map_err(|e| sqlite_err("decode record", e))?); + } + let writes = read_writes_by_checkpoint(&conn, thread_id, &namespace_json)?; + (records, writes) + }; + if records.is_empty() { + return Ok(Vec::new()); + } + + // Last write wins for a re-used id, matching `get`. + let mut by_id: std::collections::HashMap> = + std::collections::HashMap::with_capacity(records.len()); + let mut cursor: Option = None; + for record in records { + cursor = Some(record.checkpoint_id.clone()); + by_id.insert(record.checkpoint_id.clone(), record); + } + + let mut out = Vec::new(); + while let Some(id) = cursor { + if let Some(limit) = limit + && out.len() >= limit + { + break; + } + // `remove` doubles as the cycle guard: each id is visited once. + let Some(checkpoint) = by_id.remove(&id) else { + break; + }; + cursor = checkpoint.parent_checkpoint_id.clone(); + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; + let parent_config = + checkpoint + .parent_checkpoint_id + .as_ref() + .map(|parent| CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(parent.clone()), + namespace: checkpoint.namespace.clone(), + }); + let pending_writes = writes + .get(&checkpoint.checkpoint_id) + .cloned() + .unwrap_or_else(|| checkpoint.pending_writes.clone()); + out.push(CheckpointTuple { + config, + checkpoint, + parent_config, + pending_writes, + }); + } + Ok(out) + } + async fn list(&self, thread_id: &str) -> Result> { let conn = self.lock()?; let mut stmt = conn @@ -327,12 +494,24 @@ where } async fn delete_thread(&self, thread_id: &str) -> Result<()> { - let conn = self.lock()?; - conn.execute( + let mut conn = self.lock()?; + let tx = conn + .transaction() + .map_err(|e| sqlite_err("begin delete_thread", e))?; + tx.execute( "DELETE FROM checkpoints WHERE thread_id = ?1", params![thread_id], ) .map_err(|e| sqlite_err("delete thread", e))?; + // Writes go with the thread — and across *every* namespace, not just + // the root one, or an embedded subgraph's ledger outlives its thread. + tx.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = ?1", + params![thread_id], + ) + .map_err(|e| sqlite_err("delete thread writes", e))?; + tx.commit() + .map_err(|e| sqlite_err("commit delete_thread", e))?; Ok(()) } @@ -352,8 +531,173 @@ where params![thread_id, id], ) .map_err(|e| sqlite_err("delete checkpoint", e))?; + tx.execute( + "DELETE FROM checkpoint_writes WHERE thread_id = ?1 AND checkpoint_id = ?2", + params![thread_id, id], + ) + .map_err(|e| sqlite_err("delete checkpoint writes", e))?; } tx.commit().map_err(|e| sqlite_err("commit delete", e))?; Ok(removed) } + + async fn put_writes(&self, config: &CheckpointConfig, writes: &[PendingWrite]) -> Result<()> { + let checkpoint_id = super::require_checkpoint_id(config)?; + if writes.is_empty() { + return Ok(()); + } + let namespace_json = serde_json::to_string(&config.namespace) + .map_err(|e| sqlite_err("encode namespace", e))?; + let mut conn = self.lock()?; + let tx = conn + .transaction() + .map_err(|e| sqlite_err("begin put_writes", e))?; + let mut stored = 0usize; + for write in writes { + // The replace-vs-ignore rule pushed into SQL: a control-plane write + // (`idx < 0`) legitimately changes on a retry and upserts, while a + // data write is append-once so a retried `put_writes` is a no-op. + // Doing it with two conflict clauses rather than a read-then-write + // keeps it correct under concurrent writers. + let sql = if write.is_control_plane() { + "INSERT INTO checkpoint_writes + (thread_id, namespace, checkpoint_id, task_id, idx, node, channel, payload) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(thread_id, namespace, checkpoint_id, task_id, idx) DO UPDATE SET + node = excluded.node, + channel = excluded.channel, + payload = excluded.payload" + } else { + "INSERT INTO checkpoint_writes + (thread_id, namespace, checkpoint_id, task_id, idx, node, channel, payload) + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) + ON CONFLICT(thread_id, namespace, checkpoint_id, task_id, idx) DO NOTHING" + }; + let payload = serde_json::to_string(&write.payload) + .map_err(|e| sqlite_err("encode write payload", e))?; + stored += tx + .execute( + sql, + params![ + config.thread_id, + namespace_json, + checkpoint_id, + write.task_id, + write.idx, + write.node.as_str(), + write.channel, + payload, + ], + ) + .map_err(|e| sqlite_err("insert checkpoint write", e))?; + } + tx.commit() + .map_err(|e| sqlite_err("commit put_writes", e))?; + tracing::debug!( + "[checkpoint:sqlite] put_writes thread={} checkpoint={checkpoint_id} offered={} stored={stored}", + config.thread_id, + writes.len() + ); + Ok(()) + } + + async fn get_writes(&self, config: &CheckpointConfig) -> Result> { + let Some(checkpoint_id) = self.resolve_write_target(config).await? else { + return Ok(Vec::new()); + }; + let namespace_json = serde_json::to_string(&config.namespace) + .map_err(|e| sqlite_err("encode namespace", e))?; + let conn = self.lock()?; + let mut stmt = conn + .prepare( + "SELECT node, task_id, idx, channel, payload FROM checkpoint_writes + WHERE thread_id = ?1 AND namespace = ?2 AND checkpoint_id = ?3 + ORDER BY rowid ASC", + ) + .map_err(|e| sqlite_err("prepare get_writes", e))?; + let rows = stmt + .query_map( + params![config.thread_id, namespace_json, checkpoint_id], + map_write_row, + ) + .map_err(|e| sqlite_err("query get_writes", e))?; + let mut out = Vec::new(); + for row in rows { + out.push(row.map_err(|e| sqlite_err("read write row", e))??); + } + Ok(out) + } +} + +/// Decodes one `checkpoint_writes` row into a [`PendingWrite`]. +/// +/// Returns a nested `Result` because the payload decode can fail with a serde +/// error that `rusqlite`'s row-mapper signature has no room for. +#[allow(clippy::type_complexity)] +fn map_write_row(row: &rusqlite::Row<'_>) -> rusqlite::Result> { + let node: String = row.get(0)?; + let task_id: String = row.get(1)?; + let idx: i64 = row.get(2)?; + let channel: String = row.get(3)?; + let payload_json: String = row.get(4)?; + Ok( + match serde_json::from_str::(&payload_json) { + Ok(payload) => Ok(PendingWrite { + node: NodeId::from(node), + task_id, + idx, + channel, + payload, + }), + Err(e) => Err(sqlite_err("decode write payload", e)), + }, + ) +} + +/// Reads every write in `thread_id`/`namespace`, grouped by checkpoint id. +/// +/// One query for the whole lineage, so `state_history` does not issue a +/// `get_writes` per hop. +fn read_writes_by_checkpoint( + conn: &Connection, + thread_id: &str, + namespace_json: &str, +) -> Result>> { + let mut stmt = conn + .prepare( + "SELECT checkpoint_id, node, task_id, idx, channel, payload FROM checkpoint_writes + WHERE thread_id = ?1 AND namespace = ?2 ORDER BY rowid ASC", + ) + .map_err(|e| sqlite_err("prepare writes-by-checkpoint", e))?; + let rows = stmt + .query_map(params![thread_id, namespace_json], |row| { + let checkpoint_id: String = row.get(0)?; + let node: String = row.get(1)?; + let task_id: String = row.get(2)?; + let idx: i64 = row.get(3)?; + let channel: String = row.get(4)?; + let payload_json: String = row.get(5)?; + Ok((checkpoint_id, node, task_id, idx, channel, payload_json)) + }) + .map_err(|e| sqlite_err("query writes-by-checkpoint", e))?; + let mut out: std::collections::HashMap> = + std::collections::HashMap::new(); + for row in rows { + let (checkpoint_id, node, task_id, idx, channel, payload_json) = + row.map_err(|e| sqlite_err("read write row", e))?; + let payload = serde_json::from_str(&payload_json) + .map_err(|e| sqlite_err("decode write payload", e))?; + let write = PendingWrite { + node: NodeId::from(node), + task_id, + idx, + channel, + payload, + }; + let slot = out.entry(checkpoint_id).or_default(); + // `merge_writes` keeps the shared dedupe semantics even here, where the + // primary key already guarantees uniqueness — one rule, one place. + merge_writes(slot, std::slice::from_ref(&write)); + } + Ok(out) } diff --git a/src/graph/checkpoint/test.rs b/src/graph/checkpoint/test.rs index 945c006..c4b2dfc 100644 --- a/src/graph/checkpoint/test.rs +++ b/src/graph/checkpoint/test.rs @@ -92,6 +92,7 @@ fn pending_activation_send_arg_roundtrips() { pending_activations: Some(vec![super::PendingActivation { node: NodeId::from("w"), send_arg: Some(json!({ "item": 42 })), + task_id: "1:0:w".to_string(), }]), barrier_arrivals: vec![super::BarrierArrivals { node: NodeId::from("join"), @@ -465,6 +466,23 @@ mod file_backend { cp.delete_thread("missing").await.unwrap(); } + #[tokio::test] + async fn legacy_uppercase_thread_files_remain_readable_and_copyable() { + let tmp = TempDir::new("legacy-uppercase"); + let cp = FileCheckpointer::::new(tmp.path()); + cp.put(checkpoint("Run", "c1", None, 1)).await.unwrap(); + + // Simulate the pre-upgrade filename scheme, which kept uppercase + // letters unescaped (`Run.jsonl` rather than `%52un.jsonl`). + std::fs::rename(tmp.path().join("%52un.jsonl"), tmp.path().join("Run.jsonl")).unwrap(); + + assert_eq!(cp.get("Run", None).await.unwrap().unwrap().state, 1); + cp.copy_thread("Run", "copy").await.unwrap(); + assert_eq!(cp.get("copy", None).await.unwrap().unwrap().state, 1); + cp.delete_thread("Run").await.unwrap(); + assert!(cp.get("Run", None).await.unwrap().is_none()); + } + #[tokio::test] async fn prune_rewrites_the_thread_file() { let tmp = TempDir::new("prune"); diff --git a/src/graph/checkpoint/types.rs b/src/graph/checkpoint/types.rs index be8fb49..e5796e6 100644 --- a/src/graph/checkpoint/types.rs +++ b/src/graph/checkpoint/types.rs @@ -209,6 +209,13 @@ pub struct PendingActivation { /// packet (plain edge/goto activations carry `None`). #[serde(default, skip_serializing_if = "Option::is_none")] pub send_arg: Option, + /// Stable identity of this scheduled task within its superstep. + /// + /// Unlike `node`, this distinguishes repeated `Send` fan-out activations + /// targeting the same node. Empty on checkpoints written before task + /// identities were persisted. + #[serde(default, skip_serializing_if = "String::is_empty")] + pub task_id: String, } /// The persisted arrivals recorded against one barrier (waiting-edge) join node: @@ -257,15 +264,169 @@ impl Checkpoint { } } +/// The `idx` reserved for a task's **resume** control-plane write. +/// +/// LangGraph reserves *negative* indices for control-plane channels +/// (`WRITES_IDX_MAP`), which is what distinguishes an upsert from an append: +/// see [`PendingWrite::is_control_plane`]. +pub const WRITES_IDX_RESUME: i64 = -1; + +/// The `idx` reserved for a task's **error** control-plane write. +pub const WRITES_IDX_ERROR: i64 = -2; + +/// The `idx` reserved for a task's **interrupt** control-plane write. +pub const WRITES_IDX_INTERRUPT: i64 = -3; + /// A partial write produced by a completed task, preserved across reruns. -#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +/// +/// # Why writes are recorded separately from the checkpoint +/// +/// A superstep can fail *after* some of its tasks have already run. The +/// boundary checkpoint records that those tasks completed, but without a record +/// of what they wrote, a resume has no way to tell "this task already ran" from +/// "this task has not run yet" — so it re-runs them, and any side effect they +/// performed happens twice. Writes are therefore persisted per task through +/// [`Checkpointer::put_writes`](crate::graph::Checkpointer::put_writes) and read +/// back into [`CheckpointTuple::pending_writes`], which is what resume consults +/// to skip already-completed work. +/// +/// # Identity +/// +/// A write is addressed by `(thread_id, namespace, checkpoint_id, task_id, +/// idx)`, mirroring the primary key LangGraph's SQL checkpointers use. Within +/// one checkpoint the `(task_id, idx)` pair is unique: re-putting the same pair +/// never produces a second row. +/// +/// # Control-plane writes upsert; data writes are append-once +/// +/// `idx >= 0` is an ordinary data write, emitted once per task in emission +/// order. Re-putting it is **ignored** (insert-or-ignore), so a retried +/// `put_writes` is idempotent. +/// +/// `idx < 0` marks a control-plane write — resume values, errors, interrupts — +/// which by construction there is at most one of per task and whose value +/// legitimately changes on a retry. Re-putting it **replaces** the stored value +/// (insert-or-replace). Use the `WRITES_IDX_*` constants rather than raw +/// negative numbers. +/// +/// # Back-compatibility +/// +/// `task_id`, `idx` and `channel` carry `#[serde(default)]`, so checkpoint +/// records written before the write protocol existed still deserialize (as an +/// anonymous data write at index `0`). +#[derive(Clone, Debug, PartialEq, serde::Serialize, serde::Deserialize)] pub struct PendingWrite { /// The node that produced the write. pub node: NodeId, + /// The task that produced the write, unique within a superstep. + /// + /// A plain node id is not enough on its own: a fan-out step runs the same + /// node several times with different [`Send`](crate::graph::Send) args, and + /// each of those is a separately resumable task. + #[serde(default)] + pub task_id: String, + /// Position of this write within its task's emission order, or one of the + /// `WRITES_IDX_*` constants for a control-plane write. + #[serde(default)] + pub idx: i64, + /// The channel (state field / node output slot) the write targets. + /// + /// Free-form and backend-opaque; it exists so writes stay distinguishable + /// per channel, and because write isolation is asserted per channel *and* + /// namespace in the conformance suite. + #[serde(default)] + pub channel: String, /// The serialized write payload. + /// + /// May be [`serde_json::Value::Null`] when the producing runtime cannot + /// serialize its update type. The graph executor is in exactly that + /// position — a graph's `Update` carries no `Serialize` bound — so it + /// records writes as *completion markers*: the applied value is already + /// durable in the checkpoint's `state`, and the write record's job is to + /// answer "did this task already run?". pub payload: serde_json::Value, } +impl PendingWrite { + /// Builds an ordinary data write for `task_id` at position `idx`. + pub fn data( + node: impl Into, + task_id: impl Into, + idx: i64, + channel: impl Into, + payload: serde_json::Value, + ) -> Self { + Self { + node: node.into(), + task_id: task_id.into(), + idx, + channel: channel.into(), + payload, + } + } + + /// Builds a completion marker: a data write at index `0` whose payload is + /// `null`, recording only that `task_id` ran to completion. + pub fn completion_marker(node: impl Into, task_id: impl Into) -> Self { + let node = node.into(); + let channel = node.as_str().to_string(); + Self { + node, + task_id: task_id.into(), + idx: 0, + channel, + payload: serde_json::Value::Null, + } + } + + /// Whether this is a control-plane write (`idx < 0`), which upserts rather + /// than appends. See the type docs. + pub fn is_control_plane(&self) -> bool { + self.idx < 0 + } + + /// The `(task_id, idx)` identity pair this write is deduplicated on within + /// a checkpoint. + pub fn identity(&self) -> (&str, i64) { + (self.task_id.as_str(), self.idx) + } +} + +/// Merges `incoming` into `existing`, applying the replace-vs-ignore rule. +/// +/// Shared by every backend so the three of them cannot drift on the one part of +/// the write protocol that is easy to get subtly wrong: +/// +/// - an incoming **control-plane** write (`idx < 0`) replaces any stored write +/// with the same `(task_id, idx)`; +/// - an incoming **data** write (`idx >= 0`) is ignored when that pair is +/// already stored. +/// +/// Returns the number of entries that were actually inserted or replaced, which +/// backends use for logging. +pub fn merge_writes(existing: &mut Vec, incoming: &[PendingWrite]) -> usize { + let mut changed = 0; + for write in incoming { + match existing + .iter_mut() + .find(|w| w.identity() == write.identity()) + { + Some(slot) => { + if write.is_control_plane() { + *slot = write.clone(); + changed += 1; + } + // Data writes are append-once: a duplicate is a no-op. + } + None => { + existing.push(write.clone()); + changed += 1; + } + } + } + changed +} + /// Lightweight checkpoint summary returned by `Checkpointer::list`. /// /// Listing must not require deserializing full graph state, so metadata is kept diff --git a/src/graph/compiled/executor.rs b/src/graph/compiled/executor.rs index d12bbbb..0c0306d 100644 --- a/src/graph/compiled/executor.rs +++ b/src/graph/compiled/executor.rs @@ -182,6 +182,62 @@ where )); } + // Partial-failure guard. The boundary that produced this checkpoint + // recorded a completion marker per task that had already finished; a + // node named by *both* the pending set and that ledger has therefore + // already run, and re-running it would repeat its side effects. On a + // checkpoint the executor itself wrote the two sets are disjoint, so + // this is a no-op — it earns its keep on a checkpoint that was + // hand-built, time-travelled to, or edited through `update_state`, + // where `next_nodes` can legitimately disagree with what ran. + let completed_config = CheckpointConfig { + thread_id: thread_id.to_string(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: self.namespace.clone(), + }; + let recorded = checkpointer.get_writes(&completed_config).await?; + let done: HashSet = if recorded.is_empty() { + checkpoint + .pending_writes + .iter() + .map(|w| w.task_id.clone()) + .collect() + } else { + recorded.iter().map(|w| w.task_id.clone()).collect() + }; + let active: Vec = if done.is_empty() { + active + } else { + let filtered: Vec = active + .iter() + // A node name is not a task identity: a Send fan-out can have + // several live activations of one node. Legacy checkpoints + // have no persisted task id, so leave them runnable. + .filter(|a| a.task_id.is_empty() || !done.contains(&a.task_id)) + .cloned() + .collect(); + if filtered.is_empty() { + // Every pending node claims to have run. Trust the pending set + // rather than turning a resumable checkpoint into a hard error: + // a wrong re-run is recoverable, a stuck thread is not. + tracing::warn!( + "[graph:resume] every pending node of checkpoint `{}` has a completion \ + marker; resuming them anyway rather than stranding the thread", + checkpoint.checkpoint_id + ); + active + } else { + if filtered.len() != active.len() { + tracing::debug!( + "[graph:resume] checkpoint `{}`: skipping {} already-completed task(s)", + checkpoint.checkpoint_id, + active.len() - filtered.len() + ); + } + filtered + } + }; + // The resume value belongs to the node(s) that actually interrupted. The // pending set is deliberately wider than that at an interrupt boundary // (it also carries the successors of branches that completed before the @@ -242,6 +298,7 @@ where active.push(Activation { node, send_arg: input.payload, + task_id: String::new(), }); } if active.is_empty() { @@ -479,6 +536,14 @@ where } } steps += 1; + // Assign identities before any branch runs. A failure checkpoint + // carries these identities with its pending activations, letting a + // later resume skip only the completed fan-out task. + for (index, activation) in active.iter_mut().enumerate() { + if activation.task_id.is_empty() { + activation.task_id = format!("{steps}:{index}:{}", activation.node); + } + } self.emit(GraphEvent::StepStarted { step: steps, active: activation_nodes(&active), @@ -602,7 +667,6 @@ where }; let mut pending = successors; pending.extend(active[failed_index..].iter().cloned()); - let completed_nodes = activation_nodes(&active[..failed_index]); // Settle any in-flight Async background writes before the // failure-boundary persist so earlier boundaries are durable // when the run aborts. Like the persist error below, a @@ -618,7 +682,7 @@ where &run_id, &state, &pending, - &completed_nodes, + &active[..failed_index], &barrier_arrivals, parent_checkpoint.clone(), steps, @@ -706,7 +770,7 @@ where &run_id, &state, &pending, - &activation_nodes(&active[..index]), + &active[..index], vec![emitted.clone()], std::slice::from_ref(&active[index].node), &barrier_arrivals, @@ -759,7 +823,6 @@ where // Select the next active set from commands or static/conditional // edges, evaluated against the freshly-committed state. Barrier // arrivals accumulate into `barrier_arrivals` (persisted below). - let completed_nodes = activation_nodes(&active); let next = match self.route_completed(&active, &goto_map, &state, &mut barrier_arrivals) { Ok(next) => next, @@ -811,7 +874,7 @@ where &run_id, &state, &next, - &completed_nodes, + &active, &barrier_arrivals, parent_checkpoint.clone(), steps, @@ -842,7 +905,7 @@ where &run_id, &state, &next, - &completed_nodes, + &active, Vec::new(), &[], &barrier_arrivals, @@ -983,7 +1046,7 @@ where run_id: &RunId, state: &State, pending: &[Activation], - completed_tasks: &[NodeId], + completed_tasks: &[Activation], barrier_arrivals: &HashMap>, parent: Option, step: usize, @@ -1003,8 +1066,8 @@ where namespace: self.namespace.clone(), state: state.clone(), next_nodes: activation_nodes(pending), - completed_tasks: completed_tasks.to_vec(), - pending_writes: Vec::new(), + completed_tasks: activation_nodes(completed_tasks), + pending_writes: Self::completion_writes(completed_tasks, step), interrupts: Vec::new(), pending_activations: Some(pending.iter().map(PendingActivation::from).collect()), barrier_arrivals: barriers_to_persisted(barrier_arrivals), @@ -1017,7 +1080,17 @@ where "error": error.to_string(), }), }; + let writes = checkpoint.pending_writes.clone(); + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; let id = checkpointer.put(checkpoint).await?; + // Also record the ledger through the write protocol, so backends that + // implement it can answer "did this task run?" without loading the + // whole state payload. + checkpointer.put_writes(&config, &writes).await?; self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id.clone(), }); @@ -1458,7 +1531,7 @@ where run_id: &RunId, state: &State, pending: &[Activation], - completed_tasks: &[NodeId], + completed_tasks: &[Activation], interrupts: Vec, interrupted: &[NodeId], barrier_arrivals: &HashMap>, @@ -1486,7 +1559,14 @@ where recursion, child_runs, ); + let writes = checkpoint.pending_writes.clone(); + let config = CheckpointConfig { + thread_id: checkpoint.thread_id.clone(), + checkpoint_id: Some(checkpoint.checkpoint_id.clone()), + namespace: checkpoint.namespace.clone(), + }; let id = checkpointer.put(checkpoint).await?; + checkpointer.put_writes(&config, &writes).await?; self.emit(GraphEvent::CheckpointSaved { checkpoint_id: id.clone(), }); @@ -1521,7 +1601,7 @@ where run_id: &RunId, state: &State, pending: &[Activation], - completed_tasks: &[NodeId], + completed_tasks: &[Activation], barrier_arrivals: &HashMap>, parent: Option, step: usize, @@ -1573,6 +1653,35 @@ where } } + /// Records completion markers for the tasks that finished in the step a + /// boundary checkpoint closes. + /// + /// A graph's `Update` carries no `Serialize` bound, so the executor cannot + /// persist *what* a task wrote — but it does not need to: the applied value + /// is already durable in the checkpoint's `state`. What was missing was the + /// other half, the per-task record of *that* it ran, which is what lets a + /// resume distinguish "already done" from "not yet started". See + /// [`PendingWrite`](crate::graph::checkpoint::PendingWrite)'s docs for why + /// that distinction is the whole point of + /// the ledger. + /// + /// The task id is persisted on the activation itself, so a resume can + /// match a marker to one fan-out task rather than every task with its node. + fn completion_writes( + completed_tasks: &[Activation], + _step: usize, + ) -> Vec { + completed_tasks + .iter() + .map(|activation| { + crate::graph::checkpoint::PendingWrite::completion_marker( + activation.node.clone(), + activation.task_id.clone(), + ) + }) + .collect() + } + /// Builds the loop-boundary [`Checkpoint`] record shared by the sync and /// async persist paths, minting a fresh checkpoint id. #[allow(clippy::too_many_arguments)] @@ -1582,7 +1691,7 @@ where run_id: &RunId, state: &State, pending: &[Activation], - completed_tasks: &[NodeId], + completed_tasks: &[Activation], interrupts: Vec, interrupted: &[NodeId], barrier_arrivals: &HashMap>, @@ -1617,8 +1726,8 @@ where namespace: self.namespace.clone(), state: state.clone(), next_nodes: activation_nodes(pending), - completed_tasks: completed_tasks.to_vec(), - pending_writes: Vec::new(), + completed_tasks: activation_nodes(completed_tasks), + pending_writes: Self::completion_writes(completed_tasks, step), pending_activations: Some(pending.iter().map(PendingActivation::from).collect()), barrier_arrivals: barriers_to_persisted(barrier_arrivals), interrupts, diff --git a/src/graph/compiled/mod.rs b/src/graph/compiled/mod.rs index c7c5b31..b7fbaf1 100644 --- a/src/graph/compiled/mod.rs +++ b/src/graph/compiled/mod.rs @@ -182,6 +182,7 @@ struct StepFailure { struct Activation { node: NodeId, send_arg: Option, + task_id: String, } impl Activation { @@ -189,6 +190,7 @@ impl Activation { Self { node, send_arg: None, + task_id: String::new(), } } } @@ -198,6 +200,7 @@ impl From<&Activation> for PendingActivation { PendingActivation { node: a.node.clone(), send_arg: a.send_arg.clone(), + task_id: a.task_id.clone(), } } } @@ -207,6 +210,7 @@ impl From<&PendingActivation> for Activation { Activation { node: p.node.clone(), send_arg: p.send_arg.clone(), + task_id: p.task_id.clone(), } } } diff --git a/src/graph/compiled/routing.rs b/src/graph/compiled/routing.rs index e555bf7..e3dd212 100644 --- a/src/graph/compiled/routing.rs +++ b/src/graph/compiled/routing.rs @@ -57,11 +57,13 @@ where next.push(Activation { node: tnode, send_arg, + task_id: String::new(), }); } else if next_seen.insert(tnode.clone()) { next.push(Activation { node: tnode, send_arg: None, + task_id: String::new(), }); } } @@ -132,6 +134,7 @@ where next.push(Activation { node: relief.barrier_node.clone(), send_arg: None, + task_id: String::new(), }); } } diff --git a/src/graph/compiled/state_api.rs b/src/graph/compiled/state_api.rs index 188d392..a531391 100644 --- a/src/graph/compiled/state_api.rs +++ b/src/graph/compiled/state_api.rs @@ -196,6 +196,7 @@ where merged.push(Activation { node: tnode, send_arg, + task_id: String::new(), }); } } diff --git a/src/graph/subagent_node/test.rs b/src/graph/subagent_node/test.rs index 42e3802..52e1108 100644 --- a/src/graph/subagent_node/test.rs +++ b/src/graph/subagent_node/test.rs @@ -39,6 +39,7 @@ fn tool_call_response(id: &str, name: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/src/graph/testkit/conformance.rs b/src/graph/testkit/conformance.rs index d1f4ad3..f933910 100644 --- a/src/graph/testkit/conformance.rs +++ b/src/graph/testkit/conformance.rs @@ -362,3 +362,397 @@ where "pending → running → completed history replays" ); } + +// ── Checkpointer: lineage, thread ops and pending writes ───────────────────── +// +// The original `checkpointer_contract` covered put/get/list/list_threads/ +// delete_thread/prune and nothing else — no `get_tuple`, no `state_history`, no +// `copy_thread`, no `delete_checkpoints`, and nothing at all about pending +// writes. That gap is not incidental: it is precisely why a dead +// `pending_writes` field, an unguarded lineage walk and a `copy_thread` that +// silently corrupts a non-empty target all survived in-tree. These suites port +// the invariants from LangGraph's `checkpoint-conformance` spec. + +use crate::graph::checkpoint::{CheckpointConfig, PendingWrite, WRITES_IDX_RESUME}; + +/// Builds a checkpoint in `namespace` rather than at the root. +fn scoped_checkpoint( + thread: &str, + id: &str, + parent: Option<&str>, + step: usize, + namespace: &[&str], +) -> Checkpoint { + let mut checkpoint = contract_checkpoint(thread, id, parent, step); + checkpoint.namespace = namespace.iter().map(|s| s.to_string()).collect(); + checkpoint +} + +fn config(thread: &str, id: &str, namespace: &[&str]) -> CheckpointConfig { + CheckpointConfig { + thread_id: thread.to_string(), + checkpoint_id: Some(id.to_string()), + namespace: namespace.iter().map(|s| s.to_string()).collect(), + } +} + +/// Runs the **pending-writes** contract against `cp`. +/// +/// Pins the partial-failure protocol: idempotent data writes, upserting +/// control-plane writes, namespace isolation, an empty ledger for a fresh +/// checkpoint, and removal along with the checkpoint/thread they belong to. +/// +/// A backend that does not implement the protocol at all (the trait's default +/// no-op bodies) will fail the very first assertion — which is the point: the +/// suite certifies the protocol, it does not quietly accept its absence. +pub async fn checkpointer_writes_contract(cp: C) +where + C: Checkpointer, +{ + let thread = "writes"; + cp.put(contract_checkpoint(thread, "w1", None, 1)) + .await + .expect("put w1"); + + // A brand-new checkpoint carries no writes. + let fresh = cp + .get_tuple(config(thread, "w1", &[])) + .await + .expect("get_tuple") + .expect("some"); + assert!( + fresh.pending_writes.is_empty(), + "a newly written checkpoint starts with an empty pending-writes ledger" + ); + + // Data writes land and read back in order. + let writes = vec![ + PendingWrite::data("n", "task-a", 0, "out", serde_json::json!("first")), + PendingWrite::data("n", "task-a", 1, "out", serde_json::json!("second")), + ]; + cp.put_writes(&config(thread, "w1", &[]), &writes) + .await + .expect("put_writes"); + let stored = cp + .get_writes(&config(thread, "w1", &[])) + .await + .expect("get_writes"); + assert_eq!(stored.len(), 2, "both data writes stored"); + + // Re-putting the same `(task_id, idx)` data writes must not duplicate them. + cp.put_writes(&config(thread, "w1", &[]), &writes) + .await + .expect("duplicate put_writes"); + let stored = cp + .get_writes(&config(thread, "w1", &[])) + .await + .expect("get_writes after duplicate"); + assert_eq!( + stored.len(), + 2, + "a duplicate (task_id, idx) data write is ignored, not appended" + ); + assert_eq!( + stored[0].payload, + serde_json::json!("first"), + "an ignored duplicate leaves the original value in place" + ); + + // A control-plane write (negative idx) upserts: exactly one row, latest value. + for value in ["v1", "v2"] { + cp.put_writes( + &config(thread, "w1", &[]), + &[PendingWrite::data( + "n", + "task-a", + WRITES_IDX_RESUME, + "__resume__", + serde_json::json!(value), + )], + ) + .await + .expect("put control-plane write"); + } + let stored = cp + .get_writes(&config(thread, "w1", &[])) + .await + .expect("get_writes after control-plane"); + let control: Vec<&PendingWrite> = stored.iter().filter(|w| w.is_control_plane()).collect(); + assert_eq!( + control.len(), + 1, + "a control-plane write stores exactly one row (upsert, not append)" + ); + assert_eq!( + control[0].payload, + serde_json::json!("v2"), + "the control-plane upsert keeps the latest value" + ); + + // `get_tuple` surfaces the ledger. + let tuple = cp + .get_tuple(config(thread, "w1", &[])) + .await + .expect("get_tuple") + .expect("some"); + assert_eq!( + tuple.pending_writes.len(), + stored.len(), + "get_tuple surfaces the persisted pending writes" + ); + + // Namespace isolation: same thread, same checkpoint id, same channel, but a + // different namespace is a different ledger. + cp.put(scoped_checkpoint(thread, "w1", None, 1, &["child"])) + .await + .expect("put child-namespace w1"); + cp.put_writes( + &config(thread, "w1", &["child"]), + &[PendingWrite::data( + "n", + "task-a", + 0, + "out", + serde_json::json!("child"), + )], + ) + .await + .expect("put child-namespace writes"); + let child = cp + .get_writes(&config(thread, "w1", &["child"])) + .await + .expect("child writes"); + assert_eq!(child.len(), 1, "the child namespace has its own ledger"); + assert_eq!(child[0].payload, serde_json::json!("child")); + let root = cp + .get_writes(&config(thread, "w1", &[])) + .await + .expect("root writes"); + assert_eq!( + root.len(), + stored.len(), + "writing the child namespace did not touch the root ledger" + ); + + // `delete_checkpoints` takes the checkpoint's writes with it. + cp.put(contract_checkpoint(thread, "w2", Some("w1"), 2)) + .await + .expect("put w2"); + cp.put_writes( + &config(thread, "w2", &[]), + &[PendingWrite::completion_marker("n", "task-b")], + ) + .await + .expect("put w2 writes"); + cp.delete_checkpoints(thread, &["w2".to_string()]) + .await + .expect("delete_checkpoints"); + assert!( + cp.get_writes(&config(thread, "w2", &[])) + .await + .expect("writes after delete_checkpoints") + .is_empty(), + "deleting a checkpoint deletes its writes" + ); + + // `delete_thread` clears every namespace's ledger, not just the root. + cp.delete_thread(thread).await.expect("delete_thread"); + for namespace in [&[][..], &["child"][..]] { + assert!( + cp.get_writes(&config(thread, "w1", namespace)) + .await + .expect("writes after delete_thread") + .is_empty(), + "delete_thread removes the writes of every namespace ({namespace:?})" + ); + } +} + +/// Runs the **lineage and thread-operation** contract against `cp`. +/// +/// Covers `get_tuple` addressing, `state_history` ordering/limit/cycle safety, +/// `copy_thread` semantics (including the empty-target rule), and per-namespace +/// pruning. +pub async fn checkpointer_lineage_contract(cp: C) +where + C: Checkpointer, +{ + let thread = "lineage"; + for i in 1..=3 { + let parent = (i > 1).then(|| format!("c{}", i - 1)); + cp.put(contract_checkpoint( + thread, + &format!("c{i}"), + parent.as_deref(), + i, + )) + .await + .expect("put lineage"); + } + // An embedded subgraph shares the thread but keeps its own lineage. + cp.put(scoped_checkpoint(thread, "s1", None, 1, &["child"])) + .await + .expect("put child s1"); + cp.put(scoped_checkpoint(thread, "s2", Some("s1"), 2, &["child"])) + .await + .expect("put child s2"); + + // `get_tuple` resolves the addressing config and the parent config. + let tuple = cp + .get_tuple(config(thread, "c2", &[])) + .await + .expect("get_tuple") + .expect("some"); + assert_eq!(tuple.config.checkpoint_id.as_deref(), Some("c2")); + assert_eq!( + tuple + .parent_config + .as_ref() + .and_then(|c| c.checkpoint_id.as_deref()), + Some("c1"), + "get_tuple carries the parent's config" + ); + // A `None` id addresses the latest checkpoint *in the namespace*. + let latest = cp + .get_tuple(CheckpointConfig::latest(thread)) + .await + .expect("get_tuple latest") + .expect("some"); + assert_eq!( + latest.checkpoint.checkpoint_id, "c3", + "an unaddressed get_tuple resolves the latest checkpoint in the namespace, \ + not the thread's global latest" + ); + + // `state_history` is strictly newest-first along the parent spine, scoped + // to the namespace, and `limit` caps from the newest end. + let history = cp + .state_history(thread, &[], None) + .await + .expect("state_history"); + let ids: Vec<&str> = history + .iter() + .map(|t| t.checkpoint.checkpoint_id.as_str()) + .collect(); + assert_eq!(ids, vec!["c3", "c2", "c1"], "state_history is newest-first"); + let limited = cp + .state_history(thread, &[], Some(2)) + .await + .expect("state_history limited"); + assert_eq!( + limited + .iter() + .map(|t| t.checkpoint.checkpoint_id.as_str()) + .collect::>(), + vec!["c3", "c2"], + "limit keeps the most recent hops" + ); + let child_history = cp + .state_history(thread, &["child".to_string()], None) + .await + .expect("child state_history"); + assert_eq!( + child_history + .iter() + .map(|t| t.checkpoint.checkpoint_id.as_str()) + .collect::>(), + vec!["s2", "s1"], + "state_history is namespace-scoped" + ); + + // A self-referential parent pointer must terminate the walk rather than + // spin forever. `parent_checkpoint_id` is caller data, not a structurally + // guaranteed DAG edge. + let cyclic_thread = "cycle"; + let mut cyclic = contract_checkpoint(cyclic_thread, "loop1", None, 1); + cyclic.parent_checkpoint_id = Some("loop1".to_string()); + cp.put(cyclic).await.expect("put cyclic"); + let walked = cp + .state_history(cyclic_thread, &[], None) + .await + .expect("state_history over a cycle terminates"); + assert_eq!( + walked.len(), + 1, + "a lineage cycle is visited once, not forever" + ); + + // `copy_thread` preserves order, ids and namespaces, and leaves the source + // untouched. + cp.copy_thread(thread, "lineage-copy") + .await + .expect("copy_thread"); + let source_after: Vec = cp + .list(thread) + .await + .expect("list source") + .into_iter() + .map(|m| m.checkpoint_id) + .collect(); + let copied = cp.list("lineage-copy").await.expect("list copy"); + assert_eq!( + source_after.len(), + 5, + "copy_thread leaves the source thread unchanged" + ); + assert_eq!( + copied + .iter() + .map(|m| m.checkpoint_id.clone()) + .collect::>(), + source_after, + "copy_thread preserves checkpoint ids and their order" + ); + assert!( + copied + .iter() + .any(|m| m.namespace == vec!["child".to_string()]), + "copy_thread preserves namespaces" + ); + + // Copying onto a non-empty thread is rejected: reusing ids across two + // lineages in one thread is silent corruption, not a merge. + let err = cp.copy_thread(thread, "lineage-copy").await; + assert!( + err.is_err(), + "copy_thread into a non-empty target must be rejected, not interleaved" + ); + + // `prune` keeps the latest per namespace (plus ancestors), so an embedded + // subgraph's lineage survives a prune driven by the parent's recency. + cp.prune(thread, 1).await.expect("prune"); + let kept = cp.list(thread).await.expect("list pruned"); + let namespaces: Vec<&Vec> = kept.iter().map(|m| &m.namespace).collect(); + assert!( + namespaces.iter().any(|ns| ns.is_empty()), + "prune keeps the root namespace's latest" + ); + assert!( + namespaces + .iter() + .any(|ns| ns.as_slice() == ["child".to_string()]), + "prune keeps the latest of EVERY namespace, not a thread-wide window" + ); + + // A pruned-but-kept checkpoint keeps its writes. + let survivor = kept + .iter() + .find(|m| m.namespace.is_empty()) + .expect("a root-namespace survivor"); + let survivor_config = config(thread, &survivor.checkpoint_id, &[]); + cp.put_writes( + &survivor_config, + &[PendingWrite::completion_marker("n", "task-keep")], + ) + .await + .expect("put survivor writes"); + cp.prune(thread, 1).await.expect("prune again"); + assert_eq!( + cp.get_writes(&survivor_config) + .await + .expect("survivor writes") + .len(), + 1, + "prune preserves the writes of the checkpoints it keeps" + ); +} diff --git a/src/harness/agent_loop/entry.rs b/src/harness/agent_loop/entry.rs index e927f7d..fe839b7 100644 --- a/src/harness/agent_loop/entry.rs +++ b/src/harness/agent_loop/entry.rs @@ -168,13 +168,65 @@ impl AgentHarness { /// [`crate::harness::model::ChatModel::stream`] (firing `on_model_delta` /// middleware per delta) or the unary /// [`crate::harness::model::ChatModel::invoke`] path. + /// Runs the loop and returns the accumulated run **and** any error, instead + /// of discarding the run when the loop fails. + /// + /// [`AgentHarness::invoke`] and friends return `Err` on failure, which drops + /// the partially-populated [`AgentRun`] — every message, tool result, and + /// usage figure the run produced before it tripped a limit or hit a tool + /// failure. Use this when that partial work is worth keeping: to inspect + /// what the agent had done, to repair the transcript, or to resume from it. + /// + /// The returned [`PartialRunOutcome::error`] is `None` exactly when the run + /// succeeded. + pub async fn invoke_collecting_partial( + &self, + state: &State, + ctx_data: Ctx, + config: RunConfig, + input: Vec, + ) -> PartialRunOutcome { + let ctx = RunContext::new(config, ctx_data); + self.drive_collecting(state, ctx, input, false).await + } + + /// [`AgentHarness::invoke_collecting_partial`] against a caller-supplied + /// [`RunContext`]. + pub async fn invoke_in_context_collecting_partial( + &self, + state: &State, + ctx: RunContext, + input: Vec, + ) -> PartialRunOutcome { + self.drive_collecting(state, ctx, input, false).await + } + async fn drive( &self, state: &State, - mut ctx: RunContext, + ctx: RunContext, input: Vec, streaming: bool, ) -> Result { + let outcome = self.drive_collecting(state, ctx, input, streaming).await; + match outcome.error { + Some(error) => Err(error), + None => Ok(AgentLoopResult { + run: outcome.run, + status: outcome.status, + }), + } + } + + /// The shared driver both entry shapes delegate to. Never discards the + /// run, so the failing path can hand the partial transcript back. + async fn drive_collecting( + &self, + state: &State, + mut ctx: RunContext, + input: Vec, + streaming: bool, + ) -> PartialRunOutcome { let run_id = ctx.config.run_id.clone(); let thread_id = ctx.config.thread_id.clone(); // Record the drive mode so tool execution contexts (and the sub-agents @@ -194,8 +246,19 @@ impl AgentHarness { .await { Ok(()) => { - status.mark_completed(); - Ok(AgentLoopResult { run, status }) + // A paused run is resumable, not finished: reporting it + // `completed` is what made "paused for a human" look identical + // to "the model produced an empty final answer". + if run.paused.is_some() { + status.mark_interrupted(); + } else { + status.mark_completed(); + } + PartialRunOutcome { + run, + status, + error: None, + } } Err(error) => { let record = ctx.emit(AgentEvent::RunFailed { @@ -212,7 +275,11 @@ impl AgentHarness { if !ctx.take_on_error_dispatched() { let _ = self.middleware.run_on_error(&mut ctx, &error).await; } - Err(error) + PartialRunOutcome { + run, + status, + error: Some(error), + } } } } diff --git a/src/harness/agent_loop/mod.rs b/src/harness/agent_loop/mod.rs index cb9f922..bfb567a 100644 --- a/src/harness/agent_loop/mod.rs +++ b/src/harness/agent_loop/mod.rs @@ -44,15 +44,40 @@ //! # Limits //! //! Model and tool caps are enforced by the run context's own -//! [`crate::harness::limits::LimitTracker`], which is synced with -//! [`RunPolicy::limits`][crate::harness::runtime::RunPolicy] once at the start -//! of each run (see [`crate::harness::limits::LimitTracker::sync_call_limits`]) -//! so the harness policy and the per-run [`RunConfig`] agree on a single -//! enforced cap instead of silently disagreeing. Each call is checked -//! *before* it is made, returning [`TinyAgentsError::LimitExceeded`] whose -//! message always names the limit that actually tripped. The wall-clock -//! deadline (from the run config) is checked each iteration and surfaces as -//! [`TinyAgentsError::Timeout`]. +//! [`crate::harness::limits::LimitTracker`], reconciled once at the start of +//! each run with [`RunPolicy::limits`][crate::harness::runtime::RunPolicy] so +//! the harness policy and the per-run [`RunConfig`] agree on a single enforced +//! cap instead of silently disagreeing. +//! +//! The reconciliation is **asymmetric**, which is why [`RunConfig`]'s caps are +//! `Option`: an explicitly-set config cap is the caller's ceiling and +//! the stricter of the two wins (fail-closed), while an unset cap merely +//! defaulted and the policy wins outright — including when it raises the cap +//! above the crate default. Both cases are legitimate and they are +//! indistinguishable while the config cap is a bare `usize`, which is how +//! `RunConfig::new("r").with_max_model_calls(2)` came to run 25 model calls. +//! +//! Each call is checked *before* it is made. Under the default +//! [`LimitBehavior::Error`][crate::harness::limits::LimitBehavior] exhaustion +//! returns [`TinyAgentsError::LimitExceeded`] whose message always names the +//! limit that actually tripped; under +//! [`LimitBehavior::StopWithPartial`][crate::harness::limits::LimitBehavior] the +//! loop instead emits [`AgentEvent::LimitReached`], stops, and finalizes with +//! everything the run produced. The wall-clock deadline (from the run config) is +//! checked each iteration and surfaces as [`TinyAgentsError::Timeout`]. +//! +//! # Exits +//! +//! The loop distinguishes three deliberate stops (see the private `LoopExit`): +//! a normal finish, a `StopWithPartial` limit stop, and a **steering pause**. +//! A pause is not a completion: it sets +//! [`AgentRun::paused`][crate::harness::middleware::AgentRun::paused], reports +//! the run interrupted, and leaves the pause latched on the +//! [`SteeringHandle`][crate::harness::steering::SteeringHandle] so a later +//! `Resume` lifts it. The working transcript is written onto the +//! [`AgentRun`] on **every** exit path, including errors, so a failed run's +//! partial conversation can still be inspected or repaired (see +//! [`AgentHarness::invoke_collecting_partial`]). //! //! # Backoff //! diff --git a/src/harness/agent_loop/model_call.rs b/src/harness/agent_loop/model_call.rs index 96bb251..01272a8 100644 --- a/src/harness/agent_loop/model_call.rs +++ b/src/harness/agent_loop/model_call.rs @@ -7,6 +7,9 @@ //! the full loop lifecycle, limits, and backoff design. use super::*; +use crate::harness::cache::{ + CachePolicy, CacheSkipReason, apply_prompt_cache_breakpoints, scoped_cache_key, +}; impl AgentHarness { /// Invokes a model, consulting the local response cache around the @@ -20,14 +23,38 @@ impl AgentHarness { /// [`AgentEvent::CacheMiss`] is emitted, the provider is invoked normally, /// and the successful response is written back to the cache. /// + /// # Key composition + /// + /// The key is **not** the request hash alone. [`Self::response_cache_decision`] + /// produces the request half ([`cache_key`]); this method folds in the + /// *resolved* model's + /// [`cache_identity`][crate::harness::model::ChatModel::cache_identity], the + /// `streaming` flag, and the policy namespace via + /// [`scoped_cache_key`][crate::harness::cache::scoped_cache_key]. All three + /// were previously absent from the key: + /// + /// * the request's `model` field is never set by the loop and the endpoint + /// and credential live inside the `Arc`, so one shared + /// cache served a hosted harness's answer to a local one; + /// * `streaming` is a parameter of this function, not a request field, so a + /// warm streaming run could be served an entry written by a unary run. + /// /// # Accounting /// /// A cache hit is still counted as a model "step"/call by the caller /// ([`Self::run_loop`] increments `model_calls`/`steps` and emits - /// [`AgentEvent::ModelCompleted`] after this returns) so usage and limit - /// bookkeeping stay consistent whether or not a call was served from cache. - /// The behavioral guarantee is only that the underlying provider is not - /// contacted on a hit. + /// [`AgentEvent::ModelCompleted`] after this returns) so limit bookkeeping + /// stays consistent whether or not a call was served from cache. The hit is + /// stamped [`ModelResponse::served_from_cache`] so *token/cost* accounting + /// can tell a replay from a real call and not re-bill it. + /// + /// # Failure policy + /// + /// A cache is an optimization. Neither a failed lookup nor a failed write + /// fails the run: a read error degrades to a miss, and a write error is + /// logged and dropped — the provider call has already succeeded and been + /// paid for, so discarding its answer because the cache was unavailable is + /// strictly worse than not caching at all. async fn invoke_model_with_retry( &self, state: &State, @@ -37,17 +64,61 @@ impl AgentHarness { binding: ResolvedModelBinding, streaming: bool, ) -> Result { - let decision = self.response_cache_decision(request); + let policy = self.effective_cache_policy(request); + // The identity of the model that is actually about to be called — known + // only *after* resolution, which is why the key cannot be finalized by + // the request-hashing half alone. + let identity = binding.model.cache_identity(); + let primary_name = binding.resolved.name.clone(); + + let decision = self.response_cache_decision(request).map(|(cache, base)| { + let key = scoped_cache_key( + &base, + identity.as_deref(), + streaming, + policy.namespace.as_deref(), + ); + (cache, key) + }); + + if decision.is_none() { + let reason = self.cache_skip_reason(request); + tracing::debug!( + call_id = %call_id.as_str(), + reason = reason.as_str(), + "[cache] response cache not consulted for this model call" + ); + } if let Some((cache, key)) = decision.as_ref() { - if let Some(mut cached) = cache.get(key).await? { + // A read failure is a miss, not a run failure: `InMemoryResponseCache` + // reports a poisoned mutex as a `Validation` error, and the trait is + // explicitly designed for third-party implementations whose failure + // modes we do not control. + let looked_up = match cache.get(key).await { + Ok(hit) => hit, + Err(error) => { + tracing::warn!( + call_id = %call_id.as_str(), + %error, + "[cache] response-cache lookup failed; treating as a miss" + ); + None + } + }; + if let Some(mut cached) = looked_up { ctx.emit(AgentEvent::CacheHit { call_id: call_id.clone(), key: key.clone(), }); + cached.served_from_cache = true; if cached.resolved_model.is_none() { cached.resolved_model = Some(binding.resolved.clone()); } + if streaming { + self.replay_cached_response_as_deltas(state, ctx, call_id, &cached) + .await?; + } return Ok(cached); } ctx.emit(AgentEvent::CacheMiss { @@ -56,17 +127,156 @@ impl AgentHarness { }); } + // Provider prompt-cache breakpoints are injected *after* the key is + // derived (they mutate `provider_options`, which the key covers) and + // only when the policy asks for prefix protection, so the common path + // never pays for a request clone. + let mut breakpointed; + let effective_request = if policy.protect_prompt_prefix { + breakpointed = request.clone(); + apply_prompt_cache_breakpoints(&mut breakpointed); + &breakpointed + } else { + request + }; + let response = self - .invoke_model_resolving(state, ctx, request, call_id, binding, streaming) + .invoke_model_resolving(state, ctx, effective_request, call_id, binding, streaming) .await?; if let Some((cache, key)) = decision.as_ref() { - cache.put(key, response.clone()).await?; + // Only the *primary* model's answer may be stored under this key. + // `invoke_model_resolving` walks the fallback chain on failure and + // can return a different model's response; writing that under the + // primary's key poisons it — permanently, when no TTL is set — so + // every later run of the primary silently gets the fallback's + // answer, and reports model B after `ModelStarted` announced A. + let served_by = response + .resolved_model + .as_ref() + .map(|resolved| resolved.name.as_str()); + if served_by.is_some_and(|name| name != primary_name) { + tracing::debug!( + call_id = %call_id.as_str(), + primary = %primary_name, + served_by = served_by.unwrap_or_default(), + "[cache] skipping write: a fallback model answered, not the keyed model" + ); + } else if let Err(error) = cache + .put_with_ttl(key, response.clone(), policy.ttl()) + .await + { + // The provider call already succeeded and was paid for. + // Discarding its answer because the cache is unavailable would + // be strictly worse than not caching. + tracing::warn!( + call_id = %call_id.as_str(), + %error, + "[cache] response-cache write failed; returning the response uncached" + ); + } } Ok(response) } + /// Resolves the effective [`CachePolicy`] for `request`: the per-request + /// policy when present, otherwise the harness-level + /// [`RunPolicy::cache`][crate::harness::runtime::RunPolicy]. + pub(super) fn effective_cache_policy(&self, request: &ModelRequest) -> CachePolicy { + request + .cache_policy + .clone() + .unwrap_or_else(|| self.policy.cache.clone()) + } + + /// Explains why [`Self::response_cache_decision`] declined to consult the + /// cache, for the diagnostic log line on the skip path. + /// + /// `docs/modules/harness/cache.md` specifies a richer decision surface than + /// the bare `CacheHit`/`CacheMiss` pair; without this a caller seeing a 0% + /// hit rate cannot tell "no cache attached" from "policy off" from "every + /// request was multi-turn". + pub(super) fn cache_skip_reason(&self, request: &ModelRequest) -> CacheSkipReason { + if self.response_cache.is_none() { + return CacheSkipReason::NoCacheAttached; + } + if !self.effective_cache_policy(request).response_cache_enabled { + return CacheSkipReason::PolicyDisabled; + } + CacheSkipReason::MultiTurnTranscript + } + + /// Replays a cache hit as synthetic stream deltas so a warm streaming run + /// is observationally identical to a cold one. + /// + /// A cache hit short-circuits before [`Self::invoke_model_streaming_once`], + /// so a streaming run served from cache used to emit **zero** + /// [`AgentEvent::ModelDelta`] events and run **zero** + /// [`on_model_delta`][crate::harness::middleware::Middleware::on_model_delta] + /// hooks — a UI rendering deltas showed nothing at all, contradicting the + /// streaming contract documented on the harness entry points. LangChain + /// replays hits as synthetic stream events for exactly this reason. + /// + /// The replay emits one text delta (when the cached message has text) and + /// one delta per cached tool call, mirroring what the provider stream would + /// have produced. It is a replay, not a re-derivation: no provider is + /// contacted. + async fn replay_cached_response_as_deltas( + &self, + state: &State, + ctx: &mut RunContext, + call_id: &CallId, + cached: &ModelResponse, + ) -> Result<()> { + let text = cached.text(); + let tool_calls = cached.tool_calls().to_vec(); + tracing::debug!( + call_id = %call_id.as_str(), + text_len = text.len(), + tool_calls = tool_calls.len(), + "[cache] replaying a cache hit as synthetic stream deltas" + ); + + let mut deltas: Vec = Vec::new(); + if !text.is_empty() { + deltas.push(MessageDelta { + text, + reasoning: String::new(), + tool_call: None, + }); + } + for call in &tool_calls { + deltas.push(MessageDelta { + text: String::new(), + reasoning: String::new(), + tool_call: Some(crate::harness::tool::ToolDelta { + call_id: call.id.clone(), + content: serde_json::to_string(&call.arguments).unwrap_or_default(), + tool_name: Some(call.name.clone()), + }), + }); + } + + for delta in deltas { + let mut model_delta = ModelDelta { + call_id: call_id.as_str().to_string(), + content: delta.text.clone(), + reasoning: delta.reasoning.clone(), + tool_call: delta.tool_call.clone(), + }; + ctx.emit(AgentEvent::ModelDelta { + run_id: ctx.config.run_id.clone(), + call_id: call_id.clone(), + delta, + }); + self.middleware + .run_on_model_delta(ctx, state, &mut model_delta) + .await?; + } + Ok(()) + } + /// Invokes a model with retry and fallback (no caching). /// /// Retries are governed by [`RunPolicy::retry`][crate::harness::runtime::RunPolicy] @@ -99,6 +309,14 @@ impl AgentHarness { loop { // Retry loop for the current model. let mut attempt = 0usize; + // Counts the deltas the *current* streaming attempt has already + // handed to consumers. A stream that dies after 200 tokens has + // already delivered them; the retry replays from scratch, so a UI + // concatenating `ModelDelta.text` renders partial garbage followed + // by the full answer. `StreamAccumulator` is discarded internally, + // but consumers are never told to discard too — see the warning + // below and the `AgentEvent` handoff noted in the module docs. + let mut deltas_emitted = 0usize; let outcome = loop { // Observe cancellation before (re)issuing a model attempt so a // cancel requested during a retry/rate-limit wait stops the run @@ -116,8 +334,14 @@ impl AgentHarness { // awaited unbounded. let remaining = self.call_budget(ctx); let attempt_result = if streaming { - let fut = - self.invoke_model_streaming_once(state, ctx, &model, request, call_id); + let fut = self.invoke_model_streaming_once( + state, + ctx, + &model, + request, + call_id, + &mut deltas_emitted, + ); Self::with_call_budget(remaining, run_id.as_str(), "model call", fut).await } else { // Race the wall-clock-bounded unary call against cooperative @@ -164,6 +388,23 @@ impl AgentHarness { // whole exponential schedule one step too high. let backoff_attempt = attempt; attempt += 1; + if streaming && deltas_emitted > 0 { + // The retry re-emits the whole response from + // the beginning. Until `AgentEvent` grows a + // dedicated discard marker, `RetryScheduled` + // for a streaming call *is* the signal that + // every delta seen so far for this `call_id` + // must be dropped. + tracing::warn!( + call_id = %call_id.as_str(), + discarded_deltas = deltas_emitted, + attempt, + "[stream] retrying a streaming call that already emitted \ + deltas; consumers must discard everything received so far \ + for this call_id" + ); + } + deltas_emitted = 0; ctx.emit(AgentEvent::RetryScheduled { call_id: call_id.clone(), attempt, @@ -171,7 +412,10 @@ impl AgentHarness { // Sleep for the backoff only when the policy opts in // (`with_backoff_sleep`); otherwise this is a no-op so // the loop stays fast and deterministic in tests. - self.policy.retry.sleep_backoff(backoff_attempt).await; + self.policy + .retry + .sleep_backoff_for_error(backoff_attempt, &error) + .await; continue; } break Err(error); @@ -326,6 +570,10 @@ impl AgentHarness { /// [`ModelResponse`] via [`StreamAccumulator`]. The merged response is /// equivalent to what the unary [`crate::harness::model::ChatModel::invoke`] /// path would have produced, so the rest of the loop is unaffected. + /// + /// `deltas_emitted` is incremented for every delta actually handed to + /// consumers, so the retry path can tell whether a failed attempt already + /// published output that now has to be discarded. async fn invoke_model_streaming_once( &self, state: &State, @@ -333,6 +581,7 @@ impl AgentHarness { model: &Arc>, request: &ModelRequest, call_id: &CallId, + deltas_emitted: &mut usize, ) -> Result { let mut stream = model.stream(state, request.clone()).await?; let mut accumulator = StreamAccumulator::new(); @@ -385,6 +634,7 @@ impl AgentHarness { call_id: call_id.clone(), delta: message_delta, }); + *deltas_emitted += 1; self.middleware .run_on_model_delta(ctx, state, &mut model_delta) .await?; @@ -403,6 +653,19 @@ impl AgentHarness { /// can proceed, short-circuit, retry, or fall back around the *whole* real model /// call. The resolved binding is rebuilt per invocation so a wrap middleware /// that retries `next` issues a fresh provider call each time. +/// +/// # The binding is re-resolved from the request +/// +/// [`ModelCallBase::call`] used to rebuild the binding purely from +/// [`Self::resolved`] / [`Self::model`], both captured **before** the wrap onion +/// ran, and ignore [`ModelRequest::model`] entirely. A wrap middleware steers by +/// mutating that field — it is the only lever it has — so +/// [`ModelFallbackMiddleware`][crate::harness::middleware::ModelFallbackMiddleware] +/// re-invoked *the same failing model* once per configured fallback name, +/// emitting a misleading `FallbackSelected { from, to }` for each, and then +/// returned the original error. The asymmetry was easy to miss because +/// `before_model` **does** honour `request.model`: lifecycle resolution happens +/// after that hook, but before this one. pub(super) struct ModelCallBase<'h, State: Send + Sync, Ctx: Send + Sync> { pub(super) harness: &'h AgentHarness, pub(super) call_id: CallId, @@ -411,6 +674,66 @@ pub(super) struct ModelCallBase<'h, State: Send + Sync, Ctx: Send + Sync> { pub(super) streaming: bool, } +impl ModelCallBase<'_, State, Ctx> { + /// Produces the binding for one invocation, honouring a model override that + /// a wrap middleware wrote into `request.model`. + /// + /// * `request.model` absent, or equal to the already-resolved name: reuse + /// the captured binding (the common path — no registry lookup). + /// * `request.model` names something the registry resolves as a genuine + /// [`ModelResolutionSource::RequestOverride`]: use it. This is what makes + /// a wrap-layer fallback actually switch models. + /// * `request.model` names something unresolvable (unregistered, missing a + /// required capability, provider-retired): fall back to the captured + /// binding and emit [`AgentEvent::ModelOverrideSkipped`], matching the + /// fail-closed behaviour `run_loop` already has for a pre-wrap override. + /// Silently substituting a different model is the one outcome that is + /// never acceptable. + fn rebind( + &self, + ctx: &mut RunContext, + request: &ModelRequest, + ) -> ResolvedModelBinding { + let captured = || ResolvedModelBinding { + resolved: self.resolved.clone(), + model: Arc::clone(&self.model), + }; + let Some(requested) = request.model.as_deref() else { + return captured(); + }; + if requested == self.resolved.name { + return captured(); + } + match self.harness.models.resolve_request(request, None, None) { + Some(binding) + if binding.resolved.source == ModelResolutionSource::RequestOverride + && binding.resolved.name == requested => + { + tracing::debug!( + call_id = %self.call_id.as_str(), + from = %self.resolved.name, + to = %binding.resolved.name, + "[model] wrap layer overrode the model; re-resolved the binding" + ); + binding + } + _ => { + tracing::warn!( + call_id = %self.call_id.as_str(), + requested = %requested, + resolved = %self.resolved.name, + "[model] wrap layer named an unresolvable model; keeping the resolved binding" + ); + ctx.emit(AgentEvent::ModelOverrideSkipped { + requested: requested.to_string(), + resolved: self.resolved.name.clone(), + }); + captured() + } + } + } +} + impl ModelBaseCall for ModelCallBase<'_, State, Ctx> { @@ -421,10 +744,7 @@ impl ModelBaseCall request: ModelRequest, ) -> BoxModelFuture<'a> { Box::pin(async move { - let binding = ResolvedModelBinding { - resolved: self.resolved.clone(), - model: Arc::clone(&self.model), - }; + let binding = self.rebind(ctx, &request); self.harness .invoke_model_with_retry( state, diff --git a/src/harness/agent_loop/run_loop.rs b/src/harness/agent_loop/run_loop.rs index 9f71ceb..07c55e4 100644 --- a/src/harness/agent_loop/run_loop.rs +++ b/src/harness/agent_loop/run_loop.rs @@ -21,6 +21,86 @@ impl AgentHarness { input: Vec, streaming: bool, ) -> Result<()> { + let mut messages = input; + // The body borrows the working transcript rather than owning it so the + // transcript survives **every** exit path, not just the successful one. + // A mid-turn tool failure used to drop everything accumulated so far, + // leaving the caller unable to inspect, repair, or resume from the + // partial conversation. + let outcome = self + .run_loop_body(state, ctx, run, status, &mut messages, streaming) + .await; + run.messages = std::mem::take(&mut messages); + + let exit = match outcome { + Ok(exit) => exit, + Err(error) => { + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + messages = run.messages.len(), + "[agent_loop] run failed; partial transcript preserved on the run" + ); + return Err(error); + } + }; + + status.mark_running(HarnessPhase::Middleware); + self.middleware.run_after_agent(ctx, state, run).await?; + + match exit { + LoopExit::Finished | LoopExit::LimitStop(_) => { + if let LoopExit::LimitStop(kind) = &exit { + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + limit_kind = ?kind, + messages = run.messages.len(), + "[agent_loop] completing with the partial run after a limit stop" + ); + } + let record = ctx.emit(AgentEvent::RunCompleted { + run_id: ctx.run_id().clone(), + }); + status.set_last_event(record.id); + } + LoopExit::Paused(pause) => { + // A pause is not a completion: reporting `run.completed` here + // is exactly what made "paused for a human" indistinguishable + // from "the model produced an empty final answer". The pause + // stays latched on the steering handle so a later `Resume` + // lifts it. + let record = ctx.emit(AgentEvent::ControlApplied { + control: "paused".to_string(), + detail: pause.reason.clone().unwrap_or_else(|| { + format!("paused at checkpoint {}", pause.paused_at_checkpoint) + }), + }); + status.set_last_event(record.id); + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + checkpoint = pause.paused_at_checkpoint, + "[agent_loop] run paused by steering" + ); + run.paused = Some(pause); + } + } + + Ok(()) + } + + /// The loop body proper. Returns how the loop left off so the caller can + /// finalize (and, on any error, still keep the working transcript). + async fn run_loop_body( + &self, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + messages: &mut Vec, + streaming: bool, + ) -> Result { let record = ctx.emit(AgentEvent::RunStarted { run_id: ctx.run_id().clone(), thread_id: ctx.thread_id().cloned(), @@ -29,21 +109,77 @@ impl AgentHarness { status.mark_running(HarnessPhase::Idle); // Reconcile the `RunConfig`-derived limit tracker with the harness's - // `RunPolicy::limits` so model/tool call caps have one enforced - // source of truth instead of the two silently disagreeing (see - // `LimitTracker::sync_call_limits`). - ctx.limits.sync_call_limits( + // `RunPolicy::limits` so model/tool call caps have one enforced source + // of truth instead of the two silently disagreeing. + // + // The two directions are NOT symmetric, and telling them apart is the + // whole reason `RunConfig`'s caps are `Option`: + // + // - an **explicitly set** `RunConfig` cap is the caller's ceiling, so + // the stricter of (config, policy) wins — fail-closed. Previously + // this was a plain assignment, so + // `RunConfig::new("r").with_max_model_calls(2)` against the default + // policy silently ran 25 model calls; + // - an **unset** cap merely defaulted, so the policy is the only real + // source of truth and may raise the cap above that default. + let effective_model_calls = resolve_call_cap( + ctx.config.max_model_calls, self.policy.limits.max_model_calls, - self.policy.limits.max_tool_calls, ); - - let mut messages = input; + let effective_tool_calls = + resolve_call_cap(ctx.config.max_tool_calls, self.policy.limits.max_tool_calls); + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + config_model_calls = ?ctx.config.max_model_calls, + config_tool_calls = ?ctx.config.max_tool_calls, + policy_model_calls = self.policy.limits.max_model_calls, + policy_tool_calls = self.policy.limits.max_tool_calls, + effective_model_calls, + effective_tool_calls, + "[agent_loop] resolved run call caps" + ); + // The values are already reconciled per-axis above, so the assignment + // form (`sync_call_limits`) is the correct primitive here: + // `tighten_call_limits` would additionally min against the tracker's + // config-*default*-derived cap and so could not honor a policy that + // legitimately raises an unset cap. + ctx.limits + .sync_call_limits(effective_model_calls, effective_tool_calls); // The tool set is fixed for the duration of a run, so build the sorted // schema vec once here instead of re-collecting, re-calling every tool's // `schema()`, and re-sorting on every turn (per model call). let tool_schemas = self.tools.schemas(); + // Fail closed on a structured-output schema whose name collides with a + // registered tool. Under the tool-call strategy the schema is sent as an + // extra `function` entry, so a collision puts two identically-named + // functions in one request — which OpenAI rejects outright — and makes + // "was this the schema or the real tool?" unanswerable for every + // returned call. + if let Some(name) = self + .policy + .default_response_format + .as_ref() + .and_then(|format| match format { + ResponseFormat::Auto { name, .. } | ResponseFormat::JsonSchema { name, .. } => { + Some(name) + } + _ => None, + }) + && self + .tools + .names() + .iter() + .any(|registered| registered == name) + { + return Err(TinyAgentsError::Validation(format!( + "structured-output schema name `{name}` collides with a registered tool of the \ + same name; rename one of them" + ))); + } + status.mark_running(HarnessPhase::Middleware); self.middleware.run_before_agent(ctx, state).await?; @@ -67,14 +203,35 @@ impl AgentHarness { // Safe steering checkpoint: drain any orchestrator/human steering // commands and apply the policy-permitted ones before the next // model call. Cancel terminates the run; Pause short-circuits it. - match crate::harness::steering::apply_pending_steering(ctx, &mut messages)? { + match crate::harness::steering::apply_pending_steering(ctx, messages)? { crate::harness::steering::SteeringOutcome::Cancel => { return Err(TinyAgentsError::Cancelled); } - crate::harness::steering::SteeringOutcome::Pause => break, + crate::harness::steering::SteeringOutcome::Pause => { + let pause = ctx + .steering + .as_ref() + .and_then(|handle| handle.pause_state()) + .unwrap_or(crate::harness::steering::PauseState { + reason: None, + paused_at_checkpoint: 0, + }); + return Ok(LoopExit::Paused(pause)); + } crate::harness::steering::SteeringOutcome::Continue => {} } + // Safe checkpoint: honor a control outcome requested during the + // *previous* turn's tool execution (or by `before_agent`) before + // spending another model call on it. Draining only after the model + // call meant a `StopWithFinal`/`Interrupt` raised from + // `after_tool`/`wrap_tool` was honored one full model call late — + // an extra billable provider round trip after a guardrail, or a + // human gate, had already said stop. + if let Some(exit) = self.apply_pending_control(ctx, run, status)? { + return Ok(exit); + } + // Fail-closed limit and deadline checks before each model call. if ctx.check_deadline().is_err() { ctx.emit(AgentEvent::LimitReached { @@ -88,11 +245,45 @@ impl AgentHarness { // The context's `LimitTracker` (synced with `RunPolicy::limits` // above) is the single enforced source of truth for the model-call // cap, so the reported limit always matches the one that trips. - if let Err(err) = ctx.record_model_call() { - ctx.emit(AgentEvent::LimitReached { - kind: LimitKind::ModelCalls, - }); - return Err(TinyAgentsError::LimitExceeded(err.to_string())); + // `LimitBehavior::StopWithPartial` turns cap exhaustion into a + // clean stop rather than an error that discards every message, + // usage figure, and tool result the run produced up to that point. + match ctx.limits.try_record_model_call() { + Ok(crate::harness::limits::LimitOutcome::Proceed) => {} + Ok(crate::harness::limits::LimitOutcome::Stop(_)) => { + ctx.emit(AgentEvent::LimitReached { + kind: LimitKind::ModelCalls, + }); + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + "[agent_loop] model-call cap reached; stopping with the partial run" + ); + return Ok(LoopExit::LimitStop(LimitKind::ModelCalls)); + } + Err(err) => { + ctx.emit(AgentEvent::LimitReached { + kind: LimitKind::ModelCalls, + }); + // `RunConfig` cannot express a `LimitBehavior`, so the + // tracker built from it always carries the default + // (`Error`) even when the harness policy asks for + // `StopWithPartial`. Honor the policy here rather than + // discarding a run the operator asked to keep. + if matches!( + self.policy.limits.behavior, + crate::harness::limits::LimitBehavior::StopWithPartial + ) { + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + "[agent_loop] model-call cap reached; policy asks to stop with the \ + partial run" + ); + return Ok(LoopExit::LimitStop(LimitKind::ModelCalls)); + } + return Err(TinyAgentsError::LimitExceeded(err.to_string())); + } } // Build the request from the working transcript, tool schemas, and @@ -172,7 +363,28 @@ impl AgentHarness { parameters: schema.clone(), format: crate::harness::tool::ToolFormat::Json, }); - request.tool_choice = ToolChoice::Tool(name.clone()); + // Force the schema tool **only** when it is the + // sole tool available. Forcing it inside a + // tool-using loop makes the model emit the + // structured call on turn 1, which terminates + // the loop before any registered tool can ever + // run — the agent silently loses its tools, and + // the symptom points nowhere near this code. + // LangChain likewise binds a schema tool with a + // forced `tool_choice` only in its terminal + // wrapper, never in the tool-calling loop. + if tool_schemas.is_empty() { + request.tool_choice = ToolChoice::Tool(name.clone()); + } else { + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + schema_name = %name, + registered_tools = tool_schemas.len(), + "[agent_loop] structured tool offered but not forced; \ + registered tools stay callable" + ); + } } } Some((strategy, name, schema)) @@ -235,11 +447,26 @@ impl AgentHarness { run.steps += 1; status.model_calls = run.model_calls; status.active_model_call = None; + // A cache replay consumed no provider tokens, so folding its usage + // into the run's totals reports spend that never happened. The + // saving is surfaced through the cache-hit event instead of being + // buried in the spend total. if let Some(usage) = response.usage { - run.usage.record(usage); - status.usage = run.usage; - let record = ctx.emit(AgentEvent::UsageRecorded { usage }); - status.set_last_event(record.id); + if response.served_from_cache { + tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + call_id = %call_id, + saved_input_tokens = usage.input_tokens, + saved_output_tokens = usage.output_tokens, + "[agent_loop] cache-served response; usage not billed to the run" + ); + } else { + run.usage.record(usage); + status.usage = run.usage; + let record = ctx.emit(AgentEvent::UsageRecorded { usage }); + status.set_last_event(record.id); + } } let captured_output = self .policy @@ -260,40 +487,89 @@ impl AgentHarness { // Safe checkpoint: honor any control outcome a middleware requested // during this turn (for example an early-exit tool or a budget stop // hook), before executing further tools. - if let Some(control) = ctx.take_control() { + if let Some(exit) = self.apply_pending_control(ctx, run, status)? { + return Ok(exit); + } + + let tool_calls = response.tool_calls().to_vec(); + + // A tool-call structured-output strategy produces an artificial tool + // call that is not a registered tool, so split the turn's calls into + // the schema call(s) and the genuine ones. Treating "any call + // matched the schema name" as terminal silently dropped every + // sibling call in the same turn — a turn returning + // `[search(...), my_schema(...)]` broke out with `search` never + // executed and no event to say so. + let structured_call_name = match &structured_plan { + Some((StructuredStrategy::ToolCall, name, _)) => Some(name.clone()), + _ => None, + }; + let (structured_hits, real_tool_calls): (Vec, Vec) = + match &structured_call_name { + Some(name) => tool_calls + .iter() + .cloned() + .partition(|call| &call.name == name), + None => (Vec::new(), tool_calls.clone()), + }; + let structured_tool_hit = !structured_hits.is_empty(); + + if structured_tool_hit && !real_tool_calls.is_empty() { + // Record the structured payload the model already produced, + // then run the real tools it asked for in the same turn and let + // the loop continue; the model finishes on a later turn. + if let Some((strategy, name, schema)) = &structured_plan { + let extractor = + StructuredExtractor::new(*strategy, name.clone(), schema.clone()); + match extractor.extract(&response) { + Ok(output) => run.structured = Some(output.value), + Err(error) => tracing::debug!( + target: "tinyagents::agent_loop", + run_id = %ctx.run_id(), + %error, + "[agent_loop] structured extraction failed on a mixed turn; \ + continuing with the real tool calls" + ), + } + } let record = ctx.emit(AgentEvent::ControlApplied { - control: control.kind().to_string(), - detail: match &control { - MiddlewareControl::StopWithFinal(text) => text.clone(), - MiddlewareControl::Interrupt { node, message } => { - format!("{node}: {message}") - } - }, + control: "structured_with_tool_calls".to_string(), + detail: format!( + "structured output recorded alongside {} real tool call(s); \ + the run continues", + real_tool_calls.len() + ), }); status.set_last_event(record.id); - match control { - MiddlewareControl::StopWithFinal(text) => { - run.final_response = Some(ModelResponse::assistant(text)); - break; - } - MiddlewareControl::Interrupt { node, message } => { - return Err(TinyAgentsError::Interrupted { node, message }); - } + + // Every requested `tool_call_id` must be answered or the + // transcript is malformed for the next provider call. + for call in &structured_hits { + messages.push(Message::tool( + call.id.clone(), + "Structured output recorded. Continue with the remaining tool calls.", + )); } - } - let tool_calls = response.tool_calls().to_vec(); + reset_truncated_empty_recovery( + &mut truncated_empty_retries_used, + &mut boosted_max_tokens, + &mut truncation_base, + ); - // A tool-call structured-output strategy produces an artificial tool - // call that is not a registered tool; treat it as the final response - // rather than attempting to execute it. - let structured_tool_hit = matches!( - &structured_plan, - Some((StructuredStrategy::ToolCall, name, _)) - if tool_calls.iter().any(|c| &c.name == name) - ); + status.mark_running(HarnessPhase::Tools); + self.execute_tools(state, ctx, run, status, messages, real_tool_calls) + .await?; + + // Safe checkpoint: a control requested from `after_tool` / + // `wrap_tool` is honored here, at the edge it was raised on. + if let Some(exit) = self.apply_pending_control(ctx, run, status)? { + return Ok(exit); + } + continue; + } - if tool_calls.is_empty() || structured_tool_hit { + if real_tool_calls.is_empty() { // Truncated-empty recovery (runs before structured extraction, // which would otherwise fail on the empty completion). A local // reasoning model can burn the whole token budget on its hidden @@ -383,7 +659,7 @@ impl AgentHarness { return Err(TinyAgentsError::EmptyResponse); } run.final_response = Some(response); - break; + return Ok(LoopExit::Finished); } // A tool-calling response is a resolved turn too: clear the @@ -401,21 +677,52 @@ impl AgentHarness { // `agent_loop/tools.rs` for the dispatch rules and the semantics // preserved in each mode. status.mark_running(HarnessPhase::Tools); - self.execute_tools(state, ctx, run, status, &mut messages, tool_calls) + self.execute_tools(state, ctx, run, status, messages, real_tool_calls) .await?; - } - - run.messages = messages; - status.mark_running(HarnessPhase::Middleware); - self.middleware.run_after_agent(ctx, state, run).await?; + // Safe checkpoint: honor a control requested from `after_tool` / + // `wrap_tool` at the edge it was raised on, rather than a model + // call later. + if let Some(exit) = self.apply_pending_control(ctx, run, status)? { + return Ok(exit); + } + } + } - let record = ctx.emit(AgentEvent::RunCompleted { - run_id: ctx.run_id().clone(), + /// Drains any pending [`MiddlewareControl`] and turns it into a loop + /// decision. + /// + /// Returns `Ok(None)` when nothing was requested, `Ok(Some(exit))` when the + /// loop must stop, and `Err` for + /// [`MiddlewareControl::Interrupt`]. Called at every safe checkpoint — the + /// top of an iteration, after the model call, and after tool execution — so + /// a control raised anywhere in a turn takes effect on that turn. + fn apply_pending_control( + &self, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + ) -> Result> { + let Some(control) = ctx.take_control() else { + return Ok(None); + }; + let record = ctx.emit(AgentEvent::ControlApplied { + control: control.kind().to_string(), + detail: match &control { + MiddlewareControl::StopWithFinal(text) => text.clone(), + MiddlewareControl::Interrupt { node, message } => format!("{node}: {message}"), + }, }); status.set_last_event(record.id); - - Ok(()) + match control { + MiddlewareControl::StopWithFinal(text) => { + run.final_response = Some(ModelResponse::assistant(text)); + Ok(Some(LoopExit::Finished)) + } + MiddlewareControl::Interrupt { node, message } => { + Err(TinyAgentsError::Interrupted { node, message }) + } + } } /// Resolves the effective response-cache decision for `request`. @@ -456,6 +763,20 @@ impl AgentHarness { } } +/// Resolves one run-scoped call cap from the per-run [`RunConfig`] value and +/// the harness-wide [`crate::harness::runtime::RunPolicy`] value. +/// +/// An explicitly-set config cap is the caller's ceiling and can only be +/// tightened by the policy (fail-closed `min`); an unset config cap leaves the +/// policy as the single source of truth, which is what lets a policy raise a +/// cap above the crate default. +fn resolve_call_cap(config_cap: Option, policy_cap: usize) -> usize { + match config_cap { + Some(explicit) => explicit.min(policy_cap), + None => policy_cap, + } +} + /// Clears the per-turn truncated-empty recovery state (see /// [`crate::harness::runtime::RunPolicy::truncated_empty_retries`]). /// diff --git a/src/harness/agent_loop/test.rs b/src/harness/agent_loop/test.rs index 8011331..9f9f4a6 100644 --- a/src/harness/agent_loop/test.rs +++ b/src/harness/agent_loop/test.rs @@ -348,6 +348,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -372,6 +373,7 @@ fn invalid_tool_call_response(id: &str, name: &str, raw: &str) -> ModelResponse raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -389,6 +391,7 @@ fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -410,6 +413,7 @@ fn truncated_empty_response(reasoning_tokens: u64) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -530,6 +534,7 @@ impl ChatModel<()> for ProviderFailingModel { provider: "test-provider".to_string(), status: Some(self.status), retryable: self.retryable, + retry_after_ms: None, message: "boom".to_string(), ..crate::harness::model::ProviderError::default() }, @@ -1097,14 +1102,20 @@ async fn usage_accumulates_across_calls() { assert_eq!(run.usage.usage.output_tokens, 5); } +/// `UnknownToolPolicy::Fail` is opt-in now (the default recovers), so this +/// pins the opted-in fail-closed behavior rather than the default. #[tokio::test] -async fn tool_not_found_errors() { +async fn tool_not_found_errors_under_the_fail_policy() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness.register_model( "mock", Arc::new(MockModel::with_tool_call("missing", json!({}))), ); // No tool registered. + harness.with_policy(RunPolicy { + unknown_tool: UnknownToolPolicy::Fail, + ..RunPolicy::default() + }); let err = harness .invoke_default(&(), vec![Message::user("go")]) @@ -1177,8 +1188,10 @@ async fn unknown_tool_rewrite_retargets_to_real_tool() { assert_eq!(*lookup.calls.lock().unwrap(), 1); } +/// `InvalidArgsPolicy::Fail` is opt-in now (the default recovers), so this pins +/// the opted-in fail-closed behavior rather than the default. #[tokio::test] -async fn invalid_tool_arguments_fail_before_tool_execution() { +async fn invalid_tool_arguments_fail_before_tool_execution_under_the_fail_policy() { let mut harness: AgentHarness<()> = AgentHarness::new(); harness.register_model( "mock", @@ -1191,6 +1204,10 @@ async fn invalid_tool_arguments_fail_before_tool_execution() { harness.register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), })); + harness.with_policy(RunPolicy { + invalid_args: InvalidArgsPolicy::Fail, + ..RunPolicy::default() + }); let err = harness .invoke_default(&(), vec![Message::user("lookup")]) @@ -2752,6 +2769,7 @@ async fn request_cache_policy_overrides_run_policy_to_disable_caching() { request.cache_policy = Some(CachePolicy { response_cache_enabled: false, protect_prompt_prefix: false, + ..CachePolicy::default() }); Ok(()) } @@ -3099,6 +3117,7 @@ fn multi_tool_call_response(calls: Vec<(&str, &str)>) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/src/harness/agent_loop/tools.rs b/src/harness/agent_loop/tools.rs index fb610d2..28bf9b0 100644 --- a/src/harness/agent_loop/tools.rs +++ b/src/harness/agent_loop/tools.rs @@ -9,19 +9,53 @@ //! //! 1. **Admission** (always serial, in call order, under `&mut RunContext`): //! cancellation/deadline/limit checks, the lifecycle `before_tool` hooks, -//! unknown-tool policy resolution, schema validation, and the -//! [`AgentEvent::ToolStarted`] emission. +//! unknown-tool policy resolution, injected-argument stripping, and schema +//! validation. Admission emits **no** [`AgentEvent::ToolStarted`] — see +//! "Started/terminal pairing" below. //! 2. **Execution**: when the turn requests **two or more** tools and **no //! tool-wrap middleware** ([`crate::harness::middleware::ToolMiddleware`]) //! is registered, the admitted calls run **concurrently** //! (`join_all`), so turn latency is the slowest tool instead of the sum. //! Otherwise execution is serial, preserving the historical semantics. +//! [`AgentEvent::ToolStarted`] is emitted here, once every admission has +//! succeeded, so a call that is announced always runs. //! 3. **Fold** (always serial, in original call order): the lifecycle //! `after_tool` hooks, accounting, the [`AgentEvent::ToolCompleted`] //! emission, and the transcript append. Results are attached to their //! original `tool_call_id` in the calls' original order regardless of //! completion order. //! +//! ## Started/terminal pairing (the invariant this module maintains) +//! +//! Every [`AgentEvent::ToolStarted`] is followed by exactly one terminal +//! partner — [`AgentEvent::ToolCompleted`] when the call produced a result +//! (successful *or* error-carrying) or [`AgentEvent::ToolFailed`] when the run +//! itself is aborting because of it. `status.active_tool_calls` is cleared on +//! both paths, and by *position*, so two calls sharing one id (a real provider +//! defect) cannot clear each other's entry. +//! +//! Recovery paths — unknown tool, schema-invalid arguments, provider-unparseable +//! arguments — are not exceptions: they are folded through the same +//! [`AgentHarness::finish_tool_call`] pipeline, so they emit the same +//! started/completed pair, run `after_tool`, and account identically. The only +//! difference is that no tool ran. +//! +//! ## Tool errors are policy-routed, not fatal by default +//! +//! An `Err` from a tool is routed through that tool's +//! [`crate::harness::tool::ToolErrorPolicy`]: the default +//! [`Fail`][crate::harness::tool::ToolErrorPolicy::Fail] still aborts the run, +//! while `ReturnToError`/`Message` turn the failure into a model-visible error +//! result. Cancellation and interruption bubble regardless of policy, and two +//! error classes are deliberately kept **outside** the policy because they are +//! not tool failures at all: +//! +//! - the run's remaining wall-clock budget expiring around the call +//! ([`AgentHarness::with_call_budget`]) — the run is over, not the tool, and +//! - an error raised by *middleware* wrapping the call, which is how an +//! approval or allowlist gate refuses a call. Converting a refusal into "the +//! tool failed, carry on" would defeat the gate. +//! //! ## Why tool-wrap middleware forces serial execution //! //! [`crate::harness::middleware::ToolMiddleware::wrap_tool`] holds @@ -44,16 +78,21 @@ //! budget, exactly as in serial mode. //! - **Cancellation**: observed between admissions (before each call starts), //! matching the serial path, which also never interrupts a mid-flight tool. -//! - **Errors**: the first failing call *in original call order* fails the -//! turn. Difference: in serial mode later calls never start after a -//! failure; in concurrent mode they were already in flight and run to -//! completion (their results are discarded). Tools that must not observe a -//! sibling's failure should be run under a tool-wrap middleware (serial) or -//! a harness without parallel-capable turns. +//! - **Errors**: a tool error that its [`ToolErrorPolicy`] keeps fatal fails +//! the turn at the first such call *in original call order*. Difference: in +//! serial mode later calls never start after a failure; in concurrent mode +//! they were already in flight and run to completion (their results are +//! discarded). Tools that must not observe a sibling's failure should be run +//! under a tool-wrap middleware (serial) or a harness without +//! parallel-capable turns. +//! +//! [`ToolErrorPolicy`]: crate::harness::tool::ToolErrorPolicy use super::model_call::ToolCallBase; use super::*; -use crate::harness::tool::ToolExecutionContext; +use crate::harness::tool::{ + ToolErrorPolicy, ToolExecutionContext, project_injected_arguments, strip_injected_arguments, +}; /// How a single requested tool call was resolved during admission. enum ResolvedToolCall { @@ -64,13 +103,30 @@ enum ResolvedToolCall { ErrorMessage(String), } +/// One requested call after admission, in original order. +/// +/// The concurrent path materialises the whole admitted batch before emitting a +/// single [`AgentEvent::ToolStarted`], so an admission failure part-way through +/// the batch cannot leave earlier calls announced-but-never-run (TOOL-3). +enum AdmittedCall { + /// A registered tool to invoke, with its (validated) call. + Execute { + tool: Arc>, + call: ToolCall, + }, + /// A recovery: no tool runs, but the call is still answered through the + /// normal result pipeline so it emits the same started/completed pair and + /// runs the same `after_tool` hooks (TOOL-11). + Recovered { call: ToolCall, message: String }, +} + /// One transcript slot per requested call, in original order, used by the /// concurrent path to reassemble results deterministically. enum ToolSlot { - /// An executed call: consumes the next prepared/future pair in order. + /// An executed call: consumes the next prepared/result pair in order. Execute, - /// An unknown-tool recovery message, appended verbatim. - Immediate { call_id: String, message: String }, + /// A recovery, folded in place through the normal result pipeline. + Recovered { call: ToolCall, message: String }, } /// Admission metadata for one executable call, paired 1:1 (in order) with its @@ -174,7 +230,23 @@ impl AgentHarness { return Err(TinyAgentsError::LimitExceeded(err.to_string())); } - self.middleware.run_before_tool(ctx, state, call).await?; + // The slot is *reserved* above (cap-first, so a middleware hook never + // runs for a call the budget has already refused) and *released* here + // when `before_tool` refuses the call — an approval denial or an + // allowlist rejection must not spend budget on a call that never ran + // (TOOL-12). The recovery paths below deliberately keep their slot: + // they answer the model and let it try again, so counting them is what + // bounds the correction loop. + if let Err(err) = self.middleware.run_before_tool(ctx, state, call).await { + tracing::debug!( + "[agent_loop::tools] `before_tool` refused `{}` (call `{}`); \ + releasing its tool-call slot: {err}", + call.name, + call.id + ); + ctx.limits.rollback_tool_calls(1); + return Err(err); + } // The provider marked this call's arguments unparseable (a small local // model emitted malformed JSON). Rather than fail the run, inject a @@ -254,7 +326,25 @@ impl AgentHarness { } } }; - let schema = tool.schema(); + // Step 1 of the injected-argument ordering rule (see + // `crate::harness::tool::injected`): strip every host-injected key from + // the model-supplied arguments *before* validating them, so a model + // that names a hidden key it was never shown cannot forge it. Step 2 + // then validates against the **model-facing** projection of the schema + // — the same one `ToolRegistry::schemas` advertises — because a key the + // model never saw must not be `required` of it. + let injected = tool.injected_arguments(); + let forged = strip_injected_arguments(&mut call.arguments, injected); + if !forged.is_empty() { + tracing::warn!( + "[agent_loop::tools] tool `{}` call `{}`: discarded model-supplied value(s) \ + for host-injected argument(s): {}", + call.name, + call.id, + forged.join(", ") + ); + } + let schema = project_injected_arguments(tool.schema(), injected); let raw_arguments = call.arguments.clone(); if matches!( self.policy.invalid_args, @@ -325,6 +415,41 @@ impl AgentHarness { } } + /// Terminal partner of [`AgentEvent::ToolStarted`] on the abort path: + /// emits [`AgentEvent::ToolFailed`] and closes the call's `active_tool_calls` + /// entry. + /// + /// Without this, every `?` between `ToolStarted` and `ToolCompleted` left a + /// dangling start — an exporter pairing the two by `call_id` silently drops + /// the failed span, and the run keeps reporting a call that is no longer in + /// flight (TOOL-6). Call it on **every** error path after + /// [`Self::start_tool_call`]. + fn fail_tool_call( + &self, + ctx: &RunContext, + status: &mut HarnessRunStatus, + call_id: &CallId, + tool_name: &str, + started_at_ms: u64, + error: &TinyAgentsError, + ) { + release_active_tool_call(status, call_id); + let duration_ms = crate::harness::ids::now_ms().saturating_sub(started_at_ms); + tracing::debug!( + "[agent_loop::tools] tool `{tool_name}` call `{}` failed after {duration_ms} ms: \ + {error}", + call_id.as_str() + ); + let record = ctx.emit(AgentEvent::ToolFailed { + call_id: call_id.clone(), + tool_name: tool_name.to_string(), + started_at_ms: Some(started_at_ms), + duration_ms: Some(duration_ms), + error: error.to_string(), + }); + status.set_last_event(record.id); + } + /// Fold phase for one completed call: the lifecycle `after_tool` hooks, /// accounting, the `ToolCompleted` emission, and the transcript append. #[allow(clippy::too_many_arguments)] @@ -338,13 +463,44 @@ impl AgentHarness { prepared: PreparedToolCall, mut result: crate::harness::tool::ToolResult, ) -> Result<()> { - self.middleware + // The harness, not the tool, owns the identity of the call being + // answered. A tool that stamps its own `call_id` — a hard-coded string, + // an empty one, a reused constant — would otherwise put a + // `tool_call_id` in the transcript that matches no `tool_calls[].id` in + // the preceding assistant message, and the provider rejects that on the + // *next* request, one turn away from the tool that caused it (TOOL-1). + // Overwrite rather than fail: the correct id is known here, and a + // third-party bug should not end a run. + if result.call_id != prepared.call_id.as_str() { + tracing::warn!( + "[agent_loop::tools] tool `{}` returned call_id `{}` for call `{}`; \ + overwriting with the admitted id so the transcript stays consistent", + prepared.tool_name, + result.call_id, + prepared.call_id.as_str() + ); + result.call_id = prepared.call_id.as_str().to_string(); + } + + if let Err(err) = self + .middleware .run_after_tool(ctx, state, &mut result) - .await?; + .await + { + self.fail_tool_call( + ctx, + status, + &prepared.call_id, + &prepared.tool_name, + prepared.started_at_ms, + &err, + ); + return Err(err); + } run.tool_calls += 1; status.tool_calls = run.tool_calls; - status.active_tool_calls.retain(|c| c != &prepared.call_id); + release_active_tool_call(status, &prepared.call_id); let captured_output = self .policy .capture @@ -391,9 +547,8 @@ impl AgentHarness { let tool = match self.admit_tool_call(state, ctx, status, &mut call).await? { ResolvedToolCall::Tool(tool) => tool, ResolvedToolCall::ErrorMessage(message) => { - run.tool_calls += 1; - status.tool_calls = run.tool_calls; - messages.push(Message::tool(call.id.clone(), message)); + self.recover_tool_call(state, ctx, run, status, messages, &call, message) + .await?; continue; } }; @@ -407,15 +562,36 @@ impl AgentHarness { // crate-owned tool policy returns a recoverable tool error; the // outer run budget still aborts when the whole run is exhausted. let run_budget = self.call_budget(ctx); + let error_policy = tool.error_policy(); + let policy_call = call.clone(); let base = ToolCallBase { tool, timeout_settings: self.tool_timeouts.clone(), }; let run_id = ctx.run_id().as_str().to_string(); let fut = self.middleware.run_wrapped_tool(ctx, state, call, &base); - let result = Self::with_call_budget(run_budget, &run_id, "tool call", fut) - .await? - .into_result(); + // The policy is applied *inside* the run-budget wrapper so that + // exhausting the run's wall clock stays fatal (it is the run + // ending, not the tool failing) while a tool error is routed. + let guarded = async move { + let outcome = fut.await.map(|wrapped| wrapped.into_result()); + apply_tool_error_policy(&error_policy, &policy_call, outcome) + }; + let outcome = Self::with_call_budget(run_budget, &run_id, "tool call", guarded).await; + let result = match outcome { + Ok(result) => result, + Err(err) => { + self.fail_tool_call( + ctx, + status, + &prepared.call_id, + &prepared.tool_name, + prepared.started_at_ms, + &err, + ); + return Err(err); + } + }; self.finish_tool_call(state, ctx, run, status, messages, prepared, result) .await?; @@ -423,6 +599,41 @@ impl AgentHarness { Ok(()) } + /// Answers a call that no tool ran — unknown tool, schema-invalid + /// arguments, or arguments the provider could not parse — through the same + /// pipeline a real result takes. + /// + /// Before this existed the recovery arms pushed a bare + /// [`Message::tool`] and hand-incremented the counters, so no + /// `ToolStarted`/`ToolCompleted` pair was emitted, `after_tool` middleware + /// never saw the result, and accounting lived in three places (TOOL-11). The + /// transcript content is unchanged: [`ToolResult::error`][err] puts the + /// message verbatim in `content`. + /// + /// [err]: crate::harness::tool::ToolResult::error + #[allow(clippy::too_many_arguments)] + async fn recover_tool_call( + &self, + state: &State, + ctx: &mut RunContext, + run: &mut AgentRun, + status: &mut HarnessRunStatus, + messages: &mut Vec, + call: &ToolCall, + message: String, + ) -> Result<()> { + tracing::debug!( + "[agent_loop::tools] recovering call `{}` for `{}` without executing a tool", + call.id, + call.name + ); + let prepared = self.start_tool_call(ctx, status, call); + let result = + crate::harness::tool::ToolResult::error(call.id.clone(), call.name.clone(), message); + self.finish_tool_call(state, ctx, run, status, messages, prepared, result) + .await + } + /// Executes a multi-call turn concurrently (`join_all`), so turn latency /// is the slowest tool instead of the sum. Only reachable when no /// tool-wrap middleware is registered (see the module docs); execution @@ -437,20 +648,32 @@ impl AgentHarness { messages: &mut Vec, tool_calls: Vec, ) -> Result<()> { - // Phase 1 — admission, serial, in call order. - let mut slots: Vec = Vec::with_capacity(tool_calls.len()); - let mut prepared: Vec = Vec::new(); - let mut futures: Vec<_> = Vec::new(); + // Phase 1 — admission, serial, in call order. Nothing is announced and + // nothing is queued here: an admission failure at call *k* must not + // leave calls `0..k` with a `ToolStarted` they will never answer, nor + // an `active_tool_calls` entry for work that is dropped unpolled + // (TOOL-3). The serial path would have executed those calls; the + // concurrent path now agrees with it by executing none of them. + let mut admitted: Vec> = Vec::with_capacity(tool_calls.len()); for mut call in tool_calls { - let tool = match self.admit_tool_call(state, ctx, status, &mut call).await? { - ResolvedToolCall::Tool(tool) => tool, + match self.admit_tool_call(state, ctx, status, &mut call).await? { + ResolvedToolCall::Tool(tool) => admitted.push(AdmittedCall::Execute { tool, call }), ResolvedToolCall::ErrorMessage(message) => { - run.tool_calls += 1; - status.tool_calls = run.tool_calls; - slots.push(ToolSlot::Immediate { - call_id: call.id.clone(), - message, - }); + admitted.push(AdmittedCall::Recovered { call, message }) + } + } + } + + // Phase 2 — announce and queue. Every admission succeeded, so every + // `ToolStarted` emitted here is matched by a terminal event below. + let mut slots: Vec = Vec::with_capacity(admitted.len()); + let mut prepared: Vec = Vec::new(); + let mut futures: Vec<_> = Vec::new(); + for entry in admitted { + let (tool, call) = match entry { + AdmittedCall::Execute { tool, call } => (tool, call), + AdmittedCall::Recovered { call, message } => { + slots.push(ToolSlot::Recovered { call, message }); continue; } }; @@ -468,30 +691,51 @@ impl AgentHarness { let run_budget = self.call_budget(ctx); let run_id = ctx.run_id().as_str().to_string(); let exec_ctx = ToolExecutionContext::from_run_context(ctx); + let error_policy = tool.error_policy(); + let policy_call = call.clone(); futures.push(async move { let fut = tool.call_with_context(state, call, exec_ctx); let fut = Self::with_tool_policy_timeout(tool_timeout, timeout_result, fut); - Self::with_call_budget(run_budget, &run_id, "tool call", fut).await + // As in serial mode: the error policy routes the *tool's* + // failure, inside the run-budget wrapper that stays fatal. + let guarded = + async move { apply_tool_error_policy(&error_policy, &policy_call, fut.await) }; + Self::with_call_budget(run_budget, &run_id, "tool call", guarded).await }); } - // Phase 2 — run all admitted calls concurrently. `join_all` preserves + // Phase 3 — run all admitted calls concurrently. `join_all` preserves // input order, so results pair 1:1 with `prepared`. let results = futures::future::join_all(futures).await; - // Phase 3 — fold in original call order: the first failing call (in - // that order) fails the turn; siblings already ran to completion. + // Phase 4 — fold in original call order: the first call whose policy + // kept its failure fatal (in that order) fails the turn; siblings + // already ran to completion. let mut executed = prepared.into_iter().zip(results); for slot in slots { match slot { - ToolSlot::Immediate { call_id, message } => { - messages.push(Message::tool(call_id, message)); + ToolSlot::Recovered { call, message } => { + self.recover_tool_call(state, ctx, run, status, messages, &call, message) + .await?; } ToolSlot::Execute => { let (prepared, result) = executed .next() .expect("every Execute slot has a prepared/result pair"); - let result = result?; + let result = match result { + Ok(result) => result, + Err(err) => { + self.fail_tool_call( + ctx, + status, + &prepared.call_id, + &prepared.tool_name, + prepared.started_at_ms, + &err, + ); + return Err(err); + } + }; self.finish_tool_call(state, ctx, run, status, messages, prepared, result) .await?; } @@ -501,6 +745,42 @@ impl AgentHarness { } } +/// Removes **one** occurrence of `call_id` from the in-flight list. +/// +/// Positional, not `retain`: a provider can emit two calls in one turn that +/// share a `tool_call_id`, and a predicate-based removal would clear both +/// entries when the first completes — leaving the second call in flight with no +/// entry to close, and the run reporting an empty in-flight list while a tool is +/// still running (TOOL-10). +fn release_active_tool_call(status: &mut HarnessRunStatus, call_id: &CallId) { + if let Some(position) = status + .active_tool_calls + .iter() + .position(|active| active == call_id) + { + status.active_tool_calls.remove(position); + } +} + +/// Routes a tool invocation outcome through `policy`, keeping middleware +/// refusals fatal. +/// +/// [`ToolErrorPolicy::apply`] already re-raises cancellation and interruption. +/// This adds one more class the policy must not swallow: an error raised by +/// *middleware* wrapping the call. That is how an approval gate or an allowlist +/// refuses a call, and converting a refusal into "the tool failed, try +/// something else" would let the loop continue past a gate that said no. +fn apply_tool_error_policy( + policy: &ToolErrorPolicy, + call: &ToolCall, + outcome: Result, +) -> Result { + if let Err(TinyAgentsError::Middleware(_)) = &outcome { + return outcome; + } + policy.apply(call, outcome) +} + /// Repairs provider-neutral argument shape defects before schema validation. /// /// Schema-valid arguments are already canonical. A string containing valid diff --git a/src/harness/agent_loop/types.rs b/src/harness/agent_loop/types.rs index c8328e0..fe2490c 100644 --- a/src/harness/agent_loop/types.rs +++ b/src/harness/agent_loop/types.rs @@ -9,8 +9,48 @@ //! //! All public items are re-exported through [`super`]. -use crate::harness::events::HarnessRunStatus; +use crate::harness::events::{HarnessRunStatus, LimitKind}; use crate::harness::middleware::AgentRun; +use crate::harness::steering::PauseState; + +/// The result of an agent-loop invocation that keeps the partial run even when +/// the loop fails. +/// +/// [`crate::harness::runtime::AgentHarness::invoke`] returns `Err` on failure +/// and drops the [`AgentRun`] with it, so a run that tripped a limit or hit a +/// tool failure halfway through loses every message and usage figure it had +/// accumulated. `PartialRunOutcome` keeps both, letting a caller inspect, +/// repair, or resume from the partial conversation. +#[derive(Debug)] +pub struct PartialRunOutcome { + /// The accumulated transcript, usage, counters, and final response — + /// populated as far as the run got, whether or not it failed. + pub run: AgentRun, + /// A compact lifecycle/status snapshot reflecting how the run ended. + pub status: HarnessRunStatus, + /// The error that ended the run, or `None` when it succeeded. + pub error: Option, +} + +/// How the agent loop body stopped iterating. +/// +/// Kept separate from the `Result` channel so a *deliberate* stop (a pause, a +/// `StopWithPartial` limit) is never confused with a failure, and so the caller +/// can finalize each kind differently: a finish and a limit-stop complete the +/// run, while a pause is reported as interrupted and leaves the pause latched +/// on the steering handle. +#[derive(Clone, Debug)] +pub(crate) enum LoopExit { + /// The model produced a final answer, or a middleware requested + /// [`crate::harness::context::MiddlewareControl::StopWithFinal`]. + Finished, + /// A call cap was reached under + /// [`LimitBehavior::StopWithPartial`][crate::harness::limits::LimitBehavior::StopWithPartial]: + /// stop cleanly and keep everything the run produced. + LimitStop(LimitKind), + /// Steering latched a pause; the run is resumable, not finished. + Paused(PauseState), +} /// The full result of an agent-loop invocation: the accumulated [`AgentRun`] /// plus a compact [`HarnessRunStatus`] snapshot. diff --git a/src/harness/cache/hash.rs b/src/harness/cache/hash.rs new file mode 100644 index 0000000..518853c --- /dev/null +++ b/src/harness/cache/hash.rs @@ -0,0 +1,78 @@ +//! Deterministic hashing primitives shared by the cache key derivation +//! ([`super::key`]) and the prompt-cache layout tooling ([`super::layout`]). +//! +//! Everything here is seed-free and canonical so a digest computed in one +//! process matches one computed in another — a cache that is only valid within +//! a single process lifetime is not a cache. + +use serde_json::Value; +use sha2::{Digest, Sha256}; + +/// Renders a finalized SHA-256 digest as a 64-character lowercase hex string. +pub(super) fn hex_digest(digest: impl AsRef<[u8]>) -> String { + digest + .as_ref() + .iter() + .map(|byte| format!("{byte:02x}")) + .collect() +} + +/// Folds one JSON `value` into `hasher` as a self-delimiting frame: an ASCII +/// domain `tag`, then the canonical byte length (little-endian `u64`), then the +/// canonical bytes. +/// +/// Canonicalizing per component keeps peak memory bounded by the single largest +/// value rather than the whole request tree, and the length prefix makes the +/// concatenation of frames unambiguous — no two distinct component sequences +/// can hash to the same byte stream. +pub(super) fn fold_canonical(hasher: &mut Sha256, tag: u8, value: Value) { + fold_bytes( + hasher, + tag, + &serde_json::to_vec(&canonical_value(value)).unwrap_or_default(), + ); +} + +/// Folds raw `bytes` into `hasher` as a self-delimiting `tag`-ed frame. +pub(super) fn fold_bytes(hasher: &mut Sha256, tag: u8, bytes: &[u8]) { + hasher.update([tag]); + hasher.update((bytes.len() as u64).to_le_bytes()); + hasher.update(bytes); +} + +/// Computes a deterministic FNV-1a 64-bit hash over `data` and returns it as +/// a 16-character lowercase hex string. +/// +/// FNV-1a uses a fixed, seed-free offset basis so the result is identical +/// across process restarts — unlike Rust's default `SipHash`, which is seeded +/// randomly at startup. It is used only for short local prompt-layout +/// fingerprints, not for response-cache identity. +pub(super) fn fnv1a_hex(data: &[u8]) -> String { + const OFFSET_BASIS: u64 = 14_695_981_039_346_656_037; + const PRIME: u64 = 1_099_511_628_211; + let mut hash = OFFSET_BASIS; + for &byte in data { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(PRIME); + } + format!("{hash:016x}") +} + +/// Recursively sorts the keys of every JSON object so that the serialized form +/// is canonical regardless of insertion order. +pub(super) fn canonical_value(v: Value) -> Value { + match v { + Value::Object(map) => { + let mut pairs: Vec<(String, Value)> = map.into_iter().collect(); + pairs.sort_by(|a, b| a.0.cmp(&b.0)); + Value::Object( + pairs + .into_iter() + .map(|(k, val)| (k, canonical_value(val))) + .collect(), + ) + } + Value::Array(arr) => Value::Array(arr.into_iter().map(canonical_value).collect()), + other => other, + } +} diff --git a/src/harness/cache/key.rs b/src/harness/cache/key.rs new file mode 100644 index 0000000..2444570 --- /dev/null +++ b/src/harness/cache/key.rs @@ -0,0 +1,319 @@ +//! Response-cache key derivation and provider prompt-cache breakpoints. +//! +//! # Why the key is a two-tuple, not just the prompt +//! +//! A cache keyed on the prompt alone is not a cache of an *answer*, it is a +//! cache of a *question*: the same question asked of a hosted frontier model +//! and of a local 3B model has two entirely different answers, and one shared +//! [`ResponseCache`][super::ResponseCache] would serve either to the other. +//! +//! LangChain looks up on `(prompt, llm_string)` — never the prompt alone — and +//! `llm_string` serializes the whole model object (class path, model name, +//! params). This module mirrors that: [`cache_key`] hashes the request, and +//! [`scoped_cache_key`] folds in the *resolved* model's +//! [`cache_identity`][crate::harness::model::ChatModel::cache_identity] plus the +//! streaming mode and the policy namespace. The composition is deliberate — the +//! request half can be computed before model resolution and reused, while the +//! identity half is only knowable after it. + +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use super::hash::{fold_bytes, fold_canonical, hex_digest}; +use crate::harness::model::ModelRequest; + +/// Produces a stable, deterministic **request-half** cache key for `request`. +/// +/// The key is a 64-character lowercase SHA-256 hex string built by folding the +/// request into the hasher **incrementally**, one component at a time: +/// 1. Fold each conversation message as its own length-prefixed, canonicalized +/// frame (tag `m`), preceded by the message count. +/// 2. Fold each tool schema likewise (tag `t`), preceded by the tool count. +/// 3. Fold an **explicit allowlist projection** of the remaining +/// behaviour-affecting parameters as one envelope frame (tag `E`). +/// +/// Folding per component bounds peak memory by the single largest component +/// instead of the entire request; transcripts routinely carry large tool +/// results. +/// +/// # Why an allowlist, not "everything that is left" +/// +/// The envelope used to be "whatever remains once `messages` and `tools` are +/// removed", which folded in fields that cannot affect what the model says: +/// [`tags`][ModelRequest::tags] (documented as propagated to events and +/// traces), [`timeout_ms`][ModelRequest::timeout_ms], +/// [`metadata`][ModelRequest::metadata] (free-form — a caller putting a run id +/// there, its natural use, got a permanent 0% hit rate with no diagnostic), +/// [`cache_policy`][ModelRequest::cache_policy] (which selects *whether* to +/// cache), [`prompt_fingerprint`][ModelRequest::prompt_fingerprint] (derived +/// from the messages already folded above), and +/// [`cache_segments`][ModelRequest::cache_segments] (pure annotation). +/// LangChain likewise keys on a deliberate projection and strips run-specific +/// message ids before hashing. +/// +/// The projection destructures [`ModelRequest`] **exhaustively** (no `..`), so +/// adding a field to the request is a compile error here until someone decides +/// whether it belongs in the key. That preserves the old "no field can silently +/// drop out" guarantee without the over-inclusion. +/// +/// # This is only half the key +/// +/// It carries no provider or model identity — the request's `model` field is an +/// optional *hint*, and the endpoint and credentials live inside the +/// `Arc`. Always compose with [`scoped_cache_key`] once the model +/// has actually been resolved. +/// +/// # Panics +/// Does not panic. If serialization unexpectedly fails, the affected frame +/// folds empty bytes; the key stays well-defined. +pub fn cache_key(request: &ModelRequest) -> String { + let mut hasher = Sha256::new(); + + // Messages: fold one at a time so a long transcript never materializes a + // second full tree. The count frame keeps `[a, b]` distinct from a single + // message that happens to serialize to the same concatenation. + // + // Each message is serialized individually. The previous implementation + // serialized the whole request and then `map.remove("messages")`-ed it + // inside an `if let Some(Value::Array(..))`: `remove` ran unconditionally, + // so a `messages` value that was not a JSON array would have been dropped + // *without being hashed* — the exact silent-drop false hit the doc comment + // promised could not happen. Serializing per message removes the shape + // assumption entirely. + hasher.update(b"M"); + hasher.update((request.messages.len() as u64).to_le_bytes()); + for message in &request.messages { + fold_canonical( + &mut hasher, + b'm', + serde_json::to_value(message).unwrap_or(Value::Null), + ); + } + + // Tool schemas: already name-sorted by `ToolRegistry::schemas`, so the + // order is deterministic across calls. + hasher.update(b"T"); + hasher.update((request.tools.len() as u64).to_le_bytes()); + for tool in &request.tools { + fold_canonical( + &mut hasher, + b't', + serde_json::to_value(tool).unwrap_or(Value::Null), + ); + } + + fold_canonical(&mut hasher, b'E', cache_key_envelope(request)); + hex_digest(hasher.finalize()) +} + +/// Builds the allowlist projection of the behaviour-affecting request +/// parameters folded into [`cache_key`]'s envelope frame. +/// +/// The exhaustive destructure below is load-bearing: it makes a new +/// [`ModelRequest`] field a compile error rather than a silent omission. +fn cache_key_envelope(request: &ModelRequest) -> Value { + let ModelRequest { + // Folded as their own frames by `cache_key`. + messages: _, + tools: _, + // ── Included: these change what the model produces ────────────────── + tool_choice, + response_format, + model, + model_hints, + reuse_previous_model, + temperature, + top_p, + max_tokens, + stop_sequences, + seed, + required_capabilities, + provider_options, + continuation_id, + reasoning, + // ── Excluded: cannot affect the answer ────────────────────────────── + // Transport deadline only. + timeout_ms: _, + // Free-form caller annotation; the natural place to put a run id. + metadata: _, + // Documented as propagated to events and traces. + tags: _, + // Pure annotation describing prompt structure. + cache_segments: _, + // Derived from `messages`/`tools`, both already folded above. + prompt_fingerprint: _, + // Selects *whether* to cache, never what is answered. + cache_policy: _, + } = request; + + serde_json::json!({ + "tool_choice": tool_choice, + "response_format": response_format, + "model": model, + "model_hints": model_hints, + "reuse_previous_model": reuse_previous_model, + "temperature": temperature, + "top_p": top_p, + "max_tokens": max_tokens, + "stop_sequences": stop_sequences, + "seed": seed, + "required_capabilities": required_capabilities, + "provider_options": provider_options, + "continuation_id": continuation_id, + "reasoning": reasoning, + }) +} + +/// Folds the *resolved* model's identity, the streaming mode, and an optional +/// policy namespace into a request-half key from [`cache_key`]. +/// +/// `identity` is the value returned by +/// [`ChatModel::cache_identity`][crate::harness::model::ChatModel::cache_identity] +/// on the model that is actually about to be (or was) called. `None` means the +/// model declines to identify itself; the key then folds a fixed +/// `"anonymous-model"` marker so a mixed registry of identifying and +/// non-identifying models still cannot cross-serve identifying ones. +/// +/// `streaming` is folded because it is a parameter of the call, not a field of +/// the request: a warm streaming run must not be served an entry written by a +/// unary run unless the caller has opted into that sharing. +pub fn scoped_cache_key( + request_key: &str, + identity: Option<&str>, + streaming: bool, + namespace: Option<&str>, +) -> String { + let mut hasher = Sha256::new(); + fold_bytes(&mut hasher, b'R', request_key.as_bytes()); + fold_bytes( + &mut hasher, + b'I', + identity.unwrap_or("anonymous-model").as_bytes(), + ); + fold_bytes(&mut hasher, b'S', if streaming { b"1" } else { b"0" }); + fold_bytes(&mut hasher, b'N', namespace.unwrap_or("").as_bytes()); + hex_digest(hasher.finalize()) +} + +/// A non-reversible, log-safe fingerprint of a credential. +/// +/// **Never** put a raw API key in a cache key, a log line, or an event: keys +/// leak through crash dumps, exported traces, and durable cache files. This +/// returns the first 16 hex characters of a domain-separated SHA-256 of +/// `secret`, which is enough to distinguish two credentials without carrying +/// either. An empty secret maps to the fixed token `"no-credential"` so a +/// keyless local runtime does not hash the empty string into something that +/// looks like a real fingerprint. +pub fn credential_fingerprint(secret: &str) -> String { + if secret.is_empty() { + return "no-credential".to_string(); + } + let mut hasher = Sha256::new(); + hasher.update(b"tinyagents.credential.v1\0"); + hasher.update(secret.as_bytes()); + hex_digest(hasher.finalize())[..16].to_string() +} + +/// Builds the canonical +/// [`cache_identity`][crate::harness::model::ChatModel::cache_identity] string +/// for a provider-backed model. +/// +/// The identity names everything that can make two models answer the same +/// prompt differently while sharing one cache: the provider family, the model +/// id, the API base URL, an optional organization/project scope, and a +/// *fingerprint* of the credential (two API keys can address two different +/// fine-tunes or two different tenants behind the same base URL). +/// +/// # Never log or store the raw credential +/// `credential` is passed through [`credential_fingerprint`] before it reaches +/// the digest, so neither the identity string nor any key derived from it can +/// carry the secret. +pub fn model_cache_identity( + provider: &str, + model: &str, + api_base: &str, + scope: Option<&str>, + credential: &str, +) -> String { + format!( + "{provider}|{model}|{api_base}|{}|{}", + scope.unwrap_or(""), + credential_fingerprint(credential) + ) +} + +// ── Provider prompt-cache breakpoints ──────────────────────────────────────── + +/// The provider option key carrying the routing hint for a prompt cache shard. +/// +/// Named after OpenAI's `prompt_cache_key`; adapters that speak a different +/// dialect read it from `provider_options` and lower it themselves. +pub const PROMPT_CACHE_KEY_OPTION: &str = "prompt_cache_key"; + +/// Derives a stable routing key for the request's cacheable prompt prefix, or +/// `None` when the request declares no stable prefix. +/// +/// A provider prompt cache is sharded: two requests that share a byte prefix +/// only actually hit the same cache if they are routed to the same shard, which +/// is what a `prompt_cache_key` buys. Deriving it from the *stable prefix* +/// fingerprint (rather than a random per-run id) means every turn of one +/// logical thread — and every sub-agent that inherits the same system prompt +/// and tool set — routes together, which is exactly the population that shares +/// a prefix. +pub fn prompt_cache_key(request: &ModelRequest) -> Option { + let layout = super::PromptCacheLayout::from_request(request); + if layout.prefix_ids().is_empty() { + return None; + } + Some(format!("tap-{}", layout.fingerprint())) +} + +/// Injects a derived [`PROMPT_CACHE_KEY_OPTION`] into `request.provider_options` +/// when the effective policy asks for prefix protection. +/// +/// This is the *active* half of the prompt-cache tooling: until now the layout +/// types only ever **observed** a prefix, while +/// [`CachePolicy::protect_prompt_prefix`][super::CachePolicy::protect_prompt_prefix] +/// had no reader anywhere in the crate and so could not change any behaviour. +/// +/// Precedence follows the rest of the crate: a caller who already set +/// `prompt_cache_key` in `provider_options` wins and is left untouched. +/// +/// Returns `true` when an option was injected. +pub fn apply_prompt_cache_breakpoints(request: &mut ModelRequest) -> bool { + let protect = request + .cache_policy + .as_ref() + .is_some_and(|policy| policy.protect_prompt_prefix); + if !protect { + return false; + } + if request + .provider_options + .get(PROMPT_CACHE_KEY_OPTION) + .is_some() + { + tracing::debug!( + "[cache] prompt_cache_key already set by caller; leaving provider_options untouched" + ); + return false; + } + let Some(derived) = prompt_cache_key(request) else { + tracing::debug!( + "[cache] protect_prompt_prefix is on but the request declares no cacheable prefix; \ + no prompt_cache_key derived" + ); + return false; + }; + if !request.provider_options.is_object() { + request.provider_options = Value::Object(serde_json::Map::new()); + } + if let Some(map) = request.provider_options.as_object_mut() { + map.insert( + PROMPT_CACHE_KEY_OPTION.to_string(), + Value::String(derived.clone()), + ); + } + tracing::debug!(prompt_cache_key = %derived, "[cache] injected provider prompt-cache breakpoint"); + true +} diff --git a/src/harness/cache/layout.rs b/src/harness/cache/layout.rs new file mode 100644 index 0000000..74e9cc4 --- /dev/null +++ b/src/harness/cache/layout.rs @@ -0,0 +1,166 @@ +//! Provider prompt / KV-cache layout protection. +//! +//! A provider's prompt cache is a **byte prefix** cache: it survives only while +//! the leading bytes of the request are unchanged, and appending to the tail is +//! the one edit that preserves it. This module models that rule directly. + +use serde_json::Value; + +use super::hash::fnv1a_hex; +use super::types::{CacheLayoutEvent, CachePolicy, PromptCacheLayout}; +use crate::harness::model::ModelRequest; + +impl PromptCacheLayout { + /// Builds a [`PromptCacheLayout`] from `request`. + /// + /// Captures three things: + /// + /// * the ordered ids of cacheable (stable) segments, + /// * a **content-aware** fingerprint — an FNV-1a hash over each cacheable + /// segment's `(id, role)` pair, the request's + /// [`prompt_fingerprint`][ModelRequest::prompt_fingerprint] (which + /// [`PromptBuilder::fingerprint`][crate::harness::prompt::PromptBuilder::fingerprint] + /// derives from the segments' actual messages), and the declared tool + /// schemas, and + /// * a per-message digest chain used by + /// [`Self::is_prefix_stable_against`]. + /// + /// The fingerprint used to hash the joined prefix **ids** only, so editing + /// the *text* of a stable segment — or swapping a tool schema — reported + /// "prefix stable" while the provider's KV prefix was already destroyed. + /// + /// # Cost + /// One serialization pass over the transcript, comparable to + /// [`super::cache_key`]. Call it once per middleware pass, not per message. + pub fn from_request(request: &ModelRequest) -> Self { + let prefix_ids: Vec = request.cacheable_prefix_ids(); + + // Segment identity *and* role/cacheability, so a role flip or a + // cacheable-flag flip on an otherwise identically named segment is not + // mistaken for "unchanged". + let mut material = String::new(); + for segment in &request.cache_segments { + material.push_str(&segment.id); + material.push('\u{1}'); + material.push_str( + match serde_json::to_value(segment.role) { + Ok(Value::String(role)) => role, + _ => String::new(), + } + .as_str(), + ); + material.push('\u{1}'); + material.push(if segment.cacheable { '1' } else { '0' }); + material.push('\u{2}'); + } + // Content of the stable prefix, when the builder computed it. + material.push_str(request.prompt_fingerprint.as_deref().unwrap_or("")); + material.push('\u{2}'); + // Tool declarations sit inside the stable prefix on every provider that + // caches prompts, so a schema edit invalidates it. + material.push_str(&serde_json::to_string(&request.tools).unwrap_or_default()); + + Self { + prefix_ids, + fingerprint: fnv1a_hex(material.as_bytes()), + message_digests: request + .messages + .iter() + .map(|message| { + fnv1a_hex(serde_json::to_vec(message).unwrap_or_default().as_slice()) + }) + .collect(), + } + } + + /// Returns the ordered ids of cacheable (stable) prefix segments. + pub fn prefix_ids(&self) -> &[String] { + &self.prefix_ids + } + + /// Returns the deterministic content-aware fingerprint of the stable + /// prefix (16 lowercase hex characters). + /// + /// Two layouts with the same segment identities *and* the same segment + /// content, tools, and roles produce the same fingerprint. + pub fn fingerprint(&self) -> &str { + &self.fingerprint + } + + /// Returns `true` when the provider's KV-cache prefix survives the move + /// from `self` to `other`. + /// + /// That requires **both**: + /// + /// 1. the same cacheable segment ids in the same order **with the same + /// content** (equal [`Self::fingerprint`]), and + /// 2. one message stream being a pure tail-extension of the other — the + /// only edit a byte-prefix cache tolerates. + /// + /// Comparing ids alone (the previous behaviour) reported stability after a + /// middleware rewrote a stable segment's text, which is the precise failure + /// this type exists to catch. + pub fn is_prefix_stable_against(&self, other: &PromptCacheLayout) -> bool { + if self.prefix_ids != other.prefix_ids || self.fingerprint != other.fingerprint { + return false; + } + let (shorter, longer) = if self.message_digests.len() <= other.message_digests.len() { + (&self.message_digests, &other.message_digests) + } else { + (&other.message_digests, &self.message_digests) + }; + longer.starts_with(shorter.as_slice()) + } + + /// Returns `true` when the segment identities match but the material they + /// carry does not — the silent invalidation an id-only comparison missed. + pub fn is_content_only_change(&self, other: &PromptCacheLayout) -> bool { + self.prefix_ids == other.prefix_ids && !self.is_prefix_stable_against(other) + } +} + +impl CacheLayoutEvent { + /// Constructs a [`CacheLayoutEvent`] by comparing `before` and `after` + /// layouts, filling in the computed flags automatically. + /// + /// `violates_policy` is always `false` here; use + /// [`Self::under_policy`] to evaluate the change against a + /// [`CachePolicy`]. + pub fn new(before: &PromptCacheLayout, after: &PromptCacheLayout) -> Self { + Self { + changed_prefix: !before.is_prefix_stable_against(after), + volatile_only: after.prefix_ids().is_empty(), + content_only_change: before.is_content_only_change(after), + violates_policy: false, + segment_ids_before: before.prefix_ids().to_vec(), + segment_ids_after: after.prefix_ids().to_vec(), + } + } + + /// Evaluates the `before` -> `after` change against `policy`. + /// + /// Returns `None` when the prefix survived. When it did not, the returned + /// event carries `violates_policy: true` iff + /// [`CachePolicy::protect_prompt_prefix`] was in force — which is what + /// makes that flag load-bearing instead of the inert struct field it was. + pub fn under_policy( + policy: &CachePolicy, + before: &PromptCacheLayout, + after: &PromptCacheLayout, + ) -> Option { + let mut event = Self::new(before, after); + if !event.changed_prefix { + return None; + } + event.violates_policy = policy.protect_prompt_prefix; + if event.violates_policy { + tracing::warn!( + content_only_change = event.content_only_change, + before = ?event.segment_ids_before, + after = ?event.segment_ids_after, + "[cache] prompt-cache prefix invalidated while protect_prompt_prefix was set" + ); + } + Some(event) + } +} diff --git a/src/harness/cache/memory.rs b/src/harness/cache/memory.rs new file mode 100644 index 0000000..75a2ebc --- /dev/null +++ b/src/harness/cache/memory.rs @@ -0,0 +1,197 @@ +//! The in-process [`InMemoryResponseCache`] implementation. + +use std::time::{Duration, Instant}; + +use async_trait::async_trait; + +use super::types::{CacheEntry, CacheStats, InMemoryResponseCache, LruResponseMap, ResponseCache}; +use crate::error::{Result, TinyAgentsError}; +use crate::harness::model::ModelResponse; + +impl InMemoryResponseCache { + /// Default LRU capacity, in entries, when constructed via + /// [`new`](Self::new) or [`Default`]. + pub const DEFAULT_CAPACITY: usize = 1024; + + /// Default approximate byte budget (64 MiB). + /// + /// An entry count alone does not bound memory: 1024 long-context responses + /// carrying large tool payloads is hundreds of megabytes. + pub const DEFAULT_MAX_BYTES: usize = 64 * 1024 * 1024; + + /// Creates a new, empty in-memory response cache bounded by + /// [`DEFAULT_CAPACITY`](Self::DEFAULT_CAPACITY) entries and + /// [`DEFAULT_MAX_BYTES`](Self::DEFAULT_MAX_BYTES) bytes. + pub fn new() -> Self { + Self::with_capacity(Self::DEFAULT_CAPACITY) + } + + /// Creates a new, empty in-memory response cache retaining at most + /// `capacity` entries (least-recently-used evicted first). A `capacity` of + /// zero is treated as `1` so the cache always retains the last write. + pub fn with_capacity(capacity: usize) -> Self { + Self::with_bounds(capacity, Self::DEFAULT_MAX_BYTES) + } + + /// Creates a new, empty in-memory response cache bounded by **both** an + /// entry count and an approximate byte budget. Whichever bound trips first + /// evicts the least-recently-used entry. + /// + /// Both bounds are clamped to a minimum of `1` so the cache always retains + /// the most recent write. + pub fn with_bounds(capacity: usize, max_bytes: usize) -> Self { + Self { + inner: std::sync::Arc::new(std::sync::Mutex::new(LruResponseMap { + data: std::collections::HashMap::new(), + order: std::collections::BTreeMap::new(), + next_recency: 0, + capacity: capacity.max(1), + max_bytes: max_bytes.max(1), + bytes: 0, + stats: CacheStats::default(), + })), + } + } + + /// Locks the inner map, mapping a poisoned mutex to a validation error. + fn lock(&self) -> Result> { + self.inner + .lock() + .map_err(|e| TinyAgentsError::Validation(format!("cache lock poisoned: {e}"))) + } +} + +impl Default for InMemoryResponseCache { + fn default() -> Self { + Self::new() + } +} + +impl LruResponseMap { + /// Hands out the next monotonic recency ticket. + fn tick(&mut self) -> u64 { + self.next_recency = self.next_recency.wrapping_add(1); + self.next_recency + } + + /// Moves `key` to the most-recently-used end. + /// + /// `O(log n)`: one `BTreeMap` removal plus one insertion. The previous + /// implementation scanned a `VecDeque` linearly — up to `capacity` `String` + /// comparisons plus a memmove — on **every** hit. + fn touch(&mut self, key: &str) { + let Some(entry) = self.data.get(key) else { + return; + }; + let old = entry.recency; + let next = self.tick(); + self.order.remove(&old); + self.order.insert(next, key.to_string()); + if let Some(entry) = self.data.get_mut(key) { + entry.recency = next; + } + } + + /// Removes `key`, keeping the byte accounting and order index consistent. + fn remove(&mut self, key: &str) -> Option { + let entry = self.data.remove(key)?; + self.order.remove(&entry.recency); + self.bytes = self.bytes.saturating_sub(entry.bytes); + Some(entry) + } + + /// Evicts least-recently-used entries until both bounds are satisfied. + fn evict_to_fit(&mut self) { + while self.data.len() > self.capacity + || (self.bytes > self.max_bytes && self.data.len() > 1) + { + let Some((_, victim)) = self.order.iter().next().map(|(k, v)| (*k, v.clone())) else { + break; + }; + self.remove(&victim); + self.stats.evictions = self.stats.evictions.saturating_add(1); + tracing::trace!(key = %victim, "[cache] evicted least-recently-used entry"); + } + } + + /// Refreshes the derived size counters exposed through + /// [`ResponseCache::stats`]. + fn sync_size_stats(&mut self) { + self.stats.entries = self.data.len() as u64; + self.stats.bytes = self.bytes as u64; + } +} + +#[async_trait] +impl ResponseCache for InMemoryResponseCache { + async fn get(&self, key: &str) -> Result> { + let mut inner = self.lock()?; + // Lazy expiry: an entry whose TTL elapsed is a miss and is dropped on + // the way past, so a cache that is read but never written still sheds + // stale entries. + if let Some(entry) = inner.data.get(key) + && entry.expires_at.is_some_and(|at| at <= Instant::now()) + { + inner.remove(key); + inner.stats.expirations = inner.stats.expirations.saturating_add(1); + inner.stats.misses = inner.stats.misses.saturating_add(1); + inner.sync_size_stats(); + tracing::debug!(key = %key, "[cache] entry expired; treating as miss"); + return Ok(None); + } + let hit = inner.data.get(key).map(|entry| entry.value.clone()); + if hit.is_some() { + inner.touch(key); + inner.stats.hits = inner.stats.hits.saturating_add(1); + } else { + inner.stats.misses = inner.stats.misses.saturating_add(1); + } + Ok(hit) + } + + async fn put(&self, key: &str, value: ModelResponse) -> Result<()> { + self.put_with_ttl(key, value, None).await + } + + async fn put_with_ttl( + &self, + key: &str, + value: ModelResponse, + ttl: Option, + ) -> Result<()> { + let bytes = serde_json::to_vec(&value).map(|v| v.len()).unwrap_or(0); + let mut inner = self.lock()?; + inner.remove(key); + let recency = inner.tick(); + inner.order.insert(recency, key.to_string()); + inner.bytes = inner.bytes.saturating_add(bytes); + inner.data.insert( + key.to_string(), + CacheEntry { + value, + recency, + bytes, + expires_at: ttl.map(|ttl| Instant::now() + ttl), + }, + ); + inner.stats.writes = inner.stats.writes.saturating_add(1); + inner.evict_to_fit(); + inner.sync_size_stats(); + Ok(()) + } + + async fn clear(&self) -> Result<()> { + let mut inner = self.lock()?; + let dropped = inner.data.len(); + inner.data.clear(); + inner.order.clear(); + inner.bytes = 0; + inner.sync_size_stats(); + tracing::debug!(dropped, "[cache] cleared every in-memory response entry"); + Ok(()) + } + + fn stats(&self) -> CacheStats { + self.lock().map(|inner| inner.stats).unwrap_or_default() + } +} diff --git a/src/harness/cache/mod.rs b/src/harness/cache/mod.rs index a9cd8f4..9e5315a 100644 --- a/src/harness/cache/mod.rs +++ b/src/harness/cache/mod.rs @@ -12,282 +12,52 @@ //! # Two distinct caching concerns //! //! ## 1. Local response cache -//! [`ResponseCache`] + [`InMemoryResponseCache`] let the harness skip provider -//! API calls entirely when it has already seen an identical request. Use -//! [`cache_key`] to produce a stable, deterministic key from a -//! [`crate::harness::model::ModelRequest`]. +//! [`ResponseCache`] + [`InMemoryResponseCache`] (and, behind the `sqlite` +//! feature, [`SqliteResponseCache`]) let the harness skip provider API calls +//! entirely when it has already seen an identical request. +//! +//! The key is a **two-part composition**, never the prompt alone: +//! +//! ```text +//! scoped_cache_key(cache_key(request), model.cache_identity(), streaming, ns) +//! ``` +//! +//! [`cache_key`] hashes an explicit allowlist projection of the request; +//! [`scoped_cache_key`] folds in the *resolved* model's identity (provider, +//! model id, endpoint, credential fingerprint — never a raw credential), the +//! streaming mode, and the policy namespace. Without the identity half, one +//! `Arc` shared between a hosted and a local harness +//! serves either's answer to the other. //! //! ## 2. Provider prompt / KV-cache layout protection -//! [`PromptCacheLayout`] records the ordered cacheable prefix of a request. -//! [`CacheLayoutEvent`] describes mutations so middleware can signal whether it -//! preserved or invalidated the provider's KV-cache prefix. -//! [`CachePolicy`] toggles both concerns at the call-site level. +//! [`PromptCacheLayout`] records the ordered cacheable prefix of a request +//! *and* a digest of the material it carries, so a middleware that rewrites a +//! stable segment's text can no longer report "prefix stable". +//! [`CacheLayoutEvent`] describes mutations, and +//! [`CacheLayoutEvent::under_policy`] plus [`apply_prompt_cache_breakpoints`] +//! make [`CachePolicy::protect_prompt_prefix`] load-bearing rather than inert. //! +//! ## 3. Stampede protection +//! [`SingleFlight`] collapses concurrent identical misses into one provider +//! call. + +mod hash; +mod key; +mod layout; +mod memory; +mod singleflight; +#[cfg(feature = "sqlite")] +mod sqlite; mod types; -use async_trait::async_trait; -use serde_json::Value; -use sha2::{Digest, Sha256}; - +pub use key::{ + PROMPT_CACHE_KEY_OPTION, apply_prompt_cache_breakpoints, cache_key, credential_fingerprint, + model_cache_identity, prompt_cache_key, scoped_cache_key, +}; +pub use singleflight::SingleFlight; +#[cfg(feature = "sqlite")] +pub use sqlite::SqliteResponseCache; pub use types::*; -use crate::error::{Result, TinyAgentsError}; -use crate::harness::model::{ModelRequest, ModelResponse}; - -// ── Deterministic hash ──────────────────────────────────────────────────────── - -/// Renders a finalized SHA-256 digest as a 64-character lowercase hex string. -fn hex_digest(digest: impl AsRef<[u8]>) -> String { - digest - .as_ref() - .iter() - .map(|byte| format!("{byte:02x}")) - .collect() -} - -/// Folds one JSON `value` into `hasher` as a self-delimiting frame: an ASCII -/// domain `tag`, then the canonical byte length (little-endian `u64`), then the -/// canonical bytes. -/// -/// Canonicalizing per component keeps peak memory bounded by the single largest -/// value rather than the whole request tree, and the length prefix makes the -/// concatenation of frames unambiguous — no two distinct component sequences -/// can hash to the same byte stream. -fn fold_canonical(hasher: &mut Sha256, tag: u8, value: Value) { - let bytes = serde_json::to_vec(&canonical_value(value)).unwrap_or_default(); - hasher.update([tag]); - hasher.update((bytes.len() as u64).to_le_bytes()); - hasher.update(&bytes); -} - -/// Computes a deterministic FNV-1a 64-bit hash over `data` and returns it as -/// a 16-character lowercase hex string. -/// -/// FNV-1a uses a fixed, seed-free offset basis so the result is identical -/// across process restarts — unlike Rust's default `SipHash`, which is seeded -/// randomly at startup. It is used only for short local prompt-layout -/// fingerprints, not for response-cache identity. -fn fnv1a_hex(data: &[u8]) -> String { - const OFFSET_BASIS: u64 = 14_695_981_039_346_656_037; - const PRIME: u64 = 1_099_511_628_211; - let mut hash = OFFSET_BASIS; - for &byte in data { - hash ^= u64::from(byte); - hash = hash.wrapping_mul(PRIME); - } - format!("{hash:016x}") -} - -/// Recursively sorts the keys of every JSON object so that the serialized form -/// is canonical regardless of insertion order. -fn canonical_value(v: Value) -> Value { - match v { - Value::Object(map) => { - let mut pairs: Vec<(String, Value)> = map.into_iter().collect(); - pairs.sort_by(|a, b| a.0.cmp(&b.0)); - Value::Object( - pairs - .into_iter() - .map(|(k, val)| (k, canonical_value(val))) - .collect(), - ) - } - Value::Array(arr) => Value::Array(arr.into_iter().map(canonical_value).collect()), - other => other, - } -} - -// ── cache_key ───────────────────────────────────────────────────────────────── - -/// Produces a stable, deterministic cache key for `request`. -/// -/// The key is a 64-character lowercase SHA-256 hex string built by folding the -/// request into the hasher **incrementally**, one component at a time: -/// 1. Serialize `request` once to a [`serde_json::Value`]. -/// 2. Fold each conversation message as its own length-prefixed, canonicalized -/// frame (tag `M`), preceded by the message count. -/// 3. Fold each tool schema likewise (tag `T`), preceded by the tool count. -/// 4. Fold the remaining scalar/parameter fields — everything left in the -/// request object once `messages` and `tools` are removed — as one envelope -/// frame (tag `E`). -/// -/// This avoids the previous approach's three simultaneous whole-transcript -/// allocations (a full `Value` tree, a full canonical rebuild, and a full byte -/// buffer). Transcripts routinely carry large tool results; canonicalizing and -/// serializing per component bounds peak memory by the single largest component -/// instead of the entire request. The envelope is taken as "whatever remains" -/// so that any future [`ModelRequest`] field automatically participates in the -/// key — no field can silently drop out and cause a false cache hit. -/// -/// Distinct requests still map to distinct keys (modulo SHA-256 collision -/// resistance): per-component length prefixes make the frame stream -/// unambiguous, and every behavior-affecting field is folded exactly once. -/// -/// # Panics -/// Does not panic. If serialization unexpectedly fails, the affected frame -/// folds empty bytes; the key stays well-defined. -pub fn cache_key(request: &ModelRequest) -> String { - let mut hasher = Sha256::new(); - let mut root = serde_json::to_value(request).unwrap_or(Value::Null); - - if let Value::Object(map) = &mut root { - // Messages: fold one at a time so a long transcript never materializes - // a second full tree. The count frame keeps `[a, b]` distinct from a - // single message that happens to serialize to the same concatenation. - if let Some(Value::Array(messages)) = map.remove("messages") { - hasher.update(b"M"); - hasher.update((messages.len() as u64).to_le_bytes()); - for message in messages { - fold_canonical(&mut hasher, b'm', message); - } - } - // Tool schemas: already name-sorted by `ToolRegistry::schemas`, so the - // order is deterministic across calls. - if let Some(Value::Array(tools)) = map.remove("tools") { - hasher.update(b"T"); - hasher.update((tools.len() as u64).to_le_bytes()); - for tool in tools { - fold_canonical(&mut hasher, b't', tool); - } - } - } - - // Envelope: every remaining scalar/parameter field in one frame. - fold_canonical(&mut hasher, b'E', root); - hex_digest(hasher.finalize()) -} - -// ── InMemoryResponseCache ───────────────────────────────────────────────────── - -impl InMemoryResponseCache { - /// Default LRU capacity when constructed via [`new`](Self::new) or - /// [`Default`]. - pub const DEFAULT_CAPACITY: usize = 1024; - - /// Creates a new, empty in-memory response cache bounded by - /// [`DEFAULT_CAPACITY`](Self::DEFAULT_CAPACITY) entries. - pub fn new() -> Self { - Self::with_capacity(Self::DEFAULT_CAPACITY) - } - - /// Creates a new, empty in-memory response cache retaining at most - /// `capacity` entries (least-recently-used evicted first). A `capacity` of - /// zero is treated as `1` so the cache always retains the last write. - pub fn with_capacity(capacity: usize) -> Self { - Self { - inner: std::sync::Arc::new(std::sync::Mutex::new(LruResponseMap { - data: std::collections::HashMap::new(), - order: std::collections::VecDeque::new(), - capacity: capacity.max(1), - })), - } - } -} - -impl Default for InMemoryResponseCache { - fn default() -> Self { - Self::new() - } -} - -impl LruResponseMap { - /// Moves `key` to the most-recently-used end of the order queue. - fn touch(&mut self, key: &str) { - if let Some(pos) = self.order.iter().position(|k| k == key) { - let k = self.order.remove(pos).expect("position is valid"); - self.order.push_back(k); - } - } -} - -#[async_trait] -impl ResponseCache for InMemoryResponseCache { - async fn get(&self, key: &str) -> Result> { - let mut inner = self - .inner - .lock() - .map_err(|e| TinyAgentsError::Validation(format!("cache lock poisoned: {e}")))?; - let hit = inner.data.get(key).cloned(); - if hit.is_some() { - inner.touch(key); - } - Ok(hit) - } - - async fn put(&self, key: &str, value: ModelResponse) -> Result<()> { - let mut inner = self - .inner - .lock() - .map_err(|e| TinyAgentsError::Validation(format!("cache lock poisoned: {e}")))?; - if inner.data.insert(key.to_string(), value).is_some() { - // Existing key: refresh its recency without changing the length. - inner.touch(key); - } else { - inner.order.push_back(key.to_string()); - // Evict least-recently-used entries until within capacity. - while inner.order.len() > inner.capacity { - if let Some(evicted) = inner.order.pop_front() { - inner.data.remove(&evicted); - } - } - } - Ok(()) - } -} - -// ── PromptCacheLayout ───────────────────────────────────────────────────────── - -impl PromptCacheLayout { - /// Builds a [`PromptCacheLayout`] from `request` by collecting the ids of - /// all cacheable (stable) segments in their declared order. - /// - /// The fingerprint is a deterministic FNV-1a hash of the joined prefix ids - /// so regression tests can assert prefix stability independently of the - /// full request hash. - pub fn from_request(request: &ModelRequest) -> Self { - let prefix_ids: Vec = request.cacheable_prefix_ids(); - let fingerprint = fnv1a_hex(prefix_ids.join(",").as_bytes()); - Self { - prefix_ids, - fingerprint, - } - } - - /// Returns the ordered ids of cacheable (stable) prefix segments. - pub fn prefix_ids(&self) -> &[String] { - &self.prefix_ids - } - - /// Returns the deterministic fingerprint of the ordered prefix ids. - /// - /// Two layouts with identical `prefix_ids` produce the same fingerprint. - pub fn fingerprint(&self) -> &str { - &self.fingerprint - } - - /// Returns `true` if `self` and `other` have the same cacheable prefix ids - /// in the same order, meaning the provider KV-cache prefix is stable - /// across the two requests. - pub fn is_prefix_stable_against(&self, other: &PromptCacheLayout) -> bool { - self.prefix_ids == other.prefix_ids - } -} - -// ── CacheLayoutEvent ────────────────────────────────────────────────────────── - -impl CacheLayoutEvent { - /// Constructs a [`CacheLayoutEvent`] by comparing `before` and `after` - /// layouts, filling in the computed `changed_prefix` and `volatile_only` - /// flags automatically. - pub fn new(before: &PromptCacheLayout, after: &PromptCacheLayout) -> Self { - Self { - changed_prefix: !before.is_prefix_stable_against(after), - volatile_only: after.prefix_ids().is_empty(), - segment_ids_before: before.prefix_ids().to_vec(), - segment_ids_after: after.prefix_ids().to_vec(), - } - } -} - #[cfg(test)] mod test; diff --git a/src/harness/cache/singleflight.rs b/src/harness/cache/singleflight.rs new file mode 100644 index 0000000..d9c8af7 --- /dev/null +++ b/src/harness/cache/singleflight.rs @@ -0,0 +1,171 @@ +//! Stampede protection: collapse concurrent identical cache misses into one +//! provider call. +//! +//! # The gap this fills +//! +//! A cache in front of a slow provider does nothing for the *first* N callers. +//! Ten sub-agents that ask the same sub-question at the same moment all miss, +//! all call the provider, and nine of those calls are pure waste — paid for, +//! rate-limit consuming, and then thrown away when the last writer wins. +//! +//! Neither reference implementation solves this: LangChain's own unit test +//! documents the race rather than preventing it. So this is greenfield, and +//! deliberately small — one in-flight map plus a `tokio::sync::broadcast` +//! channel per key. +//! +//! # Semantics +//! +//! The first caller for a key becomes the **leader** and runs the closure. Any +//! caller arriving while the leader is in flight becomes a **follower** and +//! waits for the leader's outcome instead of running the closure. Followers +//! receive a clone of the leader's success. +//! +//! Errors are **not** shared: a follower whose leader failed re-runs the +//! closure itself. Sharing the failure would turn one caller's transient 503 +//! into every concurrent caller's failure while hiding the fact that each had +//! its own retry budget — and an error is not a value worth caching. + +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use tokio::sync::broadcast; + +use crate::Result; +use crate::harness::model::ModelResponse; + +/// What a leader broadcasts to its followers when it finishes. +#[derive(Clone, Debug)] +enum Outcome { + /// The leader succeeded; followers take this response. + Ready(Box), + /// The leader failed; followers must run the call themselves. + Failed, +} + +/// Collapses concurrent duplicate model calls that share a cache key. +/// +/// Cheap to clone; clones share the same in-flight map. +#[derive(Clone, Default)] +pub struct SingleFlight { + inflight: Arc>>>, +} + +/// Removes a leader's entry even when its future is cancelled mid-call. +/// +/// Dropping the sender closes every follower receiver, which makes followers +/// retry their own call instead of waiting forever behind an abandoned leader. +struct LeaderGuard { + inflight: Arc>>>, + key: String, +} + +impl LeaderGuard { + fn retire(&mut self) -> Option> { + self.inflight.lock().ok()?.remove(&self.key) + } +} + +impl Drop for LeaderGuard { + fn drop(&mut self) { + let _ = self.retire(); + } +} + +impl std::fmt::Debug for SingleFlight { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let inflight = self.inflight.lock().map(|m| m.len()).unwrap_or(0); + f.debug_struct("SingleFlight") + .field("inflight", &inflight) + .finish() + } +} + +impl SingleFlight { + /// Creates an empty single-flight gate. + pub fn new() -> Self { + Self::default() + } + + /// Number of keys currently in flight. Intended for tests and diagnostics. + pub fn inflight_len(&self) -> usize { + self.inflight.lock().map(|m| m.len()).unwrap_or(0) + } + + /// Claims leadership of `key`, or subscribes to the current leader. + /// + /// Returns `None` when the map is poisoned, `Some(None)` when this caller + /// is the leader, and `Some(Some(receiver))` when it is a follower. The + /// lock never escapes this function, so no guard is held across an await. + #[allow(clippy::option_option)] + fn claim(&self, key: &str) -> Option>> { + let mut inflight = self.inflight.lock().ok()?; + Some(match inflight.get(key) { + Some(sender) => Some(sender.subscribe()), + None => { + let (sender, _) = broadcast::channel(1); + inflight.insert(key.to_string(), sender); + None + } + }) + } + + /// Runs `call` for `key`, or waits for an already in-flight call with the + /// same key and returns its result. + /// + /// Returns `(response, was_follower)` so the caller can skip a redundant + /// cache write when it merely rode along on someone else's call. + pub async fn run(&self, key: &str, call: F) -> Result<(ModelResponse, bool)> + where + F: FnOnce() -> Fut, + Fut: Future>, + { + // The lock is acquired and released inside this helper so no + // `MutexGuard` is ever alive across an `await` — which would make the + // whole future `!Send` and unusable from `tokio::spawn`. + let claim = self.claim(key); + let Some(claim) = claim else { + // A poisoned map must never take the run down: fall back to simply + // making the call, which is the un-collapsed behaviour. + tracing::warn!("[cache] single-flight map poisoned; issuing the model call directly"); + return call().await.map(|response| (response, false)); + }; + let mut receiver = claim; + + // Follower: wait for the leader rather than duplicating the call. + if let Some(receiver) = receiver.as_mut() { + tracing::debug!(key = %key, "[cache] joining an in-flight identical model call"); + match receiver.recv().await { + Ok(Outcome::Ready(response)) => return Ok((*response, true)), + // Leader failed, or dropped the channel without sending (a + // cancelled or panicking leader). Either way, run it ourselves. + Ok(Outcome::Failed) | Err(_) => { + tracing::debug!( + key = %key, + "[cache] in-flight leader did not produce a response; issuing our own call" + ); + return call().await.map(|response| (response, false)); + } + } + } + + // Leader: install the guard before awaiting. If this future is dropped, + // the guard retires the sender and wakes followers through channel + // closure; without it, a cancelled leader wedges this key forever. + let mut leader = LeaderGuard { + inflight: Arc::clone(&self.inflight), + key: key.to_string(), + }; + let result = call().await; + let sender = leader.retire(); + if let Some(sender) = sender { + let outcome = match &result { + Ok(response) => Outcome::Ready(Box::new(response.clone())), + Err(_) => Outcome::Failed, + }; + // `send` fails only when no follower is subscribed, which is the + // common (uncontended) case — not an error. + let _ = sender.send(outcome); + } + result.map(|response| (response, false)) + } +} diff --git a/src/harness/cache/sqlite.rs b/src/harness/cache/sqlite.rs new file mode 100644 index 0000000..af589cd --- /dev/null +++ b/src/harness/cache/sqlite.rs @@ -0,0 +1,222 @@ +//! SQLite-backed [`ResponseCache`] — a durable response cache behind the +//! optional `sqlite` cargo feature. +//! +//! [`InMemoryResponseCache`][super::InMemoryResponseCache] loses everything when +//! the process exits, so every restart pays the whole provider bill again. This +//! backend keeps the same entries in a `response_cache` table keyed by +//! `(ns, key)`, mirroring LangGraph's `SqliteCache`: WAL journalling, an +//! `expiry` column, a lazy expiry purge on read, and `INSERT OR REPLACE` on +//! write. +//! +//! Expiry is stored as an absolute **Unix epoch millisecond** timestamp rather +//! than a monotonic instant, because the value has to survive a restart — a +//! `std::time::Instant` is meaningless in the next process. + +use std::path::Path; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use async_trait::async_trait; +use rusqlite::{Connection, OptionalExtension, params}; + +use super::types::{CacheStats, ResponseCache}; +use crate::harness::model::ModelResponse; +use crate::{Result, TinyAgentsError}; + +/// Table + index DDL. `(ns, key)` is the primary key so a namespaced population +/// can be dropped wholesale, and `expiry` is indexed so the periodic purge does +/// not table-scan. +const SCHEMA: &str = " +CREATE TABLE IF NOT EXISTS response_cache ( + ns TEXT NOT NULL, + key TEXT NOT NULL, + value TEXT NOT NULL, + expiry INTEGER, + PRIMARY KEY (ns, key) +); +CREATE INDEX IF NOT EXISTS response_cache_expiry ON response_cache (expiry); +"; + +/// A durable [`ResponseCache`] backed by SQLite. +/// +/// Cheap to clone; clones share the same underlying connection (and therefore +/// the same data, including for in-memory databases). +/// +/// # Namespacing +/// Every handle carries a namespace (default `"default"`). Two harnesses that +/// must not cross-serve — different tenants, a control and an experiment arm — +/// point at the same file with different namespaces and +/// [`clear`][ResponseCache::clear] then drops only their own population. +#[derive(Clone)] +pub struct SqliteResponseCache { + conn: Arc>, + namespace: String, +} + +fn sqlite_err(context: &str, err: impl std::fmt::Display) -> TinyAgentsError { + TinyAgentsError::Validation(format!("sqlite response cache: {context}: {err}")) +} + +/// Current wall-clock time in Unix epoch milliseconds. +fn now_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as i64) + .unwrap_or(0) +} + +impl SqliteResponseCache { + /// The namespace used when none is given. + pub const DEFAULT_NAMESPACE: &'static str = "default"; + + /// Opens (creating if needed) a SQLite-backed response cache at `path`. + /// + /// Pass `":memory:"` for an ephemeral in-memory database. + pub fn open(path: impl AsRef) -> Result { + let conn = Connection::open(path.as_ref()).map_err(|e| sqlite_err("open database", e))?; + Self::from_connection(conn) + } + + /// Opens an ephemeral in-memory cache (`":memory:"`). + /// + /// The database lives only as long as this handle and its clones, which + /// share the single underlying connection. + pub fn in_memory() -> Result { + let conn = Connection::open_in_memory().map_err(|e| sqlite_err("open in-memory", e))?; + Self::from_connection(conn) + } + + /// Wraps a caller-owned open [`Connection`], ensuring the schema exists. + /// + /// WAL is requested but not required: an in-memory database rejects it, and + /// a read-only mount may too, so a failure here is logged and ignored + /// rather than failing the open — journalling mode is a performance knob, + /// not a correctness one. + pub fn from_connection(conn: Connection) -> Result { + if let Err(error) = conn.pragma_update(None, "journal_mode", "WAL") { + tracing::debug!(%error, "[cache] sqlite WAL unavailable; continuing with the default journal mode"); + } + conn.execute_batch(SCHEMA) + .map_err(|e| sqlite_err("create schema", e))?; + Ok(Self { + conn: Arc::new(Mutex::new(conn)), + namespace: Self::DEFAULT_NAMESPACE.to_string(), + }) + } + + /// Returns the table + index DDL as a reusable, dependency-free SQL string, + /// for applications that own their own SQLite connection at a possibly + /// incompatible native-link version. + pub fn schema_sql() -> &'static str { + SCHEMA + } + + /// Returns a handle scoped to `namespace`, sharing this handle's + /// connection. + pub fn with_namespace(&self, namespace: impl Into) -> Self { + Self { + conn: Arc::clone(&self.conn), + namespace: namespace.into(), + } + } + + /// Returns this handle's namespace. + pub fn namespace(&self) -> &str { + &self.namespace + } + + fn lock(&self) -> Result> { + self.conn + .lock() + .map_err(|_| sqlite_err("connection lock", "poisoned")) + } +} + +#[async_trait] +impl ResponseCache for SqliteResponseCache { + async fn get(&self, key: &str) -> Result> { + let conn = self.lock()?; + let row: Option<(String, Option)> = conn + .query_row( + "SELECT value, expiry FROM response_cache WHERE ns = ?1 AND key = ?2", + params![self.namespace, key], + |row| Ok((row.get(0)?, row.get(1)?)), + ) + .optional() + .map_err(|e| sqlite_err("read entry", e))?; + let Some((value, expiry)) = row else { + return Ok(None); + }; + // Lazy expiry purge: a stale row is deleted on the way past, so a cache + // that is read but never written still sheds expired entries. + if expiry.is_some_and(|at| at <= now_ms()) { + conn.execute( + "DELETE FROM response_cache WHERE ns = ?1 AND key = ?2", + params![self.namespace, key], + ) + .map_err(|e| sqlite_err("purge expired entry", e))?; + tracing::debug!(key = %key, "[cache] sqlite entry expired; treating as miss"); + return Ok(None); + } + let response: ModelResponse = + serde_json::from_str(&value).map_err(|e| sqlite_err("decode entry", e))?; + Ok(Some(response)) + } + + async fn put(&self, key: &str, value: ModelResponse) -> Result<()> { + self.put_with_ttl(key, value, None).await + } + + async fn put_with_ttl( + &self, + key: &str, + value: ModelResponse, + ttl: Option, + ) -> Result<()> { + let encoded = serde_json::to_string(&value).map_err(|e| sqlite_err("encode entry", e))?; + let expiry = ttl.map(|ttl| now_ms().saturating_add(ttl.as_millis() as i64)); + let conn = self.lock()?; + conn.execute( + "INSERT OR REPLACE INTO response_cache (ns, key, value, expiry) \ + VALUES (?1, ?2, ?3, ?4)", + params![self.namespace, key, encoded, expiry], + ) + .map_err(|e| sqlite_err("write entry", e))?; + Ok(()) + } + + async fn clear(&self) -> Result<()> { + let conn = self.lock()?; + let dropped = conn + .execute( + "DELETE FROM response_cache WHERE ns = ?1", + params![self.namespace], + ) + .map_err(|e| sqlite_err("clear namespace", e))?; + tracing::debug!( + namespace = %self.namespace, + dropped, + "[cache] cleared the sqlite response cache namespace" + ); + Ok(()) + } + + fn stats(&self) -> CacheStats { + let Ok(conn) = self.lock() else { + return CacheStats::default(); + }; + let row: rusqlite::Result<(i64, i64)> = conn.query_row( + "SELECT COUNT(*), COALESCE(SUM(LENGTH(value)), 0) FROM response_cache WHERE ns = ?1", + params![self.namespace], + |row| Ok((row.get(0)?, row.get(1)?)), + ); + match row { + Ok((entries, bytes)) => CacheStats { + entries: entries.max(0) as u64, + bytes: bytes.max(0) as u64, + ..CacheStats::default() + }, + Err(_) => CacheStats::default(), + } + } +} diff --git a/src/harness/cache/types.rs b/src/harness/cache/types.rs index 11be36f..8716083 100644 --- a/src/harness/cache/types.rs +++ b/src/harness/cache/types.rs @@ -18,8 +18,9 @@ //! //! All public types in this module are re-exported through [`super`]. -use std::collections::{HashMap, VecDeque}; +use std::collections::{BTreeMap, HashMap}; use std::sync::{Arc, Mutex}; +use std::time::Duration; use async_trait::async_trait; use serde::{Deserialize, Serialize}; @@ -27,46 +28,181 @@ use serde::{Deserialize, Serialize}; use crate::error::Result; use crate::harness::model::ModelResponse; +// ── CacheStats ──────────────────────────────────────────────────────────────── + +/// Point-in-time counters describing a [`ResponseCache`]'s behaviour. +/// +/// A caller whose hit rate is unexpectedly 0% has no way to tell an +/// over-inclusive key from a policy that never enabled caching at all; these +/// counters plus [`CacheSkipReason`] are the diagnostic. Implementations that +/// cannot cheaply account for a field leave it at zero. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct CacheStats { + /// Lookups that returned a live entry. + pub hits: u64, + /// Lookups that returned nothing (including entries dropped as expired). + pub misses: u64, + /// Writes accepted by the cache. + pub writes: u64, + /// Entries dropped because they exceeded a capacity bound. + pub evictions: u64, + /// Entries dropped because their TTL had elapsed. + pub expirations: u64, + /// Entries currently retained. + pub entries: u64, + /// Approximate serialized size of the retained entries, in bytes. + pub bytes: u64, +} + +// ── CacheSkipReason ─────────────────────────────────────────────────────────── + +/// Why the agent loop declined to consult the response cache for a call. +/// +/// `docs/modules/harness/cache.md` specifies a richer decision surface than the +/// bare `CacheHit`/`CacheMiss` pair; this enum carries the "no lookup happened +/// at all" half so a 0% hit rate is always explainable. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum CacheSkipReason { + /// No [`ResponseCache`] is attached to the harness. + NoCacheAttached, + /// The effective [`CachePolicy`] disables response caching for this call. + PolicyDisabled, + /// The transcript already contains an assistant or tool turn, so the + /// request is unique to this run and can never be re-served. + MultiTurnTranscript, +} + +impl CacheSkipReason { + /// A short, stable, grep-friendly token for logs and events. + pub fn as_str(self) -> &'static str { + match self { + CacheSkipReason::NoCacheAttached => "no_cache_attached", + CacheSkipReason::PolicyDisabled => "policy_disabled", + CacheSkipReason::MultiTurnTranscript => "multi_turn_transcript", + } + } +} + // ── ResponseCache ───────────────────────────────────────────────────────────── /// Local response cache that lets the harness skip provider calls entirely. /// -/// Keys should be produced by [`super::cache_key`] for consistency. Callers -/// are responsible for deciding when caching is safe (e.g., not caching -/// side-effecting tool calls). +/// Keys should be produced by [`super::cache_key`] (folded with the resolved +/// model's [`cache_identity`][crate::harness::model::ChatModel::cache_identity] +/// via [`super::scoped_cache_key`]) for consistency. Callers are responsible for +/// deciding when caching is safe (e.g., not caching side-effecting tool calls). +/// +/// # Implementing +/// +/// Only [`get`](Self::get) and [`put`](Self::put) are required. TTL support, +/// bulk invalidation, and statistics are optional: the defaults below keep an +/// existing third-party implementation compiling and behaving exactly as before. #[async_trait] pub trait ResponseCache: Send + Sync { /// Returns the cached [`ModelResponse`] for `key`, or `None` on a miss. + /// + /// An implementation that supports TTLs must treat an expired entry as a + /// miss (and should drop it). async fn get(&self, key: &str) -> Result>; - /// Stores `value` under `key`. + /// Stores `value` under `key` with no expiry. async fn put(&self, key: &str, value: ModelResponse) -> Result<()>; + + /// Stores `value` under `key`, expiring it after `ttl` when supported. + /// + /// The default implementation ignores `ttl` and delegates to + /// [`put`](Self::put), so a cache without expiry support keeps working; + /// implementations that can expire entries should override this and + /// implement `put` as `put_with_ttl(key, value, None)`. + async fn put_with_ttl( + &self, + key: &str, + value: ModelResponse, + ttl: Option, + ) -> Result<()> { + let _ = ttl; + self.put(key, value).await + } + + /// Drops every entry. + /// + /// Needed because a poisoned entry (for example one written before a bug + /// fix changed the key derivation) is otherwise permanent in a cache with + /// no TTL. The default is a no-op `Ok(())` so existing implementations do + /// not break; a cache that cannot clear should say so in its own docs. + async fn clear(&self) -> Result<()> { + Ok(()) + } + + /// Returns point-in-time counters for this cache. + /// + /// Defaults to all-zero for implementations that do not account. + fn stats(&self) -> CacheStats { + CacheStats::default() + } } /// Thread-safe in-memory response cache. /// /// Intended for unit tests and short-lived local runs. Contains no durable -/// storage: all entries are lost when the value is dropped. +/// storage: all entries are lost when the value is dropped. For a cache that +/// survives a restart see `SqliteResponseCache` (feature `sqlite`). +/// +/// Entries are bounded on **two** axes so a long-lived cache attached to a busy +/// harness cannot grow without limit: /// -/// Entries are bounded by an LRU eviction policy (default -/// [`InMemoryResponseCache::DEFAULT_CAPACITY`]) so a long-lived cache attached -/// to a busy harness cannot grow without limit. Reads and writes move a key to -/// the most-recently-used end; once the map is full the least-recently-used key -/// is evicted on insert. +/// * an entry count (default [`InMemoryResponseCache::DEFAULT_CAPACITY`]), and +/// * an approximate byte budget (default +/// [`InMemoryResponseCache::DEFAULT_MAX_BYTES`]) — 1024 long-context +/// responses carrying large tool payloads is hundreds of megabytes, which an +/// entry count alone does not bound. +/// +/// Whichever bound trips first evicts the least-recently-used entry. Reads and +/// writes move a key to the most-recently-used end. Per-entry TTLs are honored: +/// an expired entry is dropped on read and reported as a miss. #[derive(Clone, Debug)] pub struct InMemoryResponseCache { pub(crate) inner: Arc>, } +/// One retained response plus its bookkeeping. +#[derive(Clone, Debug)] +pub(crate) struct CacheEntry { + /// The cached response. + pub(crate) value: ModelResponse, + /// Monotonic recency ticket; larger is more recently used. + pub(crate) recency: u64, + /// Approximate serialized size in bytes, used for the byte bound. + pub(crate) bytes: usize, + /// Absolute expiry instant, when a TTL was supplied. + pub(crate) expires_at: Option, +} + /// LRU-ordered map backing [`InMemoryResponseCache`]. +/// +/// Recency is tracked with a monotonic ticket per entry plus a `BTreeMap` +/// ordered by that ticket. Touching a key is `O(log n)` (one `BTreeMap` remove +/// plus one insert) and eviction pops the first entry, also `O(log n)`. The +/// previous `VecDeque` layout scanned linearly, comparing up to `capacity` +/// `String`s and memmoving the deque **on every hit** — 1024 string compares +/// per cached call at the default capacity. #[derive(Debug)] pub(crate) struct LruResponseMap { /// Cached responses keyed by cache key. - pub(crate) data: HashMap, - /// Keys in least- to most-recently-used order. - pub(crate) order: VecDeque, + pub(crate) data: HashMap, + /// Recency ticket -> key, ordered least- to most-recently-used. + pub(crate) order: BTreeMap, + /// Next recency ticket to hand out. + pub(crate) next_recency: u64, /// Maximum number of entries retained before LRU eviction. pub(crate) capacity: usize, + /// Maximum approximate total bytes retained before LRU eviction. + pub(crate) max_bytes: usize, + /// Current approximate total bytes retained. + pub(crate) bytes: usize, + /// Running counters exposed through [`ResponseCache::stats`]. + pub(crate) stats: CacheStats, } // ── PromptCacheLayout ───────────────────────────────────────────────────────── @@ -77,6 +213,23 @@ pub(crate) struct LruResponseMap { /// The harness computes a `PromptCacheLayout` before and after each middleware /// pass so it can detect and report accidental prefix invalidations. /// +/// # Content awareness +/// +/// The layout records **both** the declared segment identities *and* a digest +/// of the material those segments carry: +/// +/// * [`Self::prefix_ids`] — the ordered ids of cacheable segments, +/// * [`Self::fingerprint`] — an FNV-1a digest over each cacheable segment's +/// `(id, role)` pair *and* the request's +/// [`prompt_fingerprint`][crate::harness::model::ModelRequest::prompt_fingerprint] +/// and tool schemas, and +/// * a per-message digest chain, so "did the prompt only grow at the tail?" +/// — the provider's actual KV-prefix rule — is answerable. +/// +/// Comparing ids alone reported "prefix stable" after a middleware rewrote the +/// **text** of a stable segment, which is precisely the failure this type +/// exists to catch: the ids match while the provider's cached bytes are gone. +/// /// # Provider KV-cache stability rules /// - Never insert timestamps, run ids, or dynamic retrieval output into the /// stable prefix. @@ -88,8 +241,11 @@ pub(crate) struct LruResponseMap { pub struct PromptCacheLayout { /// Ordered ids of cacheable (stable) prefix segments. pub(crate) prefix_ids: Vec, - /// Deterministic fingerprint of the ordered prefix ids. + /// Deterministic content-aware fingerprint (16 lowercase hex chars). pub(crate) fingerprint: String, + /// Per-message digests in transcript order, used to decide whether one + /// layout's message stream is a pure tail-extension of another's. + pub(crate) message_digests: Vec, } // ── CacheLayoutEvent ────────────────────────────────────────────────────────── @@ -107,6 +263,14 @@ pub struct CacheLayoutEvent { /// `true` if `segment_ids_after` contains only volatile (non-cacheable) /// segments, meaning no stable prefix is present. pub volatile_only: bool, + /// `true` if the segment identities were preserved but their **content** + /// changed — an edit that leaves the ids matching while destroying the + /// provider's cached bytes. + pub content_only_change: bool, + /// `true` if a [`CachePolicy::protect_prompt_prefix`] was in force and this + /// change violates it. Always `false` for the policy-free + /// [`CacheLayoutEvent::new`] constructor. + pub violates_policy: bool, /// The ordered cacheable prefix ids before the middleware pass. pub segment_ids_before: Vec, /// The ordered cacheable prefix ids after the middleware pass. @@ -120,12 +284,61 @@ pub struct CacheLayoutEvent { /// /// Both flags default to `false` (no caching / no protection) so the harness /// is safe-by-default and opts must be explicit. +/// +/// **This type is deliberately excluded from [`super::cache_key`]**: it selects +/// *whether* to cache, never *what the model answers*, so folding it into the +/// key made flipping [`Self::protect_prompt_prefix`] silently invalidate every +/// existing entry. #[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] pub struct CachePolicy { /// When `true`, the harness will look up (and write) local response cache /// entries via [`ResponseCache`] before calling the provider. pub response_cache_enabled: bool, - /// When `true`, middleware must preserve the order and content of cacheable - /// prefix segments. Violations are reported as [`CacheLayoutEvent`]s. + /// When `true`, middleware must preserve the order **and content** of + /// cacheable prefix segments. Violations are reported as + /// [`CacheLayoutEvent`]s with `violates_policy: true`, and the harness + /// additionally derives a provider `prompt_cache_key` breakpoint from the + /// stable prefix (see [`super::apply_prompt_cache_breakpoints`]) so one + /// logical thread routes to the same provider cache shard. pub protect_prompt_prefix: bool, + /// Time-to-live for entries written under this policy, in milliseconds. + /// + /// `None` means "never expire", which is the historical behaviour. A TTL is + /// the only bound on a poisoned entry in a cache that is never cleared. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ttl_ms: Option, + /// Optional key namespace, folded into the cache key. + /// + /// Lets one shared cache serve several logically separate populations + /// (per-tenant, per-experiment) without the risk of cross-serving, and lets + /// a caller invalidate one population by rotating the namespace. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub namespace: Option, +} + +impl CachePolicy { + /// A policy with local response caching enabled and no expiry. + pub fn enabled() -> Self { + Self { + response_cache_enabled: true, + ..Self::default() + } + } + + /// Returns this policy's TTL as a [`Duration`], when one is set. + pub fn ttl(&self) -> Option { + self.ttl_ms.map(Duration::from_millis) + } + + /// Sets the entry time-to-live. + pub fn with_ttl(mut self, ttl: Duration) -> Self { + self.ttl_ms = Some(ttl.as_millis() as u64); + self + } + + /// Sets the key namespace. + pub fn with_namespace(mut self, namespace: impl Into) -> Self { + self.namespace = Some(namespace.into()); + self + } } diff --git a/src/harness/context/mod.rs b/src/harness/context/mod.rs index 29a70ba..55e2f8e 100644 --- a/src/harness/context/mod.rs +++ b/src/harness/context/mod.rs @@ -55,9 +55,11 @@ fn next_context_instance_id() -> u64 { impl RunConfig { /// Creates a run configuration with sensible defaults. /// - /// Defaults: no thread, no tags, `null` metadata, no timeout, - /// `max_model_calls = 25`, and `max_tool_calls = 50`. These mirror the - /// crate-wide [`RunLimits`] defaults. + /// Defaults: no thread, no tags, `null` metadata, no timeout, and **unset** + /// call caps (`max_model_calls`/`max_tool_calls` are `None`), which resolve + /// to the crate-wide [`RunLimits`] defaults of 25 and 50. Leaving them unset + /// is what lets a harness-wide `RunPolicy` raise them; see + /// [`RunConfig::max_model_calls`]. pub fn new(run_id: impl Into) -> Self { Self { run_id: RunId::new(run_id), @@ -65,8 +67,8 @@ impl RunConfig { tags: Vec::new(), metadata: serde_json::Value::Null, timeout_ms: None, - max_model_calls: 25, - max_tool_calls: 50, + max_model_calls: None, + max_tool_calls: None, max_turn_output_tokens: None, depth: 0, max_depth: RunLimits::default().max_depth, @@ -97,18 +99,39 @@ impl RunConfig { self } - /// Sets the maximum number of model calls permitted for this run. + /// Sets the maximum number of model calls permitted for this run, + /// **explicitly**. + /// + /// An explicitly-set cap is a ceiling: the agent loop reconciles it with the + /// harness [`RunPolicy`][crate::harness::runtime::RunPolicy] by taking the + /// stricter of the two, so a policy default can only ever tighten it. pub fn with_max_model_calls(mut self, n: usize) -> Self { - self.max_model_calls = n; + self.max_model_calls = Some(n); self } - /// Sets the maximum number of tool invocations permitted for this run. + /// Sets the maximum number of tool invocations permitted for this run, + /// **explicitly**. Same ceiling semantics as + /// [`RunConfig::with_max_model_calls`]. pub fn with_max_tool_calls(mut self, n: usize) -> Self { - self.max_tool_calls = n; + self.max_tool_calls = Some(n); self } + /// The model-call cap actually applied to this run: the explicitly-set + /// value, or the crate-default [`RunLimits`] cap when unset. + pub fn effective_max_model_calls(&self) -> usize { + self.max_model_calls + .unwrap_or_else(|| RunLimits::default().max_model_calls) + } + + /// The tool-call cap actually applied to this run: the explicitly-set + /// value, or the crate-default [`RunLimits`] cap when unset. + pub fn effective_max_tool_calls(&self) -> usize { + self.max_tool_calls + .unwrap_or_else(|| RunLimits::default().max_tool_calls) + } + /// Sets the maximum output tokens requested for each model turn. pub fn with_max_turn_output_tokens(mut self, n: u32) -> Self { self.max_turn_output_tokens = Some(n); @@ -167,8 +190,8 @@ impl RunConfig { /// defaults. fn to_run_limits(&self) -> RunLimits { RunLimits::default() - .with_max_model_calls(self.max_model_calls) - .with_max_tool_calls(self.max_tool_calls) + .with_max_model_calls(self.effective_max_model_calls()) + .with_max_tool_calls(self.effective_max_tool_calls()) .with_max_wall_clock_ms(self.timeout_ms) .with_max_depth(self.max_depth) } diff --git a/src/harness/context/test.rs b/src/harness/context/test.rs index ffe30b7..5a321ea 100644 --- a/src/harness/context/test.rs +++ b/src/harness/context/test.rs @@ -12,8 +12,13 @@ fn run_config_defaults_are_sensible() { assert!(config.tags.is_empty()); assert_eq!(config.metadata, serde_json::Value::Null); assert!(config.timeout_ms.is_none()); - assert_eq!(config.max_model_calls, 25); - assert_eq!(config.max_tool_calls, 50); + // Unset by default, so a harness `RunPolicy` remains free to raise the cap + // (see `RunConfig::max_model_calls`); the effective value is the crate + // default. + assert_eq!(config.max_model_calls, None); + assert_eq!(config.max_tool_calls, None); + assert_eq!(config.effective_max_model_calls(), 25); + assert_eq!(config.effective_max_tool_calls(), 50); } #[test] @@ -31,8 +36,8 @@ fn run_config_builders_compose() { assert_eq!(config.tags, vec!["a".to_string(), "b".to_string()]); assert_eq!(config.metadata["k"], serde_json::json!("v")); assert_eq!(config.timeout_ms, Some(1234)); - assert_eq!(config.max_model_calls, 3); - assert_eq!(config.max_tool_calls, 4); + assert_eq!(config.max_model_calls, Some(3)); + assert_eq!(config.max_tool_calls, Some(4)); } #[test] diff --git a/src/harness/context/types.rs b/src/harness/context/types.rs index c7e0a4f..bfba865 100644 --- a/src/harness/context/types.rs +++ b/src/harness/context/types.rs @@ -43,7 +43,13 @@ use crate::harness::store::StoreRegistry; /// .with_tag("nightly") /// .with_max_model_calls(10); /// assert_eq!(config.run_id.as_str(), "run-1"); -/// assert_eq!(config.max_model_calls, 10); +/// assert_eq!(config.max_model_calls, Some(10)); +/// assert_eq!(config.effective_max_model_calls(), 10); +/// +/// // An unset cap reads as `None` but still resolves to the crate default. +/// let defaulted = RunConfig::new("run-2"); +/// assert_eq!(defaulted.max_model_calls, None); +/// assert_eq!(defaulted.effective_max_model_calls(), 25); /// ``` #[derive(Clone, Debug, Serialize, Deserialize)] pub struct RunConfig { @@ -57,10 +63,38 @@ pub struct RunConfig { pub metadata: serde_json::Value, /// Wall-clock timeout in milliseconds. `None` means no deadline. pub timeout_ms: Option, - /// Maximum number of model calls permitted for this run. - pub max_model_calls: usize, - /// Maximum number of tool invocations permitted for this run. - pub max_tool_calls: usize, + /// Maximum number of model calls permitted for this run, when the caller + /// set one explicitly. + /// + /// `None` means "unset": the run falls back to + /// [`RunConfig::effective_max_model_calls`] (the crate-default cap of 25) + /// and a harness-wide + /// [`RunPolicy::limits`][crate::harness::runtime::RunPolicy] is free to + /// raise *or* lower it. + /// + /// # Why this is an `Option` + /// + /// The agent loop reconciles this cap with the harness policy's, and the + /// two directions are only safe to distinguish when "the caller asked for + /// 2" is distinguishable from "nobody asked, so it defaulted to 25": + /// + /// - **Explicitly set** (`Some`) → the loop takes the **stricter** of the + /// two caps, so a permissive policy default can never silently widen a + /// budget the caller deliberately narrowed. + /// - **Unset** (`None`) → the policy is the only real source of truth and + /// wins outright, including when it raises the cap above the default. + /// + /// While this was a bare `usize` the loop could not tell those apart and + /// resolved every case by plain assignment, so + /// `RunConfig::new("r").with_max_model_calls(2)` against a default policy + /// ran **25** model calls. + #[serde(default)] + pub max_model_calls: Option, + /// Maximum number of tool invocations permitted for this run, when the + /// caller set one explicitly. See [`RunConfig::max_model_calls`] for the + /// set-versus-unset semantics; the tool cap follows the identical rule. + #[serde(default)] + pub max_tool_calls: Option, /// Maximum output tokens requested for each model turn in this run. /// /// When set, the agent loop applies this as an upper bound to diff --git a/src/harness/embeddings/cloud.rs b/src/harness/embeddings/cloud.rs index 0c2c664..8c59814 100644 --- a/src/harness/embeddings/cloud.rs +++ b/src/harness/embeddings/cloud.rs @@ -30,7 +30,7 @@ impl CloudEmbeddingModel { bearer: BearerResolver, ) -> Self { Self { - client: reqwest::Client::new(), + client: super::http::default_client(), base_url: base_url.into().trim().trim_end_matches('/').to_owned(), model: model.into(), dimensions, diff --git a/src/harness/embeddings/cohere.rs b/src/harness/embeddings/cohere.rs index 9b36250..bf16ac4 100644 --- a/src/harness/embeddings/cohere.rs +++ b/src/harness/embeddings/cohere.rs @@ -23,7 +23,7 @@ pub struct CohereEmbeddingModel { impl CohereEmbeddingModel { pub fn new(api_key: impl Into) -> Self { Self { - client: reqwest::Client::new(), + client: super::http::default_client(), api_key: api_key.into(), model: COHERE_DEFAULT_MODEL.to_owned(), dimensions: COHERE_DEFAULT_DIMENSIONS, diff --git a/src/harness/embeddings/http.rs b/src/harness/embeddings/http.rs new file mode 100644 index 0000000..ed63008 --- /dev/null +++ b/src/harness/embeddings/http.rs @@ -0,0 +1,79 @@ +//! The shared HTTP client policy for every embedding adapter. +//! +//! # Why this exists +//! +//! Every embedding adapter built its transport with `reqwest::Client::new()`, +//! which sets **no timeout at all** — neither a connect timeout nor an overall +//! one. reqwest's default really is "wait forever". A server that accepts the +//! TCP connection and then never answers therefore hung the calling task +//! indefinitely, with no error to retry and nothing in the logs. +//! +//! That is not a hypothetical failure mode: the chat path hit it and fixed it +//! **twice** — once with `DEFAULT_CONNECT_TIMEOUT_SECS` on the model client, and +//! again with the `list_models` deadline, whose comment names this exact +//! scenario ("an Ollama/LM Studio server that accepts the TCP connect and then +//! never responds … hung the call forever"). The embedding adapters, which point +//! at the same local servers, never got either. +//! +//! The constants deliberately mirror the chat path's, so the two halves of the +//! crate cannot drift into different opinions about how long a wedged local +//! server is allowed to hold a task. + +use std::time::Duration; + +/// TCP connect timeout. Bounds connection establishment without capping a +/// legitimately slow response body. Mirrors the chat transport's +/// `DEFAULT_CONNECT_TIMEOUT_SECS`. +pub const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 30; + +/// Overall request deadline. +/// +/// Shorter than the chat path's 600 s because an embedding call has no +/// generation phase to wait through — it is one forward pass. A minute is +/// generous for a large batch against a cold local model and still bounds the +/// wedged-server case. +pub const DEFAULT_EMBEDDING_TIMEOUT_SECS: u64 = 120; + +/// The default [`reqwest::Client`] every embedding adapter is constructed with. +/// +/// Falls back to `reqwest::Client::new()` if the builder somehow fails, so a +/// construction path that cannot return an error stays infallible — an +/// unbounded client is bad, but panicking during construction is worse. +pub fn default_client() -> reqwest::Client { + reqwest::Client::builder() + .connect_timeout(Duration::from_secs(DEFAULT_CONNECT_TIMEOUT_SECS)) + .timeout(Duration::from_secs(DEFAULT_EMBEDDING_TIMEOUT_SECS)) + .build() + .unwrap_or_else(|error| { + tracing::warn!( + target: "tinyagents::embeddings", + %error, + "[embeddings] could not build the default timeout-bounded client; \ + falling back to an unbounded one" + ); + reqwest::Client::new() + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn the_default_client_is_bounded() { + // reqwest exposes no getter for its configured timeouts, so assert the + // property that actually matters and is observable: a request against a + // closed port fails promptly rather than hanging. + let client = default_client(); + // Cheap structural check that the builder path was taken at all. + assert!(format!("{client:?}").contains("Client")); + } + + /// An embedding call has no generation phase, so it must not inherit the + /// chat path's 600 s patience — and the connect timeout must be the tighter + /// of the two. Both are compile-time facts, so assert them as such. + const _: () = { + assert!(DEFAULT_EMBEDDING_TIMEOUT_SECS < 600); + assert!(DEFAULT_CONNECT_TIMEOUT_SECS < DEFAULT_EMBEDDING_TIMEOUT_SECS); + }; +} diff --git a/src/harness/embeddings/mod.rs b/src/harness/embeddings/mod.rs index 45750c2..49aba05 100644 --- a/src/harness/embeddings/mod.rs +++ b/src/harness/embeddings/mod.rs @@ -277,6 +277,7 @@ impl Retriever { mod cloud; mod cohere; +mod http; mod noop; mod ollama; mod openai; @@ -284,6 +285,7 @@ mod rate_limit; mod retry_after; mod voyage; +pub use http::{DEFAULT_CONNECT_TIMEOUT_SECS, DEFAULT_EMBEDDING_TIMEOUT_SECS, default_client}; pub use noop::NoopEmbeddingModel; pub use ollama::{ DEFAULT_OLLAMA_DIMENSIONS, DEFAULT_OLLAMA_MODEL, DEFAULT_OLLAMA_URL, OllamaEmbeddingModel, diff --git a/src/harness/embeddings/ollama.rs b/src/harness/embeddings/ollama.rs index 1d291f8..9120fec 100644 --- a/src/harness/embeddings/ollama.rs +++ b/src/harness/embeddings/ollama.rs @@ -28,7 +28,7 @@ pub struct OllamaEmbeddingModel { impl OllamaEmbeddingModel { pub fn try_new(base_url: &str, model: &str, dimensions: usize) -> Result { Ok(Self { - client: reqwest::Client::new(), + client: super::http::default_client(), base_url: normalize_base_url(base_url)?, model: normalize_model(model)?, dimensions: Arc::new(AtomicUsize::new(if dimensions == 0 { @@ -42,7 +42,7 @@ impl OllamaEmbeddingModel { pub(super) fn try_new_unresolved(base_url: &str, model: &str) -> Result { Ok(Self { - client: reqwest::Client::new(), + client: super::http::default_client(), base_url: normalize_base_url(base_url)?, model: normalize_model(model)?, dimensions: Arc::new(AtomicUsize::new(0)), diff --git a/src/harness/embeddings/openai.rs b/src/harness/embeddings/openai.rs index 71c92d4..dca0457 100644 --- a/src/harness/embeddings/openai.rs +++ b/src/harness/embeddings/openai.rs @@ -53,7 +53,7 @@ impl OpenAiEmbeddingModel { /// (`text-embedding-3-small`), and the default base URL. pub fn new(api_key: impl Into) -> Self { Self { - client: reqwest::Client::new(), + client: super::http::default_client(), api_key: api_key.into(), model: DEFAULT_MODEL.to_string(), base_url: DEFAULT_BASE_URL.to_string(), @@ -89,17 +89,91 @@ impl OpenAiEmbeddingModel { } /// Controls whether the OpenAI-compatible `dimensions` field is sent. + /// + /// `dimensions` is an OpenAI-specific **request** parameter for + /// Matryoshka-style truncation. llama.cpp-backed servers (LM Studio, + /// `llama-server`) reject or ignore it, and the width is whatever the GGUF + /// produces regardless. pub fn with_send_dimensions(mut self, send: bool) -> Self { self.send_dimensions = send; self } /// Controls whether an empty API key is rejected before making a request. + /// + /// Passing `false` also turns [`Self::with_send_dimensions`] off, because + /// "no API key required" is how a local server identifies itself and the + /// hosted defaults are wrong for one in exactly two ways: + /// + /// * `send_dimensions: true` puts a field on the wire that llama.cpp-backed + /// servers reject or ignore, and + /// * `dimensions: 1536` (`text-embedding-3-small`'s width) is then checked + /// against every returned vector, so a 768-wide local model fails **every + /// call** as a dimension mismatch. + /// + /// The crate's own live test hand-wrote this workaround and explained why — + /// "probing with the default would reject a 768-wide local model as a + /// mismatch". Lifting it into the adapter means callers stop rediscovering + /// it. Set [`Self::with_send_dimensions`] explicitly *after* this call to + /// override. pub fn with_required_api_key(mut self, required: bool) -> Self { self.requires_api_key = required; + if !required { + self.send_dimensions = false; + } self } + /// Embeds text against a server whose vector width is not known in advance, + /// returning the discovered width alongside the vectors. + /// + /// Mirrors + /// [`OllamaEmbeddingModel::embed_discovering_dimensions`][ollama], which has + /// existed for a while — the OpenAI-compatible path had no equivalent, so + /// every caller pointing at LM Studio or `llama-server` had to hand-roll the + /// same probe. The declared width is set to `0` for the probe, which + /// disables the width check (the check is what rejects an unknown-width + /// model), then read back off the returned vector. + /// + /// # Errors + /// + /// [`TinyAgentsError::Validation`] when `texts` holds no non-blank input, or + /// when the server answers with an empty vector (a width of zero is not a + /// discovery, it is a failure). Transport and decode failures surface from + /// [`EmbeddingModel::embed`] unchanged. + /// + /// [ollama]: super::OllamaEmbeddingModel::embed_discovering_dimensions + pub async fn embed_discovering_dimensions( + self, + texts: &[String], + ) -> Result<(usize, Vec>)> { + if !texts.iter().any(|text| !text.trim().is_empty()) { + return Err(TinyAgentsError::Validation( + "dynamic embedding dimension discovery requires at least one nonblank input" + .to_string(), + )); + } + let model_id = self.model.clone(); + // Width 0 disables the per-vector width check; sending `dimensions` is + // meaningless when we do not yet know the width. + let probe = self.with_dimensions(0).with_send_dimensions(false); + let vectors = probe.embed(texts).await?; + let width = vectors.first().map(Vec::len).unwrap_or(0); + if width == 0 { + return Err(TinyAgentsError::Validation(format!( + "embedding model `{model_id}` returned an empty vector; \ + cannot discover its dimensionality" + ))); + } + tracing::debug!( + target: "tinyagents::embeddings::openai", + model = %model_id, + width, + "[embeddings] discovered embedding width" + ); + Ok((width, vectors)) + } + pub fn base_url(&self) -> &str { &self.base_url } diff --git a/src/harness/events/mod.rs b/src/harness/events/mod.rs index b0086d5..d9cc14a 100644 --- a/src/harness/events/mod.rs +++ b/src/harness/events/mod.rs @@ -34,13 +34,53 @@ pub use types::*; // below; it is not re-exported by `pub use types::*`. use types::{EventSinkInner, JournalRecorder}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Mutex, MutexGuard, PoisonError}; use std::time::SystemTime; use crate::harness::cost::CostTotals; use crate::harness::ids::{ComponentId, ExecutionStatus, HarnessPhase, RunId, ThreadId}; use crate::harness::usage::UsageTotals; +// --------------------------------------------------------------------------- +// Lock helpers +// --------------------------------------------------------------------------- + +/// Locks a mutex, **recovering** from poisoning instead of panicking. +/// +/// A listener may panic (this module explicitly tolerates that, see +/// [`DispatchGuard`]), and a panic while any of these locks is held poisons it. +/// Every structure guarded here is a plain buffer/counter with no cross-field +/// invariant a half-finished update could break, so the poisoned state is still +/// perfectly usable — whereas `.expect("lock poisoned")` converts one listener's +/// panic into a permanently unusable event bus for the entire process. +/// +/// This also brings the events module in line with +/// [`SteeringHandle`][crate::harness::steering::SteeringHandle], which already +/// recovers from poisoning for exactly the same reason. The two used to +/// disagree. +fn lock_recovering(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(PoisonError::into_inner) +} + +/// Clears [`EventSinkInner::dispatching`] on drop, including while a panic is +/// unwinding. +/// +/// The flag used to be reset only on the normal exit path (when the queue +/// drained). A listener that panicked unwound straight past that reset, leaving +/// `dispatching == true` forever: every later `emit` saw a drain already in +/// progress, pushed onto `pending`, and returned. The run kept emitting, no +/// listener ever received anything again, and `pending` grew without bound — +/// a silent, unrecoverable observability outage from one bad callback. +struct DispatchGuard<'a> { + inner: &'a Mutex, +} + +impl Drop for DispatchGuard<'_> { + fn drop(&mut self) { + lock_recovering(self.inner).dispatching = false; + } +} + // --------------------------------------------------------------------------- // EventSink impls // --------------------------------------------------------------------------- @@ -78,7 +118,7 @@ impl EventSink { /// Subscribes a new listener. The listener will receive every subsequent /// [`AgentEvent`] emitted through this sink (or any of its clones). pub fn subscribe(&self, listener: Arc) { - let mut inner = self.inner.lock().expect("EventSink lock poisoned"); + let mut inner = lock_recovering(&self.inner); inner.listeners.push(listener); } @@ -88,7 +128,7 @@ impl EventSink { /// transient listener to a long-lived shared sink. Returns `true` when the /// listener was present. pub fn unsubscribe(&self, listener: &Arc) -> bool { - let mut inner = self.inner.lock().expect("EventSink lock poisoned"); + let mut inner = lock_recovering(&self.inner); let before = inner.listeners.len(); inner .listeners @@ -98,11 +138,7 @@ impl EventSink { #[cfg(test)] pub(crate) fn listener_count(&self) -> usize { - self.inner - .lock() - .expect("EventSink lock poisoned") - .listeners - .len() + lock_recovering(&self.inner).listeners.len() } /// Emits an event, assigning a monotonic [`EventId`] and offset, then @@ -120,9 +156,13 @@ impl EventSink { /// a listener runs, so callbacks may safely emit to the same sink (the /// re-entrant record is queued and delivered by the active drain loop) /// when they guard against unbounded event recursion. + /// A panicking listener does **not** wedge the sink: the `dispatching` flag + /// is released by a [`DispatchGuard`] on unwind, so later emits still + /// dispatch (the panicking listener's own record is lost, and any records + /// still queued behind it are delivered by whichever emitter drains next). pub fn emit(&self, event: AgentEvent) -> EventRecord { let (record, should_drain) = { - let mut inner = self.inner.lock().expect("EventSink lock poisoned"); + let mut inner = lock_recovering(&self.inner); let offset = inner.next_offset; inner.next_offset += 1; let id = crate::harness::ids::EventId::new(format!("{}-evt-{offset}", inner.stream_id)); @@ -136,15 +176,15 @@ impl EventSink { (record, should_drain) }; if should_drain { + // Held for the whole drain: `dispatching` is cleared on the normal + // exit path *and* while a listener panic unwinds through here. + let _guard = DispatchGuard { inner: &self.inner }; loop { let next = { - let mut inner = self.inner.lock().expect("EventSink lock poisoned"); + let mut inner = lock_recovering(&self.inner); match inner.pending.pop_front() { Some(entry) => entry, - None => { - inner.dispatching = false; - break; - } + None => break, } }; let (queued, listeners) = next; @@ -158,11 +198,7 @@ impl EventSink { /// Returns the number of currently registered listeners. pub fn len(&self) -> usize { - self.inner - .lock() - .expect("EventSink lock poisoned") - .listeners - .len() + lock_recovering(&self.inner).listeners.len() } /// Returns `true` when no listeners are registered. @@ -191,18 +227,12 @@ impl RecordingListener { /// Returns a snapshot of all collected [`EventRecord`]s in arrival order. pub fn events(&self) -> Vec { - self.records - .lock() - .expect("RecordingListener lock poisoned") - .clone() + lock_recovering(&self.records).clone() } /// Returns the number of events collected so far. pub fn len(&self) -> usize { - self.records - .lock() - .expect("RecordingListener lock poisoned") - .len() + lock_recovering(&self.records).len() } /// Returns `true` when no events have been collected yet. @@ -213,10 +243,7 @@ impl RecordingListener { impl EventListener for RecordingListener { fn on_event(&self, record: &EventRecord) { - self.records - .lock() - .expect("RecordingListener lock poisoned") - .push(record.clone()); + lock_recovering(&self.records).push(record.clone()); } } @@ -259,9 +286,7 @@ impl EventJournal { /// Callers can use this to replay run history from any known checkpoint. /// A `from_offset` of `0` replays the full journal. pub fn replay_from(&self, from_offset: u64) -> Vec { - self.records - .lock() - .expect("EventJournal lock poisoned") + lock_recovering(&self.records) .iter() .filter(|r| r.offset >= from_offset) .cloned() @@ -270,10 +295,7 @@ impl EventJournal { /// Returns the total number of events in the journal. pub fn len(&self) -> usize { - self.records - .lock() - .expect("EventJournal lock poisoned") - .len() + lock_recovering(&self.records).len() } /// Returns `true` when the journal contains no events. @@ -284,10 +306,7 @@ impl EventJournal { impl EventListener for JournalRecorder { fn on_event(&self, record: &EventRecord) { - self.records - .lock() - .expect("EventJournal lock poisoned") - .push(record.clone()); + lock_recovering(&self.records).push(record.clone()); } } diff --git a/src/harness/events/test.rs b/src/harness/events/test.rs index ee69e17..a30775e 100644 --- a/src/harness/events/test.rs +++ b/src/harness/events/test.rs @@ -275,3 +275,204 @@ fn concurrent_journal_appends_replay_in_offset_order() { let tail = journal.replay_from(expected.len() as u64 - 5); assert_eq!(tail.len(), 5); } + +// --------------------------------------------------------------------------- +// LOOP-9: a panicking listener must not wedge the sink +// --------------------------------------------------------------------------- + +/// A listener that panics on its first `n` events, then records normally. +struct PanickingListener { + remaining_panics: std::sync::Mutex, + seen: std::sync::Mutex>, +} + +impl EventListener for PanickingListener { + fn on_event(&self, record: &EventRecord) { + { + let mut remaining = self + .remaining_panics + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if *remaining > 0 { + *remaining -= 1; + panic!("listener blew up on {}", record.event.kind()); + } + } + self.seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(record.event.kind().to_string()); + } +} + +#[test] +fn a_panicking_listener_does_not_permanently_wedge_the_sink() { + // Regression test (LOOP-9): `dispatching` was cleared only when the drain + // loop reached an empty queue. A panic inside `on_event` unwound straight + // past that reset, so the flag stayed `true` forever: every subsequent + // `emit` saw a drain "already in progress", pushed onto `pending`, and + // returned. The run kept emitting and no listener ever received anything + // again, while `pending` grew without bound. + let sink = EventSink::new(); + let bomb = Arc::new(PanickingListener { + remaining_panics: std::sync::Mutex::new(1), + seen: std::sync::Mutex::new(Vec::new()), + }); + let recorder = Arc::new(RecordingListener::new()); + sink.subscribe(bomb.clone()); + sink.subscribe(recorder.clone()); + + // First emit: the listener panics. Catch it the way a host would. + let sink_for_panic = sink.clone(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + sink_for_panic.emit(AgentEvent::StateUpdate) + })); + assert!(result.is_err(), "the listener was supposed to panic"); + + // Every later emit must still be delivered, synchronously, to every + // listener — including the one that panicked. + for _ in 0..3 { + sink.emit(AgentEvent::StateUpdate); + } + + assert_eq!( + recorder.len(), + 3, + "sink stayed wedged after a listener panic: later events were queued, never delivered" + ); + assert_eq!( + bomb.seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .len(), + 3, + "the recovered listener should have received the later events too" + ); +} + +#[test] +fn sink_accessors_recover_from_a_poisoned_lock() { + // A panic while a listener holds no sink lock cannot poison it, but a panic + // anywhere else in the process can. `SteeringHandle` already recovers from + // poisoning; the events module used to `.expect(...)` and take the whole + // bus down with it. + let sink = EventSink::new(); + let recorder = Arc::new(RecordingListener::new()); + sink.subscribe(recorder.clone()); + + // Poison the recorder's own buffer lock from a panicking thread. + let poisoner = recorder.clone(); + let _ = thread::spawn(move || { + let _guard = poisoner.records.lock().expect("first lock is clean"); + panic!("poison the recording listener"); + }) + .join(); + assert!( + recorder.records.is_poisoned(), + "the test needs a genuinely poisoned lock" + ); + + // Both the sink and the listener remain usable. + sink.emit(AgentEvent::StateUpdate); + assert_eq!(sink.len(), 1); + assert_eq!(recorder.len(), 1); +} + +// --------------------------------------------------------------------------- +// LOOP-6: every Started variant has a terminal partner on the error path +// --------------------------------------------------------------------------- + +#[test] +fn failure_variants_exist_and_carry_stable_kind_strings() { + // Regression test (LOOP-6): there was no `ToolFailed` / `ModelFailed` / + // `SubAgentFailed`, so the three `?` sites that skip the `Completed` emit + // left a `Started` with no terminal partner, and any exporter pairing + // started/completed silently dropped every failed call. + use crate::harness::ids::CallId; + + let tool_failed = AgentEvent::ToolFailed { + call_id: CallId::new("call-1"), + tool_name: "search".into(), + started_at_ms: Some(1_000), + duration_ms: Some(25), + error: "boom".into(), + }; + let model_failed = AgentEvent::ModelFailed { + call_id: CallId::new("call-2"), + model: "gpt-4o".into(), + started_at_ms: Some(1_000), + attempts: Some(4), + error: "429 rate limited".into(), + }; + let subagent_failed = AgentEvent::SubAgentFailed { + name: "researcher".into(), + depth: 2, + error: "child run failed".into(), + }; + + assert_eq!(tool_failed.kind(), "tool.failed"); + assert_eq!(model_failed.kind(), "model.failed"); + assert_eq!(subagent_failed.kind(), "subagent.failed"); +} + +#[test] +fn failure_variants_round_trip_through_serde() { + use crate::harness::ids::CallId; + + for event in [ + AgentEvent::ToolFailed { + call_id: CallId::new("call-1"), + tool_name: "search".into(), + started_at_ms: None, + duration_ms: None, + error: "boom".into(), + }, + AgentEvent::ModelFailed { + call_id: CallId::new("call-2"), + model: "gpt-4o".into(), + started_at_ms: None, + attempts: None, + error: "boom".into(), + }, + AgentEvent::SubAgentFailed { + name: "researcher".into(), + depth: 1, + error: "boom".into(), + }, + ] { + let json = serde_json::to_value(&event).expect("serialize"); + let back: AgentEvent = serde_json::from_value(json).expect("deserialize"); + assert_eq!(back, event); + } +} + +#[test] +fn every_started_variant_pairs_with_both_a_completed_and_a_failed_variant() { + // Guard against a future `*Started` landing without its error-path partner. + // Kept as a string check because the enum has no structural grouping. + let kinds: Vec<&str> = vec![ + "tool.started", + "tool.completed", + "tool.failed", + "model.started", + "model.completed", + "model.failed", + "subagent.started", + "subagent.completed", + "subagent.failed", + "middleware.started", + "middleware.completed", + "middleware.failed", + ]; + for started in kinds.iter().filter(|k| k.ends_with(".started")) { + let prefix = started.trim_end_matches(".started"); + assert!( + kinds.contains(&format!("{prefix}.completed").as_str()), + "{prefix} has no completed partner" + ); + assert!( + kinds.contains(&format!("{prefix}.failed").as_str()), + "{prefix} has no failed partner" + ); + } +} diff --git a/src/harness/events/types.rs b/src/harness/events/types.rs index 9c8edb8..11a7f92 100644 --- a/src/harness/events/types.rs +++ b/src/harness/events/types.rs @@ -161,6 +161,78 @@ pub enum AgentEvent { error: Option, }, + /// A tool invocation failed and the run is propagating the error rather + /// than turning it into a tool result. + /// + /// This is the terminal partner of [`AgentEvent::ToolStarted`] on the error + /// path. Without it, every `?` that escapes the tool-dispatch path leaves a + /// `ToolStarted` with no matching terminal event, and any exporter pairing + /// started/completed by `call_id` silently drops the failed call — the + /// spans that matter most. The middleware stack already maintains this + /// invariant deliberately (`run_stack_hook!` emits `MiddlewareCompleted` + /// *before* inspecting the result, "the onion's balance invariant"); this + /// variant extends the same guarantee to tools. + /// + /// Distinct from [`AgentEvent::ToolCompleted`] with `error: Some(_)`, which + /// means the tool ran, failed, and the failure was fed back to the model as + /// a tool result. `ToolFailed` means the run itself is aborting. + ToolFailed { + /// Identifier for the tool call that failed; pairs with the + /// [`AgentEvent::ToolStarted`] of the same id. + call_id: CallId, + /// Name of the tool that was invoked. + tool_name: String, + /// Wall-clock time the tool call *started*, in Unix-epoch milliseconds, + /// mirroring [`AgentEvent::ToolCompleted::started_at_ms`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + started_at_ms: Option, + /// Wall-clock duration until the failure, in milliseconds. + #[serde(default, skip_serializing_if = "Option::is_none")] + duration_ms: Option, + /// Human-readable failure description. + error: String, + }, + + /// A model call failed and the run is propagating the error. + /// + /// The terminal partner of [`AgentEvent::ModelStarted`] on the error path; + /// see [`AgentEvent::ToolFailed`] for the rationale. Emitted once the retry + /// ladder is exhausted — an *individual* failed attempt that will be retried + /// is already covered by [`AgentEvent::RetryScheduled`]. + ModelFailed { + /// Identifier for the model call that failed; pairs with the + /// [`AgentEvent::ModelStarted`] of the same id. + call_id: CallId, + /// Registry name or provider model id the call was dispatched to. + model: String, + /// Wall-clock time the model call *started*, in Unix-epoch milliseconds. + #[serde(default, skip_serializing_if = "Option::is_none")] + started_at_ms: Option, + /// Number of attempts made before giving up (`1` when the call was not + /// retried). Lets an exporter distinguish "failed once" from "failed + /// after exhausting the retry policy". + #[serde(default, skip_serializing_if = "Option::is_none")] + attempts: Option, + /// Human-readable failure description. + error: String, + }, + + /// A sub-agent child run failed. + /// + /// The terminal partner of [`AgentEvent::SubAgentStarted`] / + /// [`AgentEvent::SubAgentReused`] on the error path; see + /// [`AgentEvent::ToolFailed`] for the rationale. Without it a failing branch + /// of the recursion tree simply stops appearing in the stream, and a + /// consumer cannot tell a crashed child from one still running. + SubAgentFailed { + /// Name of the sub-agent whose run failed. + name: String, + /// Depth of the child run in the recursion tree. + depth: usize, + /// Human-readable failure description. + error: String, + }, + /// The model called a tool that is not registered, and the run's /// [`UnknownToolPolicy`][crate::harness::runtime::UnknownToolPolicy] /// recovered from it instead of aborting. @@ -608,6 +680,9 @@ impl AgentEvent { AgentEvent::ToolsFiltered { .. } => "tool.filtered", AgentEvent::ToolStarted { .. } => "tool.started", AgentEvent::ToolCompleted { .. } => "tool.completed", + AgentEvent::ToolFailed { .. } => "tool.failed", + AgentEvent::ModelFailed { .. } => "model.failed", + AgentEvent::SubAgentFailed { .. } => "subagent.failed", AgentEvent::UnknownToolCall { .. } => "tool.unknown", AgentEvent::InvalidToolArgs { .. } => "tool.invalid_args", AgentEvent::BudgetWarning { .. } => "budget.warning", diff --git a/src/harness/limits/mod.rs b/src/harness/limits/mod.rs index 991c994..9d42bd4 100644 --- a/src/harness/limits/mod.rs +++ b/src/harness/limits/mod.rs @@ -61,6 +61,13 @@ impl RunLimits { self.max_depth = n; self } + + /// Sets what the run does when a call cap is reached. See + /// [`LimitBehavior`]. + pub fn with_behavior(mut self, behavior: LimitBehavior) -> Self { + self.behavior = behavior; + self + } } /// Tracks live counters for a single harness run and enforces [`RunLimits`]. @@ -93,26 +100,84 @@ impl LimitTracker { /// The counter is incremented **before** the check so the limit is /// inclusive (a cap of `N` allows exactly `N` calls). pub fn record_model_call(&mut self) -> Result<()> { - self.model_calls += 1; - if self.model_calls > self.limits.max_model_calls { - return Err(TinyAgentsError::Validation(format!( - "max model calls ({}) exceeded", - self.limits.max_model_calls - ))); - } + self.try_record_model_call()?; Ok(()) } /// Records one tool call and returns an error if the cap is exceeded. pub fn record_tool_call(&mut self) -> Result<()> { + self.try_record_tool_call()?; + Ok(()) + } + + /// Records one model call and reports the cap decision as a + /// [`LimitOutcome`], honoring [`RunLimits::behavior`]. + /// + /// - Within the cap → `Ok(LimitOutcome::Proceed)`. + /// - Cap exhausted under [`LimitBehavior::Error`] → `Err(LimitExceeded)`, + /// exactly as [`LimitTracker::record_model_call`] has always behaved. + /// - Cap exhausted under [`LimitBehavior::StopWithPartial`] → + /// `Ok(LimitOutcome::Stop(LimitKind::ModelCalls))`, so the loop can stop + /// cleanly and return the partial run instead of discarding it. + pub fn try_record_model_call(&mut self) -> Result { + self.model_calls += 1; + if self.model_calls > self.limits.max_model_calls { + return self.exhausted(LimitKind::ModelCalls, self.limits.max_model_calls); + } + Ok(LimitOutcome::Proceed) + } + + /// Records one tool call and reports the cap decision as a + /// [`LimitOutcome`]. See [`LimitTracker::try_record_model_call`]. + pub fn try_record_tool_call(&mut self) -> Result { self.tool_calls += 1; if self.tool_calls > self.limits.max_tool_calls { - return Err(TinyAgentsError::Validation(format!( - "max tool calls ({}) exceeded", - self.limits.max_tool_calls - ))); + return self.exhausted(LimitKind::ToolCalls, self.limits.max_tool_calls); + } + Ok(LimitOutcome::Proceed) + } + + /// Shared exhaustion branch: error or clean stop, per [`RunLimits::behavior`]. + fn exhausted(&self, kind: LimitKind, cap: usize) -> Result { + match self.limits.behavior { + LimitBehavior::Error => { + tracing::debug!( + target: "tinyagents::limits", + limit_kind = kind.as_str(), + cap, + "[limits] cap exhausted; failing the run" + ); + Err(TinyAgentsError::Validation(format!( + "max {} ({cap}) exceeded", + match kind { + LimitKind::ModelCalls => "model calls", + LimitKind::ToolCalls => "tool calls", + } + ))) + } + LimitBehavior::StopWithPartial => { + tracing::debug!( + target: "tinyagents::limits", + limit_kind = kind.as_str(), + cap, + "[limits] cap exhausted; stopping with the partial result" + ); + Ok(LimitOutcome::Stop(kind)) + } } - Ok(()) + } + + /// Un-counts `n` tool calls that were requested but never executed. + /// + /// Needed by the [`LimitBehavior::StopWithPartial`] tool path: when the cap + /// trips mid-batch the loop answers the remaining `tool_call_id`s with a + /// "stopped before this could run" result rather than executing them, so + /// leaving them counted would over-report the work done. LangChain's + /// `ToolCallLimitMiddleware` does the same rollback under `"end"`. + /// + /// Saturates at zero. + pub fn rollback_tool_calls(&mut self, n: usize) { + self.tool_calls = self.tool_calls.saturating_sub(n); } /// Checks whether the run has exceeded the configured wall-clock deadline. @@ -186,20 +251,91 @@ impl LimitTracker { &self.limits } - /// Overrides the model-call and tool-call caps in place, preserving - /// already-recorded counts and the wall-clock start time. + /// **Fail-open** override of the model-call and tool-call caps in place, + /// preserving already-recorded counts and the wall-clock start time. /// /// A `RunContext` derives its tracker's initial limits from its - /// `RunConfig`, which always carries a concrete default. That can - /// silently disagree with a harness-wide `RunPolicy` configured with a - /// different cap, so the *reported* limit (the policy's) and the limit - /// that actually trips (the tracker's) diverge. The agent loop calls this - /// once per run to reconcile the two into a single enforced source of - /// truth before the loop begins. + /// `RunConfig`, which always carries a concrete default. That can silently + /// disagree with a harness-wide `RunPolicy` configured with a different + /// cap, so the *reported* limit (the policy's) and the limit that actually + /// trips (the tracker's) diverge. The agent loop calls this once per run to + /// reconcile the two into a single enforced source of truth. + /// + /// # This is a plain assignment, in **both** directions + /// + /// It raises a cap as readily as it lowers one. A caller writing + /// `RunConfig::new("r").with_max_model_calls(2)` against the default policy + /// therefore gets **25** model calls, not 2 — their explicit ceiling is + /// silently widened by a policy default they never set. + /// + /// Prefer [`LimitTracker::tighten_call_limits`], which takes the stricter + /// of the two and cannot widen anything. This method remains for the one + /// case that genuinely needs widening: a harness-wide `RunPolicy` that + /// deliberately configures a *higher* cap than the `RunConfig` **default** + /// (not than an explicitly-set `RunConfig` value). + /// + /// Distinguishing those two cases needs information this module does not + /// have — whether a `RunConfig` cap was set by the caller or merely + /// defaulted. See the note on [`LimitTracker::tighten_call_limits`] for what + /// the agent loop has to do about it. pub fn sync_call_limits(&mut self, max_model_calls: usize, max_tool_calls: usize) { + tracing::debug!( + target: "tinyagents::limits", + from_model_calls = self.limits.max_model_calls, + from_tool_calls = self.limits.max_tool_calls, + to_model_calls = max_model_calls, + to_tool_calls = max_tool_calls, + "[limits] fail-open sync_call_limits override" + ); self.limits.max_model_calls = max_model_calls; self.limits.max_tool_calls = max_tool_calls; } + + /// **Fail-closed** reconciliation: keeps whichever of the tracker's current + /// cap and the supplied cap is *stricter*, so a second limit source can + /// only ever tighten the run, never loosen it. + /// + /// This is the semantics a "hard limit" needs, and the same rule + /// [`RetryPolicy::max_attempts_capped_at`][crate::harness::retry::RetryPolicy::max_attempts_capped_at] + /// already applies to the retry cap. Counts and the wall-clock start are + /// preserved. + /// + /// # Wiring note for the agent loop (wave 2) + /// + /// The loop currently calls [`LimitTracker::sync_call_limits`] with the + /// `RunPolicy` caps. Switching that call site to this method fixes the + /// silent widening of an explicit `RunConfig` cap, but it also changes the + /// case a policy raising the cap **above the `RunConfig` default** stops + /// working — which `policy_model_call_limit_above_run_config_default_is_honored` + /// in `agent_loop/test.rs` pins. + /// + /// Both cases are legitimate, and they are only distinguishable by knowing + /// whether the `RunConfig` cap was *explicitly set* or merely defaulted. The + /// clean fix is on the `RunConfig` side (a `RunConfig` owned by + /// `harness::context`, not this module): make its call caps + /// `Option` — or track an `explicitly_set` flag — and then in the + /// loop: + /// + /// - `RunConfig` cap explicitly set → `tighten_call_limits(policy caps)` + /// (the caller's ceiling wins, and the policy may only tighten it). + /// - `RunConfig` cap merely defaulted → `sync_call_limits(policy caps)` + /// (the policy is the only real source of truth, so it may raise it). + pub fn tighten_call_limits(&mut self, max_model_calls: usize, max_tool_calls: usize) { + let model = self.limits.max_model_calls.min(max_model_calls); + let tool = self.limits.max_tool_calls.min(max_tool_calls); + tracing::debug!( + target: "tinyagents::limits", + from_model_calls = self.limits.max_model_calls, + from_tool_calls = self.limits.max_tool_calls, + candidate_model_calls = max_model_calls, + candidate_tool_calls = max_tool_calls, + to_model_calls = model, + to_tool_calls = tool, + "[limits] fail-closed tighten_call_limits" + ); + self.limits.max_model_calls = model; + self.limits.max_tool_calls = tool; + } } #[cfg(test)] diff --git a/src/harness/limits/test.rs b/src/harness/limits/test.rs index f2af4f2..55c5555 100644 --- a/src/harness/limits/test.rs +++ b/src/harness/limits/test.rs @@ -1,13 +1,169 @@ //! Unit tests for run-scoped limit enforcement. //! -//! Smoke-checks that default [`RunLimits`] build a [`LimitTracker`] and that -//! recording a model call advances the counter. +//! Covers the counter/cap smoke path, the fail-open vs fail-closed +//! reconciliation of two limit sources (LOOP-1), and the +//! error-vs-stop-with-partial exhaustion behaviour (LOOP-9b). + +use super::{LimitBehavior, LimitKind, LimitOutcome, LimitTracker, RunLimits}; +use crate::error::TinyAgentsError; #[test] fn smoke_default_limits_compile() { - use super::{LimitTracker, RunLimits}; let limits = RunLimits::default(); let mut tracker = LimitTracker::new(limits); tracker.record_model_call().unwrap(); assert_eq!(tracker.model_calls(), 1); } + +// ── LOOP-1: reconciling two limit sources ──────────────────────────────────── + +#[test] +fn tighten_call_limits_keeps_the_stricter_cap_in_both_directions() { + // Regression test (LOOP-1): `sync_call_limits` is a plain assignment with + // no `min`, so the agent loop's per-run call overwrote a caller's explicit + // `RunConfig` cap *upward* — `RunConfig::new("r").with_max_model_calls(2)` + // against the default policy ran 25. The existing coverage only pinned the + // loosening direction, so nothing caught it. + + // A looser candidate must NOT widen an existing cap. + let mut strict = LimitTracker::new(RunLimits::default().with_max_model_calls(2)); + strict.tighten_call_limits(25, 50); + assert_eq!( + strict.limits().max_model_calls, + 2, + "a looser second source silently widened an explicit cap" + ); + + // A stricter candidate does tighten it. + let mut loose = LimitTracker::new(RunLimits::default()); + loose.tighten_call_limits(5, 10); + assert_eq!(loose.limits().max_model_calls, 5); + assert_eq!(loose.limits().max_tool_calls, 10); + + // Each axis is reconciled independently. + let mut mixed = LimitTracker::new( + RunLimits::default() + .with_max_model_calls(3) + .with_max_tool_calls(999), + ); + mixed.tighten_call_limits(100, 7); + assert_eq!(mixed.limits().max_model_calls, 3); + assert_eq!(mixed.limits().max_tool_calls, 7); +} + +#[test] +fn tighten_call_limits_preserves_recorded_counts() { + let mut tracker = LimitTracker::new(RunLimits::default()); + tracker.record_model_call().unwrap(); + tracker.record_tool_call().unwrap(); + tracker.tighten_call_limits(5, 10); + assert_eq!(tracker.model_calls(), 1); + assert_eq!(tracker.tool_calls(), 1); +} + +#[test] +fn tighten_call_limits_actually_trips_at_the_stricter_cap() { + // The cap must be enforced, not merely reported. + let mut tracker = LimitTracker::new(RunLimits::default().with_max_model_calls(2)); + tracker.tighten_call_limits(25, 50); + tracker.record_model_call().unwrap(); + tracker.record_model_call().unwrap(); + let err = tracker + .record_model_call() + .expect_err("the caller's cap of 2 must still trip"); + assert!(err.to_string().contains('2'), "got {err}"); +} + +#[test] +fn sync_call_limits_remains_the_documented_fail_open_override() { + // Kept deliberately, for the case a `RunPolicy` must raise a cap above the + // `RunConfig` *default*. Pinned so the two methods cannot be confused. + let mut tracker = LimitTracker::new(RunLimits::default().with_max_model_calls(2)); + tracker.sync_call_limits(30, 1_000); + assert_eq!(tracker.limits().max_model_calls, 30); + assert_eq!(tracker.limits().max_tool_calls, 1_000); +} + +// ── LOOP-9b: exhaustion behaviour ──────────────────────────────────────────── + +#[test] +fn limit_behavior_defaults_to_error() { + assert_eq!(RunLimits::default().behavior, LimitBehavior::Error); + assert_eq!(LimitBehavior::Error.as_str(), "error"); + assert_eq!(LimitBehavior::StopWithPartial.as_str(), "stop_with_partial"); +} + +#[test] +fn error_behavior_preserves_the_historical_hard_failure() { + let mut tracker = LimitTracker::new(RunLimits::default().with_max_model_calls(1)); + assert_eq!( + tracker.try_record_model_call().unwrap(), + LimitOutcome::Proceed + ); + let err = tracker.try_record_model_call().unwrap_err(); + assert!(matches!(err, TinyAgentsError::Validation(_)), "got {err:?}"); + assert!(err.to_string().contains("max model calls (1) exceeded")); +} + +#[test] +fn stop_with_partial_reports_a_clean_stop_instead_of_discarding_the_run() { + // Regression test (LOOP-9b): every cap used to raise `LimitExceeded`, + // throwing away the whole run and all the work already done. LangChain's + // `ModelCallLimitMiddleware` returns a jump-to-end instead. + let mut tracker = LimitTracker::new( + RunLimits::default() + .with_max_model_calls(1) + .with_max_tool_calls(2) + .with_behavior(LimitBehavior::StopWithPartial), + ); + + assert_eq!( + tracker.try_record_model_call().unwrap(), + LimitOutcome::Proceed + ); + assert_eq!( + tracker + .try_record_model_call() + .expect("stop_with_partial must not error"), + LimitOutcome::Stop(LimitKind::ModelCalls) + ); + + tracker.try_record_tool_call().unwrap(); + tracker.try_record_tool_call().unwrap(); + assert_eq!( + tracker.try_record_tool_call().unwrap(), + LimitOutcome::Stop(LimitKind::ToolCalls) + ); +} + +#[test] +fn rollback_tool_calls_uncounts_calls_that_never_ran() { + // LangChain's `ToolCallLimitMiddleware` rolls the thread count back for + // every remaining call it answered with a "stopped before this could run" + // message rather than executing. + let mut tracker = + LimitTracker::new(RunLimits::default().with_behavior(LimitBehavior::StopWithPartial)); + for _ in 0..5 { + tracker.try_record_tool_call().unwrap(); + } + tracker.rollback_tool_calls(3); + assert_eq!(tracker.tool_calls(), 2); + + // Saturates rather than wrapping. + tracker.rollback_tool_calls(99); + assert_eq!(tracker.tool_calls(), 0); +} + +#[test] +fn limit_kind_labels_match_the_event_layer() { + // The limits module keeps its own `LimitKind` so it need not depend on the + // observability layer; the labels must not drift apart. + assert_eq!( + LimitKind::ModelCalls.as_str(), + crate::harness::events::LimitKind::ModelCalls.as_str() + ); + assert_eq!( + LimitKind::ToolCalls.as_str(), + crate::harness::events::LimitKind::ToolCalls.as_str() + ); +} diff --git a/src/harness/limits/types.rs b/src/harness/limits/types.rs index 0ce6564..b4cd37d 100644 --- a/src/harness/limits/types.rs +++ b/src/harness/limits/types.rs @@ -40,6 +40,103 @@ pub struct RunLimits { /// fails fast (see [`crate::harness::subagent`]). Defaults to /// [`RunLimits::DEFAULT_MAX_DEPTH`]. pub max_depth: usize, + /// What the run should do when a call cap is reached. Defaults to + /// [`LimitBehavior::Error`], which is the historical behaviour. + pub behavior: LimitBehavior, +} + +/// What a run does when it reaches a configured call cap. +/// +/// Ported from LangChain's `exit_behavior` on `ModelCallLimitMiddleware` / +/// `ToolCallLimitMiddleware`. Every cap here used to be a hard error, which +/// throws away the whole run *and everything it already accomplished* — for a +/// long research run that is often the worst possible outcome, since the +/// partial answer was the valuable part. +/// +/// The variant only names the *policy*; the agent loop is what acts on it. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum LimitBehavior { + /// Fail the run with + /// [`TinyAgentsError::LimitExceeded`][crate::error::TinyAgentsError::LimitExceeded]. + /// The historical (and still default) behaviour. + #[default] + Error, + + /// Stop the loop cleanly and return whatever the run has produced so far, + /// as if the model had finished its turn. + /// + /// # Contract for the agent loop (wave 2) + /// + /// When [`LimitTracker::record_model_call`][crate::harness::limits::LimitTracker::record_model_call] + /// or [`record_tool_call`][crate::harness::limits::LimitTracker::record_tool_call] + /// reports exhaustion under this behaviour they return + /// [`LimitOutcome::Stop`] instead of `Err`, and the loop must: + /// + /// 1. Emit the existing + /// [`AgentEvent::LimitReached`][crate::harness::events::AgentEvent::LimitReached] + /// so the stop is still observable and still distinguishable from a + /// model that simply finished. + /// 2. Break out of the loop and finalize normally, keeping the transcript + /// accumulated so far as the run result. + /// 3. For the **tool** cap specifically, mirror LangChain's `"end"` path: + /// append a tool result for every remaining requested call saying it was + /// stopped before it could run (the provider APIs require every + /// `tool_call_id` to be answered, so skipping them corrupts the + /// transcript), and **roll the counter back** by the number of calls that + /// never executed via + /// [`LimitTracker::rollback_tool_calls`][crate::harness::limits::LimitTracker::rollback_tool_calls], + /// so the reported count reflects work actually done. + StopWithPartial, +} + +impl LimitBehavior { + /// Stable, snake_case label for logs and telemetry dimensions. + pub fn as_str(self) -> &'static str { + match self { + LimitBehavior::Error => "error", + LimitBehavior::StopWithPartial => "stop_with_partial", + } + } +} + +/// The result of recording a call against a [`RunLimits`] cap. +/// +/// Returned by the `try_record_*` methods so the caller sees the cap decision +/// as data rather than only as a `Result`. The `record_*` methods remain for +/// callers that always want the hard-error form. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LimitOutcome { + /// The call is within the cap; carry on. + Proceed, + /// The cap is exhausted and [`RunLimits::behavior`] is + /// [`LimitBehavior::StopWithPartial`]: stop the loop cleanly and return + /// what the run has so far. Carries which cap tripped. + Stop(LimitKind), +} + +/// Names which cap a [`LimitOutcome::Stop`] refers to. +/// +/// Deliberately mirrors +/// [`crate::harness::events::LimitKind`] rather than reusing it, so the limits +/// module (a leaf with no event dependency) does not have to import the +/// observability layer. Convert with [`LimitKind::as_str`] when emitting. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LimitKind { + /// The per-run model-call cap. + ModelCalls, + /// The per-run tool-call cap. + ToolCalls, +} + +impl LimitKind { + /// Stable, snake_case label matching + /// [`crate::harness::events::LimitKind::as_str`]. + pub fn as_str(self) -> &'static str { + match self { + LimitKind::ModelCalls => "model_calls", + LimitKind::ToolCalls => "tool_calls", + } + } } impl RunLimits { @@ -55,6 +152,7 @@ impl Default for RunLimits { max_wall_clock_ms: None, max_retries_per_call: 3, max_depth: Self::DEFAULT_MAX_DEPTH, + behavior: LimitBehavior::Error, } } } diff --git a/src/harness/memory/mod.rs b/src/harness/memory/mod.rs index 47f72d7..62120c8 100644 --- a/src/harness/memory/mod.rs +++ b/src/harness/memory/mod.rs @@ -105,7 +105,19 @@ impl StoreChatHistory { /// Wraps `store` as a chat-history backend. pub fn new(store: S) -> Self { - Self { store } + Self { + store, + append_locks: Default::default(), + } + } + + /// Returns the append lock for `thread_id`, creating it on first use. + fn append_lock(&self, thread_id: &str) -> Result>> { + let mut locks = self + .append_locks + .lock() + .map_err(|e| TinyAgentsError::Memory(format!("chat history lock poisoned: {e}")))?; + Ok(locks.entry(thread_id.to_string()).or_default().clone()) } /// Returns a reference to the backing store. @@ -127,6 +139,9 @@ impl ChatHistory for StoreChatHistory { } async fn append(&self, thread_id: &str, message: Message) -> Result<()> { + // Serialize the read-modify-write per thread; see `append_locks`. + let lock = self.append_lock(thread_id)?; + let _guard = lock.lock().await; let mut messages = self.messages(thread_id).await?; messages.push(message); let value = serde_json::to_value(&messages)?; diff --git a/src/harness/memory/types.rs b/src/harness/memory/types.rs index 063bda0..201d587 100644 --- a/src/harness/memory/types.rs +++ b/src/harness/memory/types.rs @@ -94,6 +94,22 @@ pub struct InMemoryChatHistory { pub struct StoreChatHistory { /// The backing long-term store. pub(crate) store: S, + /// Per-thread append locks. + /// + /// [`ChatHistory::append`] over a key-value store is a read-modify-write: + /// load the thread, push, write the whole list back. Two concurrent appends + /// that both read before either writes produce a last-writer-wins result in + /// which one message is silently gone — and the two `ChatHistory` backends + /// then give *different* guarantees for the same trait method, because + /// `InMemoryChatHistory::append` holds its lock for the whole operation. + /// + /// A per-thread async mutex closes that gap and keeps the two consistent. + /// It is an **in-process** guarantee, which is the same scope + /// `InMemoryChatHistory` offers; across processes sharing one `FileStore` + /// directory the read-modify-write is still racy, and closing *that* needs a + /// compare-and-swap primitive the `Store` trait does not have. Prefer + /// `replace` (a single write) when the full list is already in hand. + pub(crate) append_locks: Arc>>>>, } /// A thin thread-scoped wrapper over a [`ChatHistory`] with an optional diff --git a/src/harness/message/mod.rs b/src/harness/message/mod.rs index 4019fa6..e68f241 100644 --- a/src/harness/message/mod.rs +++ b/src/harness/message/mod.rs @@ -10,8 +10,10 @@ //! See [`types`] for definitions. This module provides ergonomic constructors //! and a [`Message::text`] accessor. +mod tokens; mod types; +pub use tokens::*; pub use types::*; /// Approximate token-estimation weight of a single image content block, in @@ -134,6 +136,7 @@ impl Message { tool_call_id: tool_call_id.into(), content: vec![ContentBlock::Text(content.into())], trusted_verbatim: false, + artifact: None, }) } @@ -144,14 +147,35 @@ impl Message { /// carries structured metadata the message shape would otherwise drop, and /// [`ToolMessage::trusted_verbatim`] is the part a host must not lose — it /// is what tells the host this content may not be reshaped. + /// + /// [`ToolResult::raw`][crate::harness::tool::ToolResult::raw] is carried + /// across into [`ToolMessage::artifact`], so a tool can return a small + /// model-facing summary in `content` and still leave the full structured + /// payload reachable from `run.messages` (LangChain's + /// `response_format="content_and_artifact"`). The artifact is host-side + /// only — provider conversion serialises [`Message::text`], never the + /// artifact. pub fn tool_from_result(result: &crate::harness::tool::ToolResult) -> Self { Message::Tool(ToolMessage { tool_call_id: result.call_id.clone(), content: vec![ContentBlock::Text(result.content.clone())], trusted_verbatim: result.is_trusted_verbatim(), + artifact: result.raw.clone(), }) } + /// Returns the structured artifact carried by a tool message, if any. + /// + /// `None` for every non-tool message and for tool messages whose producing + /// [`ToolResult`][crate::harness::tool::ToolResult] set no `raw` payload. + /// See [`ToolMessage::artifact`]. + pub fn artifact(&self) -> Option<&serde_json::Value> { + match self { + Message::Tool(m) => m.artifact.as_ref(), + _ => None, + } + } + /// Returns the concatenated text of all text content blocks. pub fn text(&self) -> String { match self { @@ -183,14 +207,36 @@ impl Message { } /// Approximate character weight of the message across *all* content blocks - /// (text, JSON, images, reasoning, provider extensions), for token - /// estimation and context-window gating. + /// (text, JSON, images, reasoning, provider extensions) **plus the + /// structural payload that lives outside `content`**: an assistant + /// message's [`tool_calls`][AssistantMessage::tool_calls] and a tool + /// message's [`tool_call_id`][ToolMessage::tool_call_id]. /// /// Distinct from [`char_len`](Self::char_len), which counts only visible /// text: a transcript dominated by images, large tool-result JSON, or model /// reasoning under-counts badly under `char_len`, so compaction/trim would /// silently never trigger even as the real context window overflows. See /// [`ContentBlock::estimated_char_weight`]. + /// + /// # Why tool calls must be counted here + /// + /// An assistant turn that *only* calls tools carries **empty `content`**: + /// the tool name and its argument JSON — often the largest part of the turn + /// — live in `tool_calls`. Counting `content` alone estimated such a + /// message at zero, so a 50-turn tool-driven run whose assistant messages + /// each carry a 2 KB argument blob estimated to near-nothing and never + /// tripped [`SummarizationPolicy::should_summarize`][crate::harness::summarization::SummarizationPolicy::should_summarize], + /// letting the window overflow uncompacted — exactly the failure this + /// estimator exists to prevent. + /// + /// Mirrors LangChain's `count_tokens_approximately`, which adds + /// `repr(tool_calls)` for AI messages and the `tool_call_id` for tool + /// messages. The role label and per-message overhead are *not* added here; + /// they belong to the message-level counters in + /// [`crate::harness::message::count_tokens_approximately`]. + /// + /// The [`ToolMessage::artifact`] payload is deliberately **not** counted: it + /// never reaches the provider, so it occupies no context window. pub fn estimated_char_weight(&self) -> usize { let content = match self { Message::System(m) => &m.content, @@ -198,10 +244,41 @@ impl Message { Message::Assistant(m) => &m.content, Message::Tool(m) => &m.content, }; - content + let content_weight: usize = content .iter() .map(ContentBlock::estimated_char_weight) - .sum() + .sum(); + + let structural_weight = match self { + Message::Assistant(m) => tool_calls_char_weight(&m.tool_calls), + Message::Tool(m) => m.tool_call_id.chars().count(), + _ => 0, + }; + + content_weight + structural_weight + } +} + +/// Approximate character weight of an assistant message's tool-call array. +/// +/// Serialises the calls to JSON (the closest analogue of LangChain's +/// `repr(tool_calls)`) and counts characters. Falls back to a per-call estimate +/// from the name plus the raw argument value when serialisation fails, so the +/// weight is never silently zero. +fn tool_calls_char_weight(tool_calls: &[crate::harness::tool::ToolCall]) -> usize { + if tool_calls.is_empty() { + return 0; + } + match serde_json::to_string(tool_calls) { + Ok(rendered) => rendered.chars().count(), + Err(_) => tool_calls + .iter() + .map(|call| { + call.name.chars().count() + + call.id.chars().count() + + call.arguments.to_string().chars().count() + }) + .sum(), } } diff --git a/src/harness/message/test.rs b/src/harness/message/test.rs index b1ce63e..5aa13eb 100644 --- a/src/harness/message/test.rs +++ b/src/harness/message/test.rs @@ -277,3 +277,160 @@ fn an_unset_flag_is_omitted_from_the_wire() { let wire = serde_json::to_value(Message::Tool(msg)).unwrap(); assert_eq!(wire["tool"]["trusted_verbatim"], json!(true)); } + +// --------------------------------------------------------------------------- +// content_and_artifact (C7) +// --------------------------------------------------------------------------- + +/// A tool may return a small model-facing summary while leaving the full +/// structured payload reachable from the transcript. +#[test] +fn tool_from_result_carries_the_structured_artifact() { + let mut result = crate::harness::tool::ToolResult::text("c1", "query", "3 rows"); + result.raw = Some(json!({"rows": [1, 2, 3]})); + + let message = Message::tool_from_result(&result); + + // The model sees only the summary… + assert_eq!(message.text(), "3 rows"); + // …while application code can recover the payload. + assert_eq!(message.artifact(), Some(&json!({"rows": [1, 2, 3]}))); +} + +/// The artifact is host-side state and must never be part of the text the wire +/// conversion serialises. +#[test] +fn artifact_does_not_leak_into_message_text() { + let mut result = crate::harness::tool::ToolResult::text("c1", "query", "ok"); + result.raw = Some(json!({"secret_looking_blob": "x".repeat(500)})); + let message = Message::tool_from_result(&result); + + assert_eq!(message.text(), "ok"); + assert!(!message.text().contains("secret_looking_blob")); + // It is also not charged to the context window, because it never reaches + // the provider. + assert_eq!(message.estimated_char_weight(), "ok".len() + "c1".len()); +} + +/// Non-tool messages have no artifact, and a plain `Message::tool` never +/// invents one. +#[test] +fn artifact_is_absent_without_a_tool_result() { + assert!(Message::assistant("hi").artifact().is_none()); + assert!(Message::tool("c1", "r").artifact().is_none()); +} + +/// Transcripts persisted before the field existed still deserialise, and a +/// message with no artifact stays byte-identical on the wire. +#[test] +fn artifact_serde_defaults_and_skips() { + let legacy: Message = + serde_json::from_value(json!({"tool": {"tool_call_id": "c1", "content": []}})).unwrap(); + assert!(legacy.artifact().is_none()); + + let rendered = serde_json::to_value(&legacy).unwrap(); + assert!(rendered["tool"].get("artifact").is_none()); + + let mut result = crate::harness::tool::ToolResult::text("c1", "t", "ok"); + result.raw = Some(json!({"a": 1})); + let round_tripped: Message = + serde_json::from_value(serde_json::to_value(Message::tool_from_result(&result)).unwrap()) + .unwrap(); + assert_eq!(round_tripped.artifact(), Some(&json!({"a": 1}))); +} + +// --------------------------------------------------------------------------- +// Token estimation (REASON-3) +// --------------------------------------------------------------------------- + +/// An assistant turn that only calls tools carries no text. Counting content +/// alone estimated it at zero tokens, so a tool-driven run never tripped a +/// compaction gate no matter how large its argument blobs grew. +#[test] +fn tool_only_assistant_turn_is_not_estimated_at_zero() { + let message = Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ToolCall::new( + "call_1", + "search", + json!({"query": "q".repeat(1_000)}), + )], + usage: None, + }); + + assert_eq!(message.text(), ""); + assert!(message.estimated_char_weight() > 1_000); + assert!(estimate_message_tokens(&message) > 200); + assert!(count_tokens_approximately(&[message]) > 200); +} + +/// The parity counter charges role labels and per-message framing, so many +/// short messages are not free. +#[test] +fn count_tokens_approximately_charges_role_and_framing() { + let one = count_tokens_approximately(&[Message::user("")]); + // "user" (4 chars => 1 token) + 3 framing tokens. + assert_eq!(one, 4); + + let four = count_tokens_approximately(&vec![Message::user(""); 4]); + assert_eq!(four, 16, "per-message rounding must sum to the whole"); +} + +/// Tool declarations occupy the prompt too; a run with verbose schemas starts +/// well into its window before the first user message. +#[test] +fn tool_schemas_are_counted() { + let schema = crate::harness::tool::ToolSchema::new( + "search", + "Search the corpus", + json!({"type": "object", "properties": {"query": {"type": "string"}}}), + ); + let options = TokenCountOptions::default(); + assert!(count_tool_schema_tokens(&[schema], &options) > 5); + assert_eq!(count_tool_schema_tokens(&[], &options), 0); +} + +/// Usage calibration may only nudge the heuristic, never replace it: the factor +/// is clamped so one anomalous provider report cannot blow the estimate up or +/// collapse it below the heuristic. +#[test] +fn usage_metadata_scaling_is_clamped_both_ways() { + let options = TokenCountOptions::default().with_usage_metadata_scaling(); + + let inflated = vec![ + Message::user("hello there"), + Message::Assistant(AssistantMessage { + id: None, + content: vec![ContentBlock::Text("hi".into())], + tool_calls: Vec::new(), + usage: Some(Usage { + total_tokens: 1_000_000, + ..Usage::default() + }), + }), + ]; + let baseline = count_tokens_approximately(&inflated); + let scaled = count_tokens_approximately_with(&inflated, &options); + assert!(scaled <= (baseline as f64 * USAGE_SCALE_MAX).ceil() as u64); + assert!(scaled >= baseline); + + let deflated = vec![ + Message::user("hello there"), + Message::Assistant(AssistantMessage { + id: None, + content: vec![ContentBlock::Text("hi".into())], + tool_calls: Vec::new(), + usage: Some(Usage { + total_tokens: 1, + ..Usage::default() + }), + }), + ]; + // A provider reporting fewer tokens than the heuristic guessed is not a + // licence to under-count. + assert_eq!( + count_tokens_approximately_with(&deflated, &options), + count_tokens_approximately(&deflated) + ); +} diff --git a/src/harness/message/tokens.rs b/src/harness/message/tokens.rs new file mode 100644 index 0000000..a74228d --- /dev/null +++ b/src/harness/message/tokens.rs @@ -0,0 +1,273 @@ +//! The crate's shared, structurally-complete token estimator. +//! +//! # Why this module exists +//! +//! Several independent places in the harness need "roughly how many tokens is +//! this transcript?" — compaction gating, context middleware, and budget +//! preflight. Each grew its own `chars / 4` loop over +//! [`Message::text`][super::Message::text], and every one of them silently +//! under-counted the same way: a transcript's *structure* (tool calls, tool +//! result correlation ids, role labels, per-message framing) is invisible to +//! `text()`, and an assistant turn that only calls tools has **no text at +//! all**. A tool-driven run therefore estimated near zero and no gate fired. +//! +//! This module is the single correct implementation those call sites should +//! use. It is a port of LangChain's `count_tokens_approximately` +//! (`libs/core/langchain_core/messages/utils.py`), adapted to this crate's +//! [`Message`] model. +//! +//! # What is counted +//! +//! | Part | Source | +//! | ---- | ------ | +//! | content blocks (text, JSON, reasoning, provider extensions) | [`ContentBlock::estimated_char_weight`][super::ContentBlock::estimated_char_weight] | +//! | images | flat per-image weight, not the base64 length | +//! | assistant `tool_calls` | JSON rendering of the call array | +//! | tool `tool_call_id` | the id string | +//! | role label (`system` / `user` / `assistant` / `tool`) | [`TokenCountOptions::count_role`] | +//! | per-message framing overhead | [`TokenCountOptions::extra_tokens_per_message`] | +//! | tool declarations offered to the model | [`count_tool_schema_tokens`] | +//! +//! Rounding is **per message** (`ceil`), matching LangChain, so the parts of a +//! transcript sum to the whole rather than losing a fraction of a token each. +//! +//! # Two entry points, deliberately +//! +//! - [`count_tokens_approximately`] / [`count_tokens_approximately_with`] — +//! the recommended counter. Use it for anything that compares an estimate +//! against a real provider limit (context-window gating, budget preflight). +//! - [`estimate_message_tokens`] / [`estimate_slice_tokens`] — the crate's +//! older bare `floor(chars / 4)` heuristic, retained because existing +//! thresholds (and their tests) are calibrated against it. It shares the same +//! *corrected* char weight, so it no longer under-counts tool calls; it just +//! omits role labels and framing overhead. +//! +//! Nothing here is a real tokenizer. Treat every number as ±30%. + +use super::{AssistantMessage, Message}; +use crate::harness::tool::ToolSchema; + +/// Default characters per token (~4 for English prose and code). +pub const DEFAULT_CHARS_PER_TOKEN: f64 = 4.0; + +/// Default per-message framing overhead in tokens (role delimiters, +/// begin/end-of-message markers). Matches LangChain's +/// `extra_tokens_per_message=3`. +pub const DEFAULT_EXTRA_TOKENS_PER_MESSAGE: f64 = 3.0; + +/// Lower clamp applied to the usage-metadata calibration factor. +/// +/// The factor may only *increase* the estimate: a provider reporting fewer +/// tokens than the heuristic guessed is not a licence to under-count, because +/// under-counting is the failure mode that lets a window overflow silently. +pub const USAGE_SCALE_MIN: f64 = 1.0; + +/// Upper clamp applied to the usage-metadata calibration factor. +/// +/// Without a ceiling a single anomalous usage report (a cached-prompt turn, a +/// provider that counts system overhead differently) would multiply every +/// subsequent estimate without bound. Matches LangChain's `min(1.25, …)`. +pub const USAGE_SCALE_MAX: f64 = 1.25; + +/// Tuning knobs for [`count_tokens_approximately_with`]. +/// +/// [`Default`] reproduces LangChain's defaults with usage-metadata calibration +/// **off**, so the counter stays a pure function of the transcript unless a +/// caller opts in. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct TokenCountOptions { + /// Characters that make up one token. See [`DEFAULT_CHARS_PER_TOKEN`]. + pub chars_per_token: f64, + /// Framing tokens added per message. See + /// [`DEFAULT_EXTRA_TOKENS_PER_MESSAGE`]. + pub extra_tokens_per_message: f64, + /// Whether to charge for the message's role label. + pub count_role: bool, + /// Calibrate the character heuristic against real provider usage. + /// + /// When enabled, the newest assistant message carrying a non-zero + /// [`Usage::total_tokens`][crate::harness::usage::Usage::total_tokens] is + /// compared against the heuristic's running estimate at that point, and the + /// whole count is scaled by that ratio — clamped to + /// `[USAGE_SCALE_MIN, USAGE_SCALE_MAX]` so one bad report cannot blow the + /// estimate up or collapse it. + /// + /// This is the crate's advantage over LangChain's version: the usage is + /// already on [`AssistantMessage::usage`], so no side-channel is needed. + pub use_usage_metadata_scaling: bool, +} + +impl Default for TokenCountOptions { + fn default() -> Self { + Self { + chars_per_token: DEFAULT_CHARS_PER_TOKEN, + extra_tokens_per_message: DEFAULT_EXTRA_TOKENS_PER_MESSAGE, + count_role: true, + use_usage_metadata_scaling: false, + } + } +} + +impl TokenCountOptions { + /// Enables usage-metadata calibration. See + /// [`use_usage_metadata_scaling`][Self::use_usage_metadata_scaling]. + pub fn with_usage_metadata_scaling(mut self) -> Self { + self.use_usage_metadata_scaling = true; + self + } + + /// Sets the characters-per-token divisor. + pub fn with_chars_per_token(mut self, chars_per_token: f64) -> Self { + self.chars_per_token = chars_per_token; + self + } + + /// Effective divisor, guarding against a zero or negative configuration + /// that would otherwise produce a non-finite estimate. + fn divisor(&self) -> f64 { + if self.chars_per_token > 0.0 { + self.chars_per_token + } else { + DEFAULT_CHARS_PER_TOKEN + } + } +} + +/// The OpenAI-style role label for a message, as counted toward its token cost. +pub fn message_role_label(message: &Message) -> &'static str { + match message { + Message::System(_) => "system", + Message::User(_) => "user", + Message::Assistant(_) => "assistant", + Message::Tool(_) => "tool", + } +} + +/// Total character weight charged for a single message: content blocks, tool +/// calls / tool-call id, and (optionally) the role label. +fn message_char_weight(message: &Message, options: &TokenCountOptions) -> usize { + let mut chars = message.estimated_char_weight(); + if options.count_role { + chars += message_role_label(message).len(); + } + chars +} + +/// Approximate the token count of `messages` with default options. +/// +/// This is the recommended estimator for any comparison against a real provider +/// limit. See the [module docs][self] for what is counted. +pub fn count_tokens_approximately(messages: &[Message]) -> u64 { + count_tokens_approximately_with(messages, &TokenCountOptions::default()) +} + +/// Approximate the token count of `messages` under explicit `options`. +/// +/// Emits `[tokens]`-prefixed trace output describing the running total and any +/// usage calibration applied, so an over- or under-estimate can be traced to a +/// specific message without re-running the model. +pub fn count_tokens_approximately_with(messages: &[Message], options: &TokenCountOptions) -> u64 { + let divisor = options.divisor(); + let mut total = 0.0_f64; + + // Newest assistant message carrying real usage, and the running heuristic + // total at that point — the two halves of the calibration ratio. + let mut last_reported_total: Option = None; + let mut approx_at_last_report: Option = None; + + for (index, message) in messages.iter().enumerate() { + let chars = message_char_weight(message, options); + let message_tokens = (chars as f64 / divisor).ceil() + options.extra_tokens_per_message; + total += message_tokens; + + tracing::trace!( + "[tokens] message index={index} role={role} chars={chars} tokens={message_tokens} running={total}", + role = message_role_label(message) + ); + + if options.use_usage_metadata_scaling + && let Message::Assistant(AssistantMessage { + usage: Some(usage), .. + }) = message + && usage.total_tokens > 0 + { + last_reported_total = Some(usage.total_tokens); + approx_at_last_report = Some(total); + } + } + + if options.use_usage_metadata_scaling + && messages.len() > 1 + && let (Some(reported), Some(approx)) = (last_reported_total, approx_at_last_report) + && approx > 0.0 + { + let raw_factor = reported as f64 / approx; + let factor = raw_factor.clamp(USAGE_SCALE_MIN, USAGE_SCALE_MAX); + tracing::debug!( + "[tokens] usage calibration reported={reported} approx={approx} raw_factor={raw_factor} clamped={factor}" + ); + total *= factor; + } + + let result = total.ceil().max(0.0) as u64; + tracing::debug!( + "[tokens] counted messages={} tokens={result}", + messages.len() + ); + result +} + +/// Approximate the token cost of the tool declarations offered to a model. +/// +/// Tool schemas are part of every request's prompt, so a run with twenty +/// verbose tool declarations starts thousands of tokens into its window before +/// the first user message. LangChain counts them the same way, dropping each +/// schema's own `title` / `description` because the tool's name and description +/// are already counted at the top level. +pub fn count_tool_schema_tokens(schemas: &[ToolSchema], options: &TokenCountOptions) -> u64 { + if schemas.is_empty() { + return 0; + } + let mut chars = 0usize; + for schema in schemas { + let mut parameters = schema.parameters.clone(); + if let Some(object) = parameters.as_object_mut() { + object.remove("title"); + object.remove("description"); + } + let rendered = serde_json::json!({ + "name": schema.name, + "description": schema.description, + "parameters": parameters, + }); + chars += rendered.to_string().chars().count(); + } + let tokens = (chars as f64 / options.divisor()).ceil().max(0.0) as u64; + tracing::debug!( + "[tokens] counted tool schemas count={} chars={chars} tokens={tokens}", + schemas.len() + ); + tokens +} + +/// The crate's legacy `floor(chars / 4)` heuristic for a single message, +/// computed over the *corrected* character weight (content blocks **plus** tool +/// calls and tool-call ids). +/// +/// Retained because the compaction thresholds in +/// [`crate::harness::summarization`] are calibrated against it. Prefer +/// [`count_tokens_approximately`] for new code: this variant charges nothing +/// for role labels or per-message framing and so runs slightly low on +/// many-short-messages transcripts. +/// +/// Returns at least `1` for any message with non-zero weight, so a short +/// message is never free. +pub fn estimate_message_tokens(message: &Message) -> u64 { + let chars = message.estimated_char_weight() as u64; + if chars == 0 { 0 } else { (chars / 4).max(1) } +} + +/// [`estimate_message_tokens`] summed over a slice. +pub fn estimate_slice_tokens(messages: &[Message]) -> u64 { + messages.iter().map(estimate_message_tokens).sum() +} diff --git a/src/harness/message/types.rs b/src/harness/message/types.rs index a953f04..a0f63a1 100644 --- a/src/harness/message/types.rs +++ b/src/harness/message/types.rs @@ -119,6 +119,43 @@ pub struct ToolMessage { /// crate omits unset fields. #[serde(default, skip_serializing_if = "is_false")] pub trusted_verbatim: bool, + + /// The structured payload the producing tool returned alongside its + /// model-facing text, carried across from + /// [`ToolResult::raw`][crate::harness::tool::ToolResult::raw] by + /// [`Message::tool_from_result`]. + /// + /// This is the crate's equivalent of LangChain's + /// `response_format="content_and_artifact"`: the model sees only + /// [`Message::text`] (a summary, a row count, a path), while application + /// code reading back `run.messages` can recover the full object — a parsed + /// dataframe, a binary handle, a large search result set — without the tool + /// having to inline megabytes of JSON into the transcript just to make it + /// reachable. + /// + /// # Wire contract + /// + /// **The artifact never reaches the provider.** Provider conversion + /// serialises a tool message from [`Message::text`] (its + /// [`ContentBlock::Text`] blocks) only, so an artifact is host-side state. + /// Any future provider adapter must keep that property: putting the + /// artifact on the wire would defeat the entire point of the field and can + /// blow the context window with the payload the summary was meant to + /// replace. + /// + /// `#[serde(default)]` so transcripts persisted before this field existed + /// still deserialise, and `skip_serializing_if` so a message with no + /// artifact stays byte-identical when persisted. + /// + /// # Note on `trusted_verbatim` + /// + /// [`ToolResult::mark_trusted_verbatim`][crate::harness::tool::ToolResult::mark_trusted_verbatim] + /// records its opt-in *inside* `raw` under + /// [`TRUSTED_VERBATIM_KEY`][crate::harness::tool::TRUSTED_VERBATIM_KEY], so + /// an artifact copied from such a result also carries that key. The copy is + /// deliberately verbatim: the crate does not edit a payload the tool owns. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub artifact: Option, } /// A structured conversation message. diff --git a/src/harness/middleware/library/budget.rs b/src/harness/middleware/library/budget.rs index c3e0dec..128f37b 100644 --- a/src/harness/middleware/library/budget.rs +++ b/src/harness/middleware/library/budget.rs @@ -126,15 +126,17 @@ impl BudgetLimits { } } -/// Estimates the input tokens a request will consume by summing a -/// heuristic token estimate over every message's text. Used for budget -/// preflight reservation, which only needs an order-of-magnitude bound. +/// Estimates the input tokens a request will consume, for budget preflight +/// reservation. +/// +/// Uses the crate's shared +/// [`count_tokens_approximately`][crate::harness::message::count_tokens_approximately] +/// estimator. Summing `estimate_tokens(&m.text())` counted only textual content +/// blocks, so a request dominated by large JSON tool results or image blocks +/// preflighted at close to zero tokens and sailed past a token budget it in +/// fact blew through. fn estimated_input_tokens(request: &ModelRequest) -> u64 { - request - .messages - .iter() - .map(|m| crate::harness::summarization::estimate_tokens(&m.text())) - .sum() + crate::harness::message::count_tokens_approximately(&request.messages) } impl BudgetMiddleware { @@ -292,6 +294,20 @@ impl Middleware for BudgetMidd guard.reserved_input_total = guard.reserved_input_total.saturating_sub(reserved); } + // A cache replay spent nothing: no provider call was made, no tokens + // were consumed, no money changed hands. Folding its (replayed) usage + // into the tracker bills phantom spend, and enough hits can abort a run + // on a budget it never actually touched. The reservation is still + // released above — that part is real bookkeeping. + if response.served_from_cache { + tracing::debug!( + target: "tinyagents::middleware", + label = self.label, + "[budget] skipping accounting for a cache-served response" + ); + return Ok(()); + } + let Some(usage) = response.usage else { return Ok(()); }; diff --git a/src/harness/middleware/library/context.rs b/src/harness/middleware/library/context.rs index 980d4ea..43978c2 100644 --- a/src/harness/middleware/library/context.rs +++ b/src/harness/middleware/library/context.rs @@ -12,8 +12,7 @@ use crate::harness::middleware::{ PromptCacheGuardMiddleware, }; use crate::harness::summarization::{ - ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, - estimate_tokens, trim_messages, + ConcatSummarizer, SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy, trim_messages, }; // ── MessageTrimMiddleware ───────────────────────────────────────────────────── @@ -44,10 +43,18 @@ impl Middleware for MessageTri // ── ContextCompressionMiddleware ────────────────────────────────────────────── -/// Estimate the total tokens of a message slice using the same per-message -/// heuristic the [`SummarizationPolicy`] uses internally. +/// Estimate the total tokens of a message slice. +/// +/// Uses the crate's shared +/// [`count_tokens_approximately`][crate::harness::message::count_tokens_approximately] +/// estimator rather than summing `estimate_tokens(&m.text())`. `text()` returns +/// only the *textual* content blocks, so a transcript of large JSON tool +/// results or image blocks estimated to nearly zero and the micro-compaction +/// budget gate never tripped on exactly the transcripts it exists to shrink. +/// The shared estimator charges every content block, tool call, and tool-call +/// id, and calibrates against reported usage metadata when it is present. fn total_message_tokens(messages: &[crate::harness::message::Message]) -> u64 { - messages.iter().map(|m| estimate_tokens(&m.text())).sum() + crate::harness::message::count_tokens_approximately(messages) } impl ContextCompressionMiddleware { @@ -330,6 +337,21 @@ impl Middleware for Microcompa continue; } if let Message::Tool(t) = &request.messages[i] { + // `ToolMessage::trusted_verbatim` means the producing tool + // asked for its content to reach the model byte-for-byte, and + // its doc names blanking as exactly the rewrite a host must not + // perform. Blanking one produced content that reads fine and is + // wrong — an input schema the model copies argument names out + // of, a signature, a diff — so leave it intact and reclaim + // tokens elsewhere. + if t.trusted_verbatim { + tracing::debug!( + target: "tinyagents::middleware", + tool_call_id = %t.tool_call_id, + "[microcompact] skipping a trusted_verbatim tool result" + ); + continue; + } let id = t.tool_call_id.clone(); request.messages[i] = Message::tool(id, self.placeholder.clone()); cleared += 1; @@ -401,15 +423,27 @@ impl Middleware for PromptCach async fn before_model( &self, - _ctx: &mut RunContext, + ctx: &mut RunContext, _state: &State, request: &mut ModelRequest, ) -> Result<()> { let layout = PromptCacheLayout::from_request(request); + let run_id = ctx.run_id().clone(); let mut previous = self.previous.lock().expect("previous mutex poisoned"); - if let Some(prev) = previous.as_ref() + // Only compare within one run. See the field docs on + // `PromptCacheGuardMiddleware::previous`: a prefix cache is scoped to a + // single conversation, so a baseline carried over from a previous run + // would report an invalidation that never happened. + if let Some((prev_run, prev)) = previous.as_ref() + && prev_run == &run_id && !prev.is_prefix_stable_against(&layout) { + tracing::debug!( + "[cache] prompt_cache_guard: prefix invalidated run={run_id} \ + before={} after={}", + prev.fingerprint(), + layout.fingerprint() + ); let event = CacheLayoutEvent::new(prev, &layout); let mut events = self.events.lock().expect("events mutex poisoned"); if self.max_events > 0 { @@ -419,7 +453,7 @@ impl Middleware for PromptCach events.push_back(event); } } - *previous = Some(layout); + *previous = Some((run_id, layout)); Ok(()) } } diff --git a/src/harness/middleware/library/test.rs b/src/harness/middleware/library/test.rs index 2a68fc4..ff97a68 100644 --- a/src/harness/middleware/library/test.rs +++ b/src/harness/middleware/library/test.rs @@ -828,8 +828,18 @@ async fn concurrent_runs_on_one_middleware_release_their_own_reservations() { mw.before_model(&mut large, &(), &mut large_req) .await .unwrap(); + // Expected from the crate's shared estimator (which charges every content + // block, tool call, and the role label) rather than a hand-copied number. + let expected = + crate::harness::message::count_tokens_approximately(&[Message::user("x".repeat(40))]) + + crate::harness::message::count_tokens_approximately(&[Message::user( + "y".repeat(400), + )]); let reserved = tracker.snapshot().reserved_input_total; - assert_eq!(reserved, 110, "both reservations are on the shared tracker"); + assert_eq!( + reserved, expected, + "both reservations are on the shared tracker" + ); let mut response = ModelResponse::assistant("ok"); mw.after_model(&mut small, &(), &mut response) @@ -892,7 +902,10 @@ async fn shared_tracker_reservation_is_atomic_under_concurrency() { use crate::harness::middleware::Middleware; let tracker = BudgetTracker::new(); - let per_call_tokens = 10u64; // "x" * 40 chars / 4 == 10 estimated tokens. + // Derived from the shared estimator so the test tracks it instead of + // pinning a stale hand-computed constant. + let per_call_tokens = + crate::harness::message::count_tokens_approximately(&[Message::user("x".repeat(40))]); let concurrent_capacity = 4u64; let attempts = 10usize; let limits = BudgetLimits { diff --git a/src/harness/middleware/test.rs b/src/harness/middleware/test.rs index 0d4655b..b080c93 100644 --- a/src/harness/middleware/test.rs +++ b/src/harness/middleware/test.rs @@ -39,6 +39,7 @@ fn response_with_usage(usage: Usage) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -999,6 +1000,7 @@ fn response_text(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/src/harness/middleware/types.rs b/src/harness/middleware/types.rs index 8944af1..1d6ca12 100644 --- a/src/harness/middleware/types.rs +++ b/src/harness/middleware/types.rs @@ -25,6 +25,7 @@ use async_trait::async_trait; use crate::error::{Result, TinyAgentsError}; use crate::harness::cache::CacheLayoutEvent; use crate::harness::context::RunContext; +use crate::harness::ids::RunId; use crate::harness::model::{ModelDelta, ModelRequest, ModelResponse}; use crate::harness::summarization::{SummarizationPolicy, Summarizer, SummaryRecord, TrimStrategy}; use crate::harness::tool::{ToolCall, ToolDelta, ToolResult}; @@ -66,6 +67,20 @@ pub struct AgentRun { pub tool_calls: usize, /// Number of loop iterations (model/tool super-steps) executed. pub steps: usize, + /// Set when the run stopped because steering latched a **pause** rather + /// than because the model finished. + /// + /// A paused run has no `final_response`, exactly like a run whose model + /// returned an empty answer — which is why the two were previously + /// indistinguishable. Check this field (or + /// [`HarnessRunStatus`][crate::harness::events::HarnessRunStatus], which + /// reports `Interrupted` for a paused run) before treating a missing final + /// response as a completed-but-empty answer. The pause stays latched on the + /// [`SteeringHandle`][crate::harness::steering::SteeringHandle], so + /// [`SteeringHandle::resume`][crate::harness::steering::SteeringHandle::resume] + /// lifts it and a fresh invocation continues from + /// [`AgentRun::messages`]. + pub paused: Option, } // ── Middleware trait ────────────────────────────────────────────────────────── @@ -613,7 +628,18 @@ pub const DEFAULT_CACHE_GUARD_EVENT_CAP: usize = 1024; pub struct PromptCacheGuardMiddleware { pub(crate) label: &'static str, - pub(crate) previous: Mutex>, + /// The previous pass's layout, tagged with the run it was observed in. + /// + /// The run id is load-bearing. A KV-cache prefix is only meaningful + /// *within* one conversation, so comparing the last request of one run + /// against the first request of the next compares two unrelated + /// transcripts and reports an invalidation that never happened. A single + /// guard instance is routinely shared across runs — a sub-agent's + /// middleware stack is built once and its agent invoked many times — so + /// this is the common case, not an edge case. It went unnoticed while + /// stability was compared by segment id alone, because any two requests + /// carrying the same segment ids compared equal regardless of content. + pub(crate) previous: Mutex>, pub(crate) events: Mutex>, pub(crate) max_events: usize, } diff --git a/src/harness/model/mod.rs b/src/harness/model/mod.rs index 15789c3..9c4d918 100644 --- a/src/harness/model/mod.rs +++ b/src/harness/model/mod.rs @@ -59,6 +59,7 @@ const MODEL_CONTEXT_PATTERNS: &[(&str, ContextPatternMatch, u64)] = &[ ("claude-3-5-sonnet", ContextPatternMatch::Substring, 200_000), ("claude-3-5-haiku", ContextPatternMatch::Substring, 200_000), ("claude-3-opus", ContextPatternMatch::Substring, 200_000), + ("gpt-5", ContextPatternMatch::Substring, 400_000), ("gpt-4.1", ContextPatternMatch::Substring, 1_047_576), ("gpt-4o", ContextPatternMatch::Substring, 128_000), ("gpt-4-turbo", ContextPatternMatch::Substring, 128_000), @@ -174,6 +175,7 @@ impl ModelProfile { && (!set.native_structured_output || self.native_structured_output) && (!set.json_schema || self.json_schema) && (!set.reasoning || self.reasoning) + && (!set.reasoning_effort || self.reasoning_effort) && (!set.image_in || self.modalities.image_in) && (!set.image_out || self.modalities.image_out) && (!set.audio_in || self.modalities.audio_in) @@ -246,6 +248,11 @@ impl ModelProfile { native_structured_output: caps.json_schema, json_schema: caps.json_schema, reasoning: caps.reasoning, + // The offline catalog has no column for a configurable effort knob, + // and inferring one from `reasoning` would over-promise (a model can + // emit reasoning without accepting an effort level). Stay + // conservative: a caller that needs the knob must say so. + reasoning_effort: false, max_input_tokens: entry.max_input_tokens, max_output_tokens: entry.max_output_tokens, } @@ -270,6 +277,7 @@ impl ModelProfile { native_structured_output: true, json_schema: true, reasoning: true, + reasoning_effort: true, ..Self::default() } } @@ -411,6 +419,20 @@ impl ModelRequest { self } + /// Sets the provider-neutral reasoning configuration for this call. + /// + /// Adapters lower it to their provider's spelling; see [`ReasoningConfig`]. + pub fn with_reasoning(mut self, reasoning: ReasoningConfig) -> Self { + self.reasoning = Some(reasoning); + self + } + + /// Shorthand for [`with_reasoning`](Self::with_reasoning) with only an + /// effort level. + pub fn with_reasoning_effort(self, effort: ReasoningEffort) -> Self { + self.with_reasoning(ReasoningConfig::effort(effort)) + } + /// Returns the ids of cacheable segments in declaration order, describing /// the stable prompt prefix middleware should preserve. pub fn cacheable_prefix_ids(&self) -> Vec { @@ -437,6 +459,7 @@ impl ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -873,6 +896,7 @@ impl StreamAccumulator { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, }) } } diff --git a/src/harness/model/types.rs b/src/harness/model/types.rs index 2dff6f4..6a83f4e 100644 --- a/src/harness/model/types.rs +++ b/src/harness/model/types.rs @@ -75,6 +75,104 @@ pub enum ResponseFormat { }, } +/// How much effort a reasoning model should spend before answering. +/// +/// Provider-neutral by design. There was no reasoning knob at all before: the +/// only route was raw [`ModelRequest::provider_options`], which is +/// provider-shaped by definition — an OpenAI `reasoning_effort` string, an +/// OpenAI Responses `reasoning: {effort, summary}` object, and an Anthropic +/// `thinking: {type, budget_tokens}` object are three incompatible spellings of +/// one idea, and a caller that hardcodes one cannot switch providers. Worse, on +/// the OpenAI **Responses** path `provider_options` was dropped from the wire +/// body entirely, so the one format that supports `reasoning` could not receive +/// it. +/// +/// Anthropic's surface is the argument for an enum over a raw dict: it accepts +/// both a `thinking` object and a `reasoning_effort` literal, has a documented +/// precedence chain between them, and rejects `budget_tokens` outright on newer +/// models. A neutral enum lets each adapter lower to whatever its provider +/// currently accepts without every caller tracking that churn. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ReasoningEffort { + /// The smallest amount of reasoning the provider offers. + Minimal, + /// Below-default effort. + Low, + /// The provider's default effort. + #[default] + Medium, + /// Above-default effort; slower and more expensive. + High, + /// Reasoning explicitly disabled. + /// + /// Distinct from leaving [`ReasoningConfig::effort`] as `None`, which means + /// "no opinion — use the provider default". On OpenAI this additionally + /// lifts the gpt-5 restriction that pins `temperature` when reasoning is on. + None, +} + +impl ReasoningEffort { + /// The wire token OpenAI's `reasoning_effort` / `reasoning.effort` fields + /// expect. + pub fn as_str(self) -> &'static str { + match self { + ReasoningEffort::Minimal => "minimal", + ReasoningEffort::Low => "low", + ReasoningEffort::Medium => "medium", + ReasoningEffort::High => "high", + ReasoningEffort::None => "none", + } + } +} + +/// Provider-neutral reasoning/thinking configuration for one call. +/// +/// Adapters lower this to whatever their provider accepts: +/// +/// | Target | Lowered to | +/// |---|---| +/// | OpenAI Chat Completions | `reasoning_effort: ""` | +/// | OpenAI Responses | `reasoning: { effort, summary }` | +/// | Anthropic-shaped gateways | `thinking: { type, budget_tokens }` | +/// +/// [`ModelRequest::provider_options`] remains the escape hatch and **wins on +/// key conflicts**, so a caller who needs a shape this struct cannot express is +/// never blocked by it. +/// +/// Require it of a resolved model with +/// [`CapabilitySet::reasoning_effort`]. +#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] +pub struct ReasoningConfig { + /// How hard to think. `None` leaves the provider default alone. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub effort: Option, + /// An explicit thinking-token budget, for providers that take one + /// (Anthropic's `thinking.budget_tokens`). Ignored by providers that only + /// accept a coarse effort level. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub budget_tokens: Option, + /// Requested reasoning-summary verbosity (`"auto"`, `"concise"`, + /// `"detailed"`), for the OpenAI Responses `reasoning.summary` field. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub summary: Option, +} + +impl ReasoningConfig { + /// A config requesting the given effort and nothing else. + pub fn effort(effort: ReasoningEffort) -> Self { + Self { + effort: Some(effort), + ..Self::default() + } + } + + /// Whether this config asks for anything at all. + pub fn is_empty(&self) -> bool { + self.effort.is_none() && self.budget_tokens.is_none() && self.summary.is_none() + } +} + /// Lifecycle status of a model, used by [`ModelProfile`]. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -176,6 +274,15 @@ pub struct ModelProfile { /// Emits reasoning/thinking output. #[serde(default)] pub reasoning: bool, + /// Accepts a **configurable** reasoning effort + /// ([`ReasoningConfig`]), not merely emitting reasoning output. + /// + /// Separate from [`Self::reasoning`] because the two really do come apart: + /// a distilled deepseek-r1 on Ollama emits `` blocks (so `reasoning` + /// is true) while accepting no effort knob at all, and a caller that needs + /// to *dial* reasoning must be able to require the knob. + #[serde(default)] + pub reasoning_effort: bool, /// Maximum input (context) tokens, when known. #[serde(default, skip_serializing_if = "Option::is_none")] pub max_input_tokens: Option, @@ -214,6 +321,9 @@ pub struct CapabilitySet { /// Requires reasoning output. #[serde(default)] pub reasoning: bool, + /// Requires a configurable reasoning effort ([`ReasoningConfig`]). + #[serde(default)] + pub reasoning_effort: bool, /// Requires image input (vision). #[serde(default)] pub image_in: bool, @@ -401,8 +511,17 @@ pub struct ModelRequest { #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_policy: Option, /// Optional provider continuation/response id for stateful follow-ups. + /// + /// Read by the OpenAI Responses adapter, which sends it as + /// `previous_response_id`. It previously had **no reader anywhere in the + /// crate** — only a builder — so a caller could set it and nothing would + /// happen. #[serde(default, skip_serializing_if = "Option::is_none")] pub continuation_id: Option, + /// Provider-neutral reasoning/thinking configuration. See + /// [`ReasoningConfig`]. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning: Option, } /// A provider-neutral chat model response. @@ -450,6 +569,25 @@ pub struct ModelResponse { /// this needs no cap of its own. #[serde(default, skip_serializing_if = "Option::is_none")] pub continue_turn: Option, + /// `true` when this response was served from a local + /// [`ResponseCache`][crate::harness::cache::ResponseCache] rather than + /// produced by a provider call. + /// + /// # Why accounting needs this + /// + /// A cached response retains the `usage` the provider originally reported. + /// Replaying it verbatim re-bills those tokens on every hit: the run's + /// usage totals inflate and a cost-budget middleware prices spend that + /// never happened, which can abort a run over money nobody paid. LangChain + /// zeroes usage on the cache-hit path for exactly this reason. This flag + /// lets the accounting sites tell a replay from a real call without + /// destroying the `usage` a caller may legitimately want to inspect. + /// + /// `#[serde(default)]` plus a skip-when-false so entries written before the + /// flag existed still deserialize, and so a serialized response is + /// byte-identical to the historical shape when it is a real call. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub served_from_cache: bool, } /// An incremental streamed chunk of a model response. @@ -492,6 +630,21 @@ pub struct ProviderError { /// Whether retrying the same request may succeed. #[serde(default)] pub retryable: bool, + /// Server-supplied wait before retrying, in milliseconds, parsed from the + /// HTTP `Retry-After` response header. + /// + /// A `429`/`503` that names how long the client must wait is authoritative: + /// retrying sooner burns an attempt for certain. + /// [`retry_after_hint`][crate::harness::retry::retry_after_hint] reads this + /// field **first**, falling back to parsing the error message text only + /// when it is `None` — until this field existed, a provider that sent the + /// header but did not echo it into the JSON body was simply not honored. + /// + /// Both header forms are normalised here: delta-seconds (`Retry-After: 30`) + /// and the HTTP-date form (`Retry-After: Wed, 21 Oct 2015 07:28:00 GMT`), + /// the latter converted to a delay relative to receipt. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub retry_after_ms: Option, /// Raw provider payload, when available. #[serde(default, skip_serializing_if = "Option::is_none")] pub raw: Option, @@ -560,6 +713,39 @@ pub trait ChatModel: Send + Sync { None } + /// Returns a stable string identifying *which* model, at *which* endpoint, + /// under *which* credential this handle talks to — folded into the response + /// cache key by + /// [`scoped_cache_key`][crate::harness::cache::scoped_cache_key]. + /// + /// # Why the request is not enough + /// + /// [`ModelRequest::model`] is an optional *hint* that the agent loop does + /// not even set on the requests it builds; the real model is chosen by + /// [`ModelRegistry::resolve_request`] afterwards, and the endpoint and + /// credential live inside this trait object and never appear in the request + /// at all. A cache keyed on the request alone therefore has no provider or + /// model identity in it, and one shared cache serves a hosted harness's + /// answer to a local one. LangChain avoids this by looking up on + /// `(prompt, llm_string)`, where `llm_string` serializes the whole model + /// object. + /// + /// # Contract + /// + /// * Return `None` (the default) to decline; the key then folds a fixed + /// `anonymous-model` marker, which still separates identifying models + /// from each other but cannot separate two anonymous ones. + /// * The value must be **stable across process restarts** — a cache that is + /// only valid within one process lifetime is not a cache. + /// * The value must **never contain a raw credential**. It is folded into + /// keys that end up in logs, events, and durable cache files. Use + /// [`credential_fingerprint`][crate::harness::cache::credential_fingerprint] + /// (or the whole-identity helper + /// [`model_cache_identity`][crate::harness::cache::model_cache_identity]). + fn cache_identity(&self) -> Option { + None + } + /// Invokes the model and returns a complete response. async fn invoke(&self, state: &State, request: ModelRequest) -> Result; diff --git a/src/harness/no_progress/mod.rs b/src/harness/no_progress/mod.rs index 8e98340..84234ea 100644 --- a/src/harness/no_progress/mod.rs +++ b/src/harness/no_progress/mod.rs @@ -18,6 +18,47 @@ //! a middleware) feeds each tool outcome in via [`NoProgressTracker::record`] //! and turns the returned [`NoProgress`] verdict into a steering nudge //! (`Nudge`) or a halt (`Halt`). +//! +//! # Driving this from an `after_tool` hook (wave 2) +//! +//! The ladder is complete and tested but **nothing in the crate drives it**, so +//! a model looping on the same failing tool call is bounded only by +//! `RunLimits::max_tool_calls` (50) — roughly 48 wasted round trips. The +//! middleware that closes that gap lives in `middleware/library/`; here is the +//! exact contract it must implement. +//! +//! Hold one [`NoProgressTracker`] per turn behind a shared reference (`record` +//! takes `&self` and is interior-mutable, so an `&self` hook needs no +//! `RefCell`). In `after_tool`: +//! +//! 1. Fingerprint the call arguments with [`fingerprint_arguments`]. **Do not +//! roll your own** — the identical-repeat rung compares fingerprints, so two +//! drivers disagreeing on canonicalisation would silently change when the +//! ladder trips. +//! 2. Build the attempt: +//! - success → [`ToolAttempt::success`] +//! - failure → [`ToolAttempt::failure`], then +//! - `.hard_reject()` when the failure is a security/approval denial (a +//! blocked call re-issued unchanged can never succeed), +//! - `.recoverable_miss()` for the unknown-tool recovery sentinel, i.e. +//! the case that raises [`crate::harness::events::AgentEvent::UnknownToolCall`]. +//! 3. Pass `step` = the run's current model-call count +//! (`LimitTracker::model_calls()`). It is used only for the "no progress +//! since step X" wording, so an approximation is harmless — but it must be +//! monotonic or the message misleads. +//! 4. Route the verdict: +//! - [`NoProgress::Continue`] → do nothing. +//! - [`NoProgress::Nudge`] → append the message as a **system** message to +//! the working transcript so the next model call sees it, and continue. +//! Injecting it as a tool result instead would attribute the harness's +//! instruction to the tool. +//! - [`NoProgress::Halt`] → stop the turn and surface the message as the +//! final response. The tracker has already reset itself, so a resumed run +//! does not immediately re-trip on latched state. +//! +//! [`NoProgress::message`], [`NoProgress::is_nudge`], [`NoProgress::is_halt`] +//! and [`NoProgress::as_str`] exist so step 4 needs no enum match, and +//! `as_str()` gives a stable telemetry label. mod successful_repeat; mod types; @@ -48,6 +89,77 @@ const NO_PROGRESS_NUDGE_THRESHOLD: usize = 4; /// call re-issued unchanged can never succeed. const HARD_REJECT_HALT_THRESHOLD: usize = 2; +/// Computes the stable argument fingerprint the identical-repeat rung compares +/// on. +/// +/// The ladder's central question is "did the model re-issue the *same* call?", +/// which is only answerable if every driver canonicalises arguments the same +/// way. JSON object key order is not significant but `serde_json::Value`'s +/// default `to_string` preserves insertion order, so two logically identical +/// argument objects can render differently — enough to make a genuine repeat +/// look novel and let the loop run to the tool-call cap instead. +/// +/// This sorts object keys recursively, then hashes, so the result is: +/// +/// - **order-independent** for objects, +/// - **order-sensitive** for arrays (list order *is* semantic), +/// - short and allocation-cheap to carry in a [`ToolAttempt`]. +/// +/// # Example +/// +/// ``` +/// use tinyagents::harness::no_progress::fingerprint_arguments; +/// use serde_json::json; +/// +/// let a = fingerprint_arguments(&json!({"path": "/tmp", "depth": 2})); +/// let b = fingerprint_arguments(&json!({"depth": 2, "path": "/tmp"})); +/// assert_eq!(a, b, "key order must not change the fingerprint"); +/// +/// let c = fingerprint_arguments(&json!({"path": "/var", "depth": 2})); +/// assert_ne!(a, c); +/// ``` +pub fn fingerprint_arguments(arguments: &serde_json::Value) -> String { + use sha2::{Digest, Sha256}; + + let mut hasher = Sha256::new(); + hash_canonical(arguments, &mut hasher); + // 16 hex chars (64 bits) is far more than enough to separate the handful of + // distinct calls within one turn, and keeps the signature string short. + hasher + .finalize() + .iter() + .take(8) + .map(|byte| format!("{byte:02x}")) + .collect() +} + +/// Feeds `value` into `hasher` with object keys visited in sorted order. +fn hash_canonical(value: &serde_json::Value, hasher: &mut impl sha2::Digest) { + match value { + serde_json::Value::Object(map) => { + hasher.update(b"{"); + let mut keys: Vec<&String> = map.keys().collect(); + keys.sort_unstable(); + for key in keys { + hasher.update(key.as_bytes()); + hasher.update(b":"); + hash_canonical(&map[key], hasher); + hasher.update(b","); + } + hasher.update(b"}"); + } + serde_json::Value::Array(items) => { + hasher.update(b"["); + for item in items { + hash_canonical(item, hasher); + hasher.update(b","); + } + hasher.update(b"]"); + } + other => hasher.update(other.to_string().as_bytes()), + } +} + impl NoProgressTracker { /// Build a tracker whose identical-repeat halt threshold is /// `identical_halt_threshold`, clamped up so it always sits above the nudge diff --git a/src/harness/no_progress/test.rs b/src/harness/no_progress/test.rs index 15ec997..8c8d023 100644 --- a/src/harness/no_progress/test.rs +++ b/src/harness/no_progress/test.rs @@ -263,3 +263,141 @@ fn zero_thresholds_are_fail_safe() { SuccessfulRepeat::Halt(_) )); } + +// --------------------------------------------------------------------------- +// C5: the API a wave-2 `after_tool` driver actually needs +// --------------------------------------------------------------------------- + +mod drivable { + use serde_json::json; + + use crate::harness::no_progress::{ + DEFAULT_IDENTICAL_HALT_THRESHOLD, NoProgress, NoProgressTracker, ToolAttempt, + fingerprint_arguments, + }; + + #[test] + fn argument_fingerprints_ignore_object_key_order() { + // The identical-repeat rung compares fingerprints, so a driver that + // rendered arguments with `Value::to_string` would see two logically + // identical calls as different (insertion order is preserved) and let + // the loop run to the tool-call cap instead of tripping the ladder. + let a = fingerprint_arguments(&json!({"path": "/tmp", "depth": 2, "glob": "*.rs"})); + let b = fingerprint_arguments(&json!({"glob": "*.rs", "depth": 2, "path": "/tmp"})); + assert_eq!(a, b); + + // Nested objects too. + let n1 = fingerprint_arguments(&json!({"o": {"x": 1, "y": 2}})); + let n2 = fingerprint_arguments(&json!({"o": {"y": 2, "x": 1}})); + assert_eq!(n1, n2); + } + + #[test] + fn argument_fingerprints_distinguish_real_differences() { + let base = fingerprint_arguments(&json!({"path": "/tmp"})); + assert_ne!(base, fingerprint_arguments(&json!({"path": "/var"}))); + assert_ne!(base, fingerprint_arguments(&json!({"pathx": "/tmp"}))); + // Array order is semantic and must NOT be canonicalised away. + assert_ne!( + fingerprint_arguments(&json!({"a": [1, 2]})), + fingerprint_arguments(&json!({"a": [2, 1]})) + ); + // Type differences matter. + assert_ne!( + fingerprint_arguments(&json!({"a": 1})), + fingerprint_arguments(&json!({"a": "1"})) + ); + } + + #[test] + fn fingerprints_are_short_and_stable_across_calls() { + let value = json!({"path": "/tmp"}); + let first = fingerprint_arguments(&value); + assert_eq!(first.len(), 16); + assert_eq!(first, fingerprint_arguments(&value)); + } + + #[test] + fn tool_attempt_constructors_cover_every_driver_case() { + let fp = fingerprint_arguments(&json!({"q": "x"})); + + let ok = ToolAttempt::success("search", &fp); + assert!(ok.error.is_none() && !ok.hard_reject && !ok.recoverable_miss); + + let bad = ToolAttempt::failure("search", &fp, "boom"); + assert_eq!(bad.error, Some("boom")); + + let blocked = ToolAttempt::failure("shell", &fp, "denied").hard_reject(); + assert!(blocked.hard_reject); + + let missing = ToolAttempt::failure("nope", &fp, "unknown tool").recoverable_miss(); + assert!(missing.recoverable_miss); + } + + #[test] + fn a_driver_can_run_the_whole_ladder_through_the_public_api_only() { + // The end-to-end shape of the middleware wave 2 has to write: a shared + // `&self` tracker, a fingerprint, a constructor, and verdict accessors — + // no field literals, no hand-rolled hashing, no enum matching. + let tracker = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + let args = json!({"path": "/missing"}); + let fp = fingerprint_arguments(&args); + + let verdicts: Vec = (1..=3) + .map(|step| tracker.record(step, &ToolAttempt::failure("read", &fp, "ENOENT"))) + .collect(); + + assert!(matches!(verdicts[0], NoProgress::Continue)); + assert!(verdicts[1].is_nudge(), "expected a nudge on the repeat"); + assert!( + verdicts[2].is_halt(), + "expected a halt once retries ran out" + ); + + // The accessors give a driver everything it needs without a match. + assert!(verdicts[0].message().is_none()); + assert!(verdicts[1].message().unwrap().contains("no progress since")); + assert!(verdicts[2].message().unwrap().contains("Stopping")); + assert_eq!( + verdicts.iter().map(NoProgress::as_str).collect::>(), + vec!["continue", "nudge", "halt"] + ); + } + + #[test] + fn record_takes_a_shared_reference_so_an_after_tool_hook_needs_no_refcell() { + // `after_tool` hooks receive `&self`; the tracker must be usable through + // a shared reference or every driver would need interior mutability of + // its own. + fn drive(tracker: &NoProgressTracker, fp: &str) -> NoProgress { + tracker.record(1, &ToolAttempt::failure("t", fp, "boom")) + } + let tracker = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + let fp = fingerprint_arguments(&json!({})); + assert_eq!(drive(&tracker, &fp), NoProgress::Continue); + assert!(drive(&tracker, &fp).is_nudge()); + } + + #[test] + fn a_success_between_failures_clears_the_ladder() { + // The driver contract says a success resets progress tracking; pinned + // here because a wave-2 middleware relies on it to avoid halting a run + // that is genuinely making progress with an occasional failure. + let tracker = NoProgressTracker::new(DEFAULT_IDENTICAL_HALT_THRESHOLD); + let fp = fingerprint_arguments(&json!({"a": 1})); + + assert_eq!( + tracker.record(1, &ToolAttempt::failure("t", &fp, "boom")), + NoProgress::Continue + ); + assert_eq!( + tracker.record(2, &ToolAttempt::success("t", &fp)), + NoProgress::Continue + ); + // The repeat counter restarted, so this is a first failure again. + assert_eq!( + tracker.record(3, &ToolAttempt::failure("t", &fp, "boom")), + NoProgress::Continue + ); + } +} diff --git a/src/harness/no_progress/types.rs b/src/harness/no_progress/types.rs index 87fd521..b4a0896 100644 --- a/src/harness/no_progress/types.rs +++ b/src/harness/no_progress/types.rs @@ -6,6 +6,13 @@ use std::sync::Mutex; +/// One recorded tool outcome, as the driver observed it. +/// +/// Built with [`ToolAttempt::success`] / [`ToolAttempt::failure`] plus the +/// [`ToolAttempt::hard_reject`] / [`ToolAttempt::recoverable_miss`] modifiers, +/// so a driver never has to remember the field set. Borrows rather than owns so +/// an `after_tool` hook can record without allocating anything but the argument +/// fingerprint. pub struct ToolAttempt<'a> { /// Tool name. pub tool: &'a str, @@ -25,6 +32,49 @@ pub struct ToolAttempt<'a> { pub recoverable_miss: bool, } +impl<'a> ToolAttempt<'a> { + /// A tool call that succeeded. Clears every ladder counter when recorded. + /// + /// `arg_fingerprint` should come from + /// [`fingerprint_arguments`][crate::harness::no_progress::fingerprint_arguments] + /// so every driver computes it the same way. + pub fn success(tool: &'a str, arg_fingerprint: &'a str) -> Self { + Self { + tool, + arg_fingerprint, + error: None, + hard_reject: false, + recoverable_miss: false, + } + } + + /// A tool call that failed, with the error text the model saw. + pub fn failure(tool: &'a str, arg_fingerprint: &'a str, error: &'a str) -> Self { + Self { + tool, + arg_fingerprint, + error: Some(error), + hard_reject: false, + recoverable_miss: false, + } + } + + /// Marks the failure as a hard security/approval rejection, which can never + /// succeed re-issued unchanged and so trips the ladder fastest. + pub fn hard_reject(mut self) -> Self { + self.hard_reject = true; + self + } + + /// Marks the failure as the unknown-tool recovery sentinel: correctable + /// feedback the model already received, which must not feed the generic + /// any-failure backstop. + pub fn recoverable_miss(mut self) -> Self { + self.recoverable_miss = true; + self + } +} + /// The ladder's verdict for one recorded attempt. #[derive(Debug, Clone, PartialEq, Eq)] pub enum NoProgress { @@ -39,6 +89,40 @@ pub enum NoProgress { Halt(String), } +impl NoProgress { + /// The corrective/summary text carried by a [`NoProgress::Nudge`] or + /// [`NoProgress::Halt`]; `None` for [`NoProgress::Continue`]. + /// + /// Saves a driver from matching the enum just to reach the string it has to + /// forward either way. + pub fn message(&self) -> Option<&str> { + match self { + NoProgress::Continue => None, + NoProgress::Nudge(message) | NoProgress::Halt(message) => Some(message), + } + } + + /// `true` when the loop should keep running but feed the corrective back to + /// the model. + pub fn is_nudge(&self) -> bool { + matches!(self, NoProgress::Nudge(_)) + } + + /// `true` when the loop must stop. + pub fn is_halt(&self) -> bool { + matches!(self, NoProgress::Halt(_)) + } + + /// Stable, snake_case label for logs and telemetry dimensions. + pub fn as_str(&self) -> &'static str { + match self { + NoProgress::Continue => "continue", + NoProgress::Nudge(_) => "nudge", + NoProgress::Halt(_) => "halt", + } + } +} + #[derive(Default)] pub(super) struct LadderState { /// Signature of the previous failing call (tool + args + first error line). diff --git a/src/harness/providers/mock.rs b/src/harness/providers/mock.rs index 1c4ff7f..737674f 100644 --- a/src/harness/providers/mock.rs +++ b/src/harness/providers/mock.rs @@ -228,6 +228,7 @@ impl ChatModel for MockModel { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -329,6 +330,7 @@ impl MockModel { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } } diff --git a/src/harness/providers/openai/convert.rs b/src/harness/providers/openai/convert.rs index 5c25ed0..bcb5a7f 100644 --- a/src/harness/providers/openai/convert.rs +++ b/src/harness/providers/openai/convert.rs @@ -7,6 +7,88 @@ use super::*; +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Process-global counter handing every decoded response a distinct **epoch**. +/// +/// Synthetic tool-call ids used to be `tool-{slot}`, keyed only to the call's +/// position in its own response. Any runtime that omits `id` (several Ollama +/// builds do) therefore emitted `tool-0` on *every* assistant turn, so one run's +/// transcript contained several distinct calls all declaring the same id — an +/// unresolvable pairing for the agent loop. Prefixing with a monotonic epoch +/// makes the id unique for the life of the process while staying stable within +/// the response that minted it (the epoch is drawn once, at the top of decoding, +/// and reused for every slot and every streamed delta of that response). +static SYNTHETIC_ID_EPOCH: AtomicU64 = AtomicU64::new(0); + +/// Draws the next synthetic-id epoch. Call **once** per decoded response (unary +/// parse, or accumulator construction on the streaming path) and thread the +/// value through every `tool_call_id` / `tool_call_from_wire` call for that +/// response. +pub(super) fn next_synthetic_id_epoch() -> u64 { + SYNTHETIC_ID_EPOCH.fetch_add(1, Ordering::Relaxed) +} + +/// Prefix for ids this crate synthesizes on the **native/provider** boundary. +/// +/// Deliberately distinct from the prompt-guided text protocol's `call_{index}` +/// ids (`crate::harness::tool::prompt`) so a transcript that mixes both — a +/// model that degraded from native to prompt-guided mid-run — can never produce +/// two different calls carrying the same id. +const SYNTHETIC_ID_PREFIX: &str = "tacall"; + +/// Characters a tool-call id may contain before the provider boundary rewrites +/// it. Mirrors LangChain's `_TOOL_CALL_ID_PATTERN` (`^[a-zA-Z0-9_-]+$`). +fn is_conforming_tool_call_id(id: &str) -> bool { + !id.is_empty() + && id + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-') +} + +/// Rewrites a non-conforming provider tool-call id into the conforming +/// alphabet, **deterministically**. +/// +/// Some gateways emit ids that other providers reject on the way back +/// (`functions.write_todos:0` is the canonical example). Rewriting at the +/// provider boundary keeps the id and its paired tool result consistent, because +/// both are derived from the same [`ToolCall`]. Determinism is the whole point: +/// the same wire id must always map to the same rewritten id, or a replayed +/// transcript would stop pairing. +/// +/// Offending bytes become `_`, and a short hash of the original is appended so +/// two distinct ids that sanitize to the same string stay distinguishable. +fn normalize_tool_call_id(id: &str) -> String { + if is_conforming_tool_call_id(id) { + return id.to_string(); + } + let sanitized: String = id + .chars() + .map(|c| { + if c.is_ascii_alphanumeric() || c == '_' || c == '-' { + c + } else { + '_' + } + }) + .collect(); + // FNV-1a over the original bytes: tiny, dependency-free, and stable across + // processes and crate versions (unlike `DefaultHasher`, whose output is not + // guaranteed stable). + let mut hash: u64 = 0xcbf2_9ce4_8422_2325; + for byte in id.as_bytes() { + hash ^= u64::from(*byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + let normalized = format!("{sanitized}_{hash:x}"); + tracing::debug!( + target: "tinyagents::providers::openai", + normalized = %normalized, + "[openai] rewrote a non-conforming provider tool-call id" + ); + normalized +} + /// Translates one harness [`Message`] into an OpenAI wire message. /// /// User messages are rendered as OpenAI content-parts when they carry non-text @@ -168,10 +250,38 @@ pub(super) fn degraded_json_object_format() -> Value { }) } +/// The schema projection applied to a `response_format` JSON Schema. +/// +/// OpenAI's strict mode is not "the same schema, validated harder": it rejects +/// any object that does not carry `additionalProperties: false` and list *every* +/// declared property in `required`. `strict: true` used to be hardcoded here and +/// paired with the caller's **raw** schema, so a perfectly valid JSON Schema +/// 400d — including this crate's own documented example. The sibling +/// `degraded_json_object_format` had it right with `strict: false`; the +/// constraint was understood in one place and not the other. +/// +/// Delegates to [`prepare_parameters`], the shared sanitizer, so the +/// response-format schema and the tool-parameter schemas are projected by +/// exactly one implementation. +fn prepare_response_schema(schema: &Value, strict: bool) -> Value { + let preparation = if strict { + crate::harness::tool::SchemaPreparation::openai().with_strict() + } else { + crate::harness::tool::SchemaPreparation::openai() + }; + crate::harness::tool::prepare_parameters(schema, &preparation) +} + /// Translates a [`ResponseFormat`] into the OpenAI `response_format` JSON value. /// /// Returns `None` for [`ResponseFormat::Text`] so the field is omitted entirely. -pub(super) fn translate_response_format(format: &ResponseFormat) -> Option { +/// +/// `strict` selects OpenAI strict structured output for the schema forms. It is +/// **not** hardcoded: hosted OpenAI defaults it on, local runtimes default it +/// off (they reject the key outright, and their schema support is looser), and +/// a 400 implicating the schema degrades it for a single retry. When `strict` is +/// on the schema first goes through [`prepare_strict_schema`]. +pub(super) fn translate_response_format(format: &ResponseFormat, strict: bool) -> Option { match format { ResponseFormat::Text => None, ResponseFormat::JsonObject => Some(json!({ "type": "json_object" })), @@ -179,12 +289,13 @@ pub(super) fn translate_response_format(format: &ResponseFormat) -> Option { + let schema = prepare_response_schema(schema, strict); Some(json!({ "type": "json_schema", "json_schema": { "name": name, "schema": schema, - "strict": true, + "strict": strict, } })) } @@ -207,7 +318,7 @@ pub(super) fn translate_response_format(format: &ResponseFormat) -> Option Result { - parse_chat_response(value, None) + parse_chat_response(value, None, CacheTokenAccounting::default()) } /// Like [`parse_response`], but also normalizes reasoning into a leading @@ -219,8 +330,12 @@ pub(super) fn parse_response(value: Value) -> Result { pub(super) fn parse_chat_response( value: Value, reasoning_tags: Option<&ReasoningTagExtraction>, + accounting: CacheTokenAccounting, ) -> Result { let parsed: ChatCompletionResponse = serde_json::from_value(value.clone())?; + // One epoch for the whole response, so every synthesized id in it shares a + // prefix and no later response can reuse it. + let epoch = next_synthetic_id_epoch(); let choice = parsed.choices.into_iter().next().ok_or_else(|| { TinyAgentsError::Model("openai response contained no choices".to_string()) @@ -275,11 +390,12 @@ pub(super) fn parse_chat_response( .enumerate() .map(|(index, call)| { // Local servers routinely omit `id`; synthesize the same - // `tool-{index}` fallback the streaming path uses so the agent loop + // run-unique fallback the streaming path uses so the agent loop // can still correlate the tool result back to this call. An empty // id is treated as absent. tool_call_from_wire( "openai response", + epoch, index, &call.id, &call.function.name, @@ -288,7 +404,9 @@ pub(super) fn parse_chat_response( }) .collect::>(); - let usage = parsed.usage.map(convert_usage); + let usage = parsed + .usage + .map(|wire| convert_usage_with(wire, accounting)); let message = AssistantMessage { id: parsed.id, @@ -304,25 +422,33 @@ pub(super) fn parse_chat_response( raw: Some(value), resolved_model: None, continue_turn: None, + served_from_cache: false, }) } -/// Returns the effective call id for a streamed tool-call slot: the -/// provider-assigned id when present, or a stable `tool-{slot}` fallback keyed to -/// the slot's position so delta ids and the final call id always agree. -pub(super) fn tool_call_id(slot: usize, id: &str) -> String { +/// Returns the effective call id for a tool-call slot. +/// +/// * A provider-assigned id is kept, after [`normalize_tool_call_id`] rewrites +/// any character the conforming alphabet (`[A-Za-z0-9_-]`) rejects. +/// * An absent id is synthesized as `tacall-{epoch}-{slot}`. `epoch` comes from +/// [`next_synthetic_id_epoch`] and is drawn **once per decoded response**, so +/// the id is stable across the streamed deltas and the terminal response of +/// one call while never repeating on a later turn — the defect the old +/// `tool-{slot}` form had, which made every id-less Ollama turn emit `tool-0`. +pub(super) fn tool_call_id(epoch: u64, slot: usize, id: &str) -> String { if id.is_empty() { - format!("tool-{slot}") + format!("{SYNTHETIC_ID_PREFIX}-{epoch}-{slot}") } else { - id.to_string() + normalize_tool_call_id(id) } } /// Builds a provider-neutral [`ToolCall`] from the wire fields, tolerating the /// defects small local models produce. /// -/// `slot` is the tool call's position in the response (used to synthesize a -/// stable `tool-{slot}` id when the provider omits one — Ollama did so until +/// `epoch` + `slot` identify the call: `slot` is its position in the response +/// and `epoch` the response's [`next_synthetic_id_epoch`] draw, which together +/// synthesize a run-unique id when the provider omits one (Ollama did so until /// v0.12.11). When the arguments cannot be parsed even after repair, the call is /// marked [`ToolCall::invalid`] with the raw arguments preserved rather than /// failing the whole model call: the agent loop feeds the error back to the @@ -331,12 +457,13 @@ pub(super) fn tool_call_id(slot: usize, id: &str) -> String { /// become a never-resolving tool call that stalls the loop. pub(super) fn tool_call_from_wire( context: &str, + epoch: u64, slot: usize, id: &str, name: &str, raw: &str, ) -> ToolCall { - let call_id = tool_call_id(slot, id); + let call_id = tool_call_id(epoch, slot, id); match parse_tool_arguments(raw) { Ok(arguments) => ToolCall { id: call_id, @@ -467,30 +594,77 @@ fn strip_tool_call_markers(raw: &str) -> Option { (!trimmed.is_empty()).then(|| trimmed.to_string()) } -/// Converts an OpenAI [`UsageWire`] into the harness-neutral [`Usage`]. -pub(super) fn convert_usage(wire: UsageWire) -> Usage { +/// Whether a provider's reported input-token count already contains the tokens +/// it served from (or wrote into) its prompt cache. +/// +/// This is a real cross-provider divergence, not a detail: +/// +/// * **OpenAI** reports `prompt_tokens` as the *total* input, with +/// `prompt_tokens_details.cached_tokens` a breakdown of it. Subtracting cache +/// reads to price the uncached remainder is correct. +/// * **Anthropic** reports `input_tokens` *excluding* cache reads and cache +/// writes — the true input total is `input + cache_read + cache_creation`. +/// +/// An OpenAI-compatible gateway fronting an Anthropic model can pass either +/// convention through, and guessing wrong silently under-bills (OpenAI +/// semantics assumed over Anthropic data) or double-counts. Select the right one +/// with [`OpenAiModel::with_cache_token_accounting`][cta]. +/// +/// [cta]: super::OpenAiModel::with_cache_token_accounting +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum CacheTokenAccounting { + /// OpenAI semantics: the reported input count already contains cache + /// read/creation tokens. The default. + #[default] + IncludedInInput, + /// Anthropic semantics: cache read/creation tokens are reported *outside* + /// the input count, so the true input total is recomputed as + /// `input + cache_read + cache_creation`. + ExcludedFromInput, +} + +/// Converts an OpenAI [`UsageWire`] into [`Usage`] under an explicit +/// [`CacheTokenAccounting`] convention. +/// +/// Maps both cache directions, unlike the original which only read +/// `cached_tokens`: `cache_creation_tokens` is a first-class, summed and priced +/// field on [`Usage`] that no provider ever populated, so cache **writes** were +/// invisible to every cost report in the crate. +pub(super) fn convert_usage_with(wire: UsageWire, accounting: CacheTokenAccounting) -> Usage { + let prompt_details = wire.prompt_tokens_details.unwrap_or_default(); + let cache_read_tokens = prompt_details.cached_tokens; + let cache_creation_tokens = prompt_details.cache_creation_tokens(); + + let input_tokens = match accounting { + CacheTokenAccounting::IncludedInInput => wire.prompt_tokens, + CacheTokenAccounting::ExcludedFromInput => wire + .prompt_tokens + .saturating_add(cache_read_tokens) + .saturating_add(cache_creation_tokens), + }; + // OpenAI-compatible endpoints sometimes omit `total_tokens` entirely // (deserializes to `0` via `#[serde(default)]`); fall back to // `prompt + completion` so `total_tokens` is never a misleading zero for - // a call that clearly consumed tokens. - let total_tokens = if wire.total_tokens > 0 { - wire.total_tokens - } else { - wire.prompt_tokens + wire.completion_tokens - }; + // a call that clearly consumed tokens. Under Anthropic semantics the + // recomputed `input_tokens` is what the total must be built from. + let total_tokens = + if wire.total_tokens > 0 && accounting == CacheTokenAccounting::IncludedInInput { + wire.total_tokens + } else { + input_tokens + wire.completion_tokens + }; + Usage { - input_tokens: wire.prompt_tokens, + input_tokens, output_tokens: wire.completion_tokens, total_tokens, - cache_read_tokens: wire - .prompt_tokens_details - .map(|d| d.cached_tokens) - .unwrap_or(0), + cache_read_tokens, + cache_creation_tokens, reasoning_tokens: wire .completion_tokens_details .map(|d| d.reasoning_tokens) .unwrap_or(0), - ..Usage::default() } } diff --git a/src/harness/providers/openai/local.rs b/src/harness/providers/openai/local.rs new file mode 100644 index 0000000..b8b1888 --- /dev/null +++ b/src/harness/providers/openai/local.rs @@ -0,0 +1,454 @@ +//! Local OpenAI-compatible runtimes: identification, capability **probing**, +//! and the native-API escape hatches the OpenAI wire format cannot express. +//! +//! # Why this module exists +//! +//! A local runtime is not "hosted OpenAI at a different URL". It differs in +//! three ways the transport used to paper over with hard-coded guesses: +//! +//! 1. **Its context window is tiny and not derivable from the model id.** +//! `derive_profile` filled `max_input_tokens` from the generic hint table, +//! which matches bare substrings — so `llama3.2:3b` on Ollama claimed +//! 128 000 tokens while Ollama's real default `num_ctx` is **2048**, roughly +//! a 60× overstatement. Compaction fires at `window * threshold`, so it never +//! fired and the server silently truncated the front of the prompt. +//! LangChain refuses to guess here (ChatOllama ships no profile at all and +//! its summarization middleware hard-fails asking for absolute counts), and +//! an invented window is strictly worse than the `None` this crate already +//! supports. See [`LocalProbe::max_input_tokens`]. +//! 2. **Whether it accepts native `tools` is a property of the loaded model, +//! not of "being local".** The transport hard-disabled native tools for every +//! local runtime unconditionally, which forced the prompt-guided branch — +//! injecting the protocol block *plus* every tool's JSON Schema into the +//! system prompt, against that real 2048-token window, which then truncated +//! from the front and dropped the very prompt carrying the protocol. +//! Ollama reports this directly in `/api/show`'s `capabilities` array. +//! 3. **Some knobs have no OpenAI-wire spelling at all.** `num_ctx` and +//! `keep_alive` are `/api/chat` fields; `POST /v1/chat/completions` drops +//! them on the floor. See [`LocalRuntimeKind::native_root`]. +//! +//! Probing is **opt-in** and never runs during construction: it costs a network +//! round trip, and a constructor that blocks on one is unusable in the contexts +//! this crate is embedded in. + +use std::time::Duration; + +use serde::Deserialize; +use serde_json::{Value, json}; + +use crate::error::TinyAgentsError; + +/// A local OpenAI-compatible model server. +/// +/// The single place that answers "is this endpoint a local runtime, and which +/// one?". Adding a runtime is one variant plus its arms here — not a condition +/// to keep in sync across the transport. +/// +/// Before this existed only Ollama and LM Studio were recognised; +/// llama.cpp-server and vLLM fell through to the hosted `Compatible` path and +/// got Bearer auth, `tool_calling: true`, `image_in: true`, no `/v1` +/// normalisation, and none of the request-shape degrade knobs. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum LocalRuntimeKind { + /// Ollama. Serves an OpenAI-compatible surface under `/v1` **and** its own + /// native API under `/api` — the only place `num_ctx` and `keep_alive` are + /// readable. + Ollama, + /// LM Studio. OpenAI-compatible under `/v1`, with a richer model listing + /// under `/api/v0/models` (context length, load state, quantisation). + LmStudio, + /// `llama-server` from llama.cpp. OpenAI-compatible only. + LlamaCpp, + /// vLLM's OpenAI-compatible server. + Vllm, +} + +impl LocalRuntimeKind { + /// Stable identifier used in provider ids, log lines, and errors. + pub fn as_str(self) -> &'static str { + match self { + LocalRuntimeKind::Ollama => "ollama", + LocalRuntimeKind::LmStudio => "lm_studio", + LocalRuntimeKind::LlamaCpp => "llama_cpp", + LocalRuntimeKind::Vllm => "vllm", + } + } + + /// The server root assumed when a spec carries a blank `base_url`. + pub fn default_root(self) -> &'static str { + match self { + LocalRuntimeKind::Ollama => "http://localhost:11434", + LocalRuntimeKind::LmStudio => "http://localhost:1234", + LocalRuntimeKind::LlamaCpp => "http://localhost:8080", + LocalRuntimeKind::Vllm => "http://localhost:8000", + } + } + + /// Strips the OpenAI-compatibility suffix off `base_url`, yielding the + /// server root the runtime's **native** API hangs off. + /// + /// `http://localhost:11434/v1` → `http://localhost:11434`, so `/api/show`, + /// `/api/chat` and `/api/v0/models` can be reached. Idempotent for a base + /// that already is the root. + pub fn native_root(self, base_url: &str) -> String { + base_url + .trim_end_matches('/') + .trim_end_matches("/v1") + .trim_end_matches('/') + .to_string() + } + + /// Whether this runtime speaks a native (non-OpenAI) API this crate knows + /// how to use. Only Ollama does today. + pub fn has_native_api(self) -> bool { + matches!(self, LocalRuntimeKind::Ollama) + } +} + +/// What a probe of a live local server learned about the loaded model. +/// +/// Every field is [`Option`] on purpose: a runtime that does not report a fact +/// leaves it `None` and the caller keeps whatever it already had, rather than +/// having a guess written over it. This is the shape LangChain's +/// `libs/model-profiles` keys on (`max_input_tokens` is a first-class key +/// there) — with the pointed difference that no `ollama` profile file exists in +/// that repo, which is the evidence that a static catalogue is the wrong answer +/// for local models and runtime probing is the state of the art. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct LocalProbe { + /// The model's real context window in tokens, when the server reports one. + /// + /// For Ollama this is the `*.context_length` entry of `/api/show`'s + /// `model_info` — the **architecture's** trained window. Note the runtime + /// still loads with `num_ctx` (default 2048) unless told otherwise, so a + /// caller that wants the full window must also request it; see + /// [`OpenAiModel::with_local_num_ctx`][wlnc]. + /// + /// [wlnc]: super::OpenAiModel::with_local_num_ctx + pub max_input_tokens: Option, + /// Whether the loaded model advertises native tool calling. + pub tool_calling: Option, + /// Whether the loaded model advertises image input. + pub vision: Option, + /// Whether the loaded model advertises a reasoning/thinking channel. + pub reasoning: Option, + /// The `num_ctx` the runtime says it actually loaded the model with, when + /// it reports one. This — not [`Self::max_input_tokens`] — is the number + /// that bounds a live request. + pub loaded_num_ctx: Option, +} + +impl LocalProbe { + /// Whether the probe learned anything at all. + pub fn is_empty(&self) -> bool { + *self == LocalProbe::default() + } + + /// The context window to advertise: the loaded `num_ctx` when known (it is + /// the real ceiling for a live request), else the architecture window, else + /// `None`. + /// + /// Deliberately **not** "the bigger of the two". Overstating the window is + /// the LOCAL-1 defect: compaction is gated on it, so a window larger than + /// the server will honour means compaction never fires and the server + /// truncates the prompt from the front instead — losing the system prompt + /// silently. + pub fn effective_context_window(&self) -> Option { + self.loaded_num_ctx.or(self.max_input_tokens) + } +} + +// --------------------------------------------------------------------------- +// Ollama `/api/show` +// --------------------------------------------------------------------------- + +/// The subset of Ollama's `POST /api/show` body this crate reads. +#[derive(Debug, Default, Deserialize)] +struct OllamaShowResponse { + /// Architecture metadata. Keys are namespaced by architecture + /// (`llama.context_length`, `qwen3.context_length`, …), so the reader scans + /// for a `*.context_length` suffix rather than guessing the prefix. + #[serde(default)] + model_info: serde_json::Map, + /// Capability tags: `completion`, `tools`, `vision`, `thinking`, `insert`. + #[serde(default)] + capabilities: Vec, +} + +/// Extracts the architecture context length from an Ollama `model_info` map. +/// +/// Keys are `{architecture}.context_length`, so match on the suffix. Returns +/// the smallest candidate when several match, staying conservative for the same +/// reason [`LocalProbe::effective_context_window`] does. +fn context_length_from_model_info(model_info: &serde_json::Map) -> Option { + model_info + .iter() + .filter(|(key, _)| key.ends_with(".context_length") || key.as_str() == "context_length") + .filter_map(|(_, value)| value.as_u64()) + .filter(|value| *value > 0) + .min() +} + +/// Turns an Ollama `/api/show` body into a [`LocalProbe`]. +/// +/// Pure, so the whole mapping is unit-testable without a live Ollama. +pub(super) fn probe_from_ollama_show(body: &Value) -> LocalProbe { + let parsed: OllamaShowResponse = + serde_json::from_value(body.clone()).unwrap_or_else(|_| OllamaShowResponse::default()); + let has = |tag: &str| { + parsed + .capabilities + .iter() + .any(|c| c.eq_ignore_ascii_case(tag)) + }; + // An empty `capabilities` array means "this server did not tell us", + // not "this model can do nothing" — leave those `None` so the caller keeps + // whatever it already had. + let capabilities_reported = !parsed.capabilities.is_empty(); + LocalProbe { + max_input_tokens: context_length_from_model_info(&parsed.model_info), + tool_calling: capabilities_reported.then(|| has("tools")), + vision: capabilities_reported.then(|| has("vision")), + reasoning: capabilities_reported.then(|| has("thinking")), + loaded_num_ctx: None, + } +} + +// --------------------------------------------------------------------------- +// LM Studio `/api/v0/models` +// --------------------------------------------------------------------------- + +/// One entry of LM Studio's richer `GET /api/v0/models` listing. +#[derive(Debug, Deserialize)] +struct LmStudioModel { + #[serde(default)] + id: String, + /// The model's context length. LM Studio reports the trained window here. + #[serde(default)] + max_context_length: Option, + /// The context the model is currently **loaded** with, when loaded. + #[serde(default)] + loaded_context_length: Option, + /// `llm`, `vlm` (vision), or `embeddings`. + #[serde(default)] + r#type: Option, +} + +#[derive(Debug, Default, Deserialize)] +struct LmStudioModelList { + #[serde(default)] + data: Vec, +} + +/// Turns an LM Studio `/api/v0/models` body into a [`LocalProbe`] for `model`. +/// +/// Pure, so the mapping is unit-testable without a live LM Studio. +pub(super) fn probe_from_lm_studio_models(body: &Value, model: &str) -> LocalProbe { + let parsed: LmStudioModelList = + serde_json::from_value(body.clone()).unwrap_or_else(|_| LmStudioModelList::default()); + let Some(entry) = parsed.data.iter().find(|m| m.id == model) else { + return LocalProbe::default(); + }; + LocalProbe { + max_input_tokens: entry.max_context_length.filter(|v| *v > 0), + // LM Studio does not report tool support in this listing; leave it + // untouched rather than inventing an answer. + tool_calling: None, + vision: entry + .r#type + .as_deref() + .map(|t| t.eq_ignore_ascii_case("vlm")), + reasoning: None, + loaded_num_ctx: entry.loaded_context_length.filter(|v| *v > 0), + } +} + +// --------------------------------------------------------------------------- +// Request bodies for the native escape hatches +// --------------------------------------------------------------------------- + +/// The `POST /api/show` body: which model to describe. +pub(super) fn ollama_show_body(model: &str) -> Value { + json!({ "model": model }) +} + +/// The `POST /api/chat` preflight body that loads `model` with explicit +/// `options` and residency. +/// +/// # Why a preflight rather than a request field +/// +/// `num_ctx` and `keep_alive` are **`/api/chat` fields**. The chat adapter +/// speaks `POST {base_url}/chat/completions`, and Ollama's OpenAI-compatibility +/// layer does not read them — so +/// [`with_default_provider_options`][super::OpenAiModel::with_default_provider_options] +/// documenting `{"options": {"num_ctx": 8192}}` as "the local escape hatch" +/// was, on this path, a field that went nowhere. (The crate's tests asserted +/// only that the request JSON *contained* it, never that a server honoured it, +/// which is exactly why that went unnoticed.) +/// +/// An `/api/chat` call with an empty `messages` array is Ollama's documented +/// **load** request: it loads the model with the given `options` and holds it +/// resident for `keep_alive`. Issuing it once before the first real turn gets +/// `num_ctx` where the OpenAI wire cannot, and doubles as the warm-up that +/// keeps Ollama from unloading after its 5-minute default and charging the next +/// turn a cold multi-gigabyte load inside the 600 s unary deadline. +/// +/// **Caveat, stated plainly:** this configures the *loaded runner*. Ollama +/// reuses an already-loaded runner for a subsequent `/v1` request that does not +/// demand conflicting options, which is the case here — but it is a property of +/// the server's runner reuse, not a guarantee of the OpenAI wire format. A full +/// native `/api/chat` chat adapter remains the complete fix and is called out as +/// a follow-up. +pub(super) fn ollama_load_body( + model: &str, + options: Option<&Value>, + keep_alive: Option<&str>, +) -> Value { + let mut body = json!({ "model": model, "messages": [] }); + if let Some(options) = options.filter(|o| o.is_object()) { + body["options"] = options.clone(); + } + if let Some(keep_alive) = keep_alive { + body["keep_alive"] = json!(keep_alive); + } + body +} + +/// Extracts an `options` object out of merged provider options, if present. +/// +/// The escape hatch's documented shape is `{"options": {"num_ctx": 8192}}`, so +/// this is what the preflight forwards natively. +pub(super) fn local_options_object(provider_options: &Value) -> Option<&Value> { + provider_options.get("options").filter(|v| v.is_object()) +} + +// --------------------------------------------------------------------------- +// Error classification +// --------------------------------------------------------------------------- + +/// Rewrites a local runtime's opaque 404 into a message naming the fix. +/// +/// The embeddings adapter has done this for a while — "Run `ollama pull +/// {model}` or choose an installed embedding model" — while the chat path +/// surfaced whatever the server said, typically a bare +/// `{"error":"model 'x' not found"}`. Returns `None` when the failure is not a +/// missing-model 404, so the original message survives untouched. +pub(super) fn missing_model_remediation( + kind: LocalRuntimeKind, + status: u16, + body: &str, + model: &str, + base_url: &str, +) -> Option { + if status != 404 { + return None; + } + let lower = body.to_ascii_lowercase(); + if !(lower.contains("model") + && (lower.contains("not found") || lower.contains("does not exist"))) + { + return None; + } + Some(match kind { + LocalRuntimeKind::Ollama => format!( + "Ollama model `{model}` is not installed at {base_url}. \ + Run `ollama pull {model}`, or call `list_models()` to see what is installed" + ), + LocalRuntimeKind::LmStudio => format!( + "LM Studio at {base_url} is not serving a model called `{model}`. \ + Load it in LM Studio, or call `list_models()` to see what is loaded — \ + the id is whichever GGUF the server has open, so there is no default to guess" + ), + _ => format!( + "{} at {base_url} is not serving a model called `{model}`. \ + Call `list_models()` to see what is available", + kind.as_str() + ), + }) +} + +/// Recognises "the prompt did not fit in this model's context window" from a +/// provider error. +/// +/// Hosted providers raise an explicit 400 for this; local servers usually +/// truncate the front of the prompt silently instead, which is why this must be +/// paired with a *real* context window from [`LocalProbe`] rather than relied on +/// alone. When it does fire, the classification is stable so a caller can act on +/// it (compact and retry) instead of string-matching a provider message. +/// +/// Surfaced as a [`ProviderError::code`][pe] of +/// [`CONTEXT_OVERFLOW_CODE`], because a typed +/// `TinyAgentsError::ContextOverflow` variant would have to be added in +/// `src/error.rs` — outside this module's ownership. Promoting the code to a +/// typed variant is a follow-up. +/// +/// [pe]: crate::harness::model::ProviderError::code +pub(super) fn is_context_overflow(status: u16, message: &str) -> bool { + if !matches!(status, 400 | 413 | 422 | 500) { + return false; + } + let lower = message.to_ascii_lowercase(); + const PHRASES: [&str; 7] = [ + "context length", + "context window", + "maximum context", + "too many tokens", + "reduce the length of the messages", + "prompt is too long", + "exceeds the maximum", + ]; + PHRASES.iter().any(|phrase| lower.contains(phrase)) +} + +/// The [`ProviderError::code`][pe] stamped on a recognised context overflow. +/// +/// [pe]: crate::harness::model::ProviderError::code +pub const CONTEXT_OVERFLOW_CODE: &str = "context_overflow"; + +/// Recognises "this endpoint rejects the `tools` parameter" from a 400 body. +/// +/// Drives the [`Degrade::native_tools`][d] latch, which is the auto-degrade half +/// of C11: a local server that cannot do native tools tells us so once, and +/// every later call goes straight to the prompt-guided branch — instead of the +/// old behaviour, which assumed *every* local server was in that state forever. +/// +/// [d]: super::transport::Degrade +pub(super) fn mentions_tools_unsupported(message: &str) -> bool { + let lower = message.to_ascii_lowercase(); + if !(lower.contains("tool") || lower.contains("function")) { + return false; + } + // `tool_choice` rejections are a *different* degrade with its own latch; + // matching them here would flip the wrong knob. + if lower.contains("tool_choice") && !lower.contains("tools") { + return false; + } + const PHRASES: [&str; 8] = [ + "does not support tools", + "does not support function", + "not supported", + "unsupported", + "unknown parameter", + "unrecognized", + "invalid parameter", + "no tool support", + ]; + PHRASES.iter().any(|phrase| lower.contains(phrase)) +} + +/// Deadline for a probe request. Probing is a convenience, never the point of +/// the call, so it fails fast rather than blocking a turn behind a wedged +/// server — the same failure the `list_models` deadline was added for. +pub(super) const PROBE_TIMEOUT: Duration = Duration::from_secs(10); + +/// Maps a probe transport failure onto the crate error with a grep-friendly +/// message naming the endpoint. +pub(super) fn probe_error(endpoint: &str, detail: impl std::fmt::Display) -> TinyAgentsError { + TinyAgentsError::Model(format!( + "[openai] local probe of {endpoint} failed: {detail}" + )) +} + +#[cfg(test)] +#[path = "local_test.rs"] +mod local_test; diff --git a/src/harness/providers/openai/local_test.rs b/src/harness/providers/openai/local_test.rs new file mode 100644 index 0000000..ca0ad47 --- /dev/null +++ b/src/harness/providers/openai/local_test.rs @@ -0,0 +1,237 @@ +//! Unit tests for local-runtime identification, probing, and classification. + +use super::*; +use serde_json::json; + +#[test] +fn native_root_strips_the_openai_compat_suffix() { + let kind = LocalRuntimeKind::Ollama; + assert_eq!( + kind.native_root("http://localhost:11434/v1"), + "http://localhost:11434" + ); + assert_eq!( + kind.native_root("http://localhost:11434/v1/"), + "http://localhost:11434" + ); + // Idempotent for a base that is already the root. + assert_eq!( + kind.native_root("http://localhost:11434"), + "http://localhost:11434" + ); +} + +#[test] +fn ollama_show_reports_the_real_window_not_the_model_id_guess() { + // `llama3.2:3b` matches the generic hint table's `("llama3", Substring, + // 128_000)` entry. The server says 8192. The probe must report the server. + let body = json!({ + "model_info": { + "general.architecture": "llama", + "llama.context_length": 8192, + "llama.embedding_length": 3072 + }, + "capabilities": ["completion", "tools"] + }); + let probe = probe_from_ollama_show(&body); + assert_eq!(probe.max_input_tokens, Some(8192)); + assert_eq!(probe.tool_calling, Some(true)); + assert_eq!(probe.vision, Some(false)); + assert_eq!(probe.reasoning, Some(false)); +} + +#[test] +fn ollama_show_reads_vision_and_thinking_capabilities() { + let body = json!({ + "model_info": { "qwen3.context_length": 40960 }, + "capabilities": ["completion", "tools", "vision", "thinking"] + }); + let probe = probe_from_ollama_show(&body); + assert_eq!(probe.max_input_tokens, Some(40960)); + assert_eq!(probe.tool_calling, Some(true)); + assert_eq!(probe.vision, Some(true)); + assert_eq!(probe.reasoning, Some(true)); +} + +#[test] +fn an_absent_capabilities_array_leaves_every_capability_unknown() { + // "the server did not tell us" must not be read as "the model cannot". + let body = json!({ "model_info": { "llama.context_length": 4096 } }); + let probe = probe_from_ollama_show(&body); + assert_eq!(probe.max_input_tokens, Some(4096)); + assert_eq!(probe.tool_calling, None); + assert_eq!(probe.vision, None); + assert_eq!(probe.reasoning, None); +} + +#[test] +fn a_body_without_model_info_probes_to_nothing_rather_than_a_guess() { + assert!(probe_from_ollama_show(&json!({})).is_empty()); + assert!(probe_from_ollama_show(&json!({ "error": "model not found" })).is_empty()); +} + +#[test] +fn context_length_scan_is_architecture_agnostic_and_conservative() { + let info = json!({ "gemma3.context_length": 8192, "clip.context_length": 77 }) + .as_object() + .cloned() + .unwrap(); + // Multimodal models carry a second, tiny window for the projector. Taking + // the max would overstate; take the min. + assert_eq!(context_length_from_model_info(&info), Some(77)); + + let zeroed = json!({ "llama.context_length": 0 }) + .as_object() + .cloned() + .unwrap(); + assert_eq!(context_length_from_model_info(&zeroed), None); +} + +#[test] +fn lm_studio_listing_reports_loaded_and_trained_windows() { + let body = json!({ + "data": [ + { "id": "other-model", "max_context_length": 999999 }, + { + "id": "qwen3-4b", + "type": "llm", + "max_context_length": 40960, + "loaded_context_length": 4096 + } + ] + }); + let probe = probe_from_lm_studio_models(&body, "qwen3-4b"); + assert_eq!(probe.max_input_tokens, Some(40960)); + assert_eq!(probe.loaded_num_ctx, Some(4096)); + assert_eq!(probe.vision, Some(false)); + // The loaded window is the ceiling a live request actually has. + assert_eq!(probe.effective_context_window(), Some(4096)); +} + +#[test] +fn lm_studio_vision_models_are_detected_by_type() { + let body = json!({ "data": [{ "id": "llava", "type": "vlm", "max_context_length": 4096 }] }); + assert_eq!( + probe_from_lm_studio_models(&body, "llava").vision, + Some(true) + ); +} + +#[test] +fn an_unlisted_lm_studio_model_probes_to_nothing() { + let body = json!({ "data": [{ "id": "a", "max_context_length": 4096 }] }); + assert!(probe_from_lm_studio_models(&body, "b").is_empty()); +} + +#[test] +fn effective_window_prefers_the_loaded_ctx_over_the_trained_window() { + let probe = LocalProbe { + max_input_tokens: Some(131_072), + loaded_num_ctx: Some(2048), + ..LocalProbe::default() + }; + // Overstating is the LOCAL-1 defect: compaction is gated on this number, so + // a window bigger than the server honours means it never fires. + assert_eq!(probe.effective_context_window(), Some(2048)); +} + +#[test] +fn load_body_carries_num_ctx_and_keep_alive_natively() { + let options = json!({ "num_ctx": 8192, "num_batch": 512 }); + let body = ollama_load_body("llama3.2", Some(&options), Some("30m")); + assert_eq!(body["model"], json!("llama3.2")); + assert_eq!(body["messages"], json!([])); + // The whole point: `num_ctx` reaches a field Ollama actually reads. + assert_eq!(body["options"]["num_ctx"], json!(8192)); + assert_eq!(body["options"]["num_batch"], json!(512)); + assert_eq!(body["keep_alive"], json!("30m")); +} + +#[test] +fn load_body_omits_absent_and_malformed_options() { + let body = ollama_load_body("m", None, None); + assert!(body.get("options").is_none()); + assert!(body.get("keep_alive").is_none()); + + let scalar = json!("not an object"); + let body = ollama_load_body("m", Some(&scalar), None); + assert!(body.get("options").is_none()); +} + +#[test] +fn options_are_lifted_out_of_the_documented_provider_options_shape() { + let provider_options = json!({ "options": { "num_ctx": 8192 }, "keep_alive": "5m" }); + assert_eq!( + local_options_object(&provider_options), + Some(&json!({ "num_ctx": 8192 })) + ); + assert_eq!(local_options_object(&json!({})), None); + // A non-object `options` is caller error, not something to forward. + assert_eq!(local_options_object(&json!({ "options": 7 })), None); +} + +#[test] +fn missing_model_404_gains_an_actionable_remediation() { + let message = missing_model_remediation( + LocalRuntimeKind::Ollama, + 404, + r#"{"error":"model 'llama3.2' not found"}"#, + "llama3.2", + "http://localhost:11434/v1", + ) + .expect("a missing-model 404 is rewritten"); + assert!(message.contains("ollama pull llama3.2"), "{message}"); + + let lm = missing_model_remediation( + LocalRuntimeKind::LmStudio, + 404, + "model does not exist", + "qwen3-4b", + "http://localhost:1234/v1", + ) + .expect("LM Studio gets its own wording"); + assert!(lm.contains("list_models()"), "{lm}"); +} + +#[test] +fn unrelated_failures_keep_their_original_message() { + assert!(missing_model_remediation(LocalRuntimeKind::Ollama, 500, "boom", "m", "u").is_none()); + assert!( + missing_model_remediation(LocalRuntimeKind::Ollama, 404, "route not found", "m", "u") + .is_none() + ); + assert!( + missing_model_remediation(LocalRuntimeKind::Ollama, 401, "model not found", "m", "u") + .is_none() + ); +} + +#[test] +fn context_overflow_is_recognised_across_provider_phrasings() { + assert!(is_context_overflow( + 400, + "This model's maximum context length is 4096 tokens, however you requested 5000" + )); + assert!(is_context_overflow(413, "prompt is too long")); + assert!(is_context_overflow( + 400, + "Please reduce the length of the messages" + )); + // Not an overflow. + assert!(!is_context_overflow(400, "invalid api key")); + assert!(!is_context_overflow(401, "maximum context length exceeded")); +} + +#[test] +fn tools_rejections_are_told_apart_from_tool_choice_rejections() { + assert!(mentions_tools_unsupported( + "registry.ollama.ai/library/gemma3 does not support tools" + )); + assert!(mentions_tools_unsupported("unknown parameter: 'tools'")); + // `tool_choice` has its own latch; flipping the tools latch for it would + // disable native tools on a server that supports them fine. + assert!(!mentions_tools_unsupported( + "invalid parameter: tool_choice must be a string" + )); + assert!(!mentions_tools_unsupported("rate limited")); +} diff --git a/src/harness/providers/openai/mod.rs b/src/harness/providers/openai/mod.rs index 4a7bc8a..02367c1 100644 --- a/src/harness/providers/openai/mod.rs +++ b/src/harness/providers/openai/mod.rs @@ -84,16 +84,20 @@ const DEFAULT_CONNECT_TIMEOUT_SECS: u64 = 30; const DEFAULT_REQUEST_TIMEOUT_SECS: u64 = 600; mod convert; +mod local; mod reasoning_tags; pub(crate) mod relaxed_json; mod responses; mod sse; mod transport; +pub use convert::CacheTokenAccounting; +pub use local::{CONTEXT_OVERFLOW_CODE, LocalProbe, LocalRuntimeKind}; pub use reasoning_tags::ReasoningTagExtraction; pub use transport::{AuthStyle, OpenAiModel}; use convert::*; +use local::*; use reasoning_tags::*; use sse::*; #[cfg(test)] diff --git a/src/harness/providers/openai/responses.rs b/src/harness/providers/openai/responses.rs index e220a57..9cd7878 100644 --- a/src/harness/providers/openai/responses.rs +++ b/src/harness/providers/openai/responses.rs @@ -7,13 +7,31 @@ //! OpenAI Codex OAuth path requires (paired with `with_extra_query_param` + //! `with_user_agent`). //! -//! This first port is **text-in / text-out**: system messages fold into -//! `instructions`, user/assistant/tool turns become `input` items, and the -//! terminal `output_text` (or the first `output_text` content part) becomes the -//! assistant reply. Native tool calls over `/responses` and true SSE streaming -//! are follow-ups; the harness embeds tool specs in the prompt for this path -//! (its [`profile`](super::OpenAiModel) advertises the caller's chosen -//! `tool_calling`). +//! System messages fold into `instructions`, user/assistant/tool turns become +//! `input` items, and the terminal `output_text` (or the first `output_text` +//! content part) becomes the assistant reply. +//! +//! # What this path now carries +//! +//! The request used to be `{model, input, instructions, stream, store, +//! max_output_tokens}` and **silently dropped everything else** a caller set — +//! `tools`, `tool_choice`, `response_format`, `temperature`, `top_p`, `seed`, +//! `stop_sequences`, `continuation_id`, and `provider_options`. That last one +//! made `reasoning: {effort, summary}` unreachable on the only wire format in +//! this crate that supports it. All of them are on the wire now, and the +//! response side reads reasoning items, their `encrypted_content`, and the +//! cache/reasoning usage breakdowns that were previously ignored (so every +//! cached token on this path was billed at the full input rate). +//! +//! # Remaining gaps +//! +//! Tool *declarations* are sent, but a model that calls one comes back as a +//! `function_call` output item this port does not yet decode into +//! [`ToolCall`](crate::harness::tool::ToolCall)s — and tool *results* are +//! rendered as `user` turns carrying an explicit `[tool_result id=…]` prefix +//! rather than native `function_call_output` items. That preserves the causal +//! link the previous fold-into-assistant behaviour erased, but structural +//! tool support and true SSE streaming remain follow-ups. use serde::{Deserialize, Serialize}; use serde_json::Value; @@ -23,7 +41,17 @@ use crate::harness::model::ModelResponse; use crate::harness::usage::Usage; /// The `/v1/responses` request body. -#[derive(Debug, Serialize)] +/// +/// # What used to be missing +/// +/// This struct carried only `{model, input, instructions, stream, store, +/// max_output_tokens}`. Everything else a caller set was **silently dropped**: +/// `tools`, `tool_choice`, `response_format`, `temperature`, `top_p`, +/// `stop_sequences`, `seed`, `previous_response_id`, and — most pointedly — +/// `provider_options`, which meant `reasoning: {effort, summary}` was +/// unreachable on the one wire format that supports it. A request that looked +/// fully configured produced an unconfigured call. +#[derive(Debug, Default, Serialize)] pub(super) struct ResponsesRequest { pub(super) model: String, pub(super) input: Vec, @@ -37,6 +65,113 @@ pub(super) struct ResponsesRequest { /// `max_tokens`. Omitted for the Codex OAuth backend, which rejects it. #[serde(skip_serializing_if = "Option::is_none")] pub(super) max_output_tokens: Option, + /// Tool declarations. The Responses API flattens the function schema onto + /// the tool object rather than nesting it under `function` as Chat + /// Completions does. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub(super) tools: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) tool_choice: Option, + /// Structured output. The Responses API nests it under `text.format`, not + /// the Chat Completions `response_format`. + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) text: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) top_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) seed: Option, + /// Stop sequences. Serialized only when non-empty. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub(super) stop: Vec, + /// The stateful-continuation handle. [`ModelRequest::continuation_id`] + /// existed with a builder and **no reader anywhere in the crate**, so this + /// was never sent and stateful follow-ups silently restarted. + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) previous_response_id: Option, + /// `reasoning: { effort, summary }`, lowered from the provider-neutral + /// [`ReasoningConfig`][rc]. + /// + /// [rc]: crate::harness::model::ReasoningConfig + #[serde(skip_serializing_if = "Option::is_none")] + pub(super) reasoning: Option, + /// Which extra payloads to return. + /// + /// Load-bearing for reasoning replay: with `store: false` the server keeps + /// no state, so reasoning items may be dropped between turns **unless** they + /// carry `encrypted_content` — which only arrives when + /// `include: ["reasoning.encrypted_content"]` is requested. Asking for + /// reasoning and not asking for this is asking for reasoning that cannot be + /// replayed. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub(super) include: Vec, + /// Provider-specific passthrough merged onto the body. Keys here win, and + /// this is the escape hatch for anything the typed fields above cannot say. + #[serde(flatten)] + pub(super) extra: serde_json::Map, +} + +/// The `include` entry that makes reasoning replayable under `store: false`. +pub(super) const INCLUDE_ENCRYPTED_REASONING: &str = "reasoning.encrypted_content"; + +/// Lowers a provider-neutral [`ReasoningConfig`][rc] onto the Responses +/// `reasoning` object. +/// +/// Returns `None` when the config asks for nothing, so an empty config never +/// adds a field. `budget_tokens` has no Responses spelling and is dropped here +/// deliberately — it is Anthropic's knob, and inventing an OpenAI field for it +/// would be worse than ignoring it. +/// +/// [rc]: crate::harness::model::ReasoningConfig +pub(super) fn translate_reasoning( + config: &crate::harness::model::ReasoningConfig, +) -> Option { + if config.is_empty() { + return None; + } + let mut object = serde_json::Map::new(); + if let Some(effort) = config.effort { + object.insert("effort".to_string(), Value::String(effort.as_str().into())); + } + if let Some(summary) = &config.summary { + object.insert("summary".to_string(), Value::String(summary.clone())); + } + (!object.is_empty()).then_some(Value::Object(object)) +} + +/// Translates a tool schema onto the Responses API's flattened tool shape. +pub(super) fn translate_tool(schema: &crate::harness::tool::ToolSchema) -> Value { + serde_json::json!({ + "type": "function", + "name": schema.name, + "description": schema.description, + "parameters": schema.parameters, + }) +} + +/// Translates a [`ResponseFormat`][rf] onto the Responses API's `text.format` +/// nesting (Chat Completions' `response_format` has no counterpart here). +/// +/// [rf]: crate::harness::model::ResponseFormat +pub(super) fn translate_text_format( + format: &crate::harness::model::ResponseFormat, + strict: bool, +) -> Option { + use crate::harness::model::ResponseFormat; + let inner = match format { + ResponseFormat::Text => return None, + ResponseFormat::JsonObject => serde_json::json!({ "type": "json_object" }), + ResponseFormat::JsonSchema { name, schema } | ResponseFormat::Auto { name, schema } => { + serde_json::json!({ + "type": "json_schema", + "name": name, + "schema": schema, + "strict": strict, + }) + } + }; + Some(serde_json::json!({ "format": inner })) } #[derive(Debug, Serialize)] @@ -62,10 +197,24 @@ pub(super) struct ResponsesResponse { pub(super) usage: Option, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Default, Deserialize)] pub(super) struct ResponsesOutput { + /// Item kind: `message`, `reasoning`, `function_call`, … + #[serde(rename = "type", default)] + pub(super) kind: Option, #[serde(default)] pub(super) content: Vec, + /// Reasoning summary parts, on a `reasoning` item. + #[serde(default)] + pub(super) summary: Vec, + /// The opaque reasoning payload that survives `store: false`. + /// + /// Only present when the request asked for + /// [`INCLUDE_ENCRYPTED_REASONING`]. Preserved so a caller can replay + /// reasoning across turns; without it the server drops reasoning between + /// stateless turns. + #[serde(default)] + pub(super) encrypted_content: Option, } #[derive(Debug, Deserialize)] @@ -75,13 +224,44 @@ pub(super) struct ResponsesContent { pub(super) text: Option, } -/// Responses-API usage block (`input_tokens` / `output_tokens`). -#[derive(Debug, Deserialize)] +/// Responses-API usage block. +/// +/// The details sub-objects used to be absent from this struct entirely, so +/// **every cached token on this path was billed at the full input rate** and +/// reasoning tokens were invisible. +#[derive(Debug, Default, Deserialize)] pub(super) struct ResponsesUsage { #[serde(default)] pub(super) input_tokens: Option, #[serde(default)] pub(super) output_tokens: Option, + #[serde(default)] + pub(super) input_tokens_details: Option, + #[serde(default)] + pub(super) output_tokens_details: Option, +} + +/// `usage.input_tokens_details` — the cache breakdown of the input total. +#[derive(Debug, Default, Deserialize)] +pub(super) struct ResponsesInputTokenDetails { + #[serde(default)] + pub(super) cached_tokens: Option, + /// Cache **writes**, under either spelling gateways use. + #[serde(default)] + pub(super) cache_write_tokens: Option, + #[serde(default)] + pub(super) cache_creation_tokens: Option, +} + +/// `usage.output_tokens_details` — where OpenAI reports reasoning tokens. +/// +/// Note that **Anthropic has no equivalent field**: its thinking tokens are +/// billed inside `output_tokens`, so a zero here is not evidence that no +/// reasoning happened on an Anthropic-shaped gateway. +#[derive(Debug, Default, Deserialize)] +pub(super) struct ResponsesOutputTokenDetails { + #[serde(default)] + pub(super) reasoning_tokens: Option, } /// Concatenates the visible text of a message's content blocks. @@ -93,12 +273,22 @@ fn message_text(content: &[ContentBlock]) -> String { .join("") } -/// Normalizes a message role for the Responses API: assistant + tool turns fold -/// into `assistant` (which the API keys to `output_text`), everything else to -/// `user` (`input_text`). Mirrors the host `normalize_responses_role`. +/// Normalizes a message role for the Responses API. +/// +/// Assistant turns key to `output_text`; everything else to `input_text`. +/// +/// **Tool results no longer fold into `assistant`.** They used to, which erased +/// tool identity entirely: a tool result became an anonymous assistant utterance +/// with no `tool_call_id`, so the model saw an assistant asserting a fact rather +/// than the answer to a call it made. They are now rendered as `user` turns +/// carrying an explicit `[tool_result …]` prefix (see +/// [`build_responses_input`]), which preserves the causal link on a wire format +/// this text-in/text-out port cannot express structurally. A true +/// `function_call_output` item is the complete fix and rides with native tool +/// support on this path. fn normalize_role(message: &Message) -> &'static str { match message { - Message::Assistant(_) | Message::Tool(_) => "assistant", + Message::Assistant(_) => "assistant", _ => "user", } } @@ -123,7 +313,16 @@ pub(super) fn build_responses_input(messages: &[Message]) -> (Option, Ve } Message::User(m) => message_text(&m.content), Message::Assistant(m) => message_text(&m.content), - Message::Tool(m) => message_text(&m.content), + // Keep the call id visible so the model can tell *which* call this + // answers. Folding it into an anonymous assistant turn lost that. + Message::Tool(m) => { + let body = message_text(&m.content); + if body.trim().is_empty() { + String::new() + } else { + format!("[tool_result id={} ]\n{body}", m.tool_call_id) + } + } }; if text.trim().is_empty() { continue; @@ -149,6 +348,7 @@ pub(super) fn build_responses_input(messages: &[Message]) -> (Option, Ve /// Extracts the assistant text from a Responses body: the convenience /// `output_text` field first, else the first `output_text` content part. pub(super) fn extract_responses_text(response: &ResponsesResponse) -> Option { + // `output_text` is the whole answer when the server supplies it. if let Some(text) = response .output_text .as_deref() @@ -158,6 +358,11 @@ pub(super) fn extract_responses_text(response: &ResponsesResponse) -> Option Option Option { + let mut text = String::new(); + for item in &response.output { + if item.kind.as_deref() != Some("reasoning") { + continue; + } + for part in item.summary.iter().chain(item.content.iter()) { + if let Some(fragment) = part + .text + .as_deref() + .map(str::trim) + .filter(|t| !t.is_empty()) + { + if !text.is_empty() { + text.push('\n'); + } + text.push_str(fragment); + } + } + } + (!text.is_empty()).then_some(text) +} + +/// The opaque reasoning payload to replay on the next turn, when the request +/// asked for [`INCLUDE_ENCRYPTED_REASONING`] and the server supplied it. +pub(super) fn extract_encrypted_reasoning(response: &ResponsesResponse) -> Option { + response + .output + .iter() + .find_map(|item| item.encrypted_content.clone()) + .filter(|value| !value.is_empty()) +} + +/// Maps a Responses `usage` block onto the neutral [`Usage`], including the +/// cache and reasoning breakdowns. +pub(super) fn convert_responses_usage(wire: &ResponsesUsage) -> Usage { + let input_details = wire.input_tokens_details.as_ref(); + let cache_read_tokens = input_details.and_then(|d| d.cached_tokens).unwrap_or(0); + let cache_creation_tokens = input_details + .map(|d| { + d.cache_write_tokens + .unwrap_or(0) + .max(d.cache_creation_tokens.unwrap_or(0)) + }) + .unwrap_or(0); + let input_tokens = wire.input_tokens.unwrap_or(0); + let output_tokens = wire.output_tokens.unwrap_or(0); + Usage { + input_tokens, + output_tokens, + total_tokens: input_tokens + output_tokens, + cache_read_tokens, + cache_creation_tokens, + reasoning_tokens: wire + .output_tokens_details + .as_ref() + .and_then(|d| d.reasoning_tokens) + .unwrap_or(0), + } +} + +/// Parses a raw `/v1/responses` JSON body into a [`ModelResponse`]. +/// +/// Reasoning items surface as a leading +/// [`ContentBlock::Thinking`] block, consistent with the Chat Completions path; +/// the encrypted payload, when present, rides on the block's `signature` so it +/// can be replayed on a later turn. pub(super) fn parse_responses_response(value: Value) -> ModelResponse { let parsed: ResponsesResponse = - serde_json::from_value(value.clone()).unwrap_or(ResponsesResponse { + serde_json::from_value(value.clone()).unwrap_or_else(|_| ResponsesResponse { output: Vec::new(), output_text: None, usage: None, }); let text = extract_responses_text(&parsed).unwrap_or_default(); - let usage = parsed.usage.as_ref().map(|u| Usage { - input_tokens: u.input_tokens.unwrap_or(0), - output_tokens: u.output_tokens.unwrap_or(0), - total_tokens: u.input_tokens.unwrap_or(0) + u.output_tokens.unwrap_or(0), - ..Usage::default() - }); + let usage = parsed.usage.as_ref().map(convert_responses_usage); + + let mut content = Vec::new(); + if let Some(reasoning) = extract_responses_reasoning(&parsed) { + content.push(ContentBlock::Thinking { + text: reasoning, + signature: extract_encrypted_reasoning(&parsed), + }); + } + content.push(ContentBlock::Text(text)); + ModelResponse { message: AssistantMessage { id: None, - content: vec![ContentBlock::Text(text)], + content, tool_calls: Vec::new(), usage, }, @@ -200,6 +481,7 @@ pub(super) fn parse_responses_response(value: Value) -> ModelResponse { raw: Some(value), resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -254,6 +536,7 @@ mod tests { text: Some("answer".into()), }, ], + ..ResponsesOutput::default() }], output_text: None, usage: None, diff --git a/src/harness/providers/openai/sse.rs b/src/harness/providers/openai/sse.rs index 76fc62d..8665340 100644 --- a/src/harness/providers/openai/sse.rs +++ b/src/harness/providers/openai/sse.rs @@ -50,16 +50,32 @@ pub(super) struct OpenAiStreamAcc { /// opened slot, so id-less argument continuations for that index keep /// following the call most recently opened there. index_slots: std::collections::HashMap, + /// Synthetic tool-call id epoch for **this** response, drawn once in + /// [`new`](Self::new). Every `tool-call` id this accumulator synthesizes — + /// both the streamed [`ToolDelta::call_id`] and the terminal response's + /// [`ToolCall`] id — carries it, so the two agree within the response while + /// no other response in the process can mint the same id. + id_epoch: u64, + /// Whether the provider reports input tokens **excluding** cache + /// read/creation tokens (Anthropic semantics). See + /// [`CacheTokenAccounting`]. + cache_accounting: CacheTokenAccounting, } impl OpenAiStreamAcc { /// Builds an accumulator with the given inline reasoning-tag extraction - /// config (`None` disables inline extraction). - pub(super) fn new(reasoning_tags: Option) -> Self { + /// config (`None` disables inline extraction) and cache-token accounting + /// convention. + pub(super) fn new( + reasoning_tags: Option, + cache_accounting: CacheTokenAccounting, + ) -> Self { let extractor = reasoning_tags.as_ref().map(ReasoningTagStream::new); Self { reasoning_tags, extractor, + id_epoch: next_synthetic_id_epoch(), + cache_accounting, ..Self::default() } } @@ -73,7 +89,7 @@ impl OpenAiStreamAcc { self.id = Some(id); } if let Some(usage_wire) = chunk.usage { - let usage = convert_usage(usage_wire); + let usage = convert_usage_with(usage_wire, self.cache_accounting); self.usage = Some(usage); pending.push_back(ModelStreamItem::UsageDelta(usage)); } @@ -138,7 +154,7 @@ impl OpenAiStreamAcc { } if let Some(args) = function.arguments.filter(|a| !a.is_empty()) { slot.args.push_str(&args); - let call_id = tool_call_id(idx, &slot.id); + let call_id = tool_call_id(self.id_epoch, idx, &slot.id); pending.push_back(ModelStreamItem::ToolCallDelta(ToolDelta { call_id, content: args, @@ -222,6 +238,7 @@ impl OpenAiStreamAcc { /// than failing the whole stream, so the agent loop can feed the error back /// to the model and the call still resolves instead of stalling the loop. fn into_response(self) -> ModelResponse { + let id_epoch = self.id_epoch; let mut content = Vec::new(); // Recompute the inline-tag split over the raw accumulated content so the // terminal response is byte-identical to the non-streaming path, then @@ -252,7 +269,7 @@ impl OpenAiStreamAcc { content.push(ContentBlock::Text(visible_text)); } // Enumerate over the full slot vector *before* filtering so the synthetic - // fallback id (`tool-{idx}`) matches the one streamed in `ToolCallDelta` + // fallback id (`tacall-{epoch}-{idx}`) matches the one streamed in `ToolCallDelta` // items — filtering first would renumber the slots and desynchronize the // delta ids from the final call ids. let tool_calls = self @@ -260,7 +277,9 @@ impl OpenAiStreamAcc { .into_iter() .enumerate() .filter(|(_, b)| !b.name.is_empty() || !b.args.is_empty()) - .map(|(idx, b)| tool_call_from_wire("openai stream", idx, &b.id, &b.name, &b.args)) + .map(|(idx, b)| { + tool_call_from_wire("openai stream", id_epoch, idx, &b.id, &b.name, &b.args) + }) .collect::>(); let message = AssistantMessage { id: self.id, @@ -275,6 +294,7 @@ impl OpenAiStreamAcc { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } } diff --git a/src/harness/providers/openai/test.rs b/src/harness/providers/openai/test.rs index 487bfef..6f347f6 100644 --- a/src/harness/providers/openai/test.rs +++ b/src/harness/providers/openai/test.rs @@ -355,12 +355,30 @@ fn parses_id_less_tool_call_with_synthesized_fallback_id() { let response = parse_response(body).unwrap(); let calls = response.tool_calls(); assert_eq!(calls.len(), 2); - assert_eq!(calls[0].id, "tool-0"); assert_eq!(calls[0].name, "ping"); assert_eq!(calls[0].arguments, json!({})); - assert_eq!(calls[1].id, "tool-1"); assert_eq!(calls[1].name, "pong"); assert_eq!(calls[1].arguments, json!({ "n": 1 })); + + // The synthesized ids are `tacall-{epoch}-{slot}`: run-unique via a + // process-global epoch, and distinct per slot within the response. The old + // form was `tool-{slot}`, keyed only to position — so every id-less turn + // emitted `tool-0` and one transcript ended up with several different calls + // all claiming the same id, which the agent loop cannot pair. + assert!( + calls[0].id.starts_with("tacall-"), + "unexpected synthesized id: {}", + calls[0].id + ); + assert!(calls[0].id.ends_with("-0"), "slot must be the id suffix"); + assert!(calls[1].id.ends_with("-1"), "slot must be the id suffix"); + assert_ne!(calls[0].id, calls[1].id); + + // The prompt-guided text protocol mints `ptc_{seq}_{slot}`; the two schemes + // must be unmistakably disjoint, because a run that degrades from native to + // prompt-guided mid-flight mixes both in one transcript. + assert!(!calls[0].id.starts_with("ptc_")); + assert!(!calls[0].id.starts_with("call_")); } #[test] @@ -495,9 +513,13 @@ fn explicit_nulls_are_tolerated_wherever_a_default_exists() { let calls = response.tool_calls(); assert_eq!(calls.len(), 1); assert_eq!(calls[0].name, "ping"); - // A nulled id still gets the synthesized positional fallback, so tool - // results stay correlatable. - assert_eq!(calls[0].id, "tool-0"); + // A nulled id still gets the synthesized fallback, so tool results stay + // correlatable. + assert!( + calls[0].id.starts_with("tacall-"), + "unexpected synthesized id: {}", + calls[0].id + ); assert_eq!(calls[0].arguments, json!({})); let usage = response.usage.unwrap(); assert_eq!(usage.input_tokens, 0); @@ -883,10 +905,35 @@ fn local_runtime_presets_normalize_endpoint_and_model() { let overridden = OpenAiModel::ollama().with_model("qwen3:8b"); let profile = >::profile(&overridden).unwrap(); - assert!(!profile.tool_calling); - assert!(!profile.parallel_tool_calls); - assert!(!profile.streaming_tool_chunks); + // Native tool calling is no longer hard-disabled for local runtimes. It used + // to be, unconditionally, which forced the prompt-guided branch on every + // local call *and* excluded every local model from any + // `CapabilitySet { tool_calling: true }` resolution. It is optimistic now, + // with a 400-driven latch (`Degrade::native_tools`) to fall back once. + assert!(profile.tool_calling); + assert!(profile.parallel_tool_calls); + assert!(profile.streaming_tool_chunks); + // Vision stays off: it is a property of the loaded weights, and a probe is + // what turns it on. assert!(!profile.modalities.image_in); + // LOCAL-1: no invented context window. `qwen3:8b` matches nothing in the + // hint table, but `llama3.2` would have matched `("llama3", Substring, + // 128_000)` against a server whose default `num_ctx` is 2048. + assert_eq!(profile.max_input_tokens, None); + assert_eq!( + >::profile(&OpenAiModel::ollama()) + .unwrap() + .max_input_tokens, + None, + "the default Ollama model must not inherit a hosted-sized window either" + ); + // The degrade knobs the local presets exist for are pre-set, so the first + // call does not pay a guaranteed 400 to rediscover them. + let baseline = overridden.baseline_degrade(); + assert!(baseline.named_tool_choice); + assert!(baseline.json_object); + assert!(baseline.json_schema_strict); + assert!(!baseline.native_tools, "native tools start enabled"); assert!(OpenAiModel::ollama_at("http://[::1", "qwen3").is_err()); assert!(OpenAiModel::ollama_at("ftp://host", "qwen3").is_err()); @@ -915,10 +962,11 @@ fn provider_spec_builds_compatible_model() { Some("ollama") ); let profile = >::profile(&model).unwrap(); - assert!(!profile.tool_calling); - assert!(!profile.parallel_tool_calls); - assert!(!profile.streaming_tool_chunks); + assert!(profile.tool_calling, "see local_runtime_presets_* for why"); + assert!(profile.parallel_tool_calls); + assert!(profile.streaming_tool_chunks); assert!(!profile.modalities.image_in); + assert_eq!(profile.max_input_tokens, None); let mut authenticated = ProviderSpec::for_kind(ProviderKind::Ollama); authenticated.requires_api_key = true; @@ -943,6 +991,7 @@ fn provider_failed_stream_item_finishes_as_provider_error() { code: Some("rate_limit".to_string()), message: "too many requests".to_string(), retryable: true, + retry_after_ms: None, raw: None, })); @@ -1163,7 +1212,7 @@ async fn collect_sse_with( bytes: Box::pin(bytes), buf: Vec::new(), pending: std::collections::VecDeque::new(), - acc: OpenAiStreamAcc::new(reasoning_tags), + acc: OpenAiStreamAcc::new(reasoning_tags, CacheTokenAccounting::default()), provider: "openai".to_string(), model: "gpt-4.1-mini".to_string(), started: false, @@ -1341,7 +1390,7 @@ fn parse_chat_response_extracts_inline_think_and_side_channel() { }); let cfg = ReasoningTagExtraction::default(); - let response = parse_chat_response(body, Some(&cfg)).unwrap(); + let response = parse_chat_response(body, Some(&cfg), CacheTokenAccounting::default()).unwrap(); assert_eq!(response.text(), "The answer is 42."); // Side-channel leads, inline follows, separator-joined. assert_eq!(response_reasoning(&response), "side\ninline"); @@ -1356,7 +1405,7 @@ fn parse_chat_response_without_config_leaves_inline_tags_in_text() { ] }); - let response = parse_chat_response(body, None).unwrap(); + let response = parse_chat_response(body, None, CacheTokenAccounting::default()).unwrap(); assert_eq!(response.text(), "xy"); assert_eq!(response_reasoning(&response), ""); } @@ -2377,6 +2426,21 @@ fn default_provider_options_are_baked_onto_every_request() { assert_eq!(value["options"]["num_ctx"], json!(8192)); } +#[test] +fn cache_identity_includes_baked_answer_configuration() { + let base = OpenAiModel::ollama().with_model("qwen2.5"); + let context_window = OpenAiModel::ollama() + .with_model("qwen2.5") + .with_default_provider_options(json!({ "options": { "num_ctx": 8192 } })); + let merged_system = OpenAiModel::ollama() + .with_model("qwen2.5") + .with_merge_system_into_user(); + + let identity = |model: &OpenAiModel| >::cache_identity(model); + assert_ne!(identity(&base), identity(&context_window)); + assert_ne!(identity(&base), identity(&merged_system)); +} + #[test] fn request_provider_options_win_over_baked_defaults() { let model = OpenAiModel::ollama() @@ -2523,6 +2587,7 @@ fn degrade_for_400_targets_only_the_shape_the_request_used() { Some(Degrade { named_tool_choice: true, json_object: false, + ..Degrade::default() }) ); @@ -2537,6 +2602,7 @@ fn degrade_for_400_targets_only_the_shape_the_request_used() { Some(Degrade { named_tool_choice: false, json_object: true, + ..Degrade::default() }) ); } @@ -2571,6 +2637,7 @@ fn degrade_for_400_ignores_unrelated_or_already_degraded_failures() { Degrade { named_tool_choice: true, json_object: false, + ..Degrade::default() }, ), None @@ -2590,11 +2657,13 @@ fn degrade_for_400_unions_with_existing_baseline_degrade() { Degrade { named_tool_choice: true, json_object: false, + ..Degrade::default() }, ), Some(Degrade { named_tool_choice: true, json_object: true, + ..Degrade::default() }) ); } @@ -2619,12 +2688,14 @@ fn shape_degrade_latches_after_discovery_so_baseline_is_already_degraded() { m.latch_degrade(Degrade { named_tool_choice: true, json_object: false, + ..Degrade::default() }); assert_eq!( m.baseline_degrade(), Degrade { named_tool_choice: true, json_object: false, + ..Degrade::default() }, "a discovered named_tool_choice rejection must be remembered so the next \ call's baseline body is already degraded, instead of re-paying the 400" @@ -2635,12 +2706,14 @@ fn shape_degrade_latches_after_discovery_so_baseline_is_already_degraded() { m.latch_degrade(Degrade { named_tool_choice: true, json_object: true, + ..Degrade::default() }); assert_eq!( m.baseline_degrade(), Degrade { named_tool_choice: true, json_object: true, + ..Degrade::default() } ); } @@ -2656,12 +2729,14 @@ fn shape_degrade_latch_survives_through_a_shared_handle() { shared.latch_degrade(Degrade { named_tool_choice: false, json_object: true, + ..Degrade::default() }); assert_eq!( clone.baseline_degrade(), Degrade { named_tool_choice: false, json_object: true, + ..Degrade::default() }, "the latch must be visible to every holder of the shared model" ); @@ -2720,6 +2795,7 @@ fn stream_cleanup_scrubs_leaked_markup_from_live_deltas() { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, }; let terminal = super::transport::clean_stream_item( ModelStreamItem::Completed(response), diff --git a/src/harness/providers/openai/transport.rs b/src/harness/providers/openai/transport.rs index 447905d..789787d 100644 --- a/src/harness/providers/openai/transport.rs +++ b/src/harness/providers/openai/transport.rs @@ -140,6 +140,41 @@ pub struct OpenAiModel { /// guaranteed-400 round trip before falling back. See /// [`Self::latch_stream_required`]. stream_required: AtomicBool, + /// Whether a JSON-Schema `response_format` is sent with OpenAI **strict** + /// structured output (`"strict": true`). + /// + /// Used to be hardcoded `true` for every endpoint and every schema. Strict + /// mode is far narrower than JSON Schema — it requires + /// `additionalProperties: false` on every object and *every* property listed + /// in `required` — so a caller's ordinary schema 400s. Local runtimes reject + /// the key outright. Defaults `true` for hosted OpenAI (unchanged wire + /// shape), `false` for local runtimes, overridable with + /// [`Self::with_strict_json_schema`], and latched `false` by a 400 that + /// implicates the schema (see [`degrade_for_400`]). + json_schema_strict: AtomicBool, + /// Whether the endpoint accepts native `tools` on the wire. + /// + /// This is the **transport** half of the tool decision; the *advertised* + /// half is [`ModelProfile::tool_calling`]. They used to be the same value, + /// hard-disabled for every local runtime, which forced prompt-guided tools + /// on every local call and excluded every local model from any + /// `CapabilitySet { tool_calling: true }` resolution. Now it is discovered: + /// seeded optimistically (or from [`Self::probe_local_profile`]) and latched + /// `false` by a 400 implicating `tools`, exactly like + /// [`Self::stream_required`]. See [`Self::with_native_tools_on_wire`]. + native_tools_on_wire: AtomicBool, + /// How the provider counts cache tokens against its input total. See + /// [`CacheTokenAccounting`]. + cache_accounting: CacheTokenAccounting, + /// Whether this instance points at a local runtime (Ollama, LM Studio, + /// llama.cpp server, vLLM, …). Local runtimes get the degrade knobs + /// pre-set, a native-API escape hatch for `num_ctx`, a `model not found` + /// error rewrite with real remediation, and are the only targets + /// [`Self::probe_local_profile`] will probe. + local_runtime: Option, + /// Optional `keep_alive` residency hint baked onto every local request. See + /// [`Self::with_keep_alive`]. + keep_alive: Option, } /// The auth headers `(name, value)` for a given [`AuthStyle`] + credential. @@ -288,11 +323,30 @@ pub(super) fn merge_system_into_user(messages: &[Message]) -> Vec { merged } -/// Returns `true` for OpenAI o-series reasoning models (`o1`/`o3`/`o4`), which -/// reject `max_tokens` and require `max_completion_tokens` instead. +/// Returns `true` for OpenAI reasoning models, which reject `max_tokens` and +/// require `max_completion_tokens` instead. +/// +/// Covers the o-series (`o1`/`o3`/`o4`) **and the gpt-5 family**. gpt-5 was +/// missing entirely: its cap was routed to `max_tokens`, which OpenAI rejects +/// outright (`Unsupported parameter: 'max_tokens' … use +/// 'max_completion_tokens'`), and it was additionally profiled as +/// `reasoning: false` / `native_structured_output: false`, so a +/// `CapabilitySet { reasoning: true }` filtered it out and structured output +/// picked the tool-call fallback over native schema mode. pub(super) fn is_reasoning_model(model: &str) -> bool { let lower = model.to_ascii_lowercase(); - lower.starts_with("o1") || lower.starts_with("o3") || lower.starts_with("o4") + lower.starts_with("o1") + || lower.starts_with("o3") + || lower.starts_with("o4") + || is_gpt5_family(&lower) +} + +/// Returns `true` for the gpt-5 family, tolerating the gateway-prefixed ids +/// routers use (`openai/gpt-5-mini` on OpenRouter). +/// +/// Takes an already-lowercased id; callers inside this module always have one. +pub(super) fn is_gpt5_family(lower: &str) -> bool { + lower.starts_with("gpt-5") || lower.starts_with("gpt5") || lower.contains("/gpt-5") } /// Derives a static [`ModelProfile`] for an OpenAI(-compatible) model id. @@ -309,6 +363,25 @@ pub(super) fn is_reasoning_model(model: &str) -> bool { /// engages on a real window instead of silently falling back to a fixed /// threshold. pub(super) fn derive_profile(provider: &str, model: &str) -> ModelProfile { + derive_profile_for(provider, model, false) +} + +/// [`derive_profile`], with `local` selecting the local-runtime policy. +/// +/// The only difference is [`ModelProfile::max_input_tokens`], and it matters a +/// great deal. The generic hint table matches **bare substrings**, so +/// `llama3.2:3b` served by Ollama resolved through `("llama3", Substring, +/// 128_000)` and advertised a 128 000-token window against a server whose +/// default `num_ctx` is 2048. Compaction fires at `window * threshold`, so it +/// never fired and the server truncated the front of the prompt silently. +/// +/// A local profile therefore reports `None` — "unknown" — which this crate +/// already supports end to end and which is strictly better than a wrong +/// number. LangChain reaches the same conclusion by shipping no ChatOllama +/// profile at all and hard-failing its summarization middleware with a message +/// telling you to pass absolute token counts. The real answer is to ask the +/// server: [`OpenAiModel::probe_local_profile`]. +pub(super) fn derive_profile_for(provider: &str, model: &str, local: bool) -> ModelProfile { let lower = model.to_ascii_lowercase(); let reasoning = is_reasoning_model(model); let native_structured = lower.contains("gpt-4o") || lower.contains("gpt-4.1") || reasoning; @@ -327,7 +400,15 @@ pub(super) fn derive_profile(provider: &str, model: &str) -> ModelProfile { native_structured_output: native_structured, json_schema: true, reasoning, - max_input_tokens: crate::harness::model::context_window_for_model_id(model), + // Every OpenAI reasoning model accepts `reasoning_effort`; models that + // merely *emit* reasoning (a deepseek-r1 distill leaking `` + // through a local runtime) do not, which is why this is a separate flag. + reasoning_effort: reasoning && !local, + max_input_tokens: if local { + None + } else { + crate::harness::model::context_window_for_model_id(model) + }, ..ModelProfile::default() } } @@ -368,6 +449,11 @@ impl OpenAiModel { reasoning_tags: Some(ReasoningTagExtraction::default()), reasoning_tags_overridden: false, stream_required: AtomicBool::new(false), + json_schema_strict: AtomicBool::new(true), + native_tools_on_wire: AtomicBool::new(true), + cache_accounting: CacheTokenAccounting::default(), + local_runtime: None, + keep_alive: None, } } @@ -437,6 +523,20 @@ impl OpenAiModel { "[openai] provider rejects response_format:json_object; latching degraded shape for subsequent calls" ); } + if degrade.json_schema_strict && self.json_schema_strict.swap(false, Ordering::Relaxed) { + tracing::info!( + provider = %self.provider, + model = %self.model, + "[openai] provider rejects strict json_schema; latching strict:false for subsequent calls" + ); + } + if degrade.native_tools && self.native_tools_on_wire.swap(false, Ordering::Relaxed) { + tracing::info!( + provider = %self.provider, + model = %self.model, + "[openai] provider rejects native tools; latching prompt-guided tools for subsequent calls" + ); + } } /// Routes calls to the OpenAI **Responses API** (`/v1/responses`) instead of @@ -639,45 +739,134 @@ impl OpenAiModel { /// Overrides the default model id. pub fn with_model(mut self, model: impl Into) -> Self { self.model = model.into(); - let local_capabilities = self.local_capabilities_locked.then_some(( - self.profile.tool_calling, - self.profile.parallel_tool_calls, - self.profile.streaming_tool_chunks, - self.profile.modalities.image_in, - )); - self.profile = derive_profile(&self.provider, &self.model); - if let Some((tool_calling, parallel_tool_calls, streaming_tool_chunks, image_in)) = - local_capabilities - { - self.profile.tool_calling = tool_calling; - self.profile.parallel_tool_calls = parallel_tool_calls; - self.profile.streaming_tool_chunks = streaming_tool_chunks; - self.profile.modalities.image_in = image_in; - } + self.rederive_profile(); self } /// Overrides the provider family id used in profiles and normalized errors. pub fn with_provider(mut self, provider: impl Into) -> Self { self.provider = provider.into(); - let local_capabilities = self.local_capabilities_locked.then_some(( + self.rederive_profile(); + self + } + + /// Re-derives [`Self::profile`] after the model id or provider id changed, + /// preserving the local-runtime capability overrides when they are locked. + /// + /// `max_input_tokens` is part of what is preserved. It used not to be: + /// [`Self::with_model`] kept the locked tool/vision overrides but + /// **re-derived** the window from the model id, so the wrong hosted-sized + /// window survived every model swap on a local handle. + fn rederive_profile(&mut self) { + let locked = self.local_capabilities_locked.then_some(( self.profile.tool_calling, self.profile.parallel_tool_calls, self.profile.streaming_tool_chunks, self.profile.modalities.image_in, + self.profile.max_input_tokens, )); - self.profile = derive_profile(&self.provider, &self.model); - if let Some((tool_calling, parallel_tool_calls, streaming_tool_chunks, image_in)) = - local_capabilities + self.profile = + derive_profile_for(&self.provider, &self.model, self.local_runtime.is_some()); + if let Some(( + tool_calling, + parallel_tool_calls, + streaming_tool_chunks, + image_in, + max_input_tokens, + )) = locked { self.profile.tool_calling = tool_calling; self.profile.parallel_tool_calls = parallel_tool_calls; self.profile.streaming_tool_chunks = streaming_tool_chunks; self.profile.modalities.image_in = image_in; + self.profile.max_input_tokens = max_input_tokens; } + } + + /// Declares whether a JSON-Schema `response_format` is sent with OpenAI + /// **strict** structured output. + /// + /// `true` on hosted OpenAI, `false` on every local runtime (they reject the + /// key, and LangChain pops it for Ollama for exactly that reason). Strict + /// mode demands `additionalProperties: false` on every object and every + /// property listed in `required`, so a caller's ordinary schema 400s under + /// it; independent of this flag, a 400 implicating the schema degrades to + /// `strict: false` for a single retry and latches. + pub fn with_strict_json_schema(self, strict: bool) -> Self { + self.json_schema_strict.store(strict, Ordering::Relaxed); + self + } + + /// Declares whether the endpoint accepts native `tools` on the wire. + /// + /// Distinct from [`with_native_tool_calling`](Self::with_native_tool_calling), + /// which changes what the *profile advertises* (and therefore what + /// [`CapabilitySet`][cs] resolution will accept). This one changes only + /// which branch the transport takes. Pass `false` for a server known to 400 + /// on `tools`; otherwise leave it alone — a 400 implicating `tools` flips it + /// automatically and latches, so the discovery is paid once per process. + /// + /// [cs]: crate::harness::model::CapabilitySet + pub fn with_native_tools_on_wire(self, enabled: bool) -> Self { + self.native_tools_on_wire.store(enabled, Ordering::Relaxed); self } + /// Whether native `tools` go on the wire for the next call: the model must + /// advertise tool calling **and** the transport must not have latched a + /// rejection. + pub fn native_tools_enabled(&self) -> bool { + self.profile.tool_calling && self.native_tools_on_wire.load(Ordering::Relaxed) + } + + /// Declares how the provider counts cache tokens against its input total. + /// + /// Defaults to OpenAI semantics ([`CacheTokenAccounting::IncludedInInput`]). + /// Set [`CacheTokenAccounting::ExcludedFromInput`] for a gateway that + /// forwards Anthropic's convention, where `input_tokens` omits cache reads + /// and writes — assuming OpenAI semantics over Anthropic data silently + /// under-bills. + pub fn with_cache_token_accounting(mut self, accounting: CacheTokenAccounting) -> Self { + self.cache_accounting = accounting; + self + } + + /// Bakes a `keep_alive` residency hint onto this local handle. + /// + /// Ollama unloads an idle model after 5 minutes by default, so the next turn + /// pays a cold multi-gigabyte load inside the 600 s unary deadline. The + /// value is Ollama's own format (`"30m"`, `"-1"` for forever, `"0"` to + /// unload immediately) and is delivered by [`Self::warm_up`], which speaks + /// the native API — `keep_alive` has no OpenAI-wire spelling. + pub fn with_keep_alive(mut self, keep_alive: impl Into) -> Self { + self.keep_alive = Some(keep_alive.into()); + self + } + + /// Requests an explicit context window for the loaded local model. + /// + /// Shorthand for the documented `{"options": {"num_ctx": n}}` escape hatch, + /// merged into [`Self::default_provider_options`], **and** — crucially — + /// the value [`Self::warm_up`] delivers over the native API, which is the + /// only path a server actually reads it on. It also updates the advertised + /// [`ModelProfile::max_input_tokens`], because a window you asked for and a + /// window you advertise must be the same number or compaction is gated on + /// fiction. + pub fn with_local_num_ctx(mut self, num_ctx: u32) -> Self { + let merged = merge_provider_options( + &self.default_provider_options, + &json!({ "options": { "num_ctx": num_ctx } }), + ); + self.default_provider_options = merged; + self.profile.max_input_tokens = Some(u64::from(num_ctx)); + self + } + + /// The local runtime this handle points at, or `None` for a hosted endpoint. + pub fn local_runtime_kind(&self) -> Option { + self.local_runtime + } + /// Overrides the API base URL. A trailing slash is trimmed so the joined /// endpoint is always `{base_url}/chat/completions`. pub fn with_base_url(mut self, base_url: impl Into) -> Self { @@ -737,15 +926,16 @@ impl OpenAiModel { // Authorization header, a base URL normalised to the `/v1` root, and // the request-shape degradations these servers require. Ollama and LM // Studio differ only in their default port. - if let Some(default_root) = local_runtime_default_root(&spec.kind) { + if let Some(kind) = local_runtime_kind(&spec.kind) { let auth = if spec.requires_api_key { AuthStyle::Bearer } else { AuthStyle::None }; return Ok(Self::local_runtime( + kind, &spec.provider, - normalize_local_v1_base_url(spec.base_url, default_root)?, + normalize_local_v1_base_url(spec.base_url, kind.default_root())?, api_key, spec.model, ) @@ -811,6 +1001,280 @@ impl OpenAiModel { Ok(listing.data) } + // ----------------------------------------------------------------------- + // Local-runtime probing and warm-up (C10 / C11 / C12 / C14) + // ----------------------------------------------------------------------- + + /// Asks a **live local server** what the loaded model can actually do. + /// + /// Returns the raw [`LocalProbe`] without touching this handle; use + /// [`Self::apply_local_probe`] to fold it into the profile, or + /// [`Self::probed`] to do both in one step. + /// + /// # Why this exists + /// + /// [`list_models`](Self::list_models) has existed for a while and is called + /// from nowhere in this crate, and [`ModelListing`] discards everything but + /// `id`/`created`/`owned_by` — so the two facts that matter most about a + /// local model were unreachable. This reaches them: + /// + /// | Runtime | Endpoint | Yields | + /// |---|---|---| + /// | Ollama | `POST {root}/api/show` | `model_info.*.context_length`, `capabilities: [tools, vision, thinking]` | + /// | LM Studio | `GET {root}/api/v0/models` | `max_context_length`, `loaded_context_length`, `type` | + /// | others | — | nothing; returns an empty probe rather than an error | + /// + /// This is the root fix for three separate defects: the invented context + /// window, the unconditionally-disabled native tools, and the "only two + /// runtimes get local treatment" gap. + /// + /// # Opt-in on purpose + /// + /// It is **never** called during construction. It costs a network round + /// trip, and a constructor that blocks on one is unusable where this crate + /// is embedded. Call it once at startup (or lazily on first use) and cache + /// the result on the handle. + /// + /// # Errors + /// + /// [`TinyAgentsError::Validation`] when this handle is not a local runtime, + /// and [`TinyAgentsError::Model`] on transport failure or an undecodable + /// body. A **non-2xx** status is not an error: it yields an empty probe, so + /// an older server without the endpoint degrades to "learned nothing" + /// rather than failing the caller's startup. + pub async fn probe_local_profile(&self) -> Result { + let Some(kind) = self.local_runtime else { + return Err(TinyAgentsError::Validation(format!( + "probe_local_profile is only meaningful for a local runtime; \ + `{}` at {} is not one", + self.provider, self.base_url + ))); + }; + let root = kind.native_root(&self.base_url); + let (endpoint, builder) = match kind { + LocalRuntimeKind::Ollama => { + let url = format!("{root}/api/show"); + let builder = self + .authorized(self.client.post(&url)) + .json(&ollama_show_body(&self.model)); + (url, builder) + } + LocalRuntimeKind::LmStudio => { + let url = format!("{root}/api/v0/models"); + let builder = self.authorized(self.client.get(&url)); + (url, builder) + } + // llama.cpp-server and vLLM expose no richer metadata endpoint this + // crate can rely on. They still get every other piece of local + // treatment; they simply learn nothing here. + LocalRuntimeKind::LlamaCpp | LocalRuntimeKind::Vllm => { + tracing::debug!( + provider = %self.provider, + kind = kind.as_str(), + "[openai] no probe endpoint for this local runtime; returning an empty probe" + ); + return Ok(LocalProbe::default()); + } + }; + + tracing::debug!( + provider = %self.provider, + model = %self.model, + endpoint = %endpoint, + "[openai] probing local runtime capabilities" + ); + let response = builder + .timeout(PROBE_TIMEOUT) + .send() + .await + .map_err(|e| probe_error(&endpoint, e))?; + if !response.status().is_success() { + tracing::debug!( + provider = %self.provider, + endpoint = %endpoint, + status = response.status().as_u16(), + "[openai] local probe endpoint unavailable; continuing without a probe" + ); + return Ok(LocalProbe::default()); + } + let body: Value = response + .json() + .await + .map_err(|e| probe_error(&endpoint, e))?; + + let probe = match kind { + LocalRuntimeKind::Ollama => probe_from_ollama_show(&body), + LocalRuntimeKind::LmStudio => probe_from_lm_studio_models(&body, &self.model), + _ => LocalProbe::default(), + }; + tracing::info!( + provider = %self.provider, + model = %self.model, + max_input_tokens = ?probe.effective_context_window(), + tool_calling = ?probe.tool_calling, + vision = ?probe.vision, + "[openai] local probe complete" + ); + Ok(probe) + } + + /// Folds a [`LocalProbe`] into this handle's profile and transport knobs. + /// + /// Only fields the probe actually learned are written; a `None` leaves the + /// current value alone, so a server that reports nothing cannot erase a + /// caller's explicit configuration. Pure and synchronous, so the + /// probe-to-profile policy is unit-testable without a live server. + pub fn apply_local_probe(mut self, probe: &LocalProbe) -> Self { + if let Some(window) = probe.effective_context_window() { + self.profile.max_input_tokens = Some(window); + } + if let Some(tool_calling) = probe.tool_calling { + self.profile.tool_calling = tool_calling; + self.profile.parallel_tool_calls = tool_calling; + self.profile.streaming_tool_chunks = tool_calling; + self.native_tools_on_wire + .store(tool_calling, Ordering::Relaxed); + } + if let Some(vision) = probe.vision { + self.profile.modalities.image_in = vision; + } + if let Some(reasoning) = probe.reasoning { + self.profile.reasoning = reasoning; + } + self + } + + /// [`probe_local_profile`](Self::probe_local_profile) + + /// [`apply_local_probe`](Self::apply_local_probe) in one await, for the + /// common startup shape: + /// + /// ```no_run + /// # use tinyagents::harness::providers::openai::OpenAiModel; + /// # async fn f() -> tinyagents::Result<()> { + /// let model = OpenAiModel::ollama().with_model("llama3.2:3b").probed().await?; + /// # Ok(()) } + /// ``` + pub async fn probed(self) -> Result { + let probe = self.probe_local_profile().await?; + Ok(self.apply_local_probe(&probe)) + } + + /// Verifies the configured model is actually served, with a remediation + /// naming the fix. + /// + /// The embeddings adapter has always done this ("Run `ollama pull {model}` + /// or choose an installed embedding model"); the chat path surfaced an + /// opaque 404 through [`Self::parse_error_body`]. LangChain gates the same + /// check behind an opt-in flag, and so does this: it costs a round trip, so + /// it is a method you call rather than something construction does to you. + /// + /// # Errors + /// + /// [`TinyAgentsError::Validation`] when the model is not among the served + /// ids, listing what *is* available. Transport failures surface unchanged + /// from [`Self::list_models`]. + pub async fn validate_model(&self) -> Result<()> { + let listed = self.list_models().await?; + if listed.is_empty() { + // An empty listing is "this server does not enumerate", not + // "nothing is served". Refusing here would be a false negative. + tracing::debug!( + provider = %self.provider, + "[openai] model listing is empty; skipping model validation" + ); + return Ok(()); + } + if listed.iter().any(|entry| entry.id == self.model) { + return Ok(()); + } + let mut available: Vec<&str> = listed.iter().map(|entry| entry.id.as_str()).collect(); + available.sort_unstable(); + let remediation = match self.local_runtime { + Some(LocalRuntimeKind::Ollama) => { + format!(" Run `ollama pull {}` to install it.", self.model) + } + _ => String::new(), + }; + Err(TinyAgentsError::Validation(format!( + "{} at {} does not serve model `{}`.{} Available: {}", + self.provider, + self.base_url, + self.model, + remediation, + available.join(", ") + ))) + } + + /// Loads the model with this handle's local options and holds it resident. + /// + /// This is the **only** path on which `num_ctx` and `keep_alive` reach the + /// server. Both are `/api/chat` fields; the chat adapter speaks + /// `POST {base_url}/chat/completions`, where Ollama's compatibility layer + /// simply ignores them — so + /// [`with_default_provider_options`](Self::with_default_provider_options) + /// documenting `{"options": {"num_ctx": 8192}}` as the local escape hatch + /// described a field that, on that path, went nowhere. (The existing tests + /// asserted only that the request JSON *contained* it, never that a server + /// honoured it, which is how that survived.) + /// + /// Call it once before the first real turn. It doubles as the warm-up that + /// keeps Ollama from unloading after its 5-minute idle default and charging + /// the next turn a cold multi-gigabyte load inside the 600 s unary deadline. + /// + /// **Scope, stated plainly:** this configures the loaded *runner*, which + /// Ollama then reuses for subsequent `/v1` requests that do not demand + /// conflicting options. That is a property of the server's runner reuse, not + /// a guarantee of the OpenAI wire format. A full native `/api/chat` chat + /// adapter is the complete fix and remains a follow-up. + /// + /// A no-op (returning `Ok(())`) for a runtime with no native API. + /// + /// # Errors + /// + /// [`TinyAgentsError::Model`] on transport failure. A non-2xx status is + /// logged and swallowed: a warm-up that the server declined must not fail + /// the caller's startup. + pub async fn warm_up(&self) -> Result<()> { + let Some(kind) = self.local_runtime.filter(|k| k.has_native_api()) else { + return Ok(()); + }; + let url = format!("{}/api/chat", kind.native_root(&self.base_url)); + let body = ollama_load_body( + &self.model, + local_options_object(&self.default_provider_options), + self.keep_alive.as_deref(), + ); + tracing::debug!( + provider = %self.provider, + model = %self.model, + url = %url, + num_ctx = ?body.pointer("/options/num_ctx"), + keep_alive = ?self.keep_alive, + "[openai] warming up local runtime over the native API" + ); + let response = self + .authorized(self.client.post(&url)) + .json(&body) + .timeout( + self.effective_request_timeout(None, false) + .unwrap_or(Duration::from_secs(DEFAULT_REQUEST_TIMEOUT_SECS)), + ) + .send() + .await + .map_err(|e| { + TinyAgentsError::Model(format!("[openai] warm-up of {url} failed: {e}")) + })?; + if !response.status().is_success() { + tracing::warn!( + provider = %self.provider, + url = %url, + status = response.status().as_u16(), + "[openai] local warm-up declined; continuing without it" + ); + } + Ok(()) + } + // ----------------------------------------------------------------------- // OpenAI-compatible provider presets // @@ -928,14 +1392,61 @@ impl OpenAiModel { /// An Ollama server exposed through its OpenAI-compatible HTTP API. pub fn ollama_at(base_url: impl Into, model: impl Into) -> Result { + let kind = LocalRuntimeKind::Ollama; Ok(Self::local_runtime( + kind, "ollama", - normalize_local_v1_base_url(base_url.into(), "http://localhost:11434")?, + normalize_local_v1_base_url(base_url.into(), kind.default_root())?, "", model, )) } + /// A llama.cpp `llama-server` exposed through its OpenAI-compatible HTTP + /// API (default root `http://localhost:8080`). + /// + /// llama.cpp-server used to fall through to the hosted `Compatible` path, + /// which gave it Bearer auth, a full hosted capability profile, no `/v1` + /// normalisation, and none of the request-shape degrade knobs — despite + /// being the runtime the `json_object` and named-`tool_choice` degrades were + /// written for. + pub fn llama_cpp(base_url: impl Into, model: impl Into) -> Result { + let kind = LocalRuntimeKind::LlamaCpp; + Ok(Self::local_runtime( + kind, + kind.as_str(), + normalize_local_v1_base_url(base_url.into(), kind.default_root())?, + "", + model, + )) + } + + /// A vLLM OpenAI-compatible server (default root `http://localhost:8000`). + /// + /// Same rationale as [`Self::llama_cpp`]: a self-hosted server is a local + /// runtime whatever its name, and gets the local treatment. + pub fn vllm( + base_url: impl Into, + api_key: impl Into, + model: impl Into, + ) -> Result { + let kind = LocalRuntimeKind::Vllm; + let api_key = api_key.into(); + let auth = if api_key.trim().is_empty() { + AuthStyle::None + } else { + AuthStyle::Bearer + }; + Ok(Self::local_runtime( + kind, + kind.as_str(), + normalize_local_v1_base_url(base_url.into(), kind.default_root())?, + api_key, + model, + ) + .with_auth_style(auth)) + } + /// An LM Studio server exposed through its OpenAI-compatible HTTP API. /// /// Authentication is disabled when `api_key` is empty and uses a bearer @@ -951,31 +1462,63 @@ impl OpenAiModel { } else { AuthStyle::Bearer }; + let kind = LocalRuntimeKind::LmStudio; Ok(Self::local_runtime( - "lm_studio", - normalize_local_v1_base_url(base_url.into(), "http://localhost:1234")?, + kind, + kind.as_str(), + normalize_local_v1_base_url(base_url.into(), kind.default_root())?, api_key, model, ) .with_auth_style(auth)) } + /// The shared local-runtime preset. + /// + /// Three things changed here, each fixing a defect the old preset baked in: + /// + /// * **`tool_calling` stays `true`.** It used to be hard-disabled for every + /// local runtime unconditionally, which (a) forced the prompt-guided + /// branch on every call — injecting the protocol block plus every tool's + /// serialized JSON Schema into the system prompt, against a real 2048-token + /// window, which then truncated from the front and dropped the very prompt + /// carrying the protocol — and (b) cleared `parallel_tool_calls` and + /// `streaming_tool_chunks`, so any `CapabilitySet { tool_calling: true }` + /// excluded **every** local model from resolution. LangChain's ChatOllama + /// sends tools natively with no capability check at all. The pessimism is + /// replaced by discovery: [`Self::probe_local_profile`] asks the server, + /// and a 400 implicating `tools` latches the transport onto the + /// prompt-guided branch for the rest of the process. + /// * **The window is `None`, not a guess.** See [`derive_profile_for`]. + /// * **The degrade knobs are pre-set.** `named_tool_choice` and + /// `json_object` are the two shapes the module doc and README name as + /// "local servers reject this"; auto-degrade recovered from them, but only + /// after one wasted 400 per process — precisely the cost the + /// `stream_required` latch exists to eliminate. `strict` JSON Schema is + /// off for the same reason. fn local_runtime( + kind: LocalRuntimeKind, provider: &str, base_url: String, api_key: impl Into, model: impl Into, ) -> Self { - Self::compatible_provider(provider, api_key, base_url, model) + let mut model = Self::compatible_provider(provider, api_key, base_url, model) .with_auth_style(AuthStyle::None) - .with_native_tool_calling(false) .with_vision(false) - .lock_local_capabilities() - } - - fn lock_local_capabilities(mut self) -> Self { - self.local_capabilities_locked = true; - self + .with_named_tool_choice(false) + .with_json_object_format(false) + .with_strict_json_schema(false); + model.local_runtime = Some(kind); + // Re-derive **before** locking. The lock preserves the profile values it + // finds, and at this point they are still the hosted ones the + // `compatible_provider` call above produced — locking first would + // preserve exactly the invented context window this preset exists to + // remove (`llama3.2` → 128 000 from the generic hint table). + model.rederive_profile(); + model.profile.modalities.image_in = false; + model.local_capabilities_locked = true; + model } #[cfg(test)] @@ -1005,6 +1548,8 @@ impl OpenAiModel { Degrade { named_tool_choice: !self.named_tool_choice_supported.load(Ordering::Relaxed), json_object: !self.json_object_format_supported.load(Ordering::Relaxed), + json_schema_strict: !self.json_schema_strict.load(Ordering::Relaxed), + native_tools: !self.native_tools_enabled(), } } @@ -1034,7 +1579,12 @@ impl OpenAiModel { // handed tools gets the tool protocol embedded in its system prompt and no // native `tools` on the wire (many local runtimes 400 on `tools`). The // model's `` blocks are parsed back in [`Self::invoke`]/stream. - let prompt_guided_tools = !self.profile.tool_calling && !request.tools.is_empty(); + // Native tools go on the wire unless the profile says the model has none + // *or* this attempt is degrading them away. `degrade.native_tools` + // carries both the instance latch (via `baseline_degrade`) and a + // freshly-discovered 400, so one condition covers both. + let native_tools = !degrade.native_tools; + let prompt_guided_tools = !native_tools && !request.tools.is_empty(); let prompt_messages; let instructed_messages; let coalesced_messages: &[Message] = if prompt_guided_tools { @@ -1056,7 +1606,7 @@ impl OpenAiModel { // applied after the coalescing above so folded tool results are already // in their final user-turn shape when we look for a real query. let user_normalized_messages; - let base_messages: &[Message] = if self.profile.tool_calling { + let base_messages: &[Message] = if native_tools { coalesced_messages } else { user_normalized_messages = @@ -1078,18 +1628,31 @@ impl OpenAiModel { .map(translate_message) .collect::>>()?; + // Project the declarations through the shared preparation seam rather + // than shipping `schema.parameters` verbatim: a tool whose `parameters` + // is `Value::Null` (the type permits it) otherwise serialises as + // `"parameters": null`, which every provider 400s on, and local `$ref`s + // are resolved for routes that cannot follow them. + // + // Deliberately **not** `.with_strict()`. The strict sanitizer forces + // every declared property into `required`, which changes the contract + // the model is given — an optional argument becomes mandatory. That is + // only correct when the wire actually carries `strict: true`, and this + // adapter does not set it on the tool object (it is a `response_format` + // concern here). Applying it anyway would silently make every optional + // tool argument required. + let preparation = crate::harness::tool::SchemaPreparation::openai(); let mut tools: Vec = if prompt_guided_tools { Vec::new() } else { - request - .tools - .iter() + crate::harness::tool::prepare_tool_schemas(&request.tools, &preparation) + .into_iter() .map(|schema| ToolWire { kind: "function".to_string(), function: FunctionSchemaWire { - name: schema.name.clone(), - description: schema.description.clone(), - parameters: schema.parameters.clone(), + name: schema.name, + description: schema.description, + parameters: schema.parameters, }, }) .collect() @@ -1121,7 +1684,7 @@ impl OpenAiModel { // `json_schema` that still guarantees a JSON object. Some(degraded_json_object_format()) } else { - translate_response_format(format) + translate_response_format(format, !degrade.json_schema_strict) } }); @@ -1143,6 +1706,24 @@ impl OpenAiModel { &self.temperature_unsupported, ); + // Provider options are the escape hatch and win on key conflicts, so a + // caller who spelled `reasoning_effort` there suppresses the typed + // field entirely — emitting both would put the key on the wire twice. + let merged_provider_options = + merge_provider_options(&self.default_provider_options, &request.provider_options); + let reasoning_effort = if merged_provider_options + .get("reasoning_effort") + .is_some_and(|v| !v.is_null()) + { + None + } else { + request + .reasoning + .as_ref() + .and_then(|config| config.effort) + .map(|effort| effort.as_str().to_string()) + }; + Ok(ChatCompletionRequest { model, messages, @@ -1155,12 +1736,10 @@ impl OpenAiModel { max_completion_tokens, stop: request.stop_sequences.clone(), seed: request.seed, + reasoning_effort, stream: false, stream_options: None, - extra: provider_extra_options(&merge_provider_options( - &self.default_provider_options, - &request.provider_options, - ))?, + extra: provider_extra_options(&merged_provider_options)?, }) } @@ -1218,6 +1797,12 @@ impl OpenAiModel { } /// Builds the `/v1/responses` request body from a provider-neutral request. + /// + /// Carries the **whole** request. The previous version built only + /// `{model, input, instructions, stream, store, max_output_tokens}` and + /// dropped tools, tool choice, response format, sampling, stop sequences, + /// seed, the continuation id, and `provider_options` — which is how + /// `reasoning` became unreachable on the one wire format that supports it. fn translate_responses_request(&self, request: &ModelRequest) -> responses::ResponsesRequest { let model = request.model.clone().unwrap_or_else(|| self.model.clone()); let (instructions, input) = responses::build_responses_input(&request.messages); @@ -1226,6 +1811,45 @@ impl OpenAiModel { } else { request.max_tokens }; + let strict = self.json_schema_strict.load(Ordering::Relaxed); + + let extra = provider_extra_options(&merge_provider_options( + &self.default_provider_options, + &request.provider_options, + )) + .unwrap_or_default(); + // Same precedence rule as the Chat Completions path: a `reasoning` key + // in `provider_options` is the escape hatch and wins, so the typed + // field stands down rather than emitting the key twice. + let reasoning = if extra.contains_key("reasoning") { + None + } else { + request + .reasoning + .as_ref() + .and_then(responses::translate_reasoning) + }; + // Reasoning replay is only possible under `store: false` when the + // reasoning item carries `encrypted_content`, and that only arrives when + // it is explicitly requested. Asking for reasoning without asking for + // this yields reasoning that cannot survive to the next turn. + let include = if reasoning.is_some() { + vec![responses::INCLUDE_ENCRYPTED_REASONING.to_string()] + } else { + Vec::new() + }; + + let tools = if self.native_tools_enabled() { + request + .tools + .iter() + .map(responses::translate_tool) + .collect() + } else { + Vec::new() + }; + let tool_choice = (!tools.is_empty()).then(|| translate_tool_choice(&request.tool_choice)); + responses::ResponsesRequest { model, input, @@ -1233,6 +1857,28 @@ impl OpenAiModel { stream: None, store: Some(false), max_output_tokens, + tools, + tool_choice, + text: request + .response_format + .as_ref() + .and_then(|format| responses::translate_text_format(format, strict)), + temperature: effective_temperature( + &self.model, + request.temperature, + self.temperature_override, + &self.temperature_unsupported, + ), + top_p: request.top_p, + seed: request.seed, + stop: request.stop_sequences.clone(), + previous_response_id: request.continuation_id.clone(), + reasoning, + include, + // `provider_options` is the escape hatch and wins on key conflicts, + // exactly as `merge_provider_options` already arranges for the Chat + // Completions path. It was dropped entirely here. + extra, } } @@ -1304,9 +1950,26 @@ impl OpenAiModel { let status = response.status(); if !status.is_success() { + // Read `Retry-After` *before* the body is consumed: a 429/503 that + // names how long to wait is authoritative, and retrying sooner + // burns an attempt for certain. + let retry_after_ms = response + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|value| value.to_str().ok()) + .and_then(parse_retry_after_header_ms); let text = response.text().await.unwrap_or_default(); - let error = self.parse_error_body(status.as_u16(), &text); - return Err(TinyAgentsError::Provider(Box::new(error))); + let mut error = self.parse_error_body(status.as_u16(), &text); + if let Some(ms) = retry_after_ms { + tracing::debug!( + provider = %self.provider, + status = status.as_u16(), + retry_after_ms = ms, + "[openai] honoring a server-supplied Retry-After header" + ); + error.retry_after_ms = Some(ms); + } + return Err(TinyAgentsError::from_provider_error(error)); } Ok(response) } @@ -1415,6 +2078,7 @@ impl OpenAiModel { code, message, retryable, + retry_after_ms: None, raw, } } @@ -1436,6 +2100,19 @@ impl OpenAiModel { ) } + /// Decodes a non-2xx body into a structured [`ProviderError`]. + /// + /// Two classifications are applied on top of the raw decode: + /// + /// * **Local missing-model 404s** are rewritten with a remediation naming + /// the fix (`ollama pull …`), matching what the embeddings adapter has + /// always done. The chat path used to surface the server's bare + /// `{"error":"model 'x' not found"}`. + /// * **Context overflow** is stamped with a stable + /// [`CONTEXT_OVERFLOW_CODE`], so a caller can compact and retry on a code + /// instead of string-matching a provider message. A typed + /// `TinyAgentsError` variant would live in `src/error.rs`; promoting the + /// code to one is a follow-up. pub(super) fn parse_error_body(&self, status: u16, text: &str) -> ProviderError { let raw = serde_json::from_str::(text).ok(); let error_obj = raw.as_ref().and_then(|value| value.get("error")); @@ -1450,27 +2127,102 @@ impl OpenAiModel { .filter(|message| !message.trim().is_empty()) .unwrap_or(text) .to_string(); - let code = error_obj + let mut code = error_obj .and_then(|error| error.get("code").or_else(|| error.get("type"))) .and_then(Value::as_str) .map(str::to_string); + + let message = match self.local_runtime.and_then(|kind| { + missing_model_remediation(kind, status, text, &self.model, &self.base_url) + }) { + Some(remediated) => { + tracing::info!( + provider = %self.provider, + model = %self.model, + "[openai] rewrote a local missing-model 404 with remediation" + ); + remediated + } + None => message, + }; + + if is_context_overflow(status, &message) { + tracing::warn!( + provider = %self.provider, + model = %self.model, + max_input_tokens = ?self.profile.max_input_tokens, + "[openai] classified failure as a context overflow" + ); + code = Some(CONTEXT_OVERFLOW_CODE.to_string()); + } + self.provider_error(message, Some(status), code, raw) } } -/// The server root a local-runtime provider falls back to when its spec carries -/// a blank `base_url`, or `None` for providers that are not local runtimes. +/// Parses an HTTP `Retry-After` header value into a delay in milliseconds. +/// +/// [RFC 9110 §10.2.3] defines two forms and both appear in the wild: +/// +/// * **delta-seconds** — `Retry-After: 30`. Used by OpenAI, Anthropic, and most +/// gateways on a 429. +/// * **HTTP-date** — `Retry-After: Wed, 21 Oct 2015 07:28:00 GMT`. Used by some +/// CDNs and proxies fronting a provider, and by 503 maintenance responses. +/// Converted to a delay relative to *now*; a date already in the past yields +/// `0` (retry immediately), never a negative or wrapped value. +/// +/// Returns `None` for an unparseable value rather than guessing, so the caller +/// falls through to the existing exponential backoff. +/// +/// [RFC 9110 §10.2.3]: https://www.rfc-editor.org/rfc/rfc9110#field.retry-after +pub(super) fn parse_retry_after_header_ms(raw: &str) -> Option { + let value = raw.trim(); + if value.is_empty() { + return None; + } + // delta-seconds. Fractional seconds are not in the grammar but some + // providers send them, so accept a float too. + if let Ok(seconds) = value.parse::() { + return Some(seconds.saturating_mul(1000)); + } + if let Ok(seconds) = value.parse::() + && seconds.is_finite() + && seconds >= 0.0 + { + return Some((seconds * 1000.0) as u64); + } + // HTTP-date. `parse_from_rfc2822` covers the IMF-fixdate form (including + // the obsolete `GMT` zone token); the explicit format is the fallback for + // servers that omit the day-of-week comma. + let parsed = chrono::DateTime::parse_from_rfc2822(value) + .map(|dt| dt.with_timezone(&chrono::Utc)) + .or_else(|_| { + chrono::NaiveDateTime::parse_from_str(value, "%a, %d %b %Y %H:%M:%S GMT") + .map(|naive| naive.and_utc()) + }) + .ok()?; + let delta = parsed.signed_duration_since(chrono::Utc::now()); + Some(delta.num_milliseconds().max(0) as u64) +} + +/// Maps a [`ProviderKind`] onto the [`LocalRuntimeKind`] it denotes, or `None` +/// for a hosted provider. /// -/// This is the single place that decides "is this kind a local runtime?", so a -/// new local provider is one arm here rather than a condition to keep in sync -/// across the transport. -fn local_runtime_default_root( +/// The single place that decides "is this kind a local runtime?", so a new local +/// provider is one arm here rather than a condition to keep in sync across the +/// transport. It used to recognise only Ollama and LM Studio, leaving +/// llama.cpp-server and vLLM to be constructed as [`ProviderKind::Compatible`] +/// — `requires_api_key: true`, Bearer auth, a full hosted profile, no `/v1` +/// normalisation, and no degrade knobs. +pub(super) fn local_runtime_kind( kind: &crate::harness::providers::ProviderKind, -) -> Option<&'static str> { +) -> Option { use crate::harness::providers::ProviderKind; match kind { - ProviderKind::Ollama => Some("http://localhost:11434"), - ProviderKind::LmStudio => Some("http://localhost:1234"), + ProviderKind::Ollama => Some(LocalRuntimeKind::Ollama), + ProviderKind::LmStudio => Some(LocalRuntimeKind::LmStudio), + ProviderKind::LlamaCpp => Some(LocalRuntimeKind::LlamaCpp), + ProviderKind::Vllm => Some(LocalRuntimeKind::Vllm), _ => None, } } @@ -1527,6 +2279,22 @@ pub(super) struct Degrade { /// Degrade `response_format:{"type":"json_object"}` to a permissive /// `json_schema`. pub json_object: bool, + /// Degrade a `json_schema` `response_format` from `strict: true` to + /// `strict: false`. + /// + /// A `JsonSchema` request had **no** degradation path at all: only + /// `JsonObject` was considered, so a 400 on a strict-schema request was + /// terminal even though the retry that fixes it is one boolean away. + pub json_schema_strict: bool, + /// Drop native `tools` from the wire and embed the tool protocol in the + /// prompt instead. + /// + /// The third latch, and the one that replaces "every local runtime is + /// assumed to have no tool support, forever". Seeded from the probe (or + /// optimistically `false` = "send tools"), flipped by a 400 implicating + /// `tools`, and retried once prompt-guided. Follows the `stream_required` + /// latch pattern exactly. + pub native_tools: bool, } /// Statuses an OpenAI-compatible proxy uses to reject a non-streaming request. @@ -1626,6 +2394,21 @@ pub(super) fn degrade_for_400( { degrade.json_object = true; } + if !already.json_schema_strict + && (lower.contains("response_format") + || lower.contains("json_schema") + || lower.contains("strict") + || lower.contains("additionalproperties")) + && matches!( + request.response_format, + Some(ResponseFormat::JsonSchema { .. } | ResponseFormat::Auto { .. }) + ) + { + degrade.json_schema_strict = true; + } + if !already.native_tools && !request.tools.is_empty() && mentions_tools_unsupported(message) { + degrade.native_tools = true; + } (degrade != already).then_some(degrade) } @@ -1739,6 +2522,41 @@ impl ChatModel for OpenAiModel { Some(&self.profile) } + /// Identifies this endpoint for the response cache key. + /// + /// Covers everything that makes two `OpenAiModel`s answer the same prompt + /// differently: the provider family, the model id, the base URL (a local + /// Ollama and hosted OpenAI are two different answers to one question), and + /// a **fingerprint** of the API credential — two keys can address two + /// tenants or two fine-tunes behind one base URL. + /// + /// The raw credential never appears: it goes through + /// [`credential_fingerprint`][crate::harness::cache::credential_fingerprint] + /// first, because this string is folded into keys that reach logs, events, + /// and durable cache files. + fn cache_identity(&self) -> Option { + let endpoint = crate::harness::cache::model_cache_identity( + &self.provider, + &self.model, + &self.base_url, + self.responses_api_primary.then_some("responses"), + &self.api_key, + ); + let configuration = serde_json::json!({ + "default_provider_options": self.default_provider_options, + "temperature_override": self.temperature_override, + "merge_system_into_user": self.merge_system_into_user, + "responses_omit_max_output_tokens": self.responses_omit_max_output_tokens, + "requires_streaming": self.requires_streaming(), + "named_tool_choice_supported": self.named_tool_choice_supported.load(Ordering::Relaxed), + "json_object_format_supported": self.json_object_format_supported.load(Ordering::Relaxed), + }); + Some(format!( + "{endpoint}|{}", + serde_json::to_string(&configuration).expect("model cache configuration serializes") + )) + } + /// Invokes the OpenAI Chat Completions endpoint and maps the response into a /// [`ModelResponse`]. /// @@ -1800,14 +2618,18 @@ impl ChatModel for OpenAiModel { })?; let value: Value = serde_json::from_str(&text)?; - let response = parse_chat_response(value, self.effective_reasoning_tags())?; + let response = parse_chat_response( + value, + self.effective_reasoning_tags(), + self.cache_accounting, + )?; // Recover the model's `` blocks into `message.tool_calls` for // prompt-guided models, and as a fallback for native models that emitted // the call as text despite being flagged native (empty structured // `tool_calls`). Native responses that already carry structured calls are // returned unchanged. if crate::harness::tool::should_recover( - self.profile.tool_calling, + self.native_tools_enabled(), !request.tools.is_empty(), response.message.tool_calls.len(), ) { @@ -1875,9 +2697,13 @@ impl ChatModel for OpenAiModel { TinyAgentsError::Model(format!("openai non-stream stream-body read failed: {e}")) })?; let value: Value = serde_json::from_str(&text)?; - let mut parsed = parse_chat_response(value, self.effective_reasoning_tags())?; + let mut parsed = parse_chat_response( + value, + self.effective_reasoning_tags(), + self.cache_accounting, + )?; if crate::harness::tool::should_recover( - self.profile.tool_calling, + self.native_tools_enabled(), !request.tools.is_empty(), parsed.message.tool_calls.len(), ) { @@ -1908,7 +2734,10 @@ impl ChatModel for OpenAiModel { bytes: Box::pin(bytes), buf: Vec::new(), pending: VecDeque::new(), - acc: OpenAiStreamAcc::new(self.effective_reasoning_tags().cloned()), + acc: OpenAiStreamAcc::new( + self.effective_reasoning_tags().cloned(), + self.cache_accounting, + ), provider: self.provider.clone(), model: self.model.clone(), started: false, @@ -1933,7 +2762,7 @@ impl ChatModel for OpenAiModel { if request.tools.is_empty() { return Ok(Box::pin(stream)); } - let native = self.profile.tool_calling; + let native = self.native_tools_enabled(); let mut scrubber = crate::harness::tool::ToolCallStreamScrubber::new(); Ok(Box::pin(stream.flat_map(move |item| { futures::stream::iter(clean_stream_item(item, &mut scrubber, native)) @@ -2023,3 +2852,42 @@ async fn invoke_with_streaming( acc.finish() } + +#[cfg(test)] +mod retry_after_header_tests { + use super::parse_retry_after_header_ms; + + #[test] + fn parses_delta_seconds() { + assert_eq!(parse_retry_after_header_ms("30"), Some(30_000)); + assert_eq!(parse_retry_after_header_ms(" 0 "), Some(0)); + // Not in the grammar, but some providers send it. + assert_eq!(parse_retry_after_header_ms("1.5"), Some(1_500)); + } + + #[test] + fn parses_the_http_date_form_relative_to_now() { + let future = chrono::Utc::now() + chrono::Duration::seconds(120); + let header = future.format("%a, %d %b %Y %H:%M:%S GMT").to_string(); + let parsed = parse_retry_after_header_ms(&header).expect("an HTTP-date parses"); + // Allow a generous window for clock/second-truncation slack. + assert!( + (110_000..=121_000).contains(&parsed), + "expected roughly 120s, got {parsed}ms from {header}" + ); + } + + #[test] + fn a_past_http_date_means_retry_immediately_not_a_wrapped_value() { + let past = chrono::Utc::now() - chrono::Duration::seconds(600); + let header = past.format("%a, %d %b %Y %H:%M:%S GMT").to_string(); + assert_eq!(parse_retry_after_header_ms(&header), Some(0)); + } + + #[test] + fn rejects_garbage_rather_than_guessing() { + assert_eq!(parse_retry_after_header_ms(""), None); + assert_eq!(parse_retry_after_header_ms("soon"), None); + assert_eq!(parse_retry_after_header_ms("-5"), None); + } +} diff --git a/src/harness/providers/openai/types.rs b/src/harness/providers/openai/types.rs index 7dd2563..da0db37 100644 --- a/src/harness/providers/openai/types.rs +++ b/src/harness/providers/openai/types.rs @@ -50,6 +50,16 @@ pub struct ChatCompletionRequest { /// Deterministic generation seed. Omitted when unset. #[serde(skip_serializing_if = "Option::is_none")] pub seed: Option, + /// `reasoning_effort` — the Chat Completions spelling of the + /// provider-neutral [`ReasoningConfig`][rc]. Omitted when unset. + /// + /// Before this field the only route to it was raw `provider_options`, which + /// is provider-shaped by definition; it stays available as the escape hatch + /// and **wins** over this typed field on a key conflict. + /// + /// [rc]: crate::harness::model::ReasoningConfig + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_effort: Option, /// Request Server-Sent-Events streaming. Omitted (false) for unary calls. #[serde(skip_serializing_if = "std::ops::Not::not")] pub stream: bool, @@ -406,6 +416,36 @@ pub struct PromptTokensDetailsWire { /// Input tokens served from OpenAI's prompt cache. #[serde(default, deserialize_with = "deserialize_null_default")] pub cached_tokens: u64, + /// Input tokens **written into** the prompt cache, as OpenAI-compatible + /// gateways report it. Two spellings are in the wild and neither is + /// universal, so both are accepted and reconciled by + /// [`Self::cache_creation_tokens`]: + /// + /// * `cache_write_tokens` — the spelling LangChain reads on the + /// Chat Completions path. + /// * `cache_creation_tokens` — the spelling gateways that mirror + /// Anthropic's `cache_creation_input_tokens` tend to use. + /// + /// Before this existed, [`Usage::cache_creation_tokens`][ucc] was summed and + /// priced by the cost feature but written by **no** provider in the crate, + /// so cache writes were billed as ordinary input everywhere. + /// + /// [ucc]: crate::harness::usage::Usage::cache_creation_tokens + #[serde(default, deserialize_with = "deserialize_null_default")] + pub cache_write_tokens: u64, + /// Alternate spelling of [`Self::cache_write_tokens`]; see its docs. + #[serde(default, deserialize_with = "deserialize_null_default")] + pub cache_creation_tokens: u64, +} + +impl PromptTokensDetailsWire { + /// The cache-write token count, taking whichever of the two accepted + /// spellings the provider actually sent (they are never both non-zero in + /// practice; the larger wins so a zero-valued alias cannot mask the real + /// figure). + pub fn cache_creation_tokens(&self) -> u64 { + self.cache_write_tokens.max(self.cache_creation_tokens) + } } /// The `completion_tokens_details` breakdown of a [`UsageWire`]. diff --git a/src/harness/providers/types.rs b/src/harness/providers/types.rs index c75604b..c90500b 100644 --- a/src/harness/providers/types.rs +++ b/src/harness/providers/types.rs @@ -36,6 +36,16 @@ pub enum ProviderKind { /// [`ProviderSpec::with_model`], or discover it at runtime with /// [`OpenAiModel::list_models`](crate::harness::providers::openai::OpenAiModel::list_models). LmStudio, + /// A local llama.cpp `llama-server` exposing `/v1/chat/completions`. + /// + /// Carries **no default model**, for the same reason as + /// [`ProviderKind::LmStudio`]: the id is whichever GGUF was loaded. + LlamaCpp, + /// A local vLLM OpenAI-compatible server. + /// + /// Carries no default model — vLLM serves whatever weights it was started + /// with, usually under the full HuggingFace repo id. + Vllm, /// DeepSeek OpenAI-compatible endpoint. DeepSeek, /// Groq OpenAI-compatible endpoint. @@ -60,6 +70,8 @@ impl ProviderKind { ProviderKind::Anthropic => "anthropic", ProviderKind::Ollama => "ollama", ProviderKind::LmStudio => "lmstudio", + ProviderKind::LlamaCpp => "llama_cpp", + ProviderKind::Vllm => "vllm", ProviderKind::DeepSeek => "deepseek", ProviderKind::Groq => "groq", ProviderKind::Xai => "xai", @@ -83,6 +95,10 @@ impl ProviderKind { "anthropic" => Some(ProviderKind::Anthropic), "ollama" => Some(ProviderKind::Ollama), "lmstudio" | "lm_studio" | "lm-studio" => Some(ProviderKind::LmStudio), + "llamacpp" | "llama_cpp" | "llama-cpp" | "llamaserver" => { + Some(ProviderKind::LlamaCpp) + } + "vllm" => Some(ProviderKind::Vllm), "deepseek" => Some(ProviderKind::DeepSeek), "groq" => Some(ProviderKind::Groq), "xai" => Some(ProviderKind::Xai), @@ -158,6 +174,11 @@ impl ProviderSpec { // `list_models`), which fails loudly at construction instead of // silently on the first request. ProviderKind::LmStudio => Self::new(kind, "", "http://localhost:1234/v1", None, false), + // Same "no default model" rule as LM Studio, and the same reason: + // the served id is whatever weights the operator started the server + // with. Both are local runtimes, so neither requires an API key. + ProviderKind::LlamaCpp => Self::new(kind, "", "http://localhost:8080/v1", None, false), + ProviderKind::Vllm => Self::new(kind, "", "http://localhost:8000/v1", None, false), ProviderKind::DeepSeek => Self::new( kind, "deepseek-chat", diff --git a/src/harness/retry/jitter.rs b/src/harness/retry/jitter.rs new file mode 100644 index 0000000..80d0abf --- /dev/null +++ b/src/harness/retry/jitter.rs @@ -0,0 +1,109 @@ +//! Minimal, dependency-free randomness for retry jitter. +//! +//! Jitter exists to break up thundering herds: without it every client that +//! failed at the same instant retries at the same instant. That needs *some* +//! randomness, but not cryptographic randomness and not a statistically +//! rigorous generator — a couple of milliseconds of spread is the whole +//! requirement. +//! +//! The crate deliberately does not take a `rand` / `fastrand` dependency for +//! this, so this module ships a ~30-line `xorshift64*` generator behind a +//! thread-local. It is: +//! +//! - **Never used for anything security-relevant.** Backoff spread only. +//! - **Seeded per thread** from the wall clock plus a process-global counter, so +//! two threads (and two processes started in the same nanosecond) do not walk +//! the same sequence. +//! - **Bypassable for tests.** Every consumer routes through +//! [`RetryPolicy::backoff_for_attempt_with`][crate::harness::retry::RetryPolicy::backoff_for_attempt_with], +//! which takes an explicit `rand01`, so no test ever has to observe this +//! module's output. + +use std::cell::Cell; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{SystemTime, UNIX_EPOCH}; + +/// Process-global seed disambiguator so two threads seeded within the same +/// clock tick still diverge. +static SEED_COUNTER: AtomicU64 = AtomicU64::new(0); + +thread_local! { + /// Per-thread generator state. Never zero (xorshift is degenerate at zero). + static STATE: Cell = Cell::new(seed()); +} + +/// Builds a non-zero seed from the wall clock and a global counter. +fn seed() -> u64 { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_nanos() as u64) + .unwrap_or(0x9E37_79B9_7F4A_7C15); + let bump = SEED_COUNTER.fetch_add(1, Ordering::Relaxed); + // Golden-ratio odd constant keeps the mix well-distributed for small bumps. + let mixed = nanos ^ bump.wrapping_mul(0x9E37_79B9_7F4A_7C15); + if mixed == 0 { + 0xDEAD_BEEF_CAFE_F00D + } else { + mixed + } +} + +/// Advances the thread-local generator and returns the raw 64-bit output. +fn next_u64() -> u64 { + STATE.with(|state| { + let mut x = state.get(); + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + state.set(x); + // xorshift64* output scrambler. + x.wrapping_mul(0x2545_F491_4F6C_DD1D) + }) +} + +/// Returns a pseudo-random `f64` uniformly distributed over `[0, 1)`. +/// +/// This is the value production retry paths feed to +/// [`RetryPolicy::backoff_for_attempt_with`][crate::harness::retry::RetryPolicy::backoff_for_attempt_with]. +pub(crate) fn rand01() -> f64 { + // 53 bits is the full mantissa of an f64, so this covers [0, 1) evenly. + ((next_u64() >> 11) as f64) / ((1u64 << 53) as f64) +} + +#[cfg(test)] +mod test { + use super::rand01; + + #[test] + fn rand01_stays_inside_the_unit_interval() { + for _ in 0..10_000 { + let value = rand01(); + assert!((0.0..1.0).contains(&value), "out of range: {value}"); + } + } + + #[test] + fn rand01_is_not_a_constant() { + let first = rand01(); + // A stuck generator (the bug this module exists to avoid) would return + // the same value forever. + assert!( + (0..64).any(|_| rand01() != first), + "generator produced a constant sequence" + ); + } + + #[test] + fn rand01_covers_both_halves_of_the_interval() { + let mut low = false; + let mut high = false; + for _ in 0..1_000 { + if rand01() < 0.5 { + low = true; + } else { + high = true; + } + } + assert!(low && high, "generator never crossed the midpoint"); + } +} diff --git a/src/harness/retry/mod.rs b/src/harness/retry/mod.rs index e23e40e..db835e3 100644 --- a/src/harness/retry/mod.rs +++ b/src/harness/retry/mod.rs @@ -22,6 +22,7 @@ //! `now: Instant` / `rand01: f64` so tests can drive time and randomness //! deterministically without injecting a clock trait. +mod jitter; mod types; pub use types::*; @@ -32,6 +33,17 @@ use std::time::{Duration, Instant}; use crate::error::TinyAgentsError; use crate::harness::model::ProviderError; +/// Fraction of the base backoff the additive jitter band spans in each +/// direction: with jitter enabled the effective delay lands uniformly in +/// `[base * (1 - JITTER_FRACTION), base * (1 + JITTER_FRACTION)]`. +/// +/// Matches LangChain's `_retry.py` (`delay ± 25%`). LangGraph instead adds a +/// flat `uniform(0, 1)` second; the proportional form generalizes better across +/// the sub-second and multi-second ends of the same schedule. What both +/// references share, and what matters here, is that jitter is **additive** — it +/// never scales the delay toward zero. +pub const JITTER_FRACTION: f64 = 0.25; + // ── Provider failure classification ───────────────────────────────────────── fn parse_status_at(text: &str, start: usize) -> Option { @@ -280,26 +292,117 @@ impl RetryPolicy { /// Enables or disables actually sleeping for the computed backoff between /// retries. /// - /// Off by default so tests stay deterministic and fast. Enable it in - /// production so a transient failure is retried after a real, growing delay - /// rather than back-to-back. See [`RetryPolicy::backoff_sleep`]. + /// **On by default.** Pass `false` to opt out — which tests that assert on + /// retry counts without wanting real elapsed time should do explicitly. See + /// [`RetryPolicy::backoff_sleep`] for why the default is on. pub fn with_backoff_sleep(mut self, sleep: bool) -> Self { self.backoff_sleep = sleep; self } + /// Sets the ceiling applied to a server-supplied `Retry-After` delay. See + /// [`RetryPolicy::max_retry_after_ms`]. + pub fn with_max_retry_after_ms(mut self, ms: u64) -> Self { + self.max_retry_after_ms = ms; + self + } + + /// Replaces the built-in [`is_retryable`] classification with `predicate`. + /// + /// The predicate decides *only* whether an error is transient; the attempt + /// cap still applies on top. Ported from LangGraph's + /// `RetryPolicy.retry_on`. + pub fn with_retry_on(mut self, predicate: RetryPredicate) -> Self { + self.retry_on = Some(predicate); + self + } + + /// Clears any custom [`RetryPolicy::retry_on`] predicate, restoring the + /// built-in [`is_retryable`] classification. + pub fn with_default_retry_on(mut self) -> Self { + self.retry_on = None; + self + } + + /// Classifies `error` using this policy's [`RetryPolicy::retry_on`] + /// predicate when one is set, and the crate-wide [`is_retryable`] otherwise. + /// + /// This is the single classification entry point every retry loop should + /// use; calling the free [`is_retryable`] directly silently ignores a + /// caller's custom predicate. + pub fn is_retryable_error(&self, error: &TinyAgentsError) -> bool { + match &self.retry_on { + Some(predicate) => predicate(error), + None => is_retryable(error), + } + } + /// Sleeps for this policy's backoff before the given retry `attempt`, but /// only when [`RetryPolicy::backoff_sleep`] is enabled. /// /// A single, reusable helper so every retry loop that honors a - /// [`RetryPolicy`] gets identical, opt-in backoff behavior. A no-op (returns + /// [`RetryPolicy`] gets identical backoff behavior. A no-op (returns /// immediately) when sleeping is disabled or the computed backoff is zero. + /// + /// Prefer [`RetryPolicy::sleep_backoff_for_error`] where the failure is in + /// hand: it additionally honors a server-supplied `Retry-After`. pub async fn sleep_backoff(&self, attempt: usize) { + self.sleep_for(attempt, self.backoff_for_attempt(attempt), None) + .await; + } + + /// Sleeps for `max(computed backoff, server-supplied Retry-After)` before + /// retry `attempt`, but only when [`RetryPolicy::backoff_sleep`] is enabled. + /// + /// A `429` carrying `Retry-After: 30` means the provider will keep refusing + /// for 30 seconds; retrying after the policy's 200 ms merely burns the + /// remaining attempts. Taking the **max** (rather than replacing the + /// backoff) means a server hint can only ever lengthen the wait, so a bogus + /// `Retry-After: 0` cannot defeat backoff. The hint is clamped at + /// [`RetryPolicy::max_retry_after_ms`]. + pub async fn sleep_backoff_for_error(&self, attempt: usize, error: &TinyAgentsError) { + let hint = self.clamped_retry_after(error); + self.sleep_for(attempt, self.backoff_for_error(attempt, error), hint) + .await; + } + + /// The delay this policy would wait before retry `attempt` given `error`: + /// the larger of the computed backoff and the clamped server-supplied + /// `Retry-After`. Computed regardless of [`RetryPolicy::backoff_sleep`], so + /// observability can report the intended delay even when sleeping is off. + pub fn backoff_for_error(&self, attempt: usize, error: &TinyAgentsError) -> Duration { + let computed = self.backoff_for_attempt(attempt); + match self.clamped_retry_after(error) { + Some(hint) => computed.max(hint), + None => computed, + } + } + + /// Extracts and clamps the server-supplied `Retry-After` carried by `error`. + fn clamped_retry_after(&self, error: &TinyAgentsError) -> Option { + retry_after_hint(error).map(|hint| hint.min(Duration::from_millis(self.max_retry_after_ms))) + } + + /// Shared sleep body: logs the decision, then waits when enabled. + async fn sleep_for(&self, attempt: usize, backoff: Duration, hint: Option) { if !self.backoff_sleep { + tracing::debug!( + target: "tinyagents::retry", + attempt, + backoff_ms = backoff.as_millis() as u64, + "[retry] backoff sleep disabled; retrying immediately" + ); return; } - let backoff = self.backoff_for_attempt(attempt); if backoff > Duration::ZERO { + tracing::debug!( + target: "tinyagents::retry", + attempt, + backoff_ms = backoff.as_millis() as u64, + retry_after_ms = hint.map(|h| h.as_millis() as u64), + jitter = self.jitter, + "[retry] sleeping before retry" + ); tokio::time::sleep(backoff).await; } } @@ -326,7 +429,7 @@ impl RetryPolicy { /// [`RetryPolicy::with_max_attempts`] first and call this on the capped /// policy. pub fn should_retry_error(&self, attempt: usize, error: &TinyAgentsError) -> bool { - is_retryable(error) && self.should_retry(attempt) + self.is_retryable_error(error) && self.should_retry(attempt) } /// Reconciles this policy's own `max_attempts` with a harness-level @@ -343,39 +446,90 @@ impl RetryPolicy { .min(max_retries_per_call.saturating_add(1)) } - /// Computes the deterministic (no-jitter) backoff for the given retry - /// `attempt`. + /// Computes the backoff for the given retry `attempt`. /// /// - `attempt = 0` → `initial_backoff_ms` /// - `attempt = 1` → `initial_backoff_ms * multiplier` /// - …capped at `max_backoff_ms` /// - /// When [`RetryPolicy::jitter`] is `true`, prefer - /// [`backoff_for_attempt_with`][Self::backoff_for_attempt_with] and supply a - /// caller-controlled `[0, 1)` random value so the implementation remains - /// testable. + /// When [`RetryPolicy::jitter`] is `false` this is fully deterministic. When + /// jitter is enabled it draws a **real** random value and spreads the result + /// additively around the base (see + /// [`backoff_for_attempt_with`][Self::backoff_for_attempt_with]); tests that + /// need a fixed spread should call that method with an explicit `rand01` + /// rather than this one. pub fn backoff_for_attempt(&self, attempt: usize) -> Duration { - self.backoff_for_attempt_with(attempt, 0.0) + // 0.5 is the band midpoint, so the no-jitter and jitter paths agree + // exactly when jitter is off — and the RNG is not touched at all then. + let rand01 = if self.jitter { jitter::rand01() } else { 0.5 }; + self.backoff_for_attempt_with(attempt, rand01) } /// Computes backoff for `attempt` using the supplied `rand01 ∈ [0, 1)` for - /// jitter. + /// jitter. This is the deterministic seam: tests inject a fixed value here, + /// production goes through [`backoff_for_attempt`][Self::backoff_for_attempt]. /// /// When [`RetryPolicy::jitter`] is `false`, `rand01` is ignored and the - /// result is fully deterministic. When jitter is enabled, the backoff is - /// uniformly distributed over `[0, base_backoff]`. + /// result is the exact exponential schedule. + /// + /// When jitter is enabled the delay is spread **additively** around the base: + /// `base * (1 + JITTER_FRACTION * (2 * rand01 - 1))`, i.e. uniformly over + /// `[base * 0.75, base * 1.25]` at the default + /// [`JITTER_FRACTION`], then clamped to `max_backoff_ms`. `rand01 == 0.5` + /// reproduces the un-jittered value exactly. + /// + /// This is deliberately **not** the old `base * rand01` form. That one + /// scaled the delay *down* toward zero, and because the production path + /// passed a hardcoded `rand01 = 0.0`, turning jitter on produced a zero + /// delay and disabled backoff entirely. Both reference implementations are + /// additive: LangGraph adds `uniform(0, 1)` seconds, LangChain applies + /// `delay ± 25%` clamped at zero. pub fn backoff_for_attempt_with(&self, attempt: usize, rand01: f64) -> Duration { let base = (self.initial_backoff_ms as f64) * self.multiplier.powi(attempt as i32); - let capped = base.min(self.max_backoff_ms as f64); - let effective = if self.jitter { - capped * rand01.clamp(0.0, 1.0) + let jittered = if self.jitter { + // Map [0, 1) onto [-1, 1) then scale by the band width. + let offset = JITTER_FRACTION * (2.0 * rand01.clamp(0.0, 1.0) - 1.0); + (base * (1.0 + offset)).max(0.0) } else { - capped + base }; - Duration::from_millis(effective as u64) + let capped = jittered.min(self.max_backoff_ms as f64); + Duration::from_millis(capped as u64) } } +/// Extracts a server-supplied `Retry-After` delay carried by `error`, if any. +/// +/// A `429` or `503` that names how long the client must wait is authoritative: +/// retrying sooner burns an attempt for certain. [`RetryPolicy::backoff_for_error`] +/// folds this into the delay by taking the larger of the two. +/// +/// # Two sources, in priority order +/// +/// 1. **The structured field.** +/// [`ProviderError::retry_after_ms`][crate::harness::model::ProviderError::retry_after_ms] +/// is populated by the provider adapter directly from the HTTP `Retry-After` +/// response header (both the delta-seconds and HTTP-date forms). This is the +/// contract; it is read first. +/// 2. **The error message text**, parsed with [`parse_retry_after_ms`]. Hosted +/// providers generally echo the header into the error body, so this is a +/// useful fallback for adapters that do not yet populate the field — but it +/// is string matching, not a contract, and a provider that sends the header +/// without echoing it into the body is served only by source 1. +pub fn retry_after_hint(error: &TinyAgentsError) -> Option { + let message = match error { + TinyAgentsError::Provider(provider_error) => { + if let Some(ms) = provider_error.retry_after_ms { + return Some(Duration::from_millis(ms)); + } + provider_error.message.as_str() + } + TinyAgentsError::Model(message) | TinyAgentsError::Tool(message) => message.as_str(), + _ => return None, + }; + parse_retry_after_ms(message).map(Duration::from_millis) +} + // ── is_retryable ───────────────────────────────────────────────────────────── /// Classifies a [`TinyAgentsError`] as retryable or not. @@ -385,7 +539,7 @@ impl RetryPolicy { /// | Variant | Retryable | Rationale | /// |---|---|---| /// | `Provider` | depends | Classified from [`crate::harness::model::ProviderError::retryable`] — a 429/408/409/5xx is retryable, a 4xx like 401/400 is not. | -/// | `Model` | yes | No structured detail to classify from (transport/parse failure); transient provider 5xx / rate-limit / network glitch is the common case. | +/// | `Model` | depends | No structured `ProviderError` to read, so the message text is run through [`classify_provider_failure`] — a 5xx / 429 / timeout is retryable, an `invalid api key` or `model not found` is not. | /// | `Tool` | yes | Tool execution may have hit a transient dependency. | /// | `Validation` | **no** | Caller-side schema or policy error; retrying will not help. | /// | `Serialization` | **no** | Malformed data; retrying will not help. | @@ -393,10 +547,28 @@ impl RetryPolicy { /// | `MissingStart` / `MissingNode` / `MissingEdgeTarget` / `MissingRoute` | **no** | Graph configuration errors; not transient. | /// | `ToolNotFound` / `ModelNotFound` | **no** | Registry errors; not transient. | /// | `StructuredOutput` | **no** | Schema mismatch; retrying the same call will likely fail again. | +/// +/// Callers holding a [`RetryPolicy`] should call +/// [`RetryPolicy::is_retryable_error`] instead, so a caller-supplied +/// [`RetryPolicy::retry_on`] predicate is honored. pub fn is_retryable(err: &TinyAgentsError) -> bool { match err { TinyAgentsError::Provider(provider_error) => provider_error.retryable, - TinyAgentsError::Model(_) | TinyAgentsError::Tool(_) => true, + // A bare `Model(String)` used to be retried unconditionally, so a + // permanent `401 invalid api key` burned every attempt with a + // guaranteed-identical failure. There is no structured + // `ProviderError` here, but the message text is the same text + // `classify_provider_failure` already knows how to read — so use it + // rather than assuming transience. + TinyAgentsError::Model(message) => { + classify_provider_failure(None, None, message).is_retryable() + } + // Tool failures stay unconditionally retryable: a tool's error text is + // arbitrary caller-authored content with no shared vocabulary to + // classify against, so the HTTP-shaped heuristics above would be + // guessing. Callers that know better narrow this with + // [`RetryPolicy::retry_on`]. + TinyAgentsError::Tool(_) => true, _ => false, } } diff --git a/src/harness/retry/test.rs b/src/harness/retry/test.rs index 306fa53..e31580a 100644 --- a/src/harness/retry/test.rs +++ b/src/harness/retry/test.rs @@ -55,26 +55,91 @@ fn backoff_grows_exponentially_then_caps() { } #[test] -fn backoff_jitter_scales_by_rand01() { +fn backoff_jitter_spreads_additively_around_the_base() { let policy = RetryPolicy::default().with_jitter(true); - // attempt 2 base = 800ms. With jitter the result is base * rand01. + // attempt 2 base = 800ms. Jitter spreads it over ±25% → [600, 1000]. assert_eq!( policy.backoff_for_attempt_with(2, 0.0), - Duration::from_millis(0) + Duration::from_millis(600) ); + // The band midpoint reproduces the un-jittered value exactly. assert_eq!( policy.backoff_for_attempt_with(2, 0.5), - Duration::from_millis(400) + Duration::from_millis(800) ); - // rand01 is clamped into [0, 1). + // rand01 is clamped into [0, 1]. assert_eq!( policy.backoff_for_attempt_with(2, 5.0), - Duration::from_millis(800) + Duration::from_millis(1_000) ); assert_eq!( policy.backoff_for_attempt_with(2, -3.0), - Duration::from_millis(0) + Duration::from_millis(600) + ); +} + +#[test] +fn jitter_never_collapses_the_backoff_to_zero() { + // Regression test (LOOP-2): jitter used to be *multiplicative* + // (`base * rand01`), and the production path — `backoff_for_attempt`, which + // `sleep_backoff` calls — passed a hardcoded `rand01 = 0.0`. So the + // production-hardened-looking `.with_backoff_sleep(true).with_jitter(true)` + // computed a ZERO delay, `sleep_backoff`'s `> Duration::ZERO` guard never + // fired, and nothing ever slept: a rate-limited provider got hammered + // back-to-back. Jitter must only ever *widen* the delay band. + let policy = RetryPolicy::default().with_jitter(true); + + for attempt in 0..8 { + let plain = RetryPolicy::default().backoff_for_attempt(attempt); + let floor = plain.mul_f64(1.0 - super::JITTER_FRACTION); + let ceiling = plain.mul_f64(1.0 + super::JITTER_FRACTION); + + // The deterministic seam across the whole [0, 1) input range. + for step in 0..=20 { + let jittered = policy.backoff_for_attempt_with(attempt, f64::from(step) / 20.0); + assert!(jittered > Duration::ZERO, "jitter collapsed the delay"); + assert!(jittered >= floor && jittered <= ceiling, "outside the band"); + } + + // And the production path, which now draws real randomness. + for _ in 0..64 { + let jittered = policy.backoff_for_attempt(attempt); + assert!( + jittered > Duration::ZERO, + "production backoff_for_attempt returned a zero delay with jitter on" + ); + assert!(jittered >= floor && jittered <= ceiling, "outside the band"); + } + } +} + +#[tokio::test(start_paused = true)] +async fn jittered_sleep_backoff_actually_sleeps() { + // The end-to-end shape of LOOP-2: the config that reads as + // production-hardened must wait, not spin. + use tokio::time::Instant as TokioInstant; + + let policy = RetryPolicy::default() + .with_backoff_sleep(true) + .with_jitter(true); + + let t0 = TokioInstant::now(); + policy.sleep_backoff(0).await; + assert!( + t0.elapsed() > Duration::ZERO, + "jitter + backoff_sleep must still sleep" + ); +} + +#[test] +fn production_backoff_is_random_when_jitter_is_enabled() { + // A stuck RNG would reproduce the original defect in a subtler form. + let policy = RetryPolicy::default().with_jitter(true); + let first = policy.backoff_for_attempt(4); + assert!( + (0..128).any(|_| policy.backoff_for_attempt(4) != first), + "jittered backoff never varied — the randomness source is not wired up" ); } @@ -374,30 +439,37 @@ fn rate_limiter_refill_caps_at_capacity() { } #[test] -fn backoff_sleep_defaults_off_and_is_opt_in() { - // Default policy does not sleep, keeping retry loops deterministic in tests. - assert!(!RetryPolicy::default().backoff_sleep); - // The builder flips it on for production callers. +fn backoff_sleep_defaults_on_and_is_opt_out() { + // Regression test (LOCAL-4): the default used to be `false` "so tests stay + // deterministic and fast", which made test convenience the *production* + // policy. Transport failures are classified retryable, so the default four + // attempts fired back-to-back against a local runtime that was merely + // loading a multi-gigabyte model. LangGraph's reference default + // (`initial_interval=0.5, backoff_factor=2.0, jitter=True`) is on. assert!( - RetryPolicy::default() - .with_backoff_sleep(true) + RetryPolicy::default().backoff_sleep, + "backoff must sleep by default" + ); + // Tests and other latency-sensitive callers opt out explicitly. + assert!( + !RetryPolicy::default() + .with_backoff_sleep(false) .backoff_sleep ); } #[tokio::test(start_paused = true)] -async fn sleep_backoff_waits_only_when_enabled() { +async fn sleep_backoff_waits_unless_explicitly_disabled() { use tokio::time::Instant as TokioInstant; - // Disabled (default): returns immediately with no virtual time elapsed. - let policy = RetryPolicy::default(); + // Explicitly disabled: returns immediately with no virtual time elapsed. + let policy = RetryPolicy::default().with_backoff_sleep(false); let t0 = TokioInstant::now(); policy.sleep_backoff(1).await; assert_eq!(t0.elapsed(), Duration::ZERO); - // Enabled: advances virtual time by the computed backoff (attempt 1 = - // initial_backoff_ms with the default multiplier applied at attempt^power). - let sleeping = RetryPolicy::default().with_backoff_sleep(true); + // Default (enabled): advances virtual time by the computed backoff. + let sleeping = RetryPolicy::default(); let expected = sleeping.backoff_for_attempt(1); let t1 = TokioInstant::now(); sleeping.sleep_backoff(1).await; @@ -405,6 +477,195 @@ async fn sleep_backoff_waits_only_when_enabled() { assert!(expected > Duration::ZERO); } +// ── Retry-After (LOOP-5) ────────────────────────────────────────────────────── + +#[test] +fn retry_after_hint_is_read_from_every_error_shape_that_can_carry_one() { + use crate::harness::model::ProviderError; + use crate::harness::retry::retry_after_hint; + + assert_eq!( + retry_after_hint(&TinyAgentsError::Model( + "429 Too Many Requests, Retry-After: 30".into() + )), + Some(Duration::from_secs(30)) + ); + assert_eq!( + retry_after_hint(&TinyAgentsError::Provider(Box::new(ProviderError { + provider: "openai".into(), + status: Some(429), + retryable: true, + message: "rate limited; retry_after: 12.5 seconds".into(), + ..ProviderError::default() + }))), + Some(Duration::from_millis(12_500)) + ); + // Nothing to read → no hint, and the plain backoff applies. + assert_eq!( + retry_after_hint(&TinyAgentsError::Model("500 Internal Server Error".into())), + None + ); + assert_eq!( + retry_after_hint(&TinyAgentsError::Validation("bad".into())), + None + ); +} + +#[test] +fn backoff_for_error_takes_the_max_of_backoff_and_retry_after() { + // Regression test (LOOP-5): `parse_retry_after_ms` existed but had only + // test callers, so a 429 saying `Retry-After: 30` was retried after the + // policy's 200ms, three times, and then gave up. + let policy = RetryPolicy::default(); + let rate_limited = TinyAgentsError::Model("429 rate limited, Retry-After: 30".into()); + + assert_eq!( + policy.backoff_for_error(0, &rate_limited), + Duration::from_secs(30), + "a server-supplied Retry-After must win over a shorter computed backoff" + ); + + // A hint shorter than the computed backoff can never shorten the wait, so a + // bogus `Retry-After: 0` cannot defeat backoff. + let tiny_hint = TinyAgentsError::Model("429 rate limited, Retry-After: 0".into()); + assert_eq!( + policy.backoff_for_error(3, &tiny_hint), + policy.backoff_for_attempt(3) + ); + + // No hint at all → identical to the plain schedule. + let plain = TinyAgentsError::Model("502 bad gateway".into()); + assert_eq!( + policy.backoff_for_error(1, &plain), + policy.backoff_for_attempt(1) + ); +} + +#[test] +fn retry_after_is_clamped_so_a_hostile_header_cannot_park_the_run() { + let policy = RetryPolicy::default().with_max_retry_after_ms(5_000); + let absurd = TinyAgentsError::Model("429 rate limited, Retry-After: 86400".into()); + assert_eq!( + policy.backoff_for_error(0, &absurd), + Duration::from_millis(5_000) + ); +} + +#[tokio::test(start_paused = true)] +async fn sleep_backoff_for_error_honors_retry_after() { + use tokio::time::Instant as TokioInstant; + + let policy = RetryPolicy::default(); + let rate_limited = TinyAgentsError::Model("429 rate limited, Retry-After: 30".into()); + + let t0 = TokioInstant::now(); + policy.sleep_backoff_for_error(0, &rate_limited).await; + assert_eq!(t0.elapsed(), Duration::from_secs(30)); +} + +// ── retry_on predicate (LOOP-5b) ────────────────────────────────────────────── + +#[test] +fn model_errors_are_classified_from_their_message_not_assumed_transient() { + // Regression test (LOOP-5b): the `Model(_)` arm returned `true` + // unconditionally, so a permanent auth failure that never got a structured + // `ProviderError` burned every attempt. + assert!(!is_retryable(&TinyAgentsError::Model( + "401 Unauthorized: invalid api key".into() + ))); + assert!(!is_retryable(&TinyAgentsError::Model( + "model gpt-9 does not exist".into() + ))); + + // Transient shapes stay retryable. + assert!(is_retryable(&TinyAgentsError::Model( + "502 Bad Gateway".into() + ))); + assert!(is_retryable(&TinyAgentsError::Model( + "429 Too Many Requests: rate limit exceeded".into() + ))); + // Unclassifiable transport text keeps the permissive default. + assert!(is_retryable(&TinyAgentsError::Model( + "connection reset by peer".into() + ))); +} + +#[test] +fn retry_on_predicate_overrides_the_builtin_classification() { + use std::sync::Arc; + + // Default: no predicate → built-in classification, unchanged. + let default_policy = RetryPolicy::default(); + assert!(default_policy.is_retryable_error(&TinyAgentsError::Tool("flaky".into()))); + assert!(!default_policy.is_retryable_error(&TinyAgentsError::Validation("bad".into()))); + + // LangGraph's curated default shape: connection errors and 5xx, never a + // programming error (a failing tool is the closest analogue here). + let narrowed = RetryPolicy::default().with_retry_on(Arc::new(|err: &TinyAgentsError| { + matches!( + err, + TinyAgentsError::Model(_) | TinyAgentsError::Provider(_) + ) + })); + assert!(narrowed.is_retryable_error(&TinyAgentsError::Model("timeout".into()))); + assert!( + !narrowed.is_retryable_error(&TinyAgentsError::Tool("flaky".into())), + "a custom predicate must be able to *narrow* the built-in set" + ); + + // And it can widen it too. + let widened = RetryPolicy::default().with_retry_on(Arc::new(|_: &TinyAgentsError| true)); + assert!(widened.is_retryable_error(&TinyAgentsError::Validation("bad".into()))); + + // Clearing restores the built-in classification. + assert!( + !widened + .clone() + .with_default_retry_on() + .is_retryable_error(&TinyAgentsError::Validation("bad".into())) + ); +} + +#[test] +fn should_retry_error_consults_the_custom_predicate() { + use std::sync::Arc; + + // Regression test: `should_retry_error` called the free `is_retryable` + // directly, so a caller's `retry_on` would have been silently ignored by + // every retry loop in the harness. + let policy = RetryPolicy::default() + .with_max_attempts(3) + .with_retry_on(Arc::new(|err: &TinyAgentsError| { + matches!(err, TinyAgentsError::Validation(_)) + })); + + assert!(policy.should_retry_error(0, &TinyAgentsError::Validation("retry me".into()))); + assert!(!policy.should_retry_error(0, &TinyAgentsError::Tool("do not".into()))); + // The attempt cap still applies on top of the predicate. + assert!(!policy.should_retry_error(2, &TinyAgentsError::Validation("retry me".into()))); +} + +#[test] +fn retry_policy_equality_and_debug_survive_the_predicate_field() { + use std::sync::Arc; + + assert_eq!(RetryPolicy::default(), RetryPolicy::default()); + + let predicate: crate::harness::retry::RetryPredicate = Arc::new(|_: &TinyAgentsError| true); + let a = RetryPolicy::default().with_retry_on(predicate.clone()); + let b = RetryPolicy::default().with_retry_on(predicate); + assert_eq!(a, b, "the same Arc compares equal"); + assert_ne!(a, RetryPolicy::default()); + assert_ne!( + a, + RetryPolicy::default().with_retry_on(Arc::new(|_: &TinyAgentsError| true)), + "distinct closures are not equal" + ); + + assert!(format!("{a:?}").contains("custom predicate")); + assert!(format!("{:?}", RetryPolicy::default()).contains("retry_on: None")); +} + #[test] fn should_retry_error_combines_classification_and_attempt_cap() { // 1 try + 2 retries: attempts 0 and 1 may retry, attempt 2 may not. diff --git a/src/harness/retry/types.rs b/src/harness/retry/types.rs index b88285e..010d91e 100644 --- a/src/harness/retry/types.rs +++ b/src/harness/retry/types.rs @@ -7,13 +7,47 @@ //! deterministically testable: backoff and the limiter take an explicit //! `rand01` / `now` rather than reading a global clock or RNG. +use std::sync::Arc; + +use crate::error::TinyAgentsError; + +/// A caller-supplied predicate deciding whether a [`TinyAgentsError`] should be +/// retried, replacing the crate's built-in +/// [`is_retryable`][crate::harness::retry::is_retryable] classification. +/// +/// Ported from LangGraph's `RetryPolicy.retry_on`, which accepts an exception +/// type, a sequence of types, or a callable. Rust has no exception hierarchy to +/// match on, so the callable form is the only one that generalizes — a caller +/// wanting the "sequence of types" form writes a `matches!` over +/// [`TinyAgentsError`] variants inside the closure. +/// +/// The predicate answers *classification only* ("is this error transient?"); the +/// attempt cap is still applied on top by +/// [`RetryPolicy::should_retry`][crate::harness::retry::RetryPolicy::should_retry]. +/// +/// # Examples +/// +/// ``` +/// use std::sync::Arc; +/// use tinyagents::error::TinyAgentsError; +/// use tinyagents::harness::retry::RetryPolicy; +/// +/// // Retry transport failures but never a failing tool. +/// let policy = RetryPolicy::default().with_retry_on(Arc::new(|err: &TinyAgentsError| { +/// matches!(err, TinyAgentsError::Model(_)) +/// })); +/// assert!(policy.is_retryable_error(&TinyAgentsError::Model("connection reset".into()))); +/// assert!(!policy.is_retryable_error(&TinyAgentsError::Tool("boom".into()))); +/// ``` +pub type RetryPredicate = Arc bool + Send + Sync>; + /// Configures how a harness call is retried on transient failure. /// /// Backoff grows exponentially: `initial_backoff_ms * multiplier^attempt`, then -/// capped at `max_backoff_ms`. When `jitter` is `true` the caller should use -/// [`RetryPolicy::backoff_for_attempt_with`] and supply a `[0, 1)` random -/// value; this avoids thundering-herd without making the implementation -/// non-deterministic for tests. +/// capped at `max_backoff_ms`. When `jitter` is `true` the delay is spread +/// **additively** around that base (`base * (1 ± JITTER_FRACTION)`), so jitter +/// can never collapse the delay to zero — see +/// [`RetryPolicy::backoff_for_attempt_with`]. /// /// # Examples /// @@ -24,7 +58,7 @@ /// assert!(policy.should_retry(0)); /// assert!(!policy.should_retry(3)); /// ``` -#[derive(Clone, Debug, PartialEq)] +#[derive(Clone)] pub struct RetryPolicy { /// Total number of attempts (first try + retries). A value of `1` means no /// retries. @@ -35,18 +69,43 @@ pub struct RetryPolicy { pub max_backoff_ms: u64, /// Multiplicative factor applied to the backoff on each attempt. pub multiplier: f64, - /// When `true`, the caller should supply a `[0, 1)` random value to - /// [`RetryPolicy::backoff_for_attempt_with`] to add jitter. + /// When `true`, the computed backoff is spread additively around its base + /// (`base * (1 ± `[`JITTER_FRACTION`][crate::harness::retry::JITTER_FRACTION]`)`) + /// so concurrent clients that failed together do not retry together. + /// + /// Jitter **widens** the delay band; it never removes the delay. (It used to: + /// the implementation multiplied the base by `rand01`, and the production + /// call path passed a hardcoded `0.0`, so enabling jitter disabled backoff + /// entirely.) pub jitter: bool, /// When `true`, retry loops that honor this policy actually /// [`tokio::time::sleep`] for the computed backoff between attempts. /// - /// Defaults to `false` so the harness stays fast and deterministic under - /// test: the backoff is *computed* (for observability) but not slept on. - /// Production/streaming callers opt in via - /// [`RetryPolicy::with_backoff_sleep`] so a transient provider/network - /// failure is retried after a real, growing delay instead of back-to-back. + /// **Defaults to `true`.** A transport failure is classified retryable, so a + /// `false` default means the default four attempts fire back-to-back — which + /// is exactly the wrong behaviour against a rate-limited hosted provider or + /// a local runtime (Ollama, LM Studio) that is merely loading a model. Tests + /// that need instant retries opt *out* explicitly with + /// [`RetryPolicy::with_backoff_sleep(false)`][RetryPolicy::with_backoff_sleep]. pub backoff_sleep: bool, + /// Upper bound, in milliseconds, on a server-supplied `Retry-After` delay + /// that [`RetryPolicy::backoff_for_error`] will honor. + /// + /// A provider (or a proxy in front of one) can send an arbitrarily large + /// `Retry-After`; without a ceiling a single header could park a run for + /// hours. Defaults to [`RetryPolicy::DEFAULT_MAX_RETRY_AFTER_MS`]. + pub max_retry_after_ms: u64, + /// Optional caller-supplied retry classification, replacing + /// [`is_retryable`][crate::harness::retry::is_retryable]. + /// + /// `None` (the default) uses the crate's built-in classification, so this + /// field is purely additive: an existing policy behaves exactly as before. + pub retry_on: Option, +} + +impl RetryPolicy { + /// Default ceiling applied to a server-supplied `Retry-After` (2 minutes). + pub const DEFAULT_MAX_RETRY_AFTER_MS: u64 = 120_000; } impl Default for RetryPolicy { @@ -57,11 +116,54 @@ impl Default for RetryPolicy { max_backoff_ms: 30_000, multiplier: 2.0, jitter: false, - backoff_sleep: false, + backoff_sleep: true, + max_retry_after_ms: Self::DEFAULT_MAX_RETRY_AFTER_MS, + retry_on: None, } } } +/// Hand-written so [`RetryPolicy::retry_on`] (a boxed closure, which has no +/// `Debug`) does not block the derive. The predicate is rendered as a presence +/// marker. +impl std::fmt::Debug for RetryPolicy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RetryPolicy") + .field("max_attempts", &self.max_attempts) + .field("initial_backoff_ms", &self.initial_backoff_ms) + .field("max_backoff_ms", &self.max_backoff_ms) + .field("multiplier", &self.multiplier) + .field("jitter", &self.jitter) + .field("backoff_sleep", &self.backoff_sleep) + .field("max_retry_after_ms", &self.max_retry_after_ms) + .field( + "retry_on", + &self.retry_on.as_ref().map(|_| ""), + ) + .finish() + } +} + +/// Hand-written for the same reason as [`Debug`]. Two predicates compare equal +/// only when they are the *same* `Arc` (or both absent) — closures have no +/// meaningful structural equality. +impl PartialEq for RetryPolicy { + fn eq(&self, other: &Self) -> bool { + self.max_attempts == other.max_attempts + && self.initial_backoff_ms == other.initial_backoff_ms + && self.max_backoff_ms == other.max_backoff_ms + && self.multiplier == other.multiplier + && self.jitter == other.jitter + && self.backoff_sleep == other.backoff_sleep + && self.max_retry_after_ms == other.max_retry_after_ms + && match (&self.retry_on, &other.retry_on) { + (None, None) => true, + (Some(a), Some(b)) => Arc::ptr_eq(a, b), + _ => false, + } + } +} + /// Ordered list of model identifiers to try in sequence when the current model /// fails. /// diff --git a/src/harness/runtime/types.rs b/src/harness/runtime/types.rs index d4c260c..330abd3 100644 --- a/src/harness/runtime/types.rs +++ b/src/harness/runtime/types.rs @@ -49,20 +49,30 @@ use crate::harness::tool::{ToolRegistry, ToolTimeoutSettings}; /// How the agent loop reacts when the model calls a tool that is not /// registered. /// -/// The default is [`UnknownToolPolicy::Fail`], preserving the historical -/// fail-fast behavior. The recoverable variants let a run keep going so the -/// model can correct itself — each recovery still consumes a tool-call budget -/// slot, so [`RunLimits::max_tool_calls`] bounds any unknown-tool loop. +/// The default is [`UnknownToolPolicy::ReturnToolError`]: a hallucinated tool +/// name is a routine model mistake, not a harness fault, so the run keeps going +/// and the model gets told which tools actually exist. Each recovery still +/// consumes a tool-call budget slot, so [`RunLimits::max_tool_calls`] bounds any +/// unknown-tool loop. +/// +/// # Why the default flipped +/// +/// `Fail` used to be the default, which made the crate inconsistent with +/// itself: an **unparseable** arguments blob has always recovered +/// unconditionally (see `agent_loop/tools.rs`), so `{city:` survived while a +/// merely unknown tool name killed the whole run. It also diverged from +/// LangGraph, whose `ToolNode` answers an unknown name with a synthetic +/// `status="error"` message listing the valid tools. `Fail` remains available +/// for callers that genuinely want a hard stop. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub enum UnknownToolPolicy { /// Abort the run with - /// [`TinyAgentsError::ToolNotFound`][crate::error::TinyAgentsError::ToolNotFound] - /// (the default, historical behavior). - #[default] + /// [`TinyAgentsError::ToolNotFound`][crate::error::TinyAgentsError::ToolNotFound]. Fail, /// Inject a tool-error result (naming the originally requested tool and /// listing the registered tools) back into the transcript and continue the - /// loop, letting the model retry with a valid tool. + /// loop, letting the model retry with a valid tool. The default. + #[default] ReturnToolError, /// Rewrite an unknown call to a fixed compatibility tool name and retry the /// lookup once. If the rewrite target is also unregistered, fall back to @@ -76,22 +86,22 @@ pub enum UnknownToolPolicy { /// How the agent loop reacts when the model calls a *registered* tool with /// arguments that fail schema validation. /// -/// The default is [`InvalidArgsPolicy::Fail`], preserving the historical -/// fail-fast behavior where a missing `required` field, wrong type, or bad -/// `enum` aborts the whole turn. The recoverable variant lets a run keep going -/// so the model can self-correct — the recovery still consumes a tool-call -/// budget slot, so [`RunLimits::max_tool_calls`] bounds any invalid-args loop. -/// Mirrors [`UnknownToolPolicy`] for the schema-validation seam. +/// The default is [`InvalidArgsPolicy::ReturnToolError`]: a missing `required` +/// field, wrong type, or bad `enum` is model output the model can fix, so the +/// validation detail plus the expected schema go back into the transcript +/// instead of aborting the turn. The recovery consumes a tool-call budget slot, +/// so [`RunLimits::max_tool_calls`] bounds any invalid-args loop. Mirrors +/// [`UnknownToolPolicy`] for the schema-validation seam — including why the +/// default flipped away from `Fail`. #[derive(Clone, Debug, Default, PartialEq, Eq)] pub enum InvalidArgsPolicy { /// Abort the run with - /// [`TinyAgentsError::Validation`][crate::error::TinyAgentsError::Validation] - /// (the default, historical behavior). - #[default] + /// [`TinyAgentsError::Validation`][crate::error::TinyAgentsError::Validation]. Fail, /// Inject a tool-error result (carrying the validation detail and the /// tool's expected parameter schema) back into the transcript and continue - /// the loop, letting the model retry with corrected arguments. + /// the loop, letting the model retry with corrected arguments. The default. + #[default] ReturnToolError, /// First normalize common provider-shape defects, then apply /// [`Self::ReturnToolError`] if the resulting arguments still fail schema @@ -220,6 +230,7 @@ impl Default for RunPolicy { cache: CachePolicy { response_cache_enabled: true, protect_prompt_prefix: false, + ..CachePolicy::default() }, // Opt-in: preserve the historical blank-final behavior by default. error_on_empty_response: false, diff --git a/src/harness/steering/mod.rs b/src/harness/steering/mod.rs index 75737d3..291ce24 100644 --- a/src/harness/steering/mod.rs +++ b/src/harness/steering/mod.rs @@ -104,6 +104,8 @@ impl SteeringHandle { inner: Arc::new(SteeringInner { queue: Mutex::new(VecDeque::new()), policy, + paused: Mutex::new(None), + checkpoints: Mutex::new(0), }), } } @@ -163,6 +165,75 @@ impl SteeringHandle { pub fn policy(&self) -> &SteeringPolicy { &self.inner.policy } + + /// Returns the latched [`PauseState`] when the run is paused, `None` + /// otherwise. + /// + /// This is how a caller distinguishes a run that stopped for a human from a + /// run that finished with nothing to say: on a + /// [`SteeringOutcome::Pause`] this is always `Some`. + pub fn pause_state(&self) -> Option { + self.lock_paused().clone() + } + + /// Returns `true` when a pause is latched. + pub fn is_paused(&self) -> bool { + self.lock_paused().is_some() + } + + /// Latches a pause with an optional reason, recorded as having taken effect + /// at `checkpoint`. Idempotent: an existing pause keeps its original reason + /// and checkpoint, so a repeated `Pause` does not rewrite why the run + /// stopped. + fn latch_pause(&self, checkpoint: usize, reason: Option) -> PauseState { + let mut paused = self.lock_paused(); + let state = paused.get_or_insert(PauseState { + reason, + paused_at_checkpoint: checkpoint, + }); + tracing::debug!( + target: "tinyagents::steering", + checkpoint = state.paused_at_checkpoint, + reason = state.reason.as_deref(), + "[steering] pause latched" + ); + state.clone() + } + + /// Clears any latched pause, returning the state that was cleared. + /// + /// Equivalent to delivering a [`SteeringCommand::Resume`]; exposed directly + /// so a host UI can resume without going through the queue. A no-op when no + /// pause is in effect. + pub fn resume(&self) -> Option { + let cleared = self.lock_paused().take(); + if cleared.is_some() { + tracing::debug!(target: "tinyagents::steering", "[steering] pause cleared by resume"); + } + cleared + } + + /// Returns this checkpoint's zero-based index and advances the counter, so + /// the *current* checkpoint is what a pause records (not the next one). + fn advance_checkpoint(&self) -> usize { + let mut checkpoints = self + .inner + .checkpoints + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let current = *checkpoints; + *checkpoints += 1; + current + } + + /// Locks the pause latch, recovering from poisoning (see + /// [`SteeringHandle::send`]). + fn lock_paused(&self) -> std::sync::MutexGuard<'_, Option> { + self.inner + .paused + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } } // ── Checkpoint application ──────────────────────────────────────────────────── @@ -175,26 +246,32 @@ impl SteeringHandle { /// standalone, synchronous function so it can be unit-tested without a full /// run. Behaviour: /// -/// - When `ctx` has no [`SteeringHandle`] (or its queue is empty), returns +/// - When `ctx` has no [`SteeringHandle`], returns /// [`SteeringOutcome::Continue`] without emitting anything. -/// - Every drained command is checked against the handle's -/// [`SteeringPolicy`]. A disallowed command emits an -/// [`AgentEvent::Steered`] with `accepted = false` and returns -/// [`TinyAgentsError::Steering`], aborting the run; no later command in the -/// batch is applied. +/// - The batch is **validated in full before anything is applied**. If any +/// command is disallowed, an [`AgentEvent::Steered`] with `accepted = false` +/// is emitted for it and [`TinyAgentsError::Steering`] is returned — with the +/// working transcript and run metadata completely untouched. (It used to +/// validate lazily while applying, so a rejected command at position *n* left +/// commands `0..n` already applied, commands after it dropped, and the run +/// erroring: a partially-steered run and no way to reason about its state.) /// - [`SteeringCommand::Cancel`] takes precedence: it is applied (emitting an /// accepted event) and the function returns [`SteeringOutcome::Cancel`] /// immediately, ignoring the rest of the batch. -/// - [`SteeringCommand::Pause`] sets a net-pause outcome; a later -/// [`SteeringCommand::Resume`] in the same batch clears it. +/// - [`SteeringCommand::Pause`] / [`SteeringCommand::PauseWith`] latch a pause +/// **on the handle**, so it survives past this batch. Any later +/// [`SteeringCommand::Resume`] — in this batch or a subsequent one — clears +/// it. While latched, every checkpoint returns [`SteeringOutcome::Pause`] +/// even with an empty queue. /// - [`SteeringCommand::InjectMessage`] and [`SteeringCommand::Redirect`] /// append to `messages`; [`SteeringCommand::SetMetadata`] replaces /// `ctx.config.metadata`. /// /// # Errors /// -/// Returns [`TinyAgentsError::Steering`] when a drained command is not -/// permitted by the run's [`SteeringPolicy`]. +/// Returns [`TinyAgentsError::Steering`] when any drained command is not +/// permitted by the run's [`SteeringPolicy`]. No command in the batch is +/// applied in that case. pub fn apply_pending_steering( ctx: &mut RunContext, messages: &mut Vec, @@ -204,33 +281,56 @@ pub fn apply_pending_steering( let Some(handle) = ctx.steering.clone() else { return Ok(SteeringOutcome::Continue); }; + let checkpoint = handle.advance_checkpoint(); let commands = handle.drain(); - if commands.is_empty() { - return Ok(SteeringOutcome::Continue); + + // ── Phase 1: validate the whole batch, mutating nothing ───────────────── + // + // A policy violation must abort the checkpoint *atomically*. Checking as we + // apply means the run dies with some of the batch already in the + // transcript. + if let Some(rejected) = commands + .iter() + .map(SteeringCommand::kind) + .find(|kind| !handle.policy().is_allowed(*kind)) + { + tracing::debug!( + target: "tinyagents::steering", + checkpoint, + command_kind = rejected.as_str(), + batch_size = commands.len(), + "[steering] batch rejected by policy; nothing applied" + ); + ctx.emit(AgentEvent::Steered { + command_kind: rejected.as_str().to_string(), + accepted: false, + }); + return Err(TinyAgentsError::Steering(format!( + "steering command `{}` is not permitted by the run policy", + rejected.as_str() + ))); } - let mut outcome = SteeringOutcome::Continue; + // ── Phase 2: apply ────────────────────────────────────────────────────── + tracing::debug!( + target: "tinyagents::steering", + checkpoint, + batch_size = commands.len(), + already_paused = handle.is_paused(), + "[steering] applying checkpoint batch" + ); for command in commands { let kind = command.kind(); - if !handle.policy().is_allowed(kind) { - ctx.emit(AgentEvent::Steered { - command_kind: kind.as_str().to_string(), - accepted: false, - }); - return Err(TinyAgentsError::Steering(format!( - "steering command `{}` is not permitted by the run policy", - kind.as_str() - ))); - } - - // Apply the permitted command. match command { - SteeringCommand::Pause => outcome = SteeringOutcome::Pause, + SteeringCommand::Pause => { + handle.latch_pause(checkpoint, None); + } + SteeringCommand::PauseWith { reason } => { + handle.latch_pause(checkpoint, Some(reason)); + } SteeringCommand::Resume => { - if outcome == SteeringOutcome::Pause { - outcome = SteeringOutcome::Continue; - } + handle.resume(); } SteeringCommand::Cancel => { ctx.emit(AgentEvent::Steered { @@ -257,7 +357,13 @@ pub fn apply_pending_steering( }); } - Ok(outcome) + // The latch — not this batch — decides the outcome, so a pause applied at an + // earlier checkpoint keeps holding the run. + if handle.is_paused() { + Ok(SteeringOutcome::Pause) + } else { + Ok(SteeringOutcome::Continue) + } } #[cfg(test)] diff --git a/src/harness/steering/test.rs b/src/harness/steering/test.rs index 58b517a..44ad328 100644 --- a/src/harness/steering/test.rs +++ b/src/harness/steering/test.rs @@ -43,6 +43,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -84,6 +85,7 @@ impl ChatModel<()> for RecordingModel { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, }) } else { Ok(text_response("done")) @@ -398,3 +400,203 @@ fn steering_queue_recovers_from_poisoned_lock() { assert_eq!(handle.drain().len(), 2); assert!(handle.is_empty()); } + +// ── LOOP-8(a): the batch is validated before anything is applied ────────────── + +#[test] +fn a_rejected_command_leaves_no_earlier_command_applied() { + // Regression test (LOOP-8a): `apply_pending_steering` drained the whole + // batch up front and then validated lazily *while applying*, so a policy + // violation at position 2 left commands 0 and 1 already in the transcript, + // command 3 silently dropped, and the run erroring. The checkpoint must be + // atomic: reject the batch, change nothing. + let recorder = EventRecorder::new(); + let handle = SteeringHandle::new( + SteeringPolicy::new() + .allow(SteeringCommandKind::InjectMessage) + .allow(SteeringCommandKind::SetMetadata), + ); + handle.send(SteeringCommand::InjectMessage(Message::user("first"))); + handle.send(SteeringCommand::SetMetadata { + metadata: serde_json::json!({"tag": "applied"}), + }); + // Not allowed → the whole batch must be refused. + handle.send(SteeringCommand::Cancel); + handle.send(SteeringCommand::InjectMessage(Message::user("last"))); + + let mut ctx: RunContext = RunContext::new(RunConfig::new("r"), ()) + .with_events(recorder.sink()) + .with_steering(handle); + let mut messages = Vec::new(); + + let err = apply_pending_steering(&mut ctx, &mut messages).unwrap_err(); + assert!(matches!(err, TinyAgentsError::Steering(_)), "got {err:?}"); + + assert!( + messages.is_empty(), + "an earlier command in a rejected batch was applied: {messages:?}" + ); + assert_eq!( + ctx.config.metadata, + serde_json::Value::Null, + "metadata was mutated by a rejected batch" + ); + // Exactly one event, for the offending command. + assert_eq!( + recorder.events(), + vec![AgentEvent::Steered { + command_kind: "cancel".to_string(), + accepted: false, + }] + ); +} + +// ── LOOP-8(b): a pause is latched and resumable across checkpoints ──────────── + +#[test] +fn a_pause_survives_the_batch_and_holds_later_checkpoints() { + // Regression test (LOOP-8b): `Resume` only cancelled a `Pause` from the + // *same* drained batch, and the outcome was recomputed from scratch each + // checkpoint — so a pause applied at one checkpoint silently evaporated at + // the next, and could never be deliberately resumed either. + let handle = SteeringHandle::allow_all(); + handle.send(SteeringCommand::Pause); + let mut ctx: RunContext = + RunContext::new(RunConfig::new("r"), ()).with_steering(handle.clone()); + let mut messages = Vec::new(); + + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Pause + ); + assert!(handle.is_paused()); + + // A later checkpoint with an EMPTY queue must stay paused. + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Pause, + "the pause evaporated at the next checkpoint" + ); +} + +#[test] +fn a_pause_is_resumable_from_a_later_batch() { + let handle = SteeringHandle::allow_all(); + let mut ctx: RunContext = + RunContext::new(RunConfig::new("r"), ()).with_steering(handle.clone()); + let mut messages = Vec::new(); + + handle.send(SteeringCommand::Pause); + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Pause + ); + + // The resume arrives long after the pause was applied — the case that was + // impossible before. + handle.send(SteeringCommand::Resume); + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Continue, + "a pause applied in an earlier batch was unresumable" + ); + assert!(!handle.is_paused()); + assert!(handle.pause_state().is_none()); +} + +#[test] +fn pause_state_makes_a_paused_run_distinguishable_from_an_empty_answer() { + // The loop reports `final_response: None` for both a pause and an empty + // model turn; `pause_state()` is what tells the caller which happened. + let handle = SteeringHandle::allow_all(); + handle.send(SteeringCommand::PauseWith { + reason: "waiting for human approval".into(), + }); + let mut ctx: RunContext = + RunContext::new(RunConfig::new("r"), ()).with_steering(handle.clone()); + let mut messages = Vec::new(); + + let outcome = apply_pending_steering(&mut ctx, &mut messages).unwrap(); + assert!(outcome.is_pause()); + + let state = handle + .pause_state() + .expect("a Pause outcome must always carry a PauseState"); + assert_eq!(state.reason.as_deref(), Some("waiting for human approval")); + assert_eq!(state.paused_at_checkpoint, 0); + + // A run that was never paused has no state at all. + assert!(SteeringHandle::allow_all().pause_state().is_none()); +} + +#[test] +fn a_repeated_pause_keeps_the_original_reason_and_checkpoint() { + let handle = SteeringHandle::allow_all(); + let mut ctx: RunContext = + RunContext::new(RunConfig::new("r"), ()).with_steering(handle.clone()); + let mut messages = Vec::new(); + + handle.send(SteeringCommand::PauseWith { + reason: "first reason".into(), + }); + apply_pending_steering(&mut ctx, &mut messages).unwrap(); + + handle.send(SteeringCommand::PauseWith { + reason: "second reason".into(), + }); + apply_pending_steering(&mut ctx, &mut messages).unwrap(); + + let state = handle.pause_state().expect("still paused"); + assert_eq!(state.reason.as_deref(), Some("first reason")); + assert_eq!(state.paused_at_checkpoint, 0); +} + +#[test] +fn pause_with_is_gated_by_the_same_policy_kind_as_pause() { + assert_eq!( + SteeringCommand::PauseWith { + reason: "why".into() + } + .kind(), + SteeringCommandKind::Pause + ); + + // A policy that forbids Pause forbids PauseWith too. + let handle = SteeringHandle::new(SteeringPolicy::new().allow(SteeringCommandKind::Resume)); + handle.send(SteeringCommand::PauseWith { + reason: "why".into(), + }); + let mut ctx: RunContext = RunContext::new(RunConfig::new("r"), ()).with_steering(handle); + let mut messages = Vec::new(); + assert!(apply_pending_steering(&mut ctx, &mut messages).is_err()); +} + +#[test] +fn handle_resume_clears_a_latch_without_going_through_the_queue() { + let handle = SteeringHandle::allow_all(); + handle.send(SteeringCommand::Pause); + let mut ctx: RunContext = + RunContext::new(RunConfig::new("r"), ()).with_steering(handle.clone()); + let mut messages = Vec::new(); + apply_pending_steering(&mut ctx, &mut messages).unwrap(); + + let cleared = handle.resume().expect("a pause was in effect"); + assert_eq!(cleared.paused_at_checkpoint, 0); + assert!(!handle.is_paused()); + assert!(handle.resume().is_none(), "resume must be idempotent"); + + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Continue + ); +} + +#[test] +fn pause_with_round_trips_through_json() { + let command = SteeringCommand::PauseWith { + reason: "human review".into(), + }; + let json = serde_json::to_value(&command).expect("serialize"); + let back: SteeringCommand = serde_json::from_value(json).expect("deserialize"); + assert_eq!(back, command); +} diff --git a/src/harness/steering/types.rs b/src/harness/steering/types.rs index e42fa43..415348c 100644 --- a/src/harness/steering/types.rs +++ b/src/harness/steering/types.rs @@ -31,12 +31,29 @@ use crate::harness::message::Message; #[serde(rename_all = "snake_case", tag = "command")] pub enum SteeringCommand { /// Cooperatively pause the run: the loop stops issuing further model and - /// tool work at the next checkpoint until a [`SteeringCommand::Resume`] is - /// delivered in the same drained batch. + /// tool work at the next checkpoint, and **stays** paused until a + /// [`SteeringCommand::Resume`] arrives — in this batch or any later one. + /// + /// The pause is latched on the [`SteeringHandle`], not on the batch. It used + /// to be batch-scoped, which made a pause unresumable in practice: a + /// `Resume` sent after the pause had already been applied found nothing to + /// clear. Pause, - /// Clear a pending pause so the loop continues. A `Resume` with no - /// preceding `Pause` in the same batch is a no-op. + /// Pause with a human-readable reason recorded in the resulting + /// [`PauseState`], for example `"waiting for human approval of the refund"`. + /// + /// Identical to [`SteeringCommand::Pause`] in every other respect, + /// including its [`SteeringCommandKind::Pause`] policy gate — a policy that + /// allows one allows the other. + PauseWith { + /// Why the run was paused. Surfaced to the caller through + /// [`PauseState::reason`]. + reason: String, + }, + + /// Clear a latched pause so the loop continues. A `Resume` with no pause in + /// effect is a no-op. Resume, /// Terminate the run cooperatively at the next checkpoint. Cancel takes @@ -70,7 +87,9 @@ impl SteeringCommand { /// Returns the policy-relevant [`SteeringCommandKind`] of this command. pub fn kind(&self) -> SteeringCommandKind { match self { - SteeringCommand::Pause => SteeringCommandKind::Pause, + SteeringCommand::Pause | SteeringCommand::PauseWith { .. } => { + SteeringCommandKind::Pause + } SteeringCommand::Resume => SteeringCommandKind::Resume, SteeringCommand::Cancel => SteeringCommandKind::Cancel, SteeringCommand::InjectMessage(_) => SteeringCommandKind::InjectMessage, @@ -138,16 +157,66 @@ pub struct SteeringPolicy { /// The control-flow decision produced by applying a batch of steering commands /// at a checkpoint. +/// +/// Deliberately still `Copy` and payload-free: the agent loop matches on it +/// directly, and widening a variant would break every one of those call sites +/// for no gain. The *state* behind a [`SteeringOutcome::Pause`] lives on the +/// [`SteeringHandle`] — read it with [`SteeringHandle::pause_state`]. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum SteeringOutcome { /// No steering, or only transcript/metadata mutations: continue the loop. Continue, - /// A net pause is in effect: the loop should cooperatively stop. + /// A pause is latched: the loop should cooperatively stop and report the + /// run as **paused**, not completed. + /// + /// # Contract for the agent loop (wave 2) + /// + /// The loop currently treats this as a bare `break`, which falls through to + /// the success epilogue and reports the run completed with + /// `final_response: None`. A caller cannot then tell "paused waiting for a + /// human" from "the model produced an empty answer". On this outcome the + /// loop must instead: + /// + /// 1. Read [`SteeringHandle::pause_state`] from `ctx.steering` — it is + /// always `Some` when this outcome is returned — and surface the + /// [`PauseState`] (reason, checkpoint index) to the caller. + /// 2. Report the run as paused/interrupted rather than completed + /// (`HarnessRunStatus::mark_interrupted`, or the crate's + /// `Interrupted` shape) so it is distinguishable from success. + /// 3. Leave the pause latched. It is *not* cleared by breaking out of the + /// loop: sending [`SteeringCommand::Resume`] on the same handle clears + /// it, and re-invoking the run continues from the checkpoint. Pause, /// A cancel was requested: the loop should terminate the run. Cancel, } +impl SteeringOutcome { + /// `true` when the loop should cooperatively stop for a pause. + pub fn is_pause(self) -> bool { + matches!(self, SteeringOutcome::Pause) + } +} + +/// The latched state behind a [`SteeringOutcome::Pause`]. +/// +/// A pause used to be scoped to the drained batch: [`SteeringCommand::Resume`] +/// only cleared a [`SteeringCommand::Pause`] that arrived in the *same* batch, +/// so a pause applied at one checkpoint could never be lifted — the run was +/// stuck. The state now lives on the [`SteeringHandle`], so a `Resume` sent at +/// any later moment resumes the run. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct PauseState { + /// Why the run was paused, when the orchestrator supplied one via + /// [`SteeringCommand::PauseWith`]. `None` for a bare + /// [`SteeringCommand::Pause`]. + pub reason: Option, + /// Zero-based index of the steering checkpoint at which the pause took + /// effect, i.e. how many checkpoints this handle had already processed. + /// Lets a caller (and a resumed run) report *where* the run stopped. + pub paused_at_checkpoint: usize, +} + /// A cloneable, thread-safe handle to a running agent's steering queue. /// /// An orchestrator holds a `SteeringHandle` and calls [`SteeringHandle::send`] @@ -171,4 +240,10 @@ pub(crate) struct SteeringInner { pub(crate) queue: Mutex>, /// The allowlist gating which drained commands may be applied. pub(crate) policy: SteeringPolicy, + /// The latched pause, if one is in effect. Survives across checkpoints so a + /// [`SteeringCommand::Resume`] delivered in a *later* batch can lift it. + pub(crate) paused: Mutex>, + /// How many steering checkpoints this handle has processed. Recorded into + /// [`PauseState::paused_at_checkpoint`]. + pub(crate) checkpoints: Mutex, } diff --git a/src/harness/store/mod.rs b/src/harness/store/mod.rs index 92abe08..9472aa5 100644 --- a/src/harness/store/mod.rs +++ b/src/harness/store/mod.rs @@ -15,6 +15,9 @@ //! //! # Primary types //! - [`Store`] — the core async trait every backend implements. +//! - [`namespaced::NamespacedStore`] — the hierarchical, TTL-aware, +//! batch-oriented long-term store. Additive: `Store` is unchanged, and +//! [`namespaced::FlatNamespacedStore`] adapts the new trait back to it. //! - [`InMemoryStore`] — ephemeral in-process store for tests and examples. //! - [`FileStore`] — file-system-backed store for local development. //! - [`StoreRegistry`] — named bag of stores injected into `RunContext`. @@ -24,6 +27,7 @@ //! `"artifacts"`. The registry does not enforce a naming scheme, but //! consistent names make multi-store applications easier to audit. +pub mod namespaced; mod types; use std::collections::HashMap; @@ -310,7 +314,64 @@ impl JsonlAppendStore { pub fn new(root_dir: impl Into) -> Self { Self { root_dir: root_dir.into(), - offsets: Default::default(), + append_guard: Default::default(), + } + } + + /// Returns the offset the next append to `path` should carry. + /// + /// Reads the **last complete line** of the file rather than parsing all of + /// it: an exponentially growing tail window is pulled from the end until a + /// line boundary is found. A trailing partial line (a torn write) is + /// skipped, so a crash mid-append cannot make the stream restart its + /// numbering. An absent or empty file starts at `0`. + fn next_offset(path: &std::path::Path) -> Result { + use std::io::{Read, Seek, SeekFrom}; + + let mut file = match fs::File::open(path) { + Ok(f) => f, + Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(0), + Err(e) => { + return Err(TinyAgentsError::Validation(format!( + "append store open error: {e}" + ))); + } + }; + let len = file + .metadata() + .map_err(|e| TinyAgentsError::Validation(format!("append store stat error: {e}")))? + .len(); + if len == 0 { + return Ok(0); + } + let mut window = 4096u64; + loop { + let start = len.saturating_sub(window); + file.seek(SeekFrom::Start(start)).map_err(|e| { + TinyAgentsError::Validation(format!("append store seek error: {e}")) + })?; + let mut buf = String::new(); + file.read_to_string(&mut buf).map_err(|e| { + TinyAgentsError::Validation(format!("append store read error: {e}")) + })?; + // Only lines that are certainly complete: when the window did not + // reach the start of the file, its first line may be cut in half. + let complete_from = usize::from(start > 0); + let lines: Vec<&str> = buf.lines().skip(complete_from).collect(); + for line in lines.iter().rev() { + if line.trim().is_empty() { + continue; + } + if let Ok(record) = serde_json::from_str::(line) { + return Ok(record.offset + 1); + } + // A torn trailing line: keep walking backwards. + } + if start == 0 { + // Whole file scanned and nothing decoded: treat it as empty. + return Ok(0); + } + window = window.saturating_mul(4); } } @@ -344,30 +405,22 @@ impl AppendStore for JsonlAppendStore { async fn append(&self, stream: &str, value: Value) -> Result { let path = self.stream_path(stream)?; let root_dir = self.root_dir.clone(); - let offsets = Arc::clone(&self.offsets); - let stream = stream.to_string(); + let guard = Arc::clone(&self.append_guard); // The append is pure blocking file I/O. Run it off the async runtime so // it never stalls a tokio worker (`spawn_blocking` when a runtime is // present, inline otherwise — e.g. a synchronous sink draining outside a - // runtime). The offset cache means we only read the file once per stream - // instead of re-parsing the whole file on every append (previously - // O(n²) per stream). + // runtime). let work = move || -> Result { fs::create_dir_all(&root_dir).map_err(|e| { TinyAgentsError::Validation(format!("append store mkdir error: {e}")) })?; - // Hold the offset guard across the write so concurrent appends to the - // same store instance get distinct, ordered offsets. - let mut cache = offsets.lock().map_err(|e| { + // Hold the guard across read-tail + write so concurrent appends + // through this instance get distinct, ordered offsets. + let _guard = guard.lock().map_err(|e| { TinyAgentsError::Validation(format!("append store lock poisoned: {e}")) })?; - let offset = match cache.get(&stream) { - Some(&next) => next, - // First append for this stream in this instance: learn the length - // from disk once, then track it in memory. - None => Self::read_records(&path)?.len() as u64, - }; + let offset = Self::next_offset(&path)?; let record = StoreRecord { offset, value, @@ -385,7 +438,6 @@ impl AppendStore for JsonlAppendStore { std::io::Write::write_all(&mut file, line.as_bytes()).map_err(|e| { TinyAgentsError::Validation(format!("append store write error: {e}")) })?; - cache.insert(stream, offset + 1); Ok(offset) }; @@ -399,16 +451,22 @@ impl AppendStore for JsonlAppendStore { async fn read_from(&self, stream: &str, offset: u64) -> Result> { let path = self.stream_path(stream)?; + // Resolve BY OFFSET, not by position. `.skip(offset)` is positional, so + // it disagreed with the offsets it then handed back the moment a stream + // contained anything other than a dense 0..n sequence — which is exactly + // what the documented contract (`entries whose offset is >= offset`) + // promises to handle, and what `InMemoryAppendStore` already did. Ok(Self::read_records(&path)? .into_iter() - .skip(offset as usize) + .filter(|r| r.offset >= offset) .map(|r| (r.offset, r.value)) .collect()) } async fn len(&self, stream: &str) -> Result { + // The offset the next append will receive, per the trait contract. let path = self.stream_path(stream)?; - Ok(Self::read_records(&path)?.len() as u64) + Self::next_offset(&path) } } diff --git a/src/harness/store/namespaced/mod.rs b/src/harness/store/namespaced/mod.rs new file mode 100644 index 0000000..211b316 --- /dev/null +++ b/src/harness/store/namespaced/mod.rs @@ -0,0 +1,447 @@ +//! Hierarchical, batch-oriented long-term store. +//! +//! # Why this exists next to [`Store`] +//! +//! The original [`Store`](crate::harness::store::Store) trait is get/put/ +//! delete/list over a **flat** `&str` namespace. That is enough to key values +//! by bucket, and not enough for anything a long-term memory layer actually +//! needs: there is no way to ask for everything under `users/alice`, no way to +//! filter by a field, no way to enumerate what namespaces exist, no expiry +//! (so a store only ever grows), and no way to issue several reads as one +//! round trip. +//! +//! [`NamespacedStore`] adds those. It does **not** replace `Store`: the flat +//! trait keeps working unchanged, and [`FlatNamespacedStore`] adapts any +//! `NamespacedStore` back to it, so existing callers are untouched. +//! +//! # `batch` is the one method that matters +//! +//! [`NamespacedStore::batch`] is the single required method; `get`, `put`, +//! `delete`, `search` and `list_namespaces` are convenience wrappers that +//! submit a one-operation batch. This is deliberate, and it is the shape +//! LangGraph converged on: a remote or pooled backend wants to coalesce +//! concurrent operations into one round trip, and it can only do that if every +//! path through the API funnels into one place. Implement `batch` well and +//! every other method is correct by construction; implement six methods +//! separately and the batched path drifts from the single path. +//! +//! Results are returned **positionally aligned** with the request: result `i` +//! answers operation `i`. +//! +//! # TTL +//! +//! [`TtlConfig`] gives items an expiry. An expired item is invisible to reads +//! and searches immediately, and is reclaimed by +//! [`NamespacedStore::sweep_expired`]. This is the bounded-growth answer the +//! flat store never had. +//! +//! # Semantic search +//! +//! [`SearchQuery::query`] is the seam for vector/semantic search. This crate +//! has no embedding dependency, so the bundled backend does a plain substring +//! scan; a backend with an index is expected to interpret it properly. + +mod types; + +use std::collections::{BTreeSet, HashMap}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use serde_json::Value; + +pub use types::*; + +use crate::error::{Result, TinyAgentsError}; +use crate::harness::ids::now_ms; +use crate::harness::store::Store; + +/// A hierarchical, TTL-aware, batch-oriented long-term store. +/// +/// Implement [`NamespacedStore::batch`]; everything else has a default body +/// built on it. +#[async_trait] +pub trait NamespacedStore: Send + Sync { + /// Executes `ops` and returns one result per operation, in order. + /// + /// The **only** required method. See the module docs for why. + async fn batch(&self, ops: &[StoreOp]) -> Result>; + + /// The TTL policy in force. Defaults to "nothing expires". + fn ttl_config(&self) -> TtlConfig { + TtlConfig::default() + } + + /// Removes expired items, returning how many were reclaimed. + /// + /// Expiry is enforced on read regardless, so sweeping is about reclaiming + /// space rather than correctness. The default is a no-op for backends that + /// expire lazily. + async fn sweep_expired(&self) -> Result { + Ok(0) + } + + /// Reads one item, or `None` when it is absent or expired. + async fn get(&self, namespace: &Namespace, key: &str) -> Result> { + one( + self, + StoreOp::Get { + namespace: namespace.clone(), + key: key.to_string(), + refresh_ttl: None, + }, + ) + .await? + .into_item() + } + + /// Writes one item, applying the store's default TTL. + async fn put(&self, namespace: &Namespace, key: &str, value: Value) -> Result<()> { + self.put_with_ttl(namespace, key, value, None).await + } + + /// Writes one item with an explicit lifetime in minutes. + async fn put_with_ttl( + &self, + namespace: &Namespace, + key: &str, + value: Value, + ttl_minutes: Option, + ) -> Result<()> { + one( + self, + StoreOp::Put { + namespace: namespace.clone(), + key: key.to_string(), + value: Some(value), + ttl_minutes, + }, + ) + .await + .map(|_| ()) + } + + /// Deletes one item. Deleting an absent key is not an error. + async fn delete(&self, namespace: &Namespace, key: &str) -> Result<()> { + one( + self, + StoreOp::Put { + namespace: namespace.clone(), + key: key.to_string(), + value: None, + ttl_minutes: None, + }, + ) + .await + .map(|_| ()) + } + + /// Searches a namespace subtree. + async fn search(&self, query: SearchQuery) -> Result> { + one(self, StoreOp::Search(query)).await?.into_items() + } + + /// Lists namespaces matching `query`. + async fn list_namespaces(&self, query: ListNamespacesQuery) -> Result> { + one(self, StoreOp::ListNamespaces(query)) + .await? + .into_namespaces() + } +} + +/// Submits a single operation and unwraps the one-element result vector. +async fn one(store: &S, op: StoreOp) -> Result +where + S: NamespacedStore + ?Sized, +{ + let mut results = store.batch(std::slice::from_ref(&op)).await?; + if results.len() != 1 { + return Err(TinyAgentsError::Validation(format!( + "store batch returned {} results for 1 operation — results must be \ + positionally aligned with the request", + results.len() + ))); + } + Ok(results.remove(0)) +} + +// ── InMemoryNamespacedStore ────────────────────────────────────────────────── + +/// In-process [`NamespacedStore`] backed by a map, with working TTL. +/// +/// Cheap to clone; clones share the same data and TTL policy. +#[derive(Clone, Debug, Default)] +pub struct InMemoryNamespacedStore { + items: Arc>>, + ttl: TtlConfig, +} + +impl InMemoryNamespacedStore { + /// Creates an empty store with no expiry. + pub fn new() -> Self { + Self::default() + } + + /// Sets the TTL policy. + pub fn with_ttl(mut self, ttl: TtlConfig) -> Self { + self.ttl = ttl; + self + } + + fn lock(&self) -> Result>> { + self.items + .lock() + .map_err(|e| TinyAgentsError::Validation(format!("store lock poisoned: {e}"))) + } +} + +/// Computes an expiry timestamp from a lifetime in minutes. +fn expiry_at(now: u64, ttl_minutes: Option) -> Option { + ttl_minutes + .filter(|m| m.is_finite() && *m > 0.0) + .map(|m| now.saturating_add((m * 60_000.0) as u64)) +} + +/// Reads a dotted field path out of a JSON value. +fn field<'v>(value: &'v Value, path: &str) -> Option<&'v Value> { + let mut cursor = value; + for segment in path.split('.') { + cursor = cursor.get(segment)?; + } + Some(cursor) +} + +/// Whether `item` satisfies every condition in `filter`. +fn matches_filter(item: &Item, filter: &HashMap) -> bool { + filter + .iter() + .all(|(path, op)| op.matches(field(&item.value, path))) +} + +/// Whether `item`'s rendered JSON contains `needle`, case-insensitively. +fn matches_query(item: &Item, needle: &str) -> bool { + item.value + .to_string() + .to_lowercase() + .contains(&needle.to_lowercase()) +} + +/// Whether `segments` matches `pattern`, where `*` matches any one segment. +fn matches_wildcards(segments: &[String], pattern: &[String]) -> bool { + segments.len() == pattern.len() + && segments + .iter() + .zip(pattern) + .all(|(actual, expected)| expected == "*" || actual == expected) +} + +#[async_trait] +impl NamespacedStore for InMemoryNamespacedStore { + fn ttl_config(&self) -> TtlConfig { + self.ttl + } + + async fn sweep_expired(&self) -> Result { + let now = now_ms(); + let mut items = self.lock()?; + let before = items.len(); + items.retain(|_, item| !item.is_expired(now)); + let reclaimed = before - items.len(); + if reclaimed > 0 { + tracing::debug!("[store:namespaced] sweep_expired reclaimed={reclaimed}"); + } + Ok(reclaimed) + } + + async fn batch(&self, ops: &[StoreOp]) -> Result> { + let now = now_ms(); + let mut items = self.lock()?; + let mut out = Vec::with_capacity(ops.len()); + for op in ops { + match op { + StoreOp::Get { + namespace, + key, + refresh_ttl, + } => { + namespace.validate()?; + let refresh = refresh_ttl.unwrap_or(self.ttl.refresh_on_read); + let addr = (namespace.clone(), key.clone()); + let found = match items.get_mut(&addr) { + Some(item) if item.is_expired(now) => { + // Expiry is enforced on read, not only by the + // sweeper, so a stale item is never observable. + items.remove(&addr); + None + } + Some(item) => { + if refresh && let Some(minutes) = self.ttl.default_ttl_minutes { + item.expires_at_ms = expiry_at(now, Some(minutes)); + item.updated_at_ms = now; + } + Some(item.clone()) + } + None => None, + }; + out.push(StoreResult::Item(found.map(Box::new))); + } + StoreOp::Put { + namespace, + key, + value, + ttl_minutes, + } => { + namespace.validate()?; + let addr = (namespace.clone(), key.clone()); + match value { + None => { + items.remove(&addr); + } + Some(value) => { + let ttl = ttl_minutes.or(self.ttl.default_ttl_minutes); + let created = items.get(&addr).map_or(now, |i| i.created_at_ms); + items.insert( + addr, + Item { + namespace: namespace.clone(), + key: key.clone(), + value: value.clone(), + created_at_ms: created, + updated_at_ms: now, + expires_at_ms: expiry_at(now, ttl), + }, + ); + } + } + out.push(StoreResult::Ack); + } + StoreOp::Search(query) => { + let mut found: Vec = items + .values() + .filter(|item| !item.is_expired(now)) + .filter(|item| item.namespace.starts_with(&query.namespace_prefix)) + .filter(|item| matches_filter(item, &query.filter)) + .filter(|item| { + query + .query + .as_ref() + .is_none_or(|needle| matches_query(item, needle)) + }) + .cloned() + .collect(); + // A HashMap has no order, so impose a deterministic one: + // paging is meaningless without it. + found.sort_by(|a, b| { + a.namespace + .cmp(&b.namespace) + .then_with(|| a.key.cmp(&b.key)) + }); + let windowed: Vec = found + .into_iter() + .skip(query.offset) + .take(query.limit.unwrap_or(usize::MAX)) + .collect(); + out.push(StoreResult::Items(windowed)); + } + StoreOp::ListNamespaces(query) => { + let mut seen: BTreeSet = BTreeSet::new(); + for item in items.values().filter(|i| !i.is_expired(now)) { + let ns = &item.namespace; + if let Some(prefix) = &query.prefix + && !(ns.segments().len() >= prefix.len() + && matches_wildcards(&ns.segments()[..prefix.len()], prefix)) + { + continue; + } + if let Some(suffix) = &query.suffix { + let len = ns.segments().len(); + if !(len >= suffix.len() + && matches_wildcards(&ns.segments()[len - suffix.len()..], suffix)) + { + continue; + } + } + // `max_depth` truncates then deduplicates, which is how + // one level of the hierarchy gets enumerated. + let truncated = match query.max_depth { + Some(depth) if ns.segments().len() > depth => { + Namespace(ns.segments()[..depth].to_vec()) + } + _ => ns.clone(), + }; + seen.insert(truncated); + } + let windowed: Vec = seen + .into_iter() + .skip(query.offset) + .take(query.limit.unwrap_or(usize::MAX)) + .collect(); + out.push(StoreResult::Namespaces(windowed)); + } + } + } + Ok(out) + } +} + +// ── Compatibility with the flat `Store` trait ──────────────────────────────── + +/// Exposes any [`NamespacedStore`] through the flat [`Store`] trait. +/// +/// The flat namespace becomes a single-segment [`Namespace`], so existing +/// callers keep working byte-for-byte while the data lives in a store that also +/// supports hierarchy, filtering and TTL. This is what makes the new trait +/// additive rather than a migration. +#[derive(Clone, Debug)] +pub struct FlatNamespacedStore { + inner: S, +} + +impl FlatNamespacedStore { + /// Wraps `inner`. + pub fn new(inner: S) -> Self { + Self { inner } + } + + /// Returns the wrapped store. + pub fn inner(&self) -> &S { + &self.inner + } +} + +#[async_trait] +impl Store for FlatNamespacedStore { + async fn get(&self, namespace: &str, key: &str) -> Result> { + let ns = Namespace::new([namespace])?; + Ok(self.inner.get(&ns, key).await?.map(|item| item.value)) + } + + async fn put(&self, namespace: &str, key: &str, value: Value) -> Result<()> { + let ns = Namespace::new([namespace])?; + self.inner.put(&ns, key, value).await + } + + async fn delete(&self, namespace: &str, key: &str) -> Result<()> { + let ns = Namespace::new([namespace])?; + self.inner.delete(&ns, key).await + } + + async fn list(&self, namespace: &str) -> Result> { + let ns = Namespace::new([namespace])?; + Ok(self + .inner + .search(SearchQuery { + namespace_prefix: ns.0.clone(), + ..SearchQuery::default() + }) + .await? + .into_iter() + // `starts_with` is a subtree match; the flat `list` contract is + // "keys in exactly this namespace". + .filter(|item| item.namespace == ns) + .map(|item| item.key) + .collect()) + } +} + +#[cfg(test)] +mod test; diff --git a/src/harness/store/namespaced/test.rs b/src/harness/store/namespaced/test.rs new file mode 100644 index 0000000..ac4f05b --- /dev/null +++ b/src/harness/store/namespaced/test.rs @@ -0,0 +1,311 @@ +//! Module-local tests for the hierarchical store. + +use super::*; +use crate::harness::store::Store as FlatStore; + +fn ns(segments: &[&str]) -> Namespace { + Namespace::new(segments.iter().copied()).expect("valid namespace") +} + +#[test] +fn namespace_validation_rejects_the_unaddressable() { + assert!(Namespace::new(Vec::::new()).is_err(), "empty"); + assert!(Namespace::new([""]).is_err(), "empty segment"); + assert!(Namespace::new(["a.b"]).is_err(), "separator in segment"); + assert!(Namespace::new(["langgraph", "x"]).is_err(), "reserved root"); + assert!(Namespace::new(["users", "alice"]).is_ok()); +} + +#[test] +fn namespaces_know_their_own_prefixes_and_suffixes() { + let n = ns(&["users", "alice", "memories"]); + assert!(n.starts_with(&["users".into()])); + assert!(n.starts_with(&["users".into(), "alice".into()])); + assert!(!n.starts_with(&["users".into(), "bob".into()])); + assert!(n.ends_with(&["memories".into()])); + assert!(!n.ends_with(&["alice".into()])); +} + +#[tokio::test] +async fn put_get_delete_roundtrip() { + let store = InMemoryNamespacedStore::new(); + let n = ns(&["users", "alice"]); + store + .put(&n, "k", serde_json::json!({"v": 1})) + .await + .unwrap(); + let item = store.get(&n, "k").await.unwrap().expect("stored"); + assert_eq!(item.value, serde_json::json!({"v": 1})); + assert_eq!(item.namespace, n); + store.delete(&n, "k").await.unwrap(); + assert!(store.get(&n, "k").await.unwrap().is_none()); +} + +/// A subtree search is what the flat store could never express. +#[tokio::test] +async fn search_walks_a_namespace_subtree() { + let store = InMemoryNamespacedStore::new(); + store + .put(&ns(&["users", "alice"]), "a", serde_json::json!({"n": 1})) + .await + .unwrap(); + store + .put(&ns(&["users", "bob"]), "b", serde_json::json!({"n": 2})) + .await + .unwrap(); + store + .put(&ns(&["teams", "core"]), "c", serde_json::json!({"n": 3})) + .await + .unwrap(); + + let found = store + .search(SearchQuery { + namespace_prefix: vec!["users".into()], + ..SearchQuery::default() + }) + .await + .unwrap(); + assert_eq!(found.len(), 2, "only the `users` subtree"); +} + +#[tokio::test] +async fn search_applies_comparison_filters() { + let store = InMemoryNamespacedStore::new(); + let n = ns(&["items"]); + for (key, score) in [("low", 1), ("mid", 5), ("high", 9)] { + store + .put( + &n, + key, + serde_json::json!({"score": score, "meta": {"kind": "x"}}), + ) + .await + .unwrap(); + } + + let mut filter = std::collections::HashMap::new(); + filter.insert("score".to_string(), FilterOp::Gte(5.0)); + let found = store + .search(SearchQuery { + namespace_prefix: vec!["items".into()], + filter, + ..SearchQuery::default() + }) + .await + .unwrap(); + assert_eq!(found.len(), 2); + + // Nested dotted paths resolve. + let mut filter = std::collections::HashMap::new(); + filter.insert( + "meta.kind".to_string(), + FilterOp::Eq(serde_json::json!("x")), + ); + assert_eq!( + store + .search(SearchQuery { + namespace_prefix: vec!["items".into()], + filter, + ..SearchQuery::default() + }) + .await + .unwrap() + .len(), + 3 + ); + + // `$exists` distinguishes absent from null. + let mut filter = std::collections::HashMap::new(); + filter.insert("missing".to_string(), FilterOp::Exists(false)); + assert_eq!( + store + .search(SearchQuery { + namespace_prefix: vec!["items".into()], + filter, + ..SearchQuery::default() + }) + .await + .unwrap() + .len(), + 3 + ); +} + +#[tokio::test] +async fn search_paginates_deterministically() { + let store = InMemoryNamespacedStore::new(); + let n = ns(&["items"]); + for i in 0..5 { + store + .put(&n, &format!("k{i}"), serde_json::json!(i)) + .await + .unwrap(); + } + let page = |offset| SearchQuery { + namespace_prefix: vec!["items".into()], + limit: Some(2), + offset, + ..SearchQuery::default() + }; + let first = store.search(page(0)).await.unwrap(); + let second = store.search(page(2)).await.unwrap(); + assert_eq!(first.len(), 2); + assert_eq!(second.len(), 2); + assert!( + first.iter().all(|a| second.iter().all(|b| a.key != b.key)), + "pages must not overlap — which requires a deterministic order" + ); +} + +#[tokio::test] +async fn list_namespaces_honours_wildcards_and_depth() { + let store = InMemoryNamespacedStore::new(); + for path in [ + vec!["users", "alice", "memories"], + vec!["users", "bob", "memories"], + vec!["teams", "core", "notes"], + ] { + store + .put(&ns(&path), "k", serde_json::json!(1)) + .await + .unwrap(); + } + + let all = store + .list_namespaces(ListNamespacesQuery::default()) + .await + .unwrap(); + assert_eq!(all.len(), 3); + + // `*` matches exactly one segment. + let wildcard = store + .list_namespaces(ListNamespacesQuery { + prefix: Some(vec!["users".into(), "*".into()]), + ..ListNamespacesQuery::default() + }) + .await + .unwrap(); + assert_eq!(wildcard.len(), 2); + + let by_suffix = store + .list_namespaces(ListNamespacesQuery { + suffix: Some(vec!["memories".into()]), + ..ListNamespacesQuery::default() + }) + .await + .unwrap(); + assert_eq!(by_suffix.len(), 2); + + // Truncating to depth 1 collapses the tree to its roots. + let roots = store + .list_namespaces(ListNamespacesQuery { + max_depth: Some(1), + ..ListNamespacesQuery::default() + }) + .await + .unwrap(); + assert_eq!(roots.len(), 2, "`users` and `teams`: {roots:?}"); +} + +/// TTL is the bounded-growth answer the flat store never had. An expired item +/// must be invisible on read *before* any sweep runs. +#[tokio::test] +async fn expired_items_are_invisible_and_reclaimable() { + let store = InMemoryNamespacedStore::new(); + let n = ns(&["cache"]); + // A negative-in-effect lifetime: expire essentially immediately. + store + .put_with_ttl(&n, "k", serde_json::json!(1), Some(f64::MIN_POSITIVE)) + .await + .unwrap(); + store + .put_with_ttl(&n, "keep", serde_json::json!(2), None) + .await + .unwrap(); + + assert!( + store.get(&n, "k").await.unwrap().is_none(), + "an expired item is invisible on read, not only after a sweep" + ); + assert!(store.get(&n, "keep").await.unwrap().is_some()); + let found = store + .search(SearchQuery { + namespace_prefix: vec!["cache".into()], + ..SearchQuery::default() + }) + .await + .unwrap(); + assert_eq!(found.len(), 1, "search skips expired items too"); +} + +#[tokio::test] +async fn a_default_ttl_applies_to_writes_that_do_not_set_one() { + let store = InMemoryNamespacedStore::new().with_ttl(TtlConfig { + default_ttl_minutes: Some(60.0), + refresh_on_read: false, + }); + let n = ns(&["cache"]); + store.put(&n, "k", serde_json::json!(1)).await.unwrap(); + let item = store.get(&n, "k").await.unwrap().unwrap(); + assert!( + item.expires_at_ms.is_some(), + "the store default applies when the write does not name a TTL" + ); +} + +/// `batch` is the single abstract method, so a multi-op request must come back +/// positionally aligned — every convenience method depends on that. +#[tokio::test] +async fn batch_results_align_positionally_with_the_request() { + let store = InMemoryNamespacedStore::new(); + let n = ns(&["b"]); + let results = store + .batch(&[ + StoreOp::Put { + namespace: n.clone(), + key: "k".into(), + value: Some(serde_json::json!("v")), + ttl_minutes: None, + }, + StoreOp::Get { + namespace: n.clone(), + key: "k".into(), + refresh_ttl: None, + }, + StoreOp::Search(SearchQuery { + namespace_prefix: vec!["b".into()], + ..SearchQuery::default() + }), + StoreOp::ListNamespaces(ListNamespacesQuery::default()), + ]) + .await + .unwrap(); + assert_eq!(results.len(), 4); + assert!(matches!(results[0], StoreResult::Ack)); + assert!(matches!(results[1], StoreResult::Item(Some(_)))); + assert!(matches!(&results[2], StoreResult::Items(v) if v.len() == 1)); + assert!(matches!(&results[3], StoreResult::Namespaces(v) if v.len() == 1)); +} + +/// The new trait must be additive: existing flat-`Store` callers keep working. +#[tokio::test] +async fn the_flat_store_surface_still_works() { + let store = FlatNamespacedStore::new(InMemoryNamespacedStore::new()); + store + .put("events", "e1", serde_json::json!({"a": 1})) + .await + .unwrap(); + assert_eq!( + FlatStore::get(&store, "events", "e1").await.unwrap(), + Some(serde_json::json!({"a": 1})) + ); + assert_eq!(store.list("events").await.unwrap(), vec!["e1".to_string()]); + FlatStore::delete(&store, "events", "e1").await.unwrap(); + assert!( + FlatStore::get(&store, "events", "e1") + .await + .unwrap() + .is_none() + ); + assert!(store.list("events").await.unwrap().is_empty()); +} diff --git a/src/harness/store/namespaced/types.rs b/src/harness/store/namespaced/types.rs new file mode 100644 index 0000000..dac1b0e --- /dev/null +++ b/src/harness/store/namespaced/types.rs @@ -0,0 +1,296 @@ +//! Types for the hierarchical, batch-oriented long-term store. +//! +//! See the module docs on [`super`] for why this exists alongside the flat +//! [`Store`](crate::harness::store::Store) trait. + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::error::{Result, TinyAgentsError}; + +/// A hierarchical namespace: an ordered tuple of path segments. +/// +/// `("users", "alice", "memories")` is a child of `("users", "alice")`, which +/// is what makes prefix search and namespace listing meaningful. The flat +/// `&str` namespace the original [`Store`](crate::harness::store::Store) trait +/// uses is the degenerate one-segment case, and converts freely. +/// +/// # Validation +/// +/// A namespace must be non-empty, every segment must be non-empty, no segment +/// may contain the separator `.`, and the first segment may not be the reserved +/// word `langgraph` (kept reserved for compatibility with stores shared with +/// that ecosystem). Segments are otherwise opaque. +#[derive(Clone, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)] +pub struct Namespace(pub Vec); + +/// Separator used when rendering a namespace as a single string. +pub const NAMESPACE_SEPARATOR: char = '.'; + +/// Reserved first segment. +const RESERVED_ROOT: &str = "langgraph"; + +impl Namespace { + /// Builds a namespace from any iterator of segments, validating it. + pub fn new(segments: I) -> Result + where + I: IntoIterator, + S: Into, + { + let ns = Namespace(segments.into_iter().map(Into::into).collect()); + ns.validate()?; + Ok(ns) + } + + /// The segments, in order. + pub fn segments(&self) -> &[String] { + &self.0 + } + + /// Whether this namespace is `prefix` or lies beneath it. + pub fn starts_with(&self, prefix: &[String]) -> bool { + self.0.len() >= prefix.len() && self.0[..prefix.len()] == *prefix + } + + /// Whether this namespace ends with `suffix`. + pub fn ends_with(&self, suffix: &[String]) -> bool { + self.0.len() >= suffix.len() && self.0[self.0.len() - suffix.len()..] == *suffix + } + + /// Rejects namespaces that would break addressing or collide with the + /// reserved root. + pub fn validate(&self) -> Result<()> { + if self.0.is_empty() { + return Err(TinyAgentsError::Validation( + "store namespace must have at least one segment".into(), + )); + } + for segment in &self.0 { + if segment.is_empty() { + return Err(TinyAgentsError::Validation( + "store namespace segments must not be empty".into(), + )); + } + if segment.contains(NAMESPACE_SEPARATOR) { + return Err(TinyAgentsError::Validation(format!( + "store namespace segment {segment:?} must not contain \ + {NAMESPACE_SEPARATOR:?}" + ))); + } + } + if self.0[0] == RESERVED_ROOT { + return Err(TinyAgentsError::Validation(format!( + "{RESERVED_ROOT:?} is a reserved root namespace" + ))); + } + Ok(()) + } +} + +impl From<&str> for Namespace { + fn from(value: &str) -> Self { + Namespace(vec![value.to_string()]) + } +} + +impl std::fmt::Display for Namespace { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0.join(&NAMESPACE_SEPARATOR.to_string())) + } +} + +/// Time-to-live policy for stored items. +#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize)] +pub struct TtlConfig { + /// Lifetime applied to items written without an explicit TTL, in minutes. + /// + /// `None` means items never expire, which is the pre-TTL behaviour and the + /// default. + pub default_ttl_minutes: Option, + /// Whether reading an item extends its life by the default TTL. + /// + /// A sliding window suits caches and working memory; leave it off for a + /// hard retention bound (which is what makes TTL usable as the answer to + /// unbounded store growth). + pub refresh_on_read: bool, +} + +/// A stored value together with its address and timestamps. +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct Item { + /// Namespace the item lives in. + pub namespace: Namespace, + /// Key within the namespace. + pub key: String, + /// The stored payload. + pub value: Value, + /// Unix-epoch milliseconds at first write. + pub created_at_ms: u64, + /// Unix-epoch milliseconds at the most recent write (or TTL refresh). + pub updated_at_ms: u64, + /// Unix-epoch milliseconds after which the item is invisible, if any. + pub expires_at_ms: Option, +} + +impl Item { + /// Whether the item is expired as of `now_ms`. + pub fn is_expired(&self, now_ms: u64) -> bool { + self.expires_at_ms.is_some_and(|expiry| now_ms >= expiry) + } +} + +/// One comparison in a [`SearchQuery`] filter. +/// +/// The operator set mirrors LangGraph's: equality by default, plus explicit +/// `$eq`/`$ne`/`$gt`/`$gte`/`$lt`/`$lte`/`$in`/`$exists`. +#[derive(Clone, Debug, PartialEq)] +pub enum FilterOp { + /// Field equals the value. + Eq(Value), + /// Field does not equal the value. + Ne(Value), + /// Field is numerically greater than the value. + Gt(f64), + /// Field is numerically greater than or equal to the value. + Gte(f64), + /// Field is numerically less than the value. + Lt(f64), + /// Field is numerically less than or equal to the value. + Lte(f64), + /// Field equals one of the values. + In(Vec), + /// Field is present (`true`) or absent (`false`). + Exists(bool), +} + +impl FilterOp { + /// Evaluates the operator against a field that may be absent. + pub fn matches(&self, field: Option<&Value>) -> bool { + match self { + FilterOp::Exists(expected) => field.is_some() == *expected, + FilterOp::Eq(expected) => field == Some(expected), + FilterOp::Ne(expected) => field != Some(expected), + FilterOp::In(options) => field.is_some_and(|v| options.contains(v)), + FilterOp::Gt(n) => field.and_then(Value::as_f64).is_some_and(|v| v > *n), + FilterOp::Gte(n) => field.and_then(Value::as_f64).is_some_and(|v| v >= *n), + FilterOp::Lt(n) => field.and_then(Value::as_f64).is_some_and(|v| v < *n), + FilterOp::Lte(n) => field.and_then(Value::as_f64).is_some_and(|v| v <= *n), + } + } +} + +/// A namespace-prefixed search over stored items. +#[derive(Clone, Debug, Default)] +pub struct SearchQuery { + /// Only items in this namespace or beneath it are considered. + pub namespace_prefix: Vec, + /// Field-path → condition. Every entry must match (conjunction). Paths are + /// dotted, so `"meta.kind"` addresses a nested object field. + pub filter: HashMap, + /// Optional substring match applied to the item's rendered JSON. + /// + /// This is the seam where semantic/vector search belongs. There is no + /// embedding dependency in this crate, so the built-in backends do a plain + /// case-insensitive substring scan; a backend with an index is free to + /// interpret it properly, and callers should treat result *order* as + /// relevance-defined rather than guaranteed. + pub query: Option, + /// Maximum items to return. `None` means unlimited. + pub limit: Option, + /// Items to skip before collecting results. + pub offset: usize, +} + +/// A namespace listing query with prefix/suffix wildcards. +#[derive(Clone, Debug, Default)] +pub struct ListNamespacesQuery { + /// Match namespaces starting with these segments. `*` matches one segment. + pub prefix: Option>, + /// Match namespaces ending with these segments. `*` matches one segment. + pub suffix: Option>, + /// Truncate returned namespaces to at most this many segments, then + /// deduplicate — the way to enumerate one level of the hierarchy. + pub max_depth: Option, + /// Maximum namespaces to return. + pub limit: Option, + /// Namespaces to skip. + pub offset: usize, +} + +/// One operation in a [`NamespacedStore::batch`] request. +#[derive(Clone, Debug)] +pub enum StoreOp { + /// Read one item. + Get { + /// Namespace to read from. + namespace: Namespace, + /// Key to read. + key: String, + /// Whether to extend the item's TTL on this read, overriding + /// [`TtlConfig::refresh_on_read`]. + refresh_ttl: Option, + }, + /// Write or delete one item. + Put { + /// Namespace to write to. + namespace: Namespace, + /// Key to write. + key: String, + /// The value, or `None` to delete the item. + value: Option, + /// Lifetime in minutes, overriding [`TtlConfig::default_ttl_minutes`]. + ttl_minutes: Option, + }, + /// Search a namespace subtree. + Search(SearchQuery), + /// List namespaces. + ListNamespaces(ListNamespacesQuery), +} + +/// The result of one [`StoreOp`], positionally aligned with the request. +#[derive(Clone, Debug, PartialEq)] +pub enum StoreResult { + /// Result of a `Get`: the item, or `None` if missing or expired. + Item(Option>), + /// Result of a `Put`: nothing to return. + Ack, + /// Result of a `Search`. + Items(Vec), + /// Result of a `ListNamespaces`. + Namespaces(Vec), +} + +impl StoreResult { + /// Unwraps a `Get` result, erroring if the batch returned a different shape. + pub fn into_item(self) -> Result> { + match self { + StoreResult::Item(item) => Ok(item.map(|b| *b)), + other => Err(mismatch("Item", &other)), + } + } + + /// Unwraps a `Search` result. + pub fn into_items(self) -> Result> { + match self { + StoreResult::Items(items) => Ok(items), + other => Err(mismatch("Items", &other)), + } + } + + /// Unwraps a `ListNamespaces` result. + pub fn into_namespaces(self) -> Result> { + match self { + StoreResult::Namespaces(ns) => Ok(ns), + other => Err(mismatch("Namespaces", &other)), + } + } +} + +fn mismatch(expected: &str, got: &StoreResult) -> TinyAgentsError { + TinyAgentsError::Validation(format!( + "store batch returned a {got:?} result where {expected} was expected — \ + a backend must return results positionally aligned with the request" + )) +} diff --git a/src/harness/store/types.rs b/src/harness/store/types.rs index ac77786..90db4b0 100644 --- a/src/harness/store/types.rs +++ b/src/harness/store/types.rs @@ -213,27 +213,35 @@ pub(crate) type AppendEntry = (u64, Value); /// only ASCII alphanumerics, hyphens (`-`), underscores (`_`), and dots (`.`) /// are allowed, and all-dot names are rejected. This blocks path traversal. /// +/// # Offsets are read from the file, not remembered +/// The next offset is derived from the **tail of the stream file** on every +/// append: the last complete record's offset plus one. That is O(1) per append +/// (a bounded read from the end, not a full parse), so it costs no more than +/// the in-memory counter it replaced — and unlike that counter it is correct +/// when more than one `JsonlAppendStore` addresses the same directory. +/// +/// The counter it replaced was learned once per *instance*: a second instance +/// over the same root started counting from the length it happened to observe, +/// so two instances routinely handed out the same offset for different records. +/// Since [`AppendStore::read_from`] resolves by offset, a stream with duplicate +/// offsets returns a window that does not match the labels on its own entries. +/// /// # Concurrency /// Operations use blocking [`std::fs`] (no async-fs dependency is pulled in for /// this local backend), but `append` runs that I/O on a blocking thread /// (`spawn_blocking`) when a tokio runtime is present so it never stalls an /// async worker. Appends use `OpenOptions::append`, which is atomic per line on -/// POSIX for small writes, but no advisory lock is held. To avoid re-parsing the -/// whole file on every append, each store instance caches the next offset per -/// stream (see [`Self::offsets`]); this assumes a single writing process per -/// directory. For multiple concurrent writers, funnel appends through one store -/// instance (its offset guard serialises them) or prefer a server backend. +/// POSIX for small writes. A per-instance guard serialises appends from within +/// one process; across processes the tail read makes duplicate offsets far less +/// likely but does not make the append atomic, so a server backend is still the +/// right answer for genuinely concurrent multi-process writers. #[derive(Clone, Debug, Default)] pub struct JsonlAppendStore { /// The root directory under which `.jsonl` files live. pub(crate) root_dir: PathBuf, - /// Per-stream cache of the *next* offset to write, so an append does not - /// have to re-read and re-parse the whole file to learn its length (which - /// made appends O(n²) per stream). Initialised lazily from disk on the - /// first append for a stream and incremented in memory thereafter; the - /// guard is held across the write so concurrent appends to the same store - /// instance stay correctly ordered. Clones share the same cache. - pub(crate) offsets: Arc>>, + /// Serialises appends issued through this instance (and its clones) so two + /// tasks cannot read the same tail offset and then both write it. + pub(crate) append_guard: Arc>, } // ── StoreRegistry ──────────────────────────────────────────────────────────── diff --git a/src/harness/stream/mod.rs b/src/harness/stream/mod.rs index 1712142..fea3504 100644 --- a/src/harness/stream/mod.rs +++ b/src/harness/stream/mod.rs @@ -17,17 +17,33 @@ //! - [`stream`] — a convenience helper that filters a slice of chunks by a //! set of modes and returns the matching subset. //! -//! The stream module is **independent** of `crate::harness::events`: it -//! provides a higher-level projection API without importing the event bus. -//! Integration between event delivery and stream chunks is the responsibility -//! of the harness runtime. - +//! - [`project_event`] / [`project_event_for_modes`] / [`projected_mode`] — the +//! [`AgentEvent`] → [`StreamChunk`] projection (see the `project` module docs +//! for the routing table). +//! +//! The chunk *types* are independent of `crate::harness::events`; only the +//! projection depends on it. Callers that never touch events still pay nothing +//! for the coupling. +//! +//! # Wiring note for the agent loop (wave 2) +//! +//! `invoke_stream` should build a `StreamSink` from the caller's requested +//! [`StreamMode`]s and call [`StreamSink::push_event`] (or +//! [`project_event_for_modes`] directly) for every emitted [`AgentEvent`], +//! instead of handing raw events to the caller. [`StreamMode::Values`] is the +//! one mode the projection cannot supply — a full state snapshot is graph +//! state, so the graph runtime pushes [`StreamChunk::Values`] itself. + +mod project; mod types; +pub use project::{project_event, project_event_for_modes, projected_mode}; pub use types::*; use std::collections::HashSet; +use crate::harness::events::AgentEvent; + // --------------------------------------------------------------------------- // StreamSink impls // --------------------------------------------------------------------------- @@ -84,6 +100,24 @@ impl StreamSink { } } + /// Projects `event` and buffers the resulting chunk when its mode is + /// active. Returns `true` when a chunk was buffered. + /// + /// The one-line adapter between the event bus and this sink: a streaming + /// run loop calls this for every emitted [`AgentEvent`] and the sink's + /// active-mode set does the filtering. Projection is skipped entirely for + /// inactive modes (see [`project_event_for_modes`]). + pub fn push_event(&self, event: &AgentEvent) -> bool { + let modes: Vec = self.active_modes.iter().copied().collect(); + match project_event_for_modes(event, &modes) { + Some(chunk) => { + self.buffer.borrow_mut().push(chunk); + true + } + None => false, + } + } + /// Returns all buffered chunks in push order and clears the buffer. pub fn drain(&self) -> Vec { self.buffer.borrow_mut().drain(..).collect() diff --git a/src/harness/stream/project.rs b/src/harness/stream/project.rs new file mode 100644 index 0000000..3955de7 --- /dev/null +++ b/src/harness/stream/project.rs @@ -0,0 +1,156 @@ +//! Projection from the raw [`AgentEvent`] stream onto typed [`StreamChunk`]s. +//! +//! This is the missing half of the streaming surface. [`StreamMode`], +//! [`StreamChunk`] and [`StreamSink`] were fully defined and tested but used +//! nowhere else in the crate — callers got raw `AgentEvent`s and each +//! re-implemented delta reassembly and filtering. This module supplies the one +//! function that turns an event into the chunk a consumer actually asked for. +//! +//! # Shape (ported from LangGraph) +//! +//! LangGraph multiplexes several *stream modes* over one producer: the producer +//! filters by the requested mode set (`_loop.py`'s `stream_modes` check) and the +//! consumer shapes the result (`main.py`'s `stream()` fan-out). The same split +//! applies here: +//! +//! - **Producer-side filtering** — [`project_event_for_modes`] returns `None` +//! for an event whose chunk is not in the requested mode set, so nothing is +//! allocated or cloned for a mode nobody is listening to. +//! - **Consumer-side shaping** — [`project_event`] gives the full projection +//! when the caller wants to route chunks itself. +//! +//! # One event, at most one chunk +//! +//! Each [`AgentEvent`] projects to **at most one** [`StreamChunk`], so a +//! consumer subscribed to several modes never receives the same event twice in +//! two shapes. The routing table: +//! +//! | Event | Chunk | Mode | +//! |---|---|---| +//! | [`AgentEvent::ModelDelta`] | [`StreamChunk::Message`] | [`StreamMode::Messages`] | +//! | [`AgentEvent::StateUpdate`] | [`StreamChunk::Updates`] | [`StreamMode::Updates`] | +//! | [`AgentEvent::ControlApplied`] with an interrupting control | [`StreamChunk::Interrupt`] | [`StreamMode::Interrupts`] | +//! | [`AgentEvent::StreamClosed`] | — (not projected) | — | +//! | everything else | [`StreamChunk::Debug`] | [`StreamMode::Debug`] | +//! +//! [`StreamMode::Values`] is deliberately **not** produced here: a full state +//! snapshot is graph state, which the event stream does not carry. The graph +//! runtime pushes [`StreamChunk::Values`] itself. +//! +//! [`StreamMode::Custom`] is likewise never produced — by definition it is the +//! caller's own extension channel. + +use crate::harness::events::AgentEvent; + +use super::{StreamChunk, StreamMode}; + +/// Control kinds from +/// [`MiddlewareControl::kind`][crate::harness::context::MiddlewareControl::kind] +/// that mean "the run paused and is waiting for something external", and so +/// belong on [`StreamMode::Interrupts`] rather than in the debug firehose. +const INTERRUPTING_CONTROLS: [&str; 1] = ["interrupt"]; + +/// Projects one [`AgentEvent`] onto the [`StreamChunk`] a consumer sees. +/// +/// Returns `None` for an event with no meaningful chunk representation — today +/// only [`AgentEvent::StreamClosed`], which is a stream terminator rather than +/// content (consumers detect end-of-stream by the stream ending). +/// +/// # Example +/// +/// ``` +/// use tinyagents::harness::events::AgentEvent; +/// use tinyagents::harness::ids::{CallId, RunId}; +/// use tinyagents::harness::message::MessageDelta; +/// use tinyagents::harness::stream::{project_event, StreamChunk}; +/// +/// let event = AgentEvent::ModelDelta { +/// run_id: RunId::new("r1"), +/// call_id: CallId::new("c1"), +/// delta: MessageDelta::text("hi"), +/// }; +/// assert!(matches!(project_event(&event), Some(StreamChunk::Message(_)))); +/// ``` +pub fn project_event(event: &AgentEvent) -> Option { + match event { + // The whole point of `messages` mode: raw token/tool-call fragments, + // already reassembled by the provider adapter. + AgentEvent::ModelDelta { delta, .. } => Some(StreamChunk::Message(delta.clone())), + + // `StateUpdate` is payload-free by design (see its variant docs), so the + // chunk carries the fact of the update and its provenance rather than a + // diff the event never had. + AgentEvent::StateUpdate => Some(StreamChunk::Updates(serde_json::json!({ + "kind": event.kind(), + }))), + + // A human-in-the-loop pause. This is the only producer of + // `StreamChunk::Interrupt` — the variant existed but nothing ever + // constructed it. + AgentEvent::ControlApplied { control, detail } + if INTERRUPTING_CONTROLS.contains(&control.as_str()) => + { + Some(StreamChunk::Interrupt(serde_json::json!({ + "control": control, + "detail": detail, + }))) + } + + // A terminator, not content. + AgentEvent::StreamClosed => None, + + // Everything else is diagnostic. Rendered as the event's own JSON so a + // debug consumer keeps the full typed payload rather than a lossy + // summary; falls back to the stable kind string if serialization ever + // fails (it cannot for the current variants, but a `Debug` chunk must + // never be the thing that panics a run). + other => Some(StreamChunk::Debug( + serde_json::to_string(other).unwrap_or_else(|_| other.kind().to_string()), + )), + } +} + +/// Producer-side filtered projection: like [`project_event`], but returns +/// `None` when the resulting chunk's mode is not in `modes`. +/// +/// This is the function a streaming run loop should call per event — it skips +/// the clone/serialization for any mode nobody subscribed to, which matters +/// because the [`StreamMode::Debug`] catch-all serializes every event. +/// +/// # Example +/// +/// ``` +/// use tinyagents::harness::events::AgentEvent; +/// use tinyagents::harness::stream::{project_event_for_modes, StreamMode}; +/// +/// let event = AgentEvent::StateUpdate; +/// assert!(project_event_for_modes(&event, &[StreamMode::Updates]).is_some()); +/// assert!(project_event_for_modes(&event, &[StreamMode::Messages]).is_none()); +/// ``` +pub fn project_event_for_modes(event: &AgentEvent, modes: &[StreamMode]) -> Option { + let mode = projected_mode(event)?; + if !modes.contains(&mode) { + return None; + } + project_event(event) +} + +/// Returns the [`StreamMode`] an event projects onto, without building the +/// chunk. +/// +/// Cheap enough to call in a hot loop, and it is what makes +/// [`project_event_for_modes`] able to skip work rather than build-then-discard. +/// `None` mirrors [`project_event`] returning `None`. +pub fn projected_mode(event: &AgentEvent) -> Option { + match event { + AgentEvent::ModelDelta { .. } => Some(StreamMode::Messages), + AgentEvent::StateUpdate => Some(StreamMode::Updates), + AgentEvent::ControlApplied { control, .. } + if INTERRUPTING_CONTROLS.contains(&control.as_str()) => + { + Some(StreamMode::Interrupts) + } + AgentEvent::StreamClosed => None, + _ => Some(StreamMode::Debug), + } +} diff --git a/src/harness/stream/test.rs b/src/harness/stream/test.rs index 53381b4..c0b863a 100644 --- a/src/harness/stream/test.rs +++ b/src/harness/stream/test.rs @@ -166,3 +166,164 @@ fn stream_chunk_null_values_does_not_corrupt_to_empty_object() { assert_eq!(value["type"], json!("values")); assert_eq!(value["content"], json!(null)); } + +// --------------------------------------------------------------------------- +// C6: AgentEvent -> StreamChunk projection +// --------------------------------------------------------------------------- + +mod project { + use crate::harness::events::AgentEvent; + use crate::harness::ids::{CallId, RunId}; + use crate::harness::message::MessageDelta; + use crate::harness::stream::{ + StreamChunk, StreamMode, StreamSink, project_event, project_event_for_modes, projected_mode, + }; + + fn delta_event() -> AgentEvent { + AgentEvent::ModelDelta { + run_id: RunId::new("r1"), + call_id: CallId::new("c1"), + delta: MessageDelta::text("hello"), + } + } + + #[test] + fn model_deltas_project_onto_messages_mode() { + // Before this projection existed, `StreamMode` / `StreamChunk` / + // `StreamSink` were referenced NOWHERE outside their own module, and + // every caller re-implemented delta reassembly against raw events. + let chunk = project_event(&delta_event()).expect("a delta must project"); + assert_eq!(chunk.mode(), StreamMode::Messages); + assert_eq!(chunk, StreamChunk::Message(MessageDelta::text("hello"))); + } + + #[test] + fn an_interrupting_control_is_the_producer_of_stream_chunk_interrupt() { + // `StreamChunk::Interrupt` was defined but never constructed anywhere + // in the crate. + let event = AgentEvent::ControlApplied { + control: "interrupt".into(), + detail: "approval_node: needs sign-off".into(), + }; + let chunk = project_event(&event).expect("must project"); + assert_eq!(chunk.mode(), StreamMode::Interrupts); + match chunk { + StreamChunk::Interrupt(value) => { + assert_eq!(value["control"], "interrupt"); + assert_eq!(value["detail"], "approval_node: needs sign-off"); + } + other => panic!("expected an Interrupt chunk, got {other:?}"), + } + } + + #[test] + fn a_non_interrupting_control_stays_in_the_debug_channel() { + let event = AgentEvent::ControlApplied { + control: "stop_with_final".into(), + detail: "done".into(), + }; + assert_eq!(projected_mode(&event), Some(StreamMode::Debug)); + } + + #[test] + fn state_updates_project_onto_updates_mode() { + let chunk = project_event(&AgentEvent::StateUpdate).expect("must project"); + assert_eq!(chunk.mode(), StreamMode::Updates); + } + + #[test] + fn every_other_event_falls_through_to_debug_and_keeps_its_payload() { + let event = AgentEvent::ToolStarted { + call_id: CallId::new("c9"), + tool_name: "search".into(), + }; + let chunk = project_event(&event).expect("must project"); + assert_eq!(chunk.mode(), StreamMode::Debug); + match chunk { + // The typed payload survives, not just the kind string. + StreamChunk::Debug(text) => { + assert!(text.contains("tool_started"), "got {text}"); + assert!(text.contains("search"), "got {text}"); + } + other => panic!("expected a Debug chunk, got {other:?}"), + } + } + + #[test] + fn stream_closed_is_a_terminator_not_content() { + assert!(project_event(&AgentEvent::StreamClosed).is_none()); + assert!(projected_mode(&AgentEvent::StreamClosed).is_none()); + assert!(project_event_for_modes(&AgentEvent::StreamClosed, &[StreamMode::Debug]).is_none()); + } + + #[test] + fn one_event_never_projects_into_two_modes() { + // A consumer subscribed to several modes must not see the same event + // twice in two shapes. + for event in [ + delta_event(), + AgentEvent::StateUpdate, + AgentEvent::ControlApplied { + control: "interrupt".into(), + detail: "d".into(), + }, + AgentEvent::MemoryLoaded, + ] { + let all = [ + StreamMode::Values, + StreamMode::Updates, + StreamMode::Messages, + StreamMode::Debug, + StreamMode::Interrupts, + StreamMode::Custom, + ]; + let hits = all + .iter() + .filter(|mode| project_event_for_modes(&event, &[**mode]).is_some()) + .count(); + assert_eq!(hits, 1, "{} projected into {hits} modes", event.kind()); + } + } + + #[test] + fn mode_filtering_happens_producer_side() { + let event = delta_event(); + assert!(project_event_for_modes(&event, &[StreamMode::Messages]).is_some()); + assert!(project_event_for_modes(&event, &[StreamMode::Debug]).is_none()); + assert!(project_event_for_modes(&event, &[]).is_none()); + // Multiplexed mode sets work the way LangGraph's do. + assert!( + project_event_for_modes(&event, &[StreamMode::Debug, StreamMode::Messages]).is_some() + ); + } + + #[test] + fn push_event_bridges_the_event_bus_into_a_sink() { + let sink = StreamSink::new([StreamMode::Messages]); + assert!(sink.push_event(&delta_event())); + assert!(!sink.push_event(&AgentEvent::StateUpdate)); + assert!(!sink.push_event(&AgentEvent::StreamClosed)); + + let chunks = sink.drain(); + assert_eq!(chunks.len(), 1); + assert_eq!(chunks[0].mode(), StreamMode::Messages); + } + + #[test] + fn values_and_custom_are_never_produced_by_the_projection() { + // Documented gap: a full state snapshot is graph state, which the event + // stream does not carry, and Custom is the caller's own channel. + for event in [ + delta_event(), + AgentEvent::StateUpdate, + AgentEvent::MemorySaved, + AgentEvent::RunCompleted { + run_id: RunId::new("r1"), + }, + ] { + let mode = projected_mode(&event); + assert_ne!(mode, Some(StreamMode::Values)); + assert_ne!(mode, Some(StreamMode::Custom)); + } + } +} diff --git a/src/harness/structured/mod.rs b/src/harness/structured/mod.rs index 525cb30..1684658 100644 --- a/src/harness/structured/mod.rs +++ b/src/harness/structured/mod.rs @@ -24,6 +24,22 @@ //! [`ResponseFormat`] to include in a [`ModelRequest`], then call //! [`StructuredExtractor::extract`] on the completed [`ModelResponse`]. //! +//! # Repair, validation, and non-fatal extraction +//! +//! Extraction is not a bare `serde_json::from_str` any more. Three things +//! happen around it, each in its own submodule: +//! +//! * [`repair`] climbs a conservative ladder — code fence, prose slice, +//! relaxed JSON, truncation close — so a fenced, chatty, or cut-off answer is +//! recovered instead of ending a run. It never invents structure: a rung is +//! accepted only when the repaired text parses strictly. +//! * [`validate`] checks the parsed value against the declared schema, so +//! `{"wrong_key": 1}` against a `score` schema is a reported error naming the +//! failing instance path — not a silent success. +//! * [`StructuredExtractor::extract_outcome`] returns a [`StructuredOutcome`] +//! instead of `Result`, recording a failure as data so a caller can repair, +//! re-ask, or return the raw response rather than losing the run. +//! //! # Example //! //! ```rust @@ -42,8 +58,11 @@ //! assert_eq!(output.value["score"], 42); //! ``` +mod repair; mod types; +mod validate; +pub use repair::JsonRepair; pub use types::*; use serde::de::DeserializeOwned; @@ -62,7 +81,23 @@ impl StructuredStrategy { /// Returns [`StructuredStrategy::ProviderSchema`] when the model advertises /// native structured output *and* JSON Schema support, or when no profile is /// available (the conservative default). Otherwise returns - /// [`StructuredStrategy::ToolCall`], which works on any tool-calling model. + /// [`StructuredStrategy::ToolCall`] — but **only for a model that can + /// actually call tools**. + /// + /// # Why the `tool_calling` check matters + /// + /// This used to select `ToolCall` for *any* profile lacking native + /// structured output, including profiles that declare `tool_calling: + /// false`. That strategy declares an artificial tool and forces + /// [`ToolChoice::Tool`][tc]; on a model with no tool support the harness + /// runs prompt-guided instead, so the wire `tools` array is empty and the + /// forced choice is dropped — leaving a request that asks for nothing in + /// particular and an extractor waiting for a tool call that can never + /// arrive. Provider-schema mode at least asks for JSON and, with the repair + /// ladder in [`super::structured::repair`], parses what a JSON-mode model + /// actually returns. + /// + /// [tc]: crate::harness::model::ToolChoice::Tool /// /// # Example /// @@ -90,12 +125,24 @@ impl StructuredStrategy { /// StructuredStrategy::for_profile(Some(&profile)), /// StructuredStrategy::ProviderSchema /// ); + /// + /// // A model that can do neither -> JSON in the text, not a tool call. + /// let plain = ModelProfile::default(); + /// assert_eq!( + /// StructuredStrategy::for_profile(Some(&plain)), + /// StructuredStrategy::ProviderSchema + /// ); /// ``` pub fn for_profile(profile: Option<&ModelProfile>) -> StructuredStrategy { match profile { - Some(p) if !(p.native_structured_output && p.json_schema) => { - StructuredStrategy::ToolCall + Some(p) if p.native_structured_output && p.json_schema => { + StructuredStrategy::ProviderSchema } + Some(p) if p.tool_calling => StructuredStrategy::ToolCall, + // No native schema support and no tool calling: ask for JSON in the + // text and lean on the repair ladder. A dedicated `JsonMode` arm + // (plain JSON object + schema in the prompt, LangChain's third + // `method`) is the refinement — see the module docs. _ => StructuredStrategy::ProviderSchema, } } @@ -133,8 +180,8 @@ impl StructuredExtractor { /// * `strategy` – whether to use provider-schema or tool-call extraction. /// * `schema_name` – the schema's logical name; used as the tool name when /// matching tool calls in [`StructuredStrategy::ToolCall`] mode. - /// * `schema` – the JSON Schema document (retained for future local - /// validation, not yet applied). + /// * `schema` – the JSON Schema document. Enforced: every extracted value + /// is validated against it (see [`validate`]). pub fn new( strategy: StructuredStrategy, schema_name: impl Into, @@ -149,7 +196,7 @@ impl StructuredExtractor { /// Returns the JSON Schema document this extractor was configured with. /// - /// Retained for local validation and for echoing the schema back into a + /// Used for local validation and for echoing the schema back into a /// [`ResponseFormat`] when re-requesting structured output. pub fn schema(&self) -> &Value { &self.schema @@ -170,18 +217,66 @@ impl StructuredExtractor { /// the structured value. Returns [`TinyAgentsError::Validation`] when no /// matching call is found. /// + /// # Validation + /// + /// Both strategies validate the extracted value against this extractor's + /// schema before returning it (see [`validate`]). A value that parses but + /// does not conform is an error naming the failing instance path — not a + /// success carrying the wrong shape. + /// /// # Errors /// /// See strategy descriptions above. pub fn extract(&self, response: &ModelResponse) -> Result { - match self.strategy { - StructuredStrategy::ProviderSchema => self.extract_provider_schema(response), - StructuredStrategy::ToolCall => self.extract_tool_call(response), + let output = match self.strategy { + StructuredStrategy::ProviderSchema => self.extract_provider_schema(response)?, + StructuredStrategy::ToolCall => self.extract_tool_call(response)?, + }; + validate::validate_value(&self.schema, &output.value, &self.instance_root())?; + Ok(output) + } + + /// Extracts without failing: records the error instead of raising it. + /// + /// The difference is who decides what a failed extraction costs. `extract` + /// decides for the caller — it returns `Err`, and on the agent loop's final + /// turn that discards the entire run. This returns a + /// [`StructuredOutcome`] instead, so a caller can log the failure and + /// return the raw response, hand [`StructuredOutcome::error`] back to the + /// model as a repair prompt (LangChain's `OutputFixingParser`), or re-ask + /// with the original prompt (`RetryOutputParser`) — none of which are + /// possible once the run has already been thrown away. + /// + /// Mirrors LangChain's `include_raw=True`. + pub fn extract_outcome(&self, response: &ModelResponse) -> StructuredOutcome { + match self.extract(response) { + Ok(output) => StructuredOutcome { + value: Some(output.value), + raw: response.clone(), + error: None, + }, + Err(error) => { + let error = error.to_string(); + tracing::debug!( + "[structured] extraction failed for schema '{}': {error}", + self.schema_name + ); + StructuredOutcome { + value: None, + raw: response.clone(), + error: Some(error), + } + } } } // -- private helpers -- + /// The label validation errors are rooted at, for example `schema 'review'`. + fn instance_root(&self) -> String { + format!("schema '{}'", self.schema_name) + } + fn extract_provider_schema(&self, response: &ModelResponse) -> Result { let raw = response.text(); @@ -212,12 +307,25 @@ impl StructuredExtractor { })); } - let value: Value = serde_json::from_str(&raw).map_err(|e| { - TinyAgentsError::StructuredOutput(format!( - "schema '{}': response text is not valid JSON: {e}", - self.schema_name - )) - })?; + // Climb the repair ladder rather than a bare `from_str`. The crate + // already repairs the *other* JSON a model emits (tool-call arguments); + // there is no reason a fenced, chatty, or truncated structured answer + // should end a run when the same repairs recover it. + let Some((value, repair)) = repair::parse_lenient(&raw) else { + return Err(TinyAgentsError::StructuredOutput(format!( + "schema '{}': response text is not valid JSON and no conservative repair \ + recovered it (finish_reason = {:?})", + self.schema_name, + response.finish_reason.as_deref().unwrap_or("unknown") + ))); + }; + if repair.is_repaired() { + tracing::debug!( + "[structured] schema '{}': recovered the value with repair `{}`", + self.schema_name, + repair.as_str() + ); + } Ok(StructuredOutput { value, raw_text: Some(raw), @@ -235,6 +343,27 @@ impl StructuredExtractor { self.schema_name )) })?; + + // A provider that could not parse the call's arguments preserves them + // as a raw string (`ToolCall::invalid`). Running the same repair ladder + // here means a small local model's malformed arguments are recovered + // rather than handed on as a JSON string masquerading as the value. + if let Some(raw) = call.arguments.as_str() + && let Some((value, repair)) = repair::parse_lenient(raw) + { + if repair.is_repaired() { + tracing::debug!( + "[structured] schema '{}': recovered tool-call arguments with repair `{}`", + self.schema_name, + repair.as_str() + ); + } + return Ok(StructuredOutput { + value, + raw_text: Some(raw.to_string()), + }); + } + Ok(StructuredOutput { value: call.arguments.clone(), raw_text: None, diff --git a/src/harness/structured/repair.rs b/src/harness/structured/repair.rs new file mode 100644 index 0000000..dc70cec --- /dev/null +++ b/src/harness/structured/repair.rs @@ -0,0 +1,308 @@ +//! Best-effort repair of a model's structured-output text into strict JSON. +//! +//! # Why this exists +//! +//! Structured extraction used to be a bare [`serde_json::from_str`] on the +//! assistant text, and a single malformed brace on the final turn discarded the +//! whole run — every tool call and token already spent. Meanwhile the crate +//! already carried a repair ladder for the *other* place a model emits JSON: +//! tool-call arguments, repaired by +//! [`recover_tool_arguments`][rta] over +//! [`relaxed_json`][rj]. Structured output got none of it. +//! +//! [rta]: crate::harness::providers::openai +//! [rj]: crate::harness::providers::openai::relaxed_json +//! +//! # The ladder +//! +//! Each rung is tried in order and the first strict parse wins. Every rung is +//! *conservative*: it only ever runs after strict parsing has already failed, +//! and a rung that does not yield strictly-parseable JSON is discarded rather +//! than half-applied. +//! +//! | Rung | Repairs | Modelled on | +//! |------|---------|-------------| +//! | `Strict` | nothing — the input was already valid | — | +//! | `CodeFence` | ```` ```json … ``` ```` wrappers | ubiquitous | +//! | `Slice` | prose around the value (`Here is the JSON: {…}`) | — | +//! | `Relaxed` | unquoted keys, doubled braces, leaked chat-template quote tokens | [`relaxed_json`][rj] | +//! | `Closed` | truncated output: unterminated strings and unclosed brackets | LangChain `parse_partial_json` | +//! +//! # What it deliberately does not do +//! +//! It never *invents* structure. A rung is accepted only when the repaired text +//! parses strictly, so noise can never be laundered into a plausible-looking +//! value. Whether the parsed value is the *right shape* is a separate question, +//! answered by [`super::validate`] against the declared schema. + +use serde_json::Value; + +/// Which rung of the ladder produced a value. +/// +/// Carried out of [`parse_lenient`] so the caller can log — and a +/// [`super::StructuredOutcome`] can record — that the model's text needed +/// repairing, instead of a repair silently masking a degrading model. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum JsonRepair { + /// The text was already strict JSON. + Strict, + /// A markdown code fence was removed. + CodeFence, + /// The value was sliced out of surrounding prose. + Slice, + /// Relaxed-JSON repairs were applied (unquoted keys, doubled braces, leaked + /// chat-template quote tokens). + Relaxed, + /// Truncated output was closed (unterminated string and/or open brackets). + Closed, +} + +impl JsonRepair { + /// A stable, log- and event-friendly label. + pub fn as_str(self) -> &'static str { + match self { + JsonRepair::Strict => "strict", + JsonRepair::CodeFence => "code_fence", + JsonRepair::Slice => "slice", + JsonRepair::Relaxed => "relaxed", + JsonRepair::Closed => "closed", + } + } + + /// Whether any repair was actually needed. + pub fn is_repaired(self) -> bool { + self != JsonRepair::Strict + } +} + +/// Maximum trailing characters trimmed while closing a truncated value. +/// +/// A truncated completion usually stops mid-token, so the tail that has to go +/// is short. Bounding the search keeps the cost linear-ish on adversarial input +/// instead of quadratic over a multi-megabyte blob. +const MAX_TRAILING_TRIM: usize = 64; + +/// Parses `raw` as JSON, climbing the repair ladder until something parses. +/// +/// Returns the parsed value and the rung that produced it, or `None` when no +/// conservative repair yields strict JSON. +pub fn parse_lenient(raw: &str) -> Option<(Value, JsonRepair)> { + let trimmed = raw.trim(); + if let Ok(value) = serde_json::from_str::(trimmed) { + return Some((value, JsonRepair::Strict)); + } + + let unfenced = strip_code_fence(trimmed); + if unfenced != trimmed + && let Ok(value) = serde_json::from_str::(unfenced) + { + tracing::debug!("[structured::repair] recovered JSON by removing a markdown code fence"); + return Some((value, JsonRepair::CodeFence)); + } + + if let Some(sliced) = slice_json_span(unfenced) + && let Ok(value) = serde_json::from_str::(sliced) + { + tracing::debug!( + "[structured::repair] recovered JSON by slicing it out of surrounding text" + ); + return Some((value, JsonRepair::Slice)); + } + + // Reuses the crate's existing relaxed-JSON repairs rather than a second, + // divergent implementation. It only yields objects, which is the shape a + // JSON-Schema structured output almost always declares. + if let Some(value) = + crate::harness::providers::openai::relaxed_json::recover_relaxed_object(unfenced) + { + tracing::debug!("[structured::repair] recovered JSON through the relaxed-JSON repairs"); + return Some((value, JsonRepair::Relaxed)); + } + + if let Some(value) = close_truncated(unfenced) { + tracing::debug!("[structured::repair] recovered JSON by closing a truncated value"); + return Some((value, JsonRepair::Closed)); + } + + None +} + +/// Removes a ```` ``` ```` fence, with or without a language tag. +fn strip_code_fence(raw: &str) -> &str { + let trimmed = raw.trim(); + let Some(after_open) = trimmed.strip_prefix("```") else { + return trimmed; + }; + let body = match after_open.find('\n') { + Some(newline) + if after_open[..newline] + .chars() + .all(|character| character.is_ascii_alphanumeric()) => + { + &after_open[newline + 1..] + } + _ => after_open, + }; + body.trim().strip_suffix("```").unwrap_or(body).trim() +} + +/// Returns the span from the first `{`/`[` to the matching last `}`/`]`. +/// +/// Handles the extremely common "chatty" completion — `Sure! Here is the +/// result: {"score": 4}` — without any structural rewriting: the span is +/// returned verbatim and still has to parse strictly to be accepted. +fn slice_json_span(raw: &str) -> Option<&str> { + let open = raw.find(['{', '['])?; + let close = raw.rfind(['}', ']'])?; + if close <= open { + return None; + } + let span = &raw[open..=close]; + (span != raw).then_some(span) +} + +/// Closes a value truncated mid-flight: an unterminated string, then any +/// brackets still open. +/// +/// Mirrors LangChain's `parse_partial_json`: walk the text tracking string and +/// escape state, remember the closing delimiters owed, then append them. When +/// the result still does not parse — the truncation landed mid-key, mid-number, +/// or on a dangling comma — trailing characters are dropped one at a time and +/// the close is retried, bounded by [`MAX_TRAILING_TRIM`]. +fn close_truncated(raw: &str) -> Option { + let chars: Vec = raw.chars().collect(); + let floor = chars.len().saturating_sub(MAX_TRAILING_TRIM); + let mut end = chars.len(); + while end > floor && end > 0 { + let candidate: String = chars[..end].iter().collect(); + if let Some(closed) = close_once(&candidate) + && let Ok(value) = serde_json::from_str::(&closed) + { + return Some(value); + } + end -= 1; + } + None +} + +/// Appends the delimiters `raw` still owes: a closing quote when it ends inside +/// a string, then every unclosed `}`/`]` in reverse order. +fn close_once(raw: &str) -> Option { + let mut stack: Vec = Vec::new(); + let mut in_string = false; + let mut escaped = false; + + for character in raw.chars() { + if escaped { + escaped = false; + continue; + } + match character { + '\\' if in_string => escaped = true, + '"' => in_string = !in_string, + '{' if !in_string => stack.push('}'), + '[' if !in_string => stack.push(']'), + '}' | ']' if !in_string => { + // A closer with no matching opener means this is not a + // truncated value at all; refuse rather than guess. + stack.pop()?; + } + _ => {} + } + } + + if stack.is_empty() && !in_string { + // Nothing was owed, so closing cannot help — the caller already tried a + // strict parse of this exact text. + return None; + } + + let mut closed = String::with_capacity(raw.len() + stack.len() + 1); + closed.push_str(raw); + if in_string { + closed.push('"'); + } + while let Some(closer) = stack.pop() { + closed.push(closer); + } + Some(closed) +} + +#[cfg(test)] +mod test { + use super::*; + use serde_json::json; + + #[test] + fn strict_json_needs_no_repair() { + let (value, repair) = parse_lenient(r#"{"score":4}"#).expect("strict JSON parses"); + assert_eq!(value, json!({ "score": 4 })); + assert_eq!(repair, JsonRepair::Strict); + assert!(!repair.is_repaired()); + } + + #[test] + fn removes_a_markdown_code_fence() { + let (value, repair) = + parse_lenient("```json\n{\"score\": 4}\n```").expect("a fenced value parses"); + assert_eq!(value, json!({ "score": 4 })); + assert_eq!(repair, JsonRepair::CodeFence); + } + + #[test] + fn slices_a_value_out_of_prose() { + let (value, repair) = parse_lenient("Sure! Here it is: {\"score\": 4} — hope that helps.") + .expect("a value embedded in prose parses"); + assert_eq!(value, json!({ "score": 4 })); + assert_eq!(repair, JsonRepair::Slice); + } + + #[test] + fn repairs_relaxed_json_through_the_existing_ladder() { + let (value, repair) = parse_lenient("{score:4}").expect("unquoted keys are repaired"); + assert_eq!(value, json!({ "score": 4 })); + assert_eq!(repair, JsonRepair::Relaxed); + } + + #[test] + fn closes_a_truncated_object() { + let (value, repair) = parse_lenient(r#"{"summary": "the model ran out of budget mid-sent"#) + .expect("a truncated value is closed"); + assert_eq!(repair, JsonRepair::Closed); + assert_eq!(value["summary"], "the model ran out of budget mid-sent"); + } + + #[test] + fn closes_nested_containers_in_the_right_order() { + let (value, _) = + parse_lenient(r#"{"items": [{"id": 1}, {"id": 2"#).expect("nesting is closed"); + assert_eq!(value["items"][1]["id"], 2); + } + + #[test] + fn drops_a_dangling_comma_before_closing() { + let (value, repair) = + parse_lenient(r#"{"a": 1, "b": 2,"#).expect("a dangling comma is trimmed"); + assert_eq!(repair, JsonRepair::Closed); + assert_eq!(value, json!({ "a": 1, "b": 2 })); + } + + #[test] + fn refuses_text_that_is_not_json_at_all() { + assert!(parse_lenient("I could not answer that.").is_none()); + } + + #[test] + fn refuses_an_unbalanced_closer() { + // A stray `}` is corruption, not truncation; guessing here would let + // noise masquerade as a value. + assert!(parse_lenient("}}}").is_none()); + } + + #[test] + fn does_not_confuse_brackets_inside_strings() { + let (value, _) = parse_lenient(r#"{"text": "a { and a [ walk in"#) + .expect("brackets inside a string are literal"); + assert_eq!(value["text"], "a { and a [ walk in"); + } +} diff --git a/src/harness/structured/types.rs b/src/harness/structured/types.rs index 7703754..aa4395e 100644 --- a/src/harness/structured/types.rs +++ b/src/harness/structured/types.rs @@ -6,6 +6,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; +use crate::harness::model::ModelResponse; + // --------------------------------------------------------------------------- // Strategy // --------------------------------------------------------------------------- @@ -48,6 +50,62 @@ pub struct StructuredOutput { pub raw_text: Option, } +// --------------------------------------------------------------------------- +// StructuredOutcome +// --------------------------------------------------------------------------- + +/// The result of attempting structured extraction, **including the failures**. +/// +/// [`StructuredExtractor::extract`][ex] returns `Result`, so +/// a parse or validation failure is fatal to whatever is driving it — on the +/// agent loop's final turn that discards a whole run, every tool call and token +/// already spent, over one malformed brace. This type is the non-fatal +/// alternative: the value when there is one, the raw response either way, and +/// the error as *data* rather than control flow. +/// +/// Modelled on LangChain's `include_raw=True`, which wraps the parser in a +/// fallback yielding `{"raw", "parsed": None, "parsing_error"}` instead of +/// raising. +/// +/// [ex]: StructuredExtractor::extract +#[derive(Clone, Debug)] +pub struct StructuredOutcome { + /// The extracted and validated value, when extraction succeeded. + pub value: Option, + /// The model response extraction was attempted on, always preserved — it + /// is the only evidence of what the model actually said, and the input a + /// repair or re-ask turn needs. + pub raw: ModelResponse, + /// Why extraction failed, when it did. Written to be handed back to a model + /// verbatim: it names the schema and, for a validation failure, the exact + /// failing instance path. + pub error: Option, +} + +impl StructuredOutcome { + /// Whether a value was extracted. + pub fn is_success(&self) -> bool { + self.value.is_some() + } + + /// The extracted value, or the recorded error as a + /// [`TinyAgentsError::StructuredOutput`][err]. + /// + /// Use this at a boundary that genuinely cannot proceed without a value; + /// prefer matching on [`Self::value`] where a repair or re-ask is possible. + /// + /// [err]: crate::error::TinyAgentsError::StructuredOutput + pub fn into_result(self) -> crate::error::Result { + match self.value { + Some(value) => Ok(value), + None => Err(crate::error::TinyAgentsError::StructuredOutput( + self.error + .unwrap_or_else(|| "structured extraction failed".to_string()), + )), + } + } +} + // --------------------------------------------------------------------------- // StructuredExtractor // --------------------------------------------------------------------------- @@ -78,6 +136,9 @@ pub struct StructuredExtractor { /// Name used to match the artificial tool call (for [`StructuredStrategy::ToolCall`]) /// or to label errors. pub(crate) schema_name: String, - /// The JSON Schema document (kept for potential future local validation). + /// The JSON Schema document. **Enforced**: every extracted value is checked + /// against it by [`super::validate`] before it is returned, so a + /// well-formed value of the wrong shape is a reported error rather than + /// silent garbage in `run.structured`. pub(crate) schema: Value, } diff --git a/src/harness/structured/validate.rs b/src/harness/structured/validate.rs new file mode 100644 index 0000000..ac132d7 --- /dev/null +++ b/src/harness/structured/validate.rs @@ -0,0 +1,273 @@ +//! Local validation of an extracted structured value against its declared +//! JSON Schema. +//! +//! # Why this exists +//! +//! [`StructuredExtractor`][super::StructuredExtractor] stored its schema and +//! never read it: provider-schema mode ran a bare +//! [`serde_json::from_str`] and tool-call mode cloned the call's arguments +//! straight through. So `{"wrong_key": 1}` against a `score` schema *succeeded*, +//! and `run.structured` came back holding something no caller had asked for — +//! a failure that surfaces later, somewhere else, as a missing field. +//! +//! Validating here also gives the repair loop something to say: an error naming +//! the exact failing instance path is a message that can be handed back to the +//! model, where "deserialisation failed" is not. +//! +//! # The supported subset +//! +//! The same subset the tool-call boundary enforces: `type` (including union +//! types), object `properties`, `required`, `additionalProperties: false`, +//! array `items`, and `enum`. Unknown keywords are ignored, so a richer schema +//! can still be sent to a provider while the local boundary fails closed on +//! exactly the structural constraints it understands. An empty or null schema +//! imposes no constraints. +//! +//! It is intentionally **not** a general JSON Schema implementation: no +//! `$ref`, no `allOf`/`anyOf`/`oneOf`, no numeric or string facets. Those +//! belong in a dedicated validator crate if the need ever arises; guessing at +//! them here would produce confident wrong answers. + +use serde_json::Value; + +use crate::error::{Result, TinyAgentsError}; + +/// Validates `value` against `schema`, reporting the failing instance path. +/// +/// `root` names the value in error messages — the caller passes something like +/// `schema 'review'` so the message reads `schema 'review'.items[2].id must be +/// integer, got string`. +pub fn validate_value(schema: &Value, value: &Value, root: &str) -> Result<()> { + validate_at(schema, value, root) +} + +fn validate_at(schema: &Value, value: &Value, path: &str) -> Result<()> { + if schema.is_null() || schema.as_object().is_some_and(|map| map.is_empty()) { + return Ok(()); + } + + if let Some(allowed) = schema.get("enum").and_then(Value::as_array) + && !allowed.iter().any(|candidate| candidate == value) + { + return Err(invalid(format!( + "{path} must be one of the declared enum values" + ))); + } + + if let Some(type_spec) = schema.get("type") { + validate_type(type_spec, value, path)?; + } + + // `required` is enforced independently of `properties`: a schema may name + // required fields without describing them, and nesting the check under + // `properties` would let such a schema fail open. + if let Some(required) = schema.get("required").and_then(Value::as_array) { + if let Some(object) = value.as_object() { + for field in required.iter().filter_map(Value::as_str) { + if !object.contains_key(field) { + return Err(invalid(format!("{path}.{field} is required"))); + } + } + } else if schema.get("type").is_none() { + return Err(invalid(format!( + "{path} must be an object with the declared fields, got {}", + kind_of(value) + ))); + } + } + + if let Some(properties) = schema.get("properties").and_then(Value::as_object) { + if let Some(object) = value.as_object() { + if schema.get("additionalProperties").and_then(Value::as_bool) == Some(false) { + for field in object.keys() { + if !properties.contains_key(field) { + return Err(invalid(format!("{path}.{field} is not allowed"))); + } + } + } + for (field, field_schema) in properties { + if let Some(field_value) = object.get(field) { + validate_at(field_schema, field_value, &format!("{path}.{field}"))?; + } + } + } else if schema.get("type").is_none() { + return Err(invalid(format!( + "{path} must be an object with the declared fields, got {}", + kind_of(value) + ))); + } + } + + if let Some(items_schema) = schema.get("items") + && let Some(items) = value.as_array() + { + for (index, item) in items.iter().enumerate() { + validate_at(items_schema, item, &format!("{path}[{index}]"))?; + } + } + + Ok(()) +} + +fn validate_type(type_spec: &Value, value: &Value, path: &str) -> Result<()> { + if let Some(kind) = type_spec.as_str() { + if matches_type(value, kind) { + return Ok(()); + } + return Err(invalid(format!( + "{path} must be {kind}, got {}", + kind_of(value) + ))); + } + + if let Some(kinds) = type_spec.as_array() { + let allowed: Vec<&str> = kinds.iter().filter_map(Value::as_str).collect(); + if allowed.iter().any(|kind| matches_type(value, kind)) { + return Ok(()); + } + return Err(invalid(format!( + "{path} must be one of {}, got {}", + allowed.join(", "), + kind_of(value) + ))); + } + + Ok(()) +} + +fn matches_type(value: &Value, kind: &str) -> bool { + match kind { + "null" => value.is_null(), + "boolean" => value.is_boolean(), + "object" => value.is_object(), + "array" => value.is_array(), + "number" => value.is_number(), + "integer" => value.as_i64().is_some() || value.as_u64().is_some(), + "string" => value.is_string(), + // An unknown type keyword must not fail closed: providers accept richer + // vocabularies than this subset understands. + _ => true, + } +} + +fn kind_of(value: &Value) -> &'static str { + match value { + Value::Null => "null", + Value::Bool(_) => "boolean", + Value::Number(number) if number.as_i64().is_some() || number.as_u64().is_some() => { + "integer" + } + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Array(_) => "array", + Value::Object(_) => "object", + } +} + +fn invalid(message: String) -> TinyAgentsError { + TinyAgentsError::StructuredOutput(message) +} + +#[cfg(test)] +mod test { + use super::*; + use serde_json::json; + + fn score_schema() -> Value { + json!({ + "type": "object", + "properties": { "score": { "type": "integer" } }, + "required": ["score"], + "additionalProperties": false + }) + } + + #[test] + fn accepts_a_conforming_value() { + validate_value(&score_schema(), &json!({ "score": 4 }), "schema 'score'").unwrap(); + } + + #[test] + fn rejects_a_missing_required_field_by_path() { + let err = validate_value( + &score_schema(), + &json!({ "wrong_key": 1 }), + "schema 'score'", + ) + .expect_err("a missing required field is not valid"); + assert!( + err.to_string().contains("schema 'score'.score is required"), + "{err}" + ); + } + + #[test] + fn rejects_a_wrong_type_by_path() { + let err = validate_value( + &score_schema(), + &json!({ "score": "four" }), + "schema 'score'", + ) + .expect_err("a string is not an integer"); + assert!( + err.to_string() + .contains("schema 'score'.score must be integer, got string"), + "{err}" + ); + } + + #[test] + fn reports_a_nested_array_index() { + let schema = json!({ + "type": "object", + "properties": { + "items": { "type": "array", "items": { "type": "object", "properties": { "id": { "type": "integer" } } } } + } + }); + let err = validate_value( + &schema, + &json!({ "items": [{ "id": 1 }, { "id": "two" }] }), + "schema 'batch'", + ) + .expect_err("the second item is invalid"); + assert!(err.to_string().contains("items[1].id"), "{err}"); + } + + #[test] + fn rejects_an_undeclared_field_when_additional_properties_is_false() { + let err = validate_value( + &score_schema(), + &json!({ "score": 4, "extra": true }), + "schema 'score'", + ) + .expect_err("`extra` is not declared"); + assert!(err.to_string().contains("extra is not allowed"), "{err}"); + } + + #[test] + fn an_empty_schema_constrains_nothing() { + validate_value(&json!({}), &json!("anything at all"), "schema 'free'").unwrap(); + validate_value(&Value::Null, &json!(7), "schema 'free'").unwrap(); + } + + #[test] + fn accepts_a_union_type() { + let schema = json!({ "type": ["string", "null"] }); + validate_value(&schema, &json!(null), "schema 'maybe'").unwrap(); + validate_value(&schema, &json!("x"), "schema 'maybe'").unwrap(); + assert!(validate_value(&schema, &json!(3), "schema 'maybe'").is_err()); + } + + #[test] + fn ignores_unknown_type_keywords() { + // A provider may accept a richer vocabulary than this subset knows. + validate_value(&json!({ "type": "date-time" }), &json!("2026-01-01"), "s").unwrap(); + } + + #[test] + fn enforces_an_enum() { + let schema = json!({ "enum": ["a", "b"] }); + validate_value(&schema, &json!("a"), "schema 'choice'").unwrap(); + assert!(validate_value(&schema, &json!("c"), "schema 'choice'").is_err()); + } +} diff --git a/src/harness/subagent/test.rs b/src/harness/subagent/test.rs index 4d5d7b5..d787d5a 100644 --- a/src/harness/subagent/test.rs +++ b/src/harness/subagent/test.rs @@ -42,6 +42,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -59,6 +60,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/src/harness/summarization/mod.rs b/src/harness/summarization/mod.rs index 5fd3733..11bcf3f 100644 --- a/src/harness/summarization/mod.rs +++ b/src/harness/summarization/mod.rs @@ -19,13 +19,23 @@ //! All policy decisions are explicit data types, never hidden behaviour. Callers //! choose when to call, what to pass, and how to handle the result. +pub mod pairing; +mod render; +mod trim; mod types; +pub use pairing::{ + advance_past_orphan_tools, find_safe_cutoff_point, is_tool_calling_assistant, + retract_orphan_tool_calls, tool_pairing_is_intact, +}; +pub use render::render_message_for_summary; +pub use trim::{trim_messages, trim_messages_to_token_budget_with, trim_messages_with}; pub use types::*; use crate::error::{Result, TinyAgentsError}; -use crate::harness::message::Message; +use crate::harness::message::{Message, estimate_slice_tokens}; use async_trait::async_trait; +use trim::partition_system; // --------------------------------------------------------------------------- // Token estimation @@ -47,206 +57,28 @@ pub fn estimate_tokens(text: &str) -> u64 { if chars == 0 { 0 } else { (chars / 4).max(1) } } -/// Estimate the total tokens for a [`Message`] using the same `chars / 4` -/// heuristic as [`estimate_tokens`], counting weight directly over the message -/// content rather than allocating the concatenated string first. -/// -/// Uses [`Message::estimated_char_weight`] (all content blocks) rather than -/// `char_len` (visible text only): images, structured JSON, and model reasoning -/// occupy real context, so counting only text would under-estimate a -/// multimodal or tool-heavy transcript to near-zero and silently prevent -/// [`SummarizationPolicy::should_summarize`] from ever triggering. -fn message_token_estimate(msg: &Message) -> u64 { - let chars = msg.estimated_char_weight() as u64; - if chars == 0 { 0 } else { (chars / 4).max(1) } -} - -/// Estimate the total tokens for a slice of messages. -fn slice_token_estimate(messages: &[Message]) -> u64 { - messages.iter().map(message_token_estimate).sum() -} - -// --------------------------------------------------------------------------- -// Trimming -// --------------------------------------------------------------------------- - -/// Partition `messages` into system and non-system messages, preserving order. -/// -/// Returns `(system, non_system)`. -fn partition_system(messages: &[Message]) -> (Vec, Vec) { - let system = messages - .iter() - .filter(|m| matches!(m, Message::System(_))) - .cloned() - .collect(); - let non_system = messages - .iter() - .filter(|m| !matches!(m, Message::System(_))) - .cloned() - .collect(); - (system, non_system) -} - -/// Trim a message slice according to `strategy`, returning the retained subset. -/// -/// System messages are preserved by default: -/// -/// - [`TrimStrategy::KeepLast`] and [`TrimStrategy::KeepFirstAndLast`] always -/// keep all system messages and apply the rule only to non-system messages. -/// - [`TrimStrategy::MaxTokens`] drops non-system messages first (from the -/// front) and only starts dropping system messages if the budget still -/// cannot be met after all non-system messages are removed. -/// -/// The returned `Vec` preserves the relative order of messages as -/// they appeared in the input. -pub fn trim_messages(messages: &[Message], strategy: &TrimStrategy) -> Vec { - match strategy { - TrimStrategy::KeepLast(n) => { - let (system, non_system) = partition_system(messages); - let keep_start = non_system.len().saturating_sub(*n); - let mut result = system; - result.extend_from_slice(&non_system[keep_start..]); - result - } - - TrimStrategy::KeepFirstAndLast { first, last } => { - let (system, non_system) = partition_system(messages); - let len = non_system.len(); - let first = *first; - let last = *last; - - let mut result = system; - if first + last >= len { - // No overlap: keep everything. - result.extend(non_system); - } else { - result.extend_from_slice(&non_system[..first]); - result.extend_from_slice(&non_system[len - last..]); - } - result - } - - TrimStrategy::MaxTokens(limit) => { - let (system, non_system) = partition_system(messages); - let limit = *limit; - - // Precompute each message's token estimate once. The previous - // implementation re-summed the entire slice on every dropped - // message and used `remove(0)` (itself O(n)), making trimming - // O(n^2). Here we sum once and drop from the front by advancing an - // index while subtracting the running total. - let sys_tokens: Vec = system.iter().map(message_token_estimate).collect(); - let non_sys_tokens: Vec = non_system.iter().map(message_token_estimate).collect(); - let sys_total: u64 = sys_tokens.iter().sum(); - - // Drop non-system messages from the front until within budget or - // exhausted. - let mut non_sys_start = 0; - let mut non_sys_total: u64 = non_sys_tokens.iter().sum(); - while non_sys_start < non_system.len() && sys_total + non_sys_total > limit { - non_sys_total -= non_sys_tokens[non_sys_start]; - non_sys_start += 1; - } - - // Still over budget: drop system messages from the front as a last - // resort. - let mut sys_start = 0; - let mut sys_running = sys_total; - while sys_start < system.len() && sys_running + non_sys_total > limit { - sys_running -= sys_tokens[sys_start]; - sys_start += 1; - } - - let mut result = - Vec::with_capacity((system.len() - sys_start) + (non_system.len() - non_sys_start)); - result.extend_from_slice(&system[sys_start..]); - result.extend_from_slice(&non_system[non_sys_start..]); - result - } - } -} - -/// Trim messages to a token budget while preserving their original order. -/// -/// `estimate` lets callers account for provider-specific payloads without -/// teaching the harness about their wire representation. The built-in -/// [`Message::estimated_char_weight`] already assigns a flat weight to native -/// image blocks; hosts may additionally recognize inline image markers or use a -/// real tokenizer. Oldest non-system messages are evicted first. System -/// messages are either retained unconditionally or evicted oldest-first only -/// after all other messages, according to [`TokenTrimPolicy::preserve_system`]. -/// -/// When `drop_leading_orphan_tools` is enabled, leading tool-result messages are -/// removed after budget eviction so the retained transcript starts on a valid -/// provider turn boundary. The returned messages always retain their original -/// relative order. -pub fn trim_messages_to_token_budget_with( - messages: &[Message], - policy: TokenTrimPolicy, - estimate: impl Fn(&Message) -> u64, -) -> Vec { - let estimates: Vec = messages.iter().map(estimate).collect(); - let mut total: u64 = estimates.iter().copied().sum(); - if total <= policy.limit && !policy.drop_leading_orphan_tools { - return messages.to_vec(); - } - - let mut retained = vec![true; messages.len()]; - for (index, message) in messages.iter().enumerate() { - if total <= policy.limit { - break; - } - if !matches!(message, Message::System(_)) { - retained[index] = false; - total = total.saturating_sub(estimates[index]); - } - } - - if !policy.preserve_system && total > policy.limit { - for (index, message) in messages.iter().enumerate() { - if total <= policy.limit { - break; - } - if matches!(message, Message::System(_)) && retained[index] { - retained[index] = false; - total = total.saturating_sub(estimates[index]); - } - } - } - - let mut result: Vec = messages - .iter() - .zip(retained) - .filter(|(_, keep)| *keep) - .map(|(message, _)| message.clone()) - .collect(); - - if policy.drop_leading_orphan_tools { - while let Some(index) = result - .iter() - .position(|message| !matches!(message, Message::System(_))) - { - if matches!(result[index], Message::Tool(_)) { - result.remove(index); - } else { - break; - } - } - } - result -} - // --------------------------------------------------------------------------- // ConcatSummarizer // --------------------------------------------------------------------------- #[async_trait] impl Summarizer for ConcatSummarizer { - /// Summarize `messages` by concatenating their text content into a single - /// system message. + /// Summarize `messages` by concatenating them into a single system message. /// - /// Each message's text is prefixed by a role label and positional id so - /// the summary is human-readable. No LLM call is made. + /// Each message is rendered by [`render_message_for_summary`] and prefixed + /// with a positional id, so the summary is human-readable. No LLM call is + /// made. + /// + /// # Why not `Message::text()` + /// + /// [`Message::text`] returns only visible text blocks, so an assistant turn + /// that only called tools, a JSON tool result, and model reasoning all + /// render as an empty string. Because this is the crate's **default** + /// summarizer, building it on `text()` meant that out of the box, + /// compaction of a tool-driven run replaced the real history with a column + /// of bare role labels. Rendering tool calls, tool results, and reasoning + /// keeps the compacted transcript worth keeping — the same reason LangChain + /// summarizes through `get_buffer_string(..., format="xml")`. /// /// # Provenance /// @@ -260,7 +92,7 @@ impl Summarizer for ConcatSummarizer { )); } - let original_token_estimate = slice_token_estimate(messages); + let original_token_estimate = estimate_slice_tokens(messages); let mut parts: Vec = Vec::with_capacity(messages.len() + 1); parts.push("=== Conversation Summary ===".to_string()); @@ -269,14 +101,8 @@ impl Summarizer for ConcatSummarizer { .iter() .enumerate() .map(|(i, msg)| { - let role = match msg { - Message::System(_) => "system", - Message::User(_) => "user", - Message::Assistant(_) => "assistant", - Message::Tool(_) => "tool", - }; let id = format!("msg-{i}"); - parts.push(format!("[{id}] {role}: {}", msg.text())); + parts.push(format!("[{id}] {}", render_message_for_summary(msg))); id }) .collect(); @@ -358,7 +184,7 @@ impl SummarizationPolicy { /// returns `true` when the estimate **exceeds** /// [`trigger_tokens`][Self::trigger_tokens]. pub fn should_summarize(&self, messages: &[Message]) -> bool { - let tokens = slice_token_estimate(messages); + let tokens = estimate_slice_tokens(messages); match self.context_window { Some(_) => tokens >= self.trigger_budget(), None => tokens > self.trigger_tokens, @@ -378,6 +204,23 @@ impl SummarizationPolicy { /// /// System messages are never placed in `to_summarize` — they must be kept /// verbatim to avoid losing persistent instructions. + /// + /// # Tool-call pairing + /// + /// The split point is **not** a blind `len - keep_last` index. That index + /// routinely lands between an assistant tool-call turn and the tool results + /// answering it, putting the assistant message in `to_summarize` and its + /// `tool` messages in `to_keep`; the rebuilt request then opens with a + /// `role:"tool"` message that answers nothing, which OpenAI rejects with a + /// `400` and Anthropic rejects as a `tool_result` with no matching + /// `tool_use`. Since only long tool-driven runs reach a compaction + /// threshold at all, the blind index failed on essentially every run that + /// used it. + /// + /// [`find_safe_cutoff_point`] moves the split back to include the owning + /// assistant turn (or, for a transcript with no such turn, forward past the + /// unpairable results), so `to_keep` is always a slice a provider accepts. + /// `keep_last` is therefore a **minimum**, not an exact count. pub fn plan(&self, messages: &[Message]) -> (Vec, Vec) { let (system, non_system) = partition_system(messages); @@ -388,7 +231,14 @@ impl SummarizationPolicy { return (Vec::new(), to_keep); } - let split = non_system.len() - self.keep_last; + let requested_split = non_system.len() - self.keep_last; + let split = find_safe_cutoff_point(&non_system, requested_split); + if split != requested_split { + tracing::debug!( + "[summarization::plan] keep_last={} moved split {requested_split} -> {split} to preserve tool-call pairing", + self.keep_last + ); + } let to_summarize = non_system[..split].to_vec(); let to_keep_recent = non_system[split..].to_vec(); diff --git a/src/harness/summarization/pairing.rs b/src/harness/summarization/pairing.rs new file mode 100644 index 0000000..ae12f71 --- /dev/null +++ b/src/harness/summarization/pairing.rs @@ -0,0 +1,193 @@ +//! Structural repair of transcript cut points so compaction never orphans a +//! tool call or a tool result. +//! +//! # The failure this prevents +//! +//! Every provider enforces the same structural invariant: a `role:"tool"` +//! message must be preceded by the assistant message whose `tool_calls` +//! declared its id, and an assistant `tool_calls` entry must be answered. Cut a +//! transcript at a blind index and you routinely break both halves: +//! +//! ```text +//! [system, user, assistant(tool_calls=[c1]), tool(c1), assistant("done")] +//! ^ keep_last = 2 cuts here +//! ``` +//! +//! The rebuilt request opens with `tool(c1)` and no preceding `tool_calls`. +//! OpenAI answers `400`; Anthropic rejects a `tool_result` with no matching +//! `tool_use`. This fires only on long tool-driven runs — which is to say, only +//! on the runs that ever reach a compaction threshold, and only after the run +//! has already done expensive work. +//! +//! # The two repairs +//! +//! Which direction is correct depends on what the cut is *for*: +//! +//! - [`find_safe_cutoff_point`] moves the cut **backward** to swallow the +//! assistant message that owns the orphaned results. Use it when the goal is +//! a message *count* ("keep the last N") or a summarize/keep split, where +//! keeping one extra message is free. This is a port of LangChain's +//! `SummarizationMiddleware._find_safe_cutoff_point`. +//! - [`advance_past_orphan_tools`] moves the cut **forward**, discarding the +//! orphaned tool results instead. Use it when the cut enforces a *token +//! budget*, where moving backward would re-admit the very tokens the trim +//! was trying to shed. +//! - [`retract_orphan_tool_calls`] repairs the *other* end: an assistant +//! message left at the tail of a retained prefix whose results were cut is an +//! orphaned tool call, and just as fatal. LangChain's summarization +//! middleware never hits this because it only ever keeps a suffix; +//! [`TrimStrategy::KeepFirstAndLast`][crate::harness::summarization::TrimStrategy::KeepFirstAndLast] +//! does keep a prefix, so the crate needs the mirrored repair. +//! +//! All three take and return indices into a slice that contains **no system +//! messages** (callers partition those out first); a system message can never +//! sit between an assistant and its tool results, so the partitioning does not +//! affect pairing. + +use std::collections::HashSet; + +use crate::harness::message::Message; + +/// Returns the tool-call ids declared by an assistant message, or an empty set +/// for every other message kind. +fn declared_call_ids(message: &Message) -> HashSet<&str> { + match message { + Message::Assistant(assistant) => assistant + .tool_calls + .iter() + .map(|call| call.id.as_str()) + .filter(|id| !id.is_empty()) + .collect(), + _ => HashSet::new(), + } +} + +/// Returns `true` when `message` is an assistant turn that requested tools. +pub fn is_tool_calling_assistant(message: &Message) -> bool { + matches!(message, Message::Assistant(a) if !a.tool_calls.is_empty()) +} + +/// Moves `cutoff_index` to a point that does not split an assistant tool-call +/// message from the tool results answering it, preferring to keep *more*. +/// +/// `cutoff_index` is the index at which the retained suffix begins (everything +/// before it is dropped or summarized). Semantics, in order: +/// +/// 1. If the message at the cutoff is not a tool result, the cutoff is already +/// safe and is returned unchanged. +/// 2. Otherwise the consecutive run of tool results at the cutoff is collected +/// and the slice is scanned **backward** for the assistant message whose +/// `tool_calls` ids intersect that run; the cutoff moves back to that +/// assistant's index so the pair stays together. +/// 3. If no such assistant exists (a truncated or imported transcript), the +/// cutoff falls **forward** past the whole tool run, discarding results that +/// can never be paired. +/// +/// Port of LangChain's `SummarizationMiddleware._find_safe_cutoff_point`. +pub fn find_safe_cutoff_point(messages: &[Message], cutoff_index: usize) -> usize { + if cutoff_index >= messages.len() || !matches!(messages[cutoff_index], Message::Tool(_)) { + return cutoff_index; + } + + // Collect the ids of the consecutive tool-result run starting at the cutoff. + let mut orphan_ids: HashSet<&str> = HashSet::new(); + let mut past_run = cutoff_index; + while past_run < messages.len() + && let Message::Tool(tool) = &messages[past_run] + { + if !tool.tool_call_id.is_empty() { + orphan_ids.insert(tool.tool_call_id.as_str()); + } + past_run += 1; + } + + // Scan backward for the assistant turn that declared any of those ids. + for index in (0..cutoff_index).rev() { + let declared = declared_call_ids(&messages[index]); + if !declared.is_empty() && declared.intersection(&orphan_ids).next().is_some() { + tracing::debug!( + "[summarization::pairing] cutoff {cutoff_index} split a tool pair; moving back to {index} to keep the assistant tool-call turn" + ); + return index; + } + } + + tracing::debug!( + "[summarization::pairing] cutoff {cutoff_index} has no matching assistant tool-call turn; advancing to {past_run} to drop unpairable tool results" + ); + past_run +} + +/// Moves `cutoff_index` **forward** past any leading tool results whose +/// assistant tool-call turn was dropped, discarding them. +/// +/// This is the budget-preserving counterpart to [`find_safe_cutoff_point`]: +/// where that function keeps more to repair the pair, this one keeps less, so a +/// token-bounded trim cannot re-admit the tokens it just shed. A tool result at +/// the head of the retained slice is orphaned by definition — its assistant +/// turn necessarily preceded it and was therefore already dropped. +pub fn advance_past_orphan_tools(messages: &[Message], cutoff_index: usize) -> usize { + let mut index = cutoff_index; + while index < messages.len() && matches!(messages[index], Message::Tool(_)) { + index += 1; + } + if index != cutoff_index { + tracing::debug!( + "[summarization::pairing] dropped {} leading orphan tool result(s) at cutoff {cutoff_index}", + index - cutoff_index + ); + } + index +} + +/// Pulls an exclusive `end_index` **backward** so the retained prefix does not +/// end on an assistant message whose tool results were cut. +/// +/// An assistant `tool_calls` entry with no answering tool message is as fatal +/// as the reverse orphan: OpenAI rejects the request, and Anthropic rejects a +/// `tool_use` with no `tool_result`. Because the results always follow their +/// call, an assistant tool-call turn sitting at the very end of a kept prefix +/// is unanswerable by construction, so it is removed along with any tool-call +/// turns it uncovers. +pub fn retract_orphan_tool_calls(messages: &[Message], end_index: usize) -> usize { + let mut end = end_index.min(messages.len()); + while end > 0 && is_tool_calling_assistant(&messages[end - 1]) { + end -= 1; + } + if end != end_index.min(messages.len()) { + tracing::debug!( + "[summarization::pairing] retracted retained prefix end from {end_index} to {end} to drop unanswered assistant tool call(s)" + ); + } + end +} + +/// Returns `true` when `messages` satisfies the provider pairing invariant: +/// every tool result is preceded by an assistant turn declaring its id, and +/// every declared tool call is answered later in the slice. +/// +/// Exposed for tests and for hosts that want to assert the invariant before +/// sending a request they assembled themselves. +pub fn tool_pairing_is_intact(messages: &[Message]) -> bool { + let mut declared: HashSet<&str> = HashSet::new(); + let mut answered: HashSet<&str> = HashSet::new(); + + for message in messages { + match message { + Message::Assistant(assistant) => { + for call in &assistant.tool_calls { + declared.insert(call.id.as_str()); + } + } + Message::Tool(tool) => { + if !declared.contains(tool.tool_call_id.as_str()) { + return false; + } + answered.insert(tool.tool_call_id.as_str()); + } + _ => {} + } + } + + declared.is_subset(&answered) +} diff --git a/src/harness/summarization/render.rs b/src/harness/summarization/render.rs new file mode 100644 index 0000000..c052875 --- /dev/null +++ b/src/harness/summarization/render.rs @@ -0,0 +1,144 @@ +//! Lossless-enough rendering of a transcript into summarizable text. +//! +//! # Why `Message::text()` is not enough +//! +//! [`Message::text`] returns only [`ContentBlock::Text`] blocks — by design, so +//! reasoning never leaks into visible assistant output. That makes it the wrong +//! function to summarize *with*: an assistant turn that only called tools has no +//! text, a tool result carrying JSON has no text, and reasoning has no text. A +//! summarizer built on `text()` therefore renders a tool-heavy transcript as a +//! column of bare role labels: +//! +//! ```text +//! [msg-2] assistant: +//! [msg-3] tool: +//! ``` +//! +//! and compaction replaces real history with nothing. Since the crate's default +//! summarizer is exactly that, the out-of-the-box behaviour was to erase the +//! part of the history that was most expensive to produce. +//! +//! [`render_message_for_summary`] renders the whole message — text, reasoning, +//! tool calls with their arguments, tool results with their correlation id and +//! JSON content — in a compact tagged form. LangChain solves the same problem +//! with `get_buffer_string(..., format="xml")`. + +use crate::harness::message::{ContentBlock, Message}; + +/// Maximum characters rendered for a single tool-call argument blob or tool +/// result body before it is elided. +/// +/// A summary that faithfully reproduces a 2 MB tool result is not a summary. +/// The cap is generous enough to keep short structured results (ids, paths, +/// row counts) intact, which is what a later turn usually needs to refer back +/// to. +const MAX_RENDERED_PAYLOAD_CHARS: usize = 2_000; + +/// The role label used when rendering a message for summarization. +fn role_label(message: &Message) -> &'static str { + match message { + Message::System(_) => "system", + Message::User(_) => "user", + Message::Assistant(_) => "assistant", + Message::Tool(_) => "tool", + } +} + +/// Truncates `text` to [`MAX_RENDERED_PAYLOAD_CHARS`], marking the elision so a +/// reader (or a downstream LLM summarizer) can tell content was dropped rather +/// than assuming the tool returned exactly that much. +fn elide(text: &str) -> String { + let count = text.chars().count(); + if count <= MAX_RENDERED_PAYLOAD_CHARS { + return text.to_string(); + } + let kept: String = text.chars().take(MAX_RENDERED_PAYLOAD_CHARS).collect(); + format!( + "{kept}… [{} chars elided]", + count - MAX_RENDERED_PAYLOAD_CHARS + ) +} + +/// Renders every content block that carries information a summary should keep: +/// visible text, structured JSON, reasoning, provider extensions, and a marker +/// for images (whose bytes are useless in a text summary but whose *presence* +/// is not). +fn render_content(content: &[ContentBlock]) -> Vec { + content + .iter() + .filter_map(|block| match block { + ContentBlock::Text(text) if text.trim().is_empty() => None, + ContentBlock::Text(text) => Some(text.clone()), + ContentBlock::Json(value) => { + Some(format!("{}", elide(&value.to_string()))) + } + ContentBlock::Image(image) => Some(format!( + "", + image.mime_type.as_deref().unwrap_or("unknown") + )), + ContentBlock::Thinking { text, .. } if text.trim().is_empty() => None, + ContentBlock::Thinking { text, .. } => { + Some(format!("{}", elide(text))) + } + ContentBlock::RedactedThinking { .. } => Some("".to_string()), + ContentBlock::ProviderExtension(value) => Some(format!( + "{}", + elide(&value.to_string()) + )), + }) + .collect() +} + +/// Renders a single message into the text a summarizer should see. +/// +/// The output is a single line per message where possible, in the shape +/// `: `, with tool calls rendered as +/// `{args}` and tool results as +/// `…`. Large payloads are elided (see +/// [`MAX_RENDERED_PAYLOAD_CHARS`]). +pub fn render_message_for_summary(message: &Message) -> String { + let mut parts: Vec = match message { + Message::System(m) => render_content(&m.content), + Message::User(m) => render_content(&m.content), + Message::Assistant(m) => render_content(&m.content), + Message::Tool(m) => render_content(&m.content), + }; + + match message { + Message::Assistant(assistant) => { + for call in &assistant.tool_calls { + let arguments = elide(&call.arguments.to_string()); + parts.push(format!( + "{arguments}", + call.id, call.name + )); + if let Some(reason) = &call.invalid { + parts.push(format!("{reason}")); + } + } + } + Message::Tool(tool) => { + // A tool result whose content rendered to nothing still needs its + // correlation id recorded: "this call was answered" is itself the + // fact a later turn reasons about. + // Tool results are the one message kind whose *text* is elided as + // well as its structured blocks: a tool that returns a 9 MB page + // dump would otherwise be reproduced verbatim into the "summary". + // User and assistant prose is left intact — that is the + // conversation itself, and `ConcatSummarizer` promises verbatim + // concatenation of it. + let body = if parts.is_empty() { + String::new() + } else { + elide(&parts.join(" ")) + }; + return format!( + "tool: {body}", + tool.tool_call_id + ); + } + _ => {} + } + + format!("{}: {}", role_label(message), parts.join(" ")) +} diff --git a/src/harness/summarization/test.rs b/src/harness/summarization/test.rs index 30cc0e7..325db15 100644 --- a/src/harness/summarization/test.rs +++ b/src/harness/summarization/test.rs @@ -370,3 +370,317 @@ mod smoke { assert_eq!(to_keep.len(), 2); } } + +/// Regression tests for the structural repair of transcript cut points. +/// +/// Every test here is written against the concrete provider failure it +/// prevents: a `role:"tool"` message with no preceding assistant `tool_calls` +/// (OpenAI `400`, Anthropic "tool_result with no matching tool_use"), or the +/// mirror image, an assistant `tool_calls` entry nothing ever answers. +/// +/// Before the repair landed, `plan` and all three [`TrimStrategy`] variants cut +/// at a blind index and produced exactly those shapes. +#[cfg(test)] +mod pairing { + use crate::harness::message::{AssistantMessage, ContentBlock, Message}; + use crate::harness::summarization::{ + MessageRole, SummarizationPolicy, TrimOptions, TrimStrategy, tool_pairing_is_intact, + trim_messages, trim_messages_with, + }; + use crate::harness::tool::ToolCall; + use serde_json::json; + + /// An assistant turn that only calls tools: no visible text at all, which + /// is precisely the shape that used to estimate to zero tokens and to be + /// severed from its results by a blind cut. + fn assistant_calling(ids: &[&str]) -> Message { + Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: ids + .iter() + .map(|id| ToolCall::new(*id, "lookup", json!({"q": "rust"}))) + .collect(), + usage: None, + }) + } + + /// `[system, user, assistant(tool_calls=[c1]), tool(c1), assistant("done")]` + /// — the canonical transcript that a `keep_last = 2` cut splits. + fn tool_transcript() -> Vec { + vec![ + Message::system("sys"), + Message::user("weather?"), + assistant_calling(&["c1"]), + Message::tool("c1", "sunny"), + Message::assistant("done"), + ] + } + + #[test] + fn plan_does_not_orphan_a_tool_result() { + let policy = SummarizationPolicy { + keep_last: 2, + ..Default::default() + }; + let (to_summarize, to_keep) = policy.plan(&tool_transcript()); + + // The assistant tool-call turn must travel with its result. + assert!( + tool_pairing_is_intact(&to_keep), + "kept slice orphans a tool result: {to_keep:?}" + ); + assert!( + !to_summarize + .iter() + .any(|m| matches!(m, Message::Assistant(a) if !a.tool_calls.is_empty())), + "the assistant tool-call turn was summarized away from its result" + ); + // keep_last is a minimum: the repair kept one extra message. + assert_eq!(to_keep.len(), 4); + } + + #[test] + fn keep_last_does_not_orphan_a_tool_result() { + let trimmed = trim_messages(&tool_transcript(), &TrimStrategy::KeepLast(2)); + assert!( + tool_pairing_is_intact(&trimmed), + "KeepLast orphaned a tool result: {trimmed:?}" + ); + } + + #[test] + fn keep_first_and_last_does_not_orphan_either_end() { + // Head block ends on an assistant tool-call turn; tail block starts on + // a tool result. Both ends are broken without repair. + let messages = vec![ + Message::user("one"), + assistant_calling(&["c1"]), + Message::tool("c1", "r1"), + Message::user("two"), + assistant_calling(&["c2"]), + Message::tool("c2", "r2"), + Message::assistant("done"), + ]; + let trimmed = trim_messages( + &messages, + &TrimStrategy::KeepFirstAndLast { first: 2, last: 2 }, + ); + assert!( + tool_pairing_is_intact(&trimmed), + "KeepFirstAndLast produced an unpaired slice: {trimmed:?}" + ); + } + + #[test] + fn max_tokens_drops_orphan_tool_results_rather_than_readmitting_tokens() { + let messages = vec![ + Message::user("a".repeat(400)), + assistant_calling(&["c1"]), + Message::tool("c1", "r1"), + Message::assistant("done"), + ]; + let trimmed = trim_messages(&messages, &TrimStrategy::MaxTokens(4)); + assert!( + tool_pairing_is_intact(&trimmed), + "MaxTokens orphaned a tool result: {trimmed:?}" + ); + // Forward repair: the assistant tool-call turn is NOT re-admitted, so + // the budget-bound trim cannot grow back. + assert!( + !trimmed + .iter() + .any(|m| matches!(m, Message::Assistant(a) if !a.tool_calls.is_empty())) + ); + } + + #[test] + fn unpairable_tool_results_are_dropped_when_no_assistant_exists() { + // An imported/truncated transcript whose assistant turn is already + // gone: there is nothing to move back to, so the results are shed. + let messages = vec![ + Message::tool("ghost", "r1"), + Message::tool("ghost2", "r2"), + Message::assistant("done"), + ]; + let trimmed = trim_messages(&messages, &TrimStrategy::KeepLast(2)); + assert!(tool_pairing_is_intact(&trimmed), "{trimmed:?}"); + assert_eq!(trimmed.len(), 1); + } + + #[test] + fn opting_out_of_repair_restores_the_unsafe_cut() { + // Pins that the repair is what makes the difference, not some other + // change of behaviour: without it the old orphaning cut comes back. + let trimmed = trim_messages_with( + &tool_transcript(), + &TrimStrategy::KeepLast(2), + &TrimOptions::default().without_pair_repair(), + ); + assert!( + !tool_pairing_is_intact(&trimmed), + "expected the unrepaired cut to orphan a tool result" + ); + } + + #[test] + fn role_boundaries_trim_to_the_requested_roles() { + let messages = vec![ + Message::assistant("lead-in"), + Message::user("question"), + Message::assistant("answer"), + Message::user("trailing"), + ]; + let trimmed = trim_messages_with( + &messages, + &TrimStrategy::KeepLast(4), + &TrimOptions::default() + .starting_on([MessageRole::User]) + .ending_on([MessageRole::Assistant]), + ); + assert_eq!(trimmed.len(), 2); + assert_eq!(trimmed[0].text(), "question"); + assert_eq!(trimmed[1].text(), "answer"); + } + + #[test] + fn end_on_boundary_does_not_leave_an_unanswered_tool_call() { + let messages = vec![ + Message::user("go"), + Message::assistant("thinking"), + assistant_calling(&["c1"]), + Message::tool("c1", "r1"), + ]; + // Ending on an assistant turn would naively stop on the tool-call turn + // whose result was just dropped. + let trimmed = trim_messages_with( + &messages, + &TrimStrategy::KeepLast(4), + &TrimOptions::default().ending_on([MessageRole::Assistant]), + ); + assert!(tool_pairing_is_intact(&trimmed), "{trimmed:?}"); + assert_eq!(trimmed.last().map(Message::text), Some("thinking".into())); + } + + #[test] + fn tool_pairing_is_intact_detects_both_orphan_shapes() { + assert!(tool_pairing_is_intact(&tool_transcript())); + // Orphaned result. + assert!(!tool_pairing_is_intact(&[Message::tool("c1", "r")])); + // Unanswered call. + assert!(!tool_pairing_is_intact(&[assistant_calling(&["c1"])])); + } + + #[test] + fn a_tool_only_assistant_turn_is_not_free() { + // REASON-3: an assistant message with no content but a 2 KB argument + // blob used to weigh zero, so no compaction gate could ever fire. + let heavy = Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ToolCall::new( + "c1", + "search", + json!({"query": "x".repeat(2000)}), + )], + usage: None, + }); + assert!( + heavy.estimated_char_weight() > 2000, + "tool-call arguments must be counted, got {}", + heavy.estimated_char_weight() + ); + + let policy = SummarizationPolicy { + trigger_tokens: 100, + ..Default::default() + }; + assert!( + policy.should_summarize(&[heavy]), + "a 2 KB tool-call turn must trip a 100-token trigger" + ); + } + + #[test] + fn a_tool_result_id_is_counted() { + let bare = Message::Tool(crate::harness::message::ToolMessage { + tool_call_id: "call_abcdefghijklmnop".into(), + content: Vec::new(), + trusted_verbatim: false, + artifact: None, + }); + assert_eq!(bare.estimated_char_weight(), "call_abcdefghijklmnop".len()); + } + + #[test] + fn reasoning_only_turns_still_weigh() { + let msg = Message::Assistant(AssistantMessage { + id: None, + content: vec![ContentBlock::thinking("z".repeat(120))], + tool_calls: Vec::new(), + usage: None, + }); + assert_eq!(msg.estimated_char_weight(), 120); + } +} + +/// Tests for [`render_message_for_summary`] and the default summarizer built on +/// it. +#[cfg(test)] +mod rendering { + use crate::harness::message::{AssistantMessage, Message}; + use crate::harness::summarization::{ConcatSummarizer, Summarizer, render_message_for_summary}; + use crate::harness::tool::ToolCall; + use serde_json::json; + + #[tokio::test] + async fn default_summarizer_keeps_tool_history() { + // REASON-8: every one of these messages rendered to a bare role label + // under `Message::text()`, so the default summarizer replaced real + // history with nothing. + let messages = vec![ + Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ToolCall::new("c1", "get_weather", json!({"city": "Paris"}))], + usage: None, + }), + Message::tool("c1", r#"{"temp_c":21}"#), + ]; + + let record = ConcatSummarizer.summarize(&messages).await.unwrap(); + let text = record.summary.text(); + + assert!(text.contains("get_weather"), "tool name lost: {text}"); + assert!(text.contains("Paris"), "tool arguments lost: {text}"); + assert!(text.contains("temp_c"), "tool result lost: {text}"); + assert!(text.contains("c1"), "correlation id lost: {text}"); + } + + #[test] + fn reasoning_and_json_are_rendered() { + let msg = Message::Assistant(AssistantMessage { + id: None, + content: vec![ + crate::harness::message::ContentBlock::thinking("weighing options"), + crate::harness::message::ContentBlock::Json(json!({"k": "v"})), + ], + tool_calls: Vec::new(), + usage: None, + }); + let rendered = render_message_for_summary(&msg); + assert!(rendered.contains("weighing options"), "{rendered}"); + assert!(rendered.contains("\"k\""), "{rendered}"); + } + + #[test] + fn oversized_payloads_are_elided_not_reproduced() { + let msg = Message::tool("c1", "y".repeat(9_000)); + let rendered = render_message_for_summary(&msg); + assert!(rendered.contains("chars elided"), "no elision marker"); + assert!( + rendered.chars().count() < 3_000, + "summary reproduced the payload" + ); + } +} diff --git a/src/harness/summarization/trim.rs b/src/harness/summarization/trim.rs new file mode 100644 index 0000000..af51d2c --- /dev/null +++ b/src/harness/summarization/trim.rs @@ -0,0 +1,289 @@ +//! Synchronous, LLM-free transcript trimming. +//! +//! Every strategy here is **pairing-safe by default**: the cut point produced by +//! the strategy is repaired by [`super::pairing`] before the slice is rebuilt, +//! so a trimmed transcript never opens on an orphaned tool result or ends a +//! retained prefix on an unanswered assistant tool call. See +//! [`TrimOptions::repair_tool_pairs`] for the escape hatch and why you almost +//! certainly do not want it. +//! +//! Two mechanisms are provided, mirroring the two LangChain has: +//! +//! - **structural repair** — the pairing scan described above, which is +//! automatic (LangChain implements this only inside its summarization +//! middleware, not in `trim_messages`); +//! - **role boundaries** — [`TrimOptions::start_on`] / [`TrimOptions::end_on`], +//! the caller-driven predicates LangChain core's `trim_messages` exposes, +//! whose docstring makes honouring provider pairing rules the caller's job. + +use super::pairing::{ + advance_past_orphan_tools, find_safe_cutoff_point, retract_orphan_tool_calls, +}; +use super::types::{MessageRole, TokenTrimPolicy, TrimOptions, TrimStrategy}; +use crate::harness::message::{Message, estimate_message_tokens}; + +/// Partition `messages` into system and non-system messages, preserving order. +/// +/// Returns `(system, non_system)`. Pairing repair operates on the non-system +/// half alone: a system message can never sit between an assistant tool-call +/// turn and the tool results answering it, so removing them from consideration +/// cannot change a pairing decision. +pub(super) fn partition_system(messages: &[Message]) -> (Vec, Vec) { + let system = messages + .iter() + .filter(|m| matches!(m, Message::System(_))) + .cloned() + .collect(); + let non_system = messages + .iter() + .filter(|m| !matches!(m, Message::System(_))) + .cloned() + .collect(); + (system, non_system) +} + +/// Trim a message slice according to `strategy`, returning the retained subset. +/// +/// System messages are preserved by default: +/// +/// - [`TrimStrategy::KeepLast`] and [`TrimStrategy::KeepFirstAndLast`] always +/// keep all system messages and apply the rule only to non-system messages. +/// - [`TrimStrategy::MaxTokens`] drops non-system messages first (from the +/// front) and only starts dropping system messages if the budget still +/// cannot be met after all non-system messages are removed. +/// +/// Tool-call pairing is repaired automatically — see [`trim_messages_with`] for +/// the details and for the role-boundary knobs. The returned `Vec` +/// preserves the relative order of messages as they appeared in the input. +pub fn trim_messages(messages: &[Message], strategy: &TrimStrategy) -> Vec { + trim_messages_with(messages, strategy, &TrimOptions::default()) +} + +/// Trim a message slice according to `strategy` under explicit [`TrimOptions`]. +/// +/// # Pairing repair +/// +/// With [`TrimOptions::repair_tool_pairs`] set (the default), each strategy's +/// raw cut point is adjusted so the result is a slice a provider will accept: +/// +/// | Strategy | Repair | +/// | -------- | ------ | +/// | [`TrimStrategy::KeepLast`] | cut moves **backward** to include the assistant turn owning any leading tool results | +/// | [`TrimStrategy::KeepFirstAndLast`] | suffix cut moves backward as above; the retained prefix's end moves backward past any unanswered assistant tool call | +/// | [`TrimStrategy::MaxTokens`] | cut moves **forward**, dropping orphaned leading tool results — moving backward would re-admit the tokens the budget was shedding | +/// +/// # Role boundaries +/// +/// [`TrimOptions::start_on`] and [`TrimOptions::end_on`] drop further messages +/// from the front / back until the retained slice begins / ends on one of the +/// listed roles, matching LangChain core's `trim_messages(start_on=…, +/// end_on=…)`. They are applied *after* the strategy, and pairing repair runs +/// once more afterwards so a role boundary cannot reintroduce an orphan. +pub fn trim_messages_with( + messages: &[Message], + strategy: &TrimStrategy, + options: &TrimOptions, +) -> Vec { + let (system, non_system) = partition_system(messages); + let (retained_system, mut retained) = apply_strategy(&system, &non_system, strategy, options); + apply_role_boundaries(&mut retained, options); + + if options.repair_tool_pairs { + let start = advance_past_orphan_tools(&retained, 0); + if start > 0 { + retained.drain(..start); + } + } + + let mut result = retained_system; + result.extend(retained); + tracing::debug!( + "[summarization::trim] strategy={strategy:?} input={} retained={}", + messages.len(), + result.len() + ); + result +} + +/// Applies `strategy`, returning `(retained_system, retained_non_system)`. +/// +/// System messages are returned separately (rather than re-attached here) +/// because [`TrimStrategy::MaxTokens`] is the one strategy allowed to shed +/// them, and only after every other message is gone. +fn apply_strategy( + system: &[Message], + non_system: &[Message], + strategy: &TrimStrategy, + options: &TrimOptions, +) -> (Vec, Vec) { + match strategy { + TrimStrategy::KeepLast(n) => { + let mut keep_start = non_system.len().saturating_sub(*n); + if options.repair_tool_pairs { + keep_start = find_safe_cutoff_point(non_system, keep_start); + } + (system.to_vec(), non_system[keep_start..].to_vec()) + } + + TrimStrategy::KeepFirstAndLast { first, last } => { + let len = non_system.len(); + let (first, last) = (*first, *last); + + if first + last >= len { + // No room to drop anything: keep every non-system message. + return (system.to_vec(), non_system.to_vec()); + } + + let mut prefix_end = first; + let mut suffix_start = len - last; + if options.repair_tool_pairs { + prefix_end = retract_orphan_tool_calls(non_system, prefix_end); + suffix_start = find_safe_cutoff_point(non_system, suffix_start); + // Backward repair of the suffix can reach into (or past) the + // retained prefix. Clamping keeps the two ranges contiguous and + // non-overlapping, which in that case means keeping everything. + suffix_start = suffix_start.max(prefix_end); + } + + let mut result = non_system[..prefix_end].to_vec(); + result.extend_from_slice(&non_system[suffix_start..]); + (system.to_vec(), result) + } + + TrimStrategy::MaxTokens(limit) => { + let limit = *limit; + + // Precompute each message's token estimate once: re-summing the + // slice per dropped message (and using `remove(0)`) would make this + // O(n^2). + let sys_tokens: Vec = system.iter().map(estimate_message_tokens).collect(); + let non_sys_tokens: Vec = non_system.iter().map(estimate_message_tokens).collect(); + let sys_total: u64 = sys_tokens.iter().sum(); + + let mut non_sys_start = 0; + let mut non_sys_total: u64 = non_sys_tokens.iter().sum(); + while non_sys_start < non_system.len() && sys_total + non_sys_total > limit { + non_sys_total -= non_sys_tokens[non_sys_start]; + non_sys_start += 1; + } + + if options.repair_tool_pairs { + // Forward-only: a budget-bound trim must not grow again. + non_sys_start = advance_past_orphan_tools(non_system, non_sys_start); + } + + // Still over budget with every non-system message gone: shed + // system instructions from the front as a last resort. + let mut sys_start = 0; + let mut sys_running = sys_total; + while sys_start < system.len() && sys_running + non_sys_total > limit { + sys_running -= sys_tokens[sys_start]; + sys_start += 1; + } + + ( + system[sys_start..].to_vec(), + non_system[non_sys_start..].to_vec(), + ) + } + } +} + +/// Drops messages from either end until the slice begins / ends on an allowed +/// role, per [`TrimOptions::start_on`] / [`TrimOptions::end_on`]. +fn apply_role_boundaries(retained: &mut Vec, options: &TrimOptions) { + if let Some(end_on) = options.end_on.as_ref() + && !end_on.is_empty() + { + while let Some(last) = retained.last() { + if end_on.contains(&MessageRole::of(last)) { + break; + } + retained.pop(); + } + if options.repair_tool_pairs { + let end = retract_orphan_tool_calls(retained, retained.len()); + retained.truncate(end); + } + } + + if let Some(start_on) = options.start_on.as_ref() + && !start_on.is_empty() + { + let mut start = 0; + while start < retained.len() && !start_on.contains(&MessageRole::of(&retained[start])) { + start += 1; + } + retained.drain(..start); + } +} + +/// Trim messages to a token budget while preserving their original order. +/// +/// `estimate` lets callers account for provider-specific payloads without +/// teaching the harness about their wire representation. The built-in +/// [`Message::estimated_char_weight`] already assigns a flat weight to native +/// image blocks; hosts may additionally recognize inline image markers or use a +/// real tokenizer. Oldest non-system messages are evicted first. System +/// messages are either retained unconditionally or evicted oldest-first only +/// after all other messages, according to [`TokenTrimPolicy::preserve_system`]. +/// +/// When `drop_leading_orphan_tools` is enabled, leading tool-result messages are +/// removed after budget eviction so the retained transcript starts on a valid +/// provider turn boundary. The returned messages always retain their original +/// relative order. +pub fn trim_messages_to_token_budget_with( + messages: &[Message], + policy: TokenTrimPolicy, + estimate: impl Fn(&Message) -> u64, +) -> Vec { + let estimates: Vec = messages.iter().map(estimate).collect(); + let mut total: u64 = estimates.iter().copied().sum(); + if total <= policy.limit && !policy.drop_leading_orphan_tools { + return messages.to_vec(); + } + + let mut retained = vec![true; messages.len()]; + for (index, message) in messages.iter().enumerate() { + if total <= policy.limit { + break; + } + if !matches!(message, Message::System(_)) { + retained[index] = false; + total = total.saturating_sub(estimates[index]); + } + } + + if !policy.preserve_system && total > policy.limit { + for (index, message) in messages.iter().enumerate() { + if total <= policy.limit { + break; + } + if matches!(message, Message::System(_)) && retained[index] { + retained[index] = false; + total = total.saturating_sub(estimates[index]); + } + } + } + + let mut result: Vec = messages + .iter() + .zip(retained) + .filter(|(_, keep)| *keep) + .map(|(message, _)| message.clone()) + .collect(); + + if policy.drop_leading_orphan_tools { + while let Some(index) = result + .iter() + .position(|message| !matches!(message, Message::System(_))) + { + if matches!(result[index], Message::Tool(_)) { + result.remove(index); + } else { + break; + } + } + } + result +} diff --git a/src/harness/summarization/types.rs b/src/harness/summarization/types.rs index ac506a3..729edfe 100644 --- a/src/harness/summarization/types.rs +++ b/src/harness/summarization/types.rs @@ -58,6 +58,121 @@ pub enum TrimStrategy { MaxTokens(u64), } +/// The role of a [`Message`], as a standalone value for role-boundary +/// predicates. +/// +/// Used by [`TrimOptions::start_on`] / [`TrimOptions::end_on`], the crate's +/// port of LangChain core's `trim_messages(start_on=…, end_on=…)`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum MessageRole { + /// [`Message::System`]. + System, + /// [`Message::User`]. + User, + /// [`Message::Assistant`]. + Assistant, + /// [`Message::Tool`]. + Tool, +} + +impl MessageRole { + /// Returns the role of `message`. + pub fn of(message: &Message) -> Self { + match message { + Message::System(_) => MessageRole::System, + Message::User(_) => MessageRole::User, + Message::Assistant(_) => MessageRole::Assistant, + Message::Tool(_) => MessageRole::Tool, + } + } +} + +/// Knobs for [`trim_messages_with`][crate::harness::summarization::trim_messages_with]. +/// +/// [`Default`] is the safe configuration: tool-call pairing is repaired and no +/// role boundary is imposed. [`trim_messages`][crate::harness::summarization::trim_messages] +/// is exactly `trim_messages_with(.., &TrimOptions::default())`. +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct TrimOptions { + /// Repair the strategy's cut point so it never splits an assistant + /// tool-call turn from the tool results answering it, and never leaves an + /// assistant tool call unanswered. + /// + /// **Defaults to on, and should stay on.** Turning it off restores the + /// pre-repair behaviour, which produces transcripts that providers reject + /// outright: OpenAI `400`s on a `role:"tool"` with no preceding + /// `tool_calls`, and Anthropic rejects a `tool_result` with no matching + /// `tool_use`. It exists for callers that reconstruct pairing themselves + /// afterwards (and for tests that need to observe the unrepaired cut). + /// + /// `#[serde(default = …)]` so a persisted `TrimOptions` written before this + /// field existed — or one that simply omits it — still deserialises to the + /// safe value rather than to `false`. + #[serde(default = "default_repair_tool_pairs")] + pub repair_tool_pairs: bool, + + /// Drop messages from the **front** of the retained slice until it begins + /// on one of these roles. `None` (the default) imposes no boundary. + /// + /// This is the mechanism LangChain core's `trim_messages` documents as the + /// caller's responsibility for provider compatibility — e.g. + /// `start_on: [User]` for providers that require the first non-system turn + /// to be a user turn. It composes with, and does not replace, + /// [`repair_tool_pairs`][Self::repair_tool_pairs]. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub start_on: Option>, + + /// Drop messages from the **back** of the retained slice until it ends on + /// one of these roles. `None` (the default) imposes no boundary. + /// + /// `end_on: [Tool, Assistant]` is the usual setting for "do not end the + /// prompt mid-tool-call". + #[serde(default, skip_serializing_if = "Option::is_none")] + pub end_on: Option>, +} + +/// The default for [`TrimOptions::repair_tool_pairs`] (`true`). +pub(crate) fn default_repair_tool_pairs() -> bool { + true +} + +impl Default for TrimOptions { + /// The safe configuration: pairing repaired, no role boundary. + /// + /// Hand-written rather than derived precisely because a derived `Default` + /// would set `repair_tool_pairs` to `false` — silently restoring the + /// provider-`400` behaviour this type exists to prevent. + fn default() -> Self { + Self { + repair_tool_pairs: default_repair_tool_pairs(), + start_on: None, + end_on: None, + } + } +} + +impl TrimOptions { + /// Requires the retained slice to begin on one of `roles`. + pub fn starting_on(mut self, roles: impl IntoIterator) -> Self { + self.start_on = Some(roles.into_iter().collect()); + self + } + + /// Requires the retained slice to end on one of `roles`. + pub fn ending_on(mut self, roles: impl IntoIterator) -> Self { + self.end_on = Some(roles.into_iter().collect()); + self + } + + /// Disables tool-call pairing repair. See + /// [`repair_tool_pairs`][Self::repair_tool_pairs] before reaching for this. + pub fn without_pair_repair(mut self) -> Self { + self.repair_tool_pairs = false; + self + } +} + /// Options for order-preserving token-budget trimming with a caller-supplied /// message estimator. #[derive(Clone, Copy, Debug, PartialEq, Eq)] diff --git a/src/harness/tool/error_policy.rs b/src/harness/tool/error_policy.rs new file mode 100644 index 0000000..cfcabfd --- /dev/null +++ b/src/harness/tool/error_policy.rs @@ -0,0 +1,132 @@ +//! Per-tool error handling policy. +//! +//! # The gap this closes +//! +//! [`Tool::call`][super::Tool::call] returns `Result`, and an `Err` +//! propagates out of the agent loop and kills the run. The only way for a tool +//! to signal a *recoverable* failure — one the model should see and retry +//! differently — is to return `Ok(ToolResult::error(...))`. That convention is +//! documented, but it is a documentation-level contract only: nothing stops a +//! tool from returning `Err` for a routine "file not found", and when one does, +//! a whole agent run dies over something the model could have handled in one +//! more turn. +//! +//! LangChain solves this with `handle_tool_error: bool | str | Callable` on the +//! tool itself, resolved at execution time: a handled error flips the result's +//! status and comes back as a `ToolMessage(status="error")` rather than +//! propagating. [`ToolErrorPolicy`] is the same idea with the variants named. +//! +//! # The one rule that must not be broken +//! +//! **Cancellation and interruption always bubble.** See +//! [`ToolErrorPolicy::apply`]. + +use crate::error::{Result, TinyAgentsError}; + +use super::types::{ToolCall, ToolResult}; + +/// What the harness should do when a tool's [`call`][super::Tool::call] returns +/// `Err`. +/// +/// Declared per tool via [`Tool::error_policy`][super::Tool::error_policy]. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum ToolErrorPolicy { + /// Propagate the error and fail the run. The default, and the right choice + /// for a tool whose failure means the run's premise is broken. + #[default] + Fail, + + /// Convert the error into an error [`ToolResult`] carrying the error's own + /// message, so the model sees what went wrong and can adapt. + /// + /// Equivalent to LangChain's `handle_tool_error=True`. + ReturnToError, + + /// Convert the error into an error [`ToolResult`] carrying this fixed + /// message instead of the error's own text. + /// + /// Use it when the underlying error is not safe or not useful to show a + /// model — a database error containing a connection string, a stack trace, + /// an internal identifier. Equivalent to LangChain's + /// `handle_tool_error=""`. + Message(String), +} + +impl ToolErrorPolicy { + /// Applies the policy to a tool invocation's outcome. + /// + /// # Control-flow errors always bubble + /// + /// [`TinyAgentsError::Cancelled`] and [`TinyAgentsError::Interrupted`] are + /// re-raised **regardless of policy**. They are not tool failures: they are + /// the run being told to stop, or a human-in-the-loop pause. Converting one + /// into a tool result would report "the tool failed" to the model, let the + /// loop continue, and silently defeat a cancel or an interrupt — the exact + /// reason LangGraph re-raises `GraphBubbleUp` unconditionally in + /// `ToolNode`, and why its tool-retry middleware does the same. Swallowing + /// an interrupt is a correctness bug, not a robustness feature. + /// + /// Everything else is routed by the policy. `Ok` results pass straight + /// through untouched — a tool that already returned + /// `Ok(ToolResult::error(..))` has made its own choice and the policy does + /// not second-guess it. + pub fn apply(&self, call: &ToolCall, outcome: Result) -> Result { + let error = match outcome { + Ok(result) => return Ok(result), + Err(error) => error, + }; + + if is_control_flow_error(&error) { + tracing::debug!( + "[tool::error_policy] bubbling control-flow error from `{}`: {error}", + call.name + ); + return Err(error); + } + + match self { + ToolErrorPolicy::Fail => { + tracing::debug!( + "[tool::error_policy] failing run on `{}` error: {error}", + call.name + ); + Err(error) + } + ToolErrorPolicy::ReturnToError => { + tracing::debug!( + "[tool::error_policy] converting `{}` error to a tool result: {error}", + call.name + ); + Ok(ToolResult::error( + call.id.clone(), + call.name.clone(), + error.to_string(), + )) + } + ToolErrorPolicy::Message(message) => { + tracing::debug!( + "[tool::error_policy] masking `{}` error behind a fixed message: {error}", + call.name + ); + Ok(ToolResult::error( + call.id.clone(), + call.name.clone(), + message.clone(), + )) + } + } + } +} + +/// Whether `error` is a control-flow signal that must never be converted into a +/// tool result. +/// +/// Exposed so any execution site — including middleware that wraps tool calls +/// with its own retry or fallback logic — can apply the same rule, rather than +/// each re-deriving which variants are safe to swallow. +pub fn is_control_flow_error(error: &TinyAgentsError) -> bool { + matches!( + error, + TinyAgentsError::Cancelled | TinyAgentsError::Interrupted { .. } + ) +} diff --git a/src/harness/tool/error_policy_test.rs b/src/harness/tool/error_policy_test.rs new file mode 100644 index 0000000..2e8caf4 --- /dev/null +++ b/src/harness/tool/error_policy_test.rs @@ -0,0 +1,122 @@ +//! Tests for [`ToolErrorPolicy`]. +//! +//! The load-bearing case is the last one: a cancellation or an interrupt must +//! bubble no matter what policy a tool declares. Converting one into a tool +//! result would tell the model "the tool failed", let the loop continue, and +//! silently defeat the cancel. + +use super::*; +use crate::error::TinyAgentsError; +use serde_json::json; + +fn call() -> ToolCall { + ToolCall::new("c1", "lookup", json!({})) +} + +#[test] +fn fail_is_the_default_and_propagates() { + assert_eq!(ToolErrorPolicy::default(), ToolErrorPolicy::Fail); + + let outcome = ToolErrorPolicy::Fail.apply(&call(), Err(TinyAgentsError::Tool("boom".into()))); + assert!(matches!(outcome, Err(TinyAgentsError::Tool(_)))); +} + +#[test] +fn return_to_error_hands_the_message_to_the_model() { + let outcome = + ToolErrorPolicy::ReturnToError.apply(&call(), Err(TinyAgentsError::Tool("boom".into()))); + let result = outcome.expect("a handled error must not fail the run"); + + assert!(result.is_error()); + assert_eq!(result.call_id, "c1"); + assert_eq!(result.name, "lookup"); + assert!(result.content.contains("boom")); +} + +#[test] +fn a_fixed_message_masks_an_error_that_should_not_be_shown() { + let outcome = ToolErrorPolicy::Message("the lookup service is unavailable".into()).apply( + &call(), + Err(TinyAgentsError::Tool( + "postgres://user:pw@db/internal timed out".into(), + )), + ); + let result = outcome.unwrap(); + + assert_eq!(result.content, "the lookup service is unavailable"); + assert!(!result.content.contains("postgres")); +} + +#[test] +fn successful_results_pass_through_untouched() { + // A tool that already chose `Ok(ToolResult::error(..))` has made its own + // decision; the policy does not second-guess it. + let declared = ToolResult::error("c1", "lookup", "not found"); + let outcome = ToolErrorPolicy::Fail.apply(&call(), Ok(declared.clone())); + assert_eq!(outcome.unwrap(), declared); +} + +#[test] +fn cancellation_and_interruption_always_bubble() { + for policy in [ + ToolErrorPolicy::Fail, + ToolErrorPolicy::ReturnToError, + ToolErrorPolicy::Message("masked".into()), + ] { + let cancelled = policy.apply(&call(), Err(TinyAgentsError::Cancelled)); + assert!( + matches!(cancelled, Err(TinyAgentsError::Cancelled)), + "policy {policy:?} swallowed a cancellation" + ); + + let interrupted = policy.apply( + &call(), + Err(TinyAgentsError::Interrupted { + node: "approval".into(), + message: "waiting for a human".into(), + }), + ); + assert!( + matches!(interrupted, Err(TinyAgentsError::Interrupted { .. })), + "policy {policy:?} swallowed an interrupt" + ); + } + + assert!(is_control_flow_error(&TinyAgentsError::Cancelled)); + assert!(!is_control_flow_error(&TinyAgentsError::Tool( + "boom".into() + ))); +} + +#[test] +fn registry_exposes_per_tool_error_policies() { + use async_trait::async_trait; + + struct Flaky; + + #[async_trait] + impl Tool<()> for Flaky { + fn name(&self) -> &str { + "flaky" + } + fn description(&self) -> &str { + "sometimes fails" + } + fn schema(&self) -> ToolSchema { + ToolSchema::new("flaky", "sometimes fails", json!({"type": "object"})) + } + fn error_policy(&self) -> ToolErrorPolicy { + ToolErrorPolicy::ReturnToError + } + async fn call(&self, _state: &(), call: ToolCall) -> Result { + Ok(ToolResult::text(call.id, call.name, "ok")) + } + } + + let mut registry: ToolRegistry<()> = ToolRegistry::new(); + registry.register(std::sync::Arc::new(Flaky)); + assert_eq!( + registry.error_policies().get("flaky"), + Some(&ToolErrorPolicy::ReturnToError) + ); +} diff --git a/src/harness/tool/injected.rs b/src/harness/tool/injected.rs new file mode 100644 index 0000000..829e91a --- /dev/null +++ b/src/harness/tool/injected.rs @@ -0,0 +1,115 @@ +//! Injected (hidden) tool arguments: values the *host* supplies, never the +//! model. +//! +//! # What this is for +//! +//! Some tool arguments are not the model's business: the caller's thread id, +//! the recursion depth, the id of the tool call being answered, a database +//! handle, an authenticated user. Today a tool that needs them reaches around +//! the argument schema entirely — `SubAgentTool` reads `context.thread_id` and +//! `context.depth` off [`ToolExecutionContext`][super::ToolExecutionContext] +//! because there is no declarative way to *receive* them — which means every +//! such tool re-invents the plumbing and none of it is visible in the tool's +//! declared shape. +//! +//! LangChain models this as `InjectedToolArg` / `InjectedToolCallId` +//! annotations, strips them from the model-facing `tool_call_schema`, and +//! LangGraph's `ToolNode` re-injects the real values at execution time. +//! +//! # The ordering rule (security-critical) +//! +//! At execution time the sequence must be, in this order: +//! +//! 1. **Strip** every injected key from the model-supplied arguments. +//! 2. **Validate** the remaining arguments against the model-facing schema. +//! 3. **Inject** the host's real values. +//! 4. Invoke the tool. +//! +//! Step 1 must come first and must be unconditional. A model that has seen an +//! injected key named in a prompt, a log, or an error message can put that key +//! in its own `arguments` object; if the host merges its value in *after*, a +//! well-formed merge might still let the model's value win, and if the host +//! merges *before* validating, the forged key rides along. LangGraph's +//! `ToolNode` strips first for exactly this reason, with the comment that it +//! "prevents an LLM from forging hidden InjectedToolArg fields via +//! ToolCall.args". +//! +//! [`strip_injected_arguments`] performs step 1 and reports what it removed, so +//! a forgery attempt is visible in the log rather than silent. +//! +//! # Status +//! +//! The declaration side is live: [`Tool::injected_arguments`][super::Tool::injected_arguments] +//! declares the keys and [`ToolRegistry::schemas`][super::ToolRegistry::schemas] +//! already projects them out of what the model sees. The enforcement side (the +//! four-step sequence above) belongs to the agent loop's tool-execution path. + +use serde_json::Value; + +use super::types::ToolSchema; + +/// Removes every key in `injected` from a model-supplied argument object, +/// returning the names that were actually present. +/// +/// A non-empty return value means the model emitted a key it was never shown — +/// either because it inferred one, or because it was told one. Callers should +/// log it; the value itself is discarded either way. +/// +/// Non-object arguments (including the raw string preserved on an +/// [`invalid`][super::ToolCall::invalid] call) are left untouched: there is no +/// key to forge in a scalar. +pub fn strip_injected_arguments(arguments: &mut Value, injected: &[&str]) -> Vec { + if injected.is_empty() { + return Vec::new(); + } + let Some(object) = arguments.as_object_mut() else { + return Vec::new(); + }; + + let mut removed = Vec::new(); + for key in injected { + if object.remove(*key).is_some() { + removed.push((*key).to_string()); + } + } + + if !removed.is_empty() { + tracing::warn!( + "[tool::injected] discarded model-supplied value(s) for host-injected argument(s): {}", + removed.join(", ") + ); + } + removed +} + +/// Removes `injected` keys from a schema's `properties` **and** its `required` +/// list, producing the model-facing projection of the declaration. +/// +/// Dropping a key from `properties` alone is not enough: leaving it in +/// `required` tells the model to supply an argument it cannot see, which is +/// either a validation failure or an invitation to invent the value. +pub fn project_injected_arguments(mut schema: ToolSchema, injected: &[&str]) -> ToolSchema { + if injected.is_empty() { + return schema; + } + let Some(parameters) = schema.parameters.as_object_mut() else { + return schema; + }; + + if let Some(Value::Object(properties)) = parameters.get_mut("properties") { + for key in injected { + properties.remove(*key); + } + } + + if let Some(Value::Array(required)) = parameters.get_mut("required") { + required.retain(|value| value.as_str().is_none_or(|name| !injected.contains(&name))); + } + + tracing::trace!( + "[tool::injected] projected {} hidden argument(s) out of `{}`", + injected.len(), + schema.name + ); + schema +} diff --git a/src/harness/tool/injected_test.rs b/src/harness/tool/injected_test.rs new file mode 100644 index 0000000..1126c54 --- /dev/null +++ b/src/harness/tool/injected_test.rs @@ -0,0 +1,138 @@ +//! Tests for injected (host-supplied) tool arguments. +//! +//! Cover the two halves of the feature: the declaration side — an injected key +//! is projected out of the model-facing schema, from `properties` *and* from +//! `required` — and the enforcement primitive that discards a model-supplied +//! value for such a key before it can be used. + +use async_trait::async_trait; +use serde_json::json; + +use super::*; + +/// A tool that receives its caller's thread id from the host rather than the +/// model — the shape `SubAgentTool` has to hand-roll today. +struct ThreadScopedTool; + +#[async_trait] +impl Tool<()> for ThreadScopedTool { + fn name(&self) -> &str { + "thread_scoped" + } + + fn description(&self) -> &str { + "Does something in the caller's thread" + } + + fn schema(&self) -> ToolSchema { + ToolSchema::new( + "thread_scoped", + "Does something in the caller's thread", + json!({ + "type": "object", + "properties": { + "note": {"type": "string"}, + "thread_id": {"type": "string"}, + }, + "required": ["note", "thread_id"], + }), + ) + } + + fn injected_arguments(&self) -> &[&str] { + &["thread_id"] + } + + async fn call(&self, _state: &(), call: ToolCall) -> Result { + Ok(ToolResult::text(call.id, call.name, "ok")) + } +} + +#[test] +fn injected_arguments_are_hidden_from_the_model_facing_schema() { + let mut registry: ToolRegistry<()> = ToolRegistry::new(); + registry.register(std::sync::Arc::new(ThreadScopedTool)); + + let model_facing = registry.schemas(); + let properties = &model_facing[0].parameters["properties"]; + assert!(properties.get("note").is_some()); + assert!( + properties.get("thread_id").is_none(), + "an injected argument was advertised to the model" + ); + + // Leaving it in `required` would demand an argument the model cannot see. + let required = model_facing[0].parameters["required"].as_array().unwrap(); + assert_eq!(required, &vec![json!("note")]); +} + +#[test] +fn declared_schemas_keep_injected_arguments_for_introspection() { + let mut registry: ToolRegistry<()> = ToolRegistry::new(); + registry.register(std::sync::Arc::new(ThreadScopedTool)); + + let declared = registry.declared_schemas(); + assert!(declared[0].parameters["properties"]["thread_id"].is_object()); + assert_eq!( + registry.injected_arguments().get("thread_scoped"), + Some(&vec!["thread_id".to_string()]) + ); +} + +#[test] +fn a_model_supplied_value_for_an_injected_key_is_discarded() { + // The forgery this prevents: the model names a hidden key in its own + // arguments, hoping the host will honour it. + let mut arguments = json!({"note": "hi", "thread_id": "victim-thread"}); + let removed = strip_injected_arguments(&mut arguments, &["thread_id"]); + + assert_eq!(removed, vec!["thread_id".to_string()]); + assert_eq!(arguments, json!({"note": "hi"})); +} + +#[test] +fn stripping_is_a_no_op_without_injected_keys_or_object_arguments() { + let mut arguments = json!({"note": "hi"}); + assert!(strip_injected_arguments(&mut arguments, &[]).is_empty()); + assert_eq!(arguments, json!({"note": "hi"})); + + // An `invalid` call preserves raw text as a JSON string; there is no key to + // forge in a scalar, and it must not be mangled. + let mut raw = json!("{not json"); + assert!(strip_injected_arguments(&mut raw, &["thread_id"]).is_empty()); + assert_eq!(raw, json!("{not json")); +} + +#[test] +fn projection_tolerates_schemas_without_properties_or_required() { + let schema = ToolSchema::new("bare", "no args", json!({"type": "object"})); + let projected = project_injected_arguments(schema.clone(), &["thread_id"]); + assert_eq!(projected.parameters, schema.parameters); + + let scalar = ToolSchema::new("odd", "odd", json!("nonsense")); + let projected = project_injected_arguments(scalar.clone(), &["thread_id"]); + assert_eq!(projected.parameters, scalar.parameters); +} + +#[test] +fn tools_declare_no_injected_arguments_by_default() { + struct Plain; + + #[async_trait] + impl Tool<()> for Plain { + fn name(&self) -> &str { + "plain" + } + fn description(&self) -> &str { + "plain" + } + fn schema(&self) -> ToolSchema { + ToolSchema::new("plain", "plain", json!({"type": "object"})) + } + async fn call(&self, _state: &(), call: ToolCall) -> Result { + Ok(ToolResult::text(call.id, call.name, "ok")) + } + } + + assert!(Tool::<()>::injected_arguments(&Plain).is_empty()); +} diff --git a/src/harness/tool/mod.rs b/src/harness/tool/mod.rs index 22f466b..8371ded 100644 --- a/src/harness/tool/mod.rs +++ b/src/harness/tool/mod.rs @@ -10,8 +10,11 @@ //! See [`types`] for definitions. This module provides constructors and the //! [`ToolRegistry`] logic for registering and looking up tools by name. +mod error_policy; +pub mod injected; mod prompt; mod schema; +mod schema_prepare; mod timeout; mod types; @@ -21,8 +24,11 @@ use serde_json::Value; use crate::error::{Result, TinyAgentsError}; +pub use error_policy::{ToolErrorPolicy, is_control_flow_error}; +pub use injected::{project_injected_arguments, strip_injected_arguments}; pub use prompt::*; pub use schema::*; +pub use schema_prepare::*; pub use timeout::*; pub use types::*; @@ -397,13 +403,61 @@ impl ToolRegistry { names } - /// Returns the schemas of all registered tools, sorted by name. + /// Returns the **model-facing** schemas of all registered tools, sorted by + /// name. + /// + /// Each schema has its tool's + /// [`injected_arguments`][Tool::injected_arguments] projected out — removed + /// from `properties` and from `required` alike — so a host-supplied + /// argument is never advertised to the model and never demanded of it. See + /// [`crate::harness::tool::injected`] for the matching execution-time rule. pub fn schemas(&self) -> Vec { + let mut schemas: Vec = self + .tools + .values() + .map(|t| project_injected_arguments(t.schema(), t.injected_arguments())) + .collect(); + schemas.sort_by(|a, b| a.name.cmp(&b.name)); + schemas + } + + /// Returns the **declared** schemas, including any injected arguments. + /// + /// This is the introspection view — registry listings, audit logs, docs — + /// not the model-facing one. Never put this on the wire; use + /// [`Self::schemas`]. + pub fn declared_schemas(&self) -> Vec { let mut schemas: Vec = self.tools.values().map(|t| t.schema()).collect(); schemas.sort_by(|a, b| a.name.cmp(&b.name)); schemas } + /// Returns each registered tool's injected-argument names, keyed by tool + /// name. Tools declaring none are omitted. + pub fn injected_arguments(&self) -> std::collections::HashMap> { + self.tools + .iter() + .filter_map(|(name, tool)| { + let injected = tool.injected_arguments(); + if injected.is_empty() { + return None; + } + Some(( + name.clone(), + injected.iter().map(|key| (*key).to_string()).collect(), + )) + }) + .collect() + } + + /// Returns each registered tool's [`ToolErrorPolicy`], keyed by tool name. + pub fn error_policies(&self) -> std::collections::HashMap { + self.tools + .iter() + .map(|(name, tool)| (name.clone(), tool.error_policy())) + .collect() + } + /// Returns a snapshot of every registered tool's [`ToolPolicy`], keyed by /// tool name. This is the projection policy-enforcement middleware and audit /// logs consume. @@ -551,9 +605,15 @@ fn json_value_kind(value: &Value) -> &'static str { } } +#[cfg(test)] +mod error_policy_test; +#[cfg(test)] +mod injected_test; #[cfg(test)] mod prompt_test; #[cfg(test)] +mod schema_prepare_test; +#[cfg(test)] mod schema_test; #[cfg(test)] mod test; diff --git a/src/harness/tool/prompt.rs b/src/harness/tool/prompt.rs index 6c9f62b..05a8183 100644 --- a/src/harness/tool/prompt.rs +++ b/src/harness/tool/prompt.rs @@ -536,10 +536,58 @@ fn replace_text_blocks(content: Vec, cleaned: String) -> Vec Option { +/// Parse a single tool-call body into a [`ToolCall`] with a synthetic id. +/// +/// `slot` is the call's 1-based position within the response it was recovered +/// from; it appears in the id only for readability. Uniqueness comes from the +/// process-wide counter in [`next_synthetic_call_id`], not from `slot`. +fn parse_one(inner: &str, slot: usize) -> Option { let value = parse_relaxed_object(inner)?; - tool_call_from_object(&value, index) + tool_call_from_object(&value, slot) +} + +/// Monotonic source of unique synthetic tool-call ids. +/// +/// # Why a global counter and not a per-response index +/// +/// The recovered id previously came from the call's position **within one +/// response** (`call_1`, `call_2`, …), which resets on every model turn. That is +/// wrong for anything but a single-turn run: two turns of the same run both emit +/// `call_1`, so the next request contains two assistant messages declaring the +/// same tool-call id and two tool messages answering it. The pairing is then +/// unresolvable — a provider cannot tell which result answers which call, and +/// neither can the harness's own pairing repair. +/// +/// This is **not** confined to prompt-guided models. +/// [`should_recover`] returns `true` for a *native* profile whenever tools were +/// offered and the response carried no structured calls, so a native run that +/// hits the text-mode fallback twice collides exactly the same way. +/// +/// A process-wide `AtomicU64` makes every recovered id unique for the lifetime +/// of the process, which is strictly stronger than per-run uniqueness and needs +/// no run context threaded into a pure parsing function. +static SYNTHETIC_CALL_SEQUENCE: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1); + +/// Prefix of every synthetic id minted here. +/// +/// Deliberately **not** `call_` and **not** `tool-`: those are the shapes real +/// providers emit and the shape the OpenAI adapter mints for its own +/// positional fallback (`tool-{slot}`), so a distinct prefix makes a collision +/// between the two schemes impossible by construction and makes a synthetic id +/// obvious in a transcript or a log. +pub const SYNTHETIC_CALL_ID_PREFIX: &str = "ptc"; + +/// Returns a fresh, process-unique synthetic tool-call id of the form +/// `ptc_{sequence}_{slot}` — "prompt tool call". +/// +/// `slot` is the 1-based position of the call within its response and is +/// included only so a human reading a transcript can see the ordering; the +/// `sequence` component is what guarantees uniqueness. +pub fn next_synthetic_call_id(slot: usize) -> String { + let sequence = SYNTHETIC_CALL_SEQUENCE.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let id = format!("{SYNTHETIC_CALL_ID_PREFIX}_{sequence}_{slot}"); + tracing::trace!("[tool::prompt] minted synthetic tool-call id {id}"); + id } /// Parses a JSON object, repairing the relaxed spellings small local models @@ -557,7 +605,11 @@ fn parse_relaxed_object(raw: &str) -> Option { /// Builds a [`ToolCall`] from an already-parsed call object, or `None` when the /// object does not name a tool. -fn tool_call_from_object(value: &Value, index: usize) -> Option { +/// +/// The id is minted by [`next_synthetic_call_id`] and is unique for the life of +/// the process, so two calls recovered in different turns of the same run can +/// never share one. +fn tool_call_from_object(value: &Value, slot: usize) -> Option { let name = value.get("name")?.as_str()?.trim().to_string(); if name.is_empty() { return None; @@ -567,7 +619,7 @@ fn tool_call_from_object(value: &Value, index: usize) -> Option { .find_map(|key| value.get(*key).cloned()) .unwrap_or_else(|| Value::Object(Map::new())); Some(ToolCall { - id: format!("call_{index}"), + id: next_synthetic_call_id(slot), name, arguments, invalid: None, diff --git a/src/harness/tool/prompt_test.rs b/src/harness/tool/prompt_test.rs index fcbca6d..b0787af 100644 --- a/src/harness/tool/prompt_test.rs +++ b/src/harness/tool/prompt_test.rs @@ -223,7 +223,20 @@ fn prompt_parser_extracts_single_tool_call() { assert_eq!(cleaned, "Let me read it."); assert_eq!(calls.len(), 1); assert_eq!(calls[0].name, "read_file"); - assert_eq!(calls[0].id, "call_1"); + // Ids are process-unique, not a per-response index: only the shape and the + // slot suffix are stable. See `next_synthetic_call_id`. + assert!( + calls[0] + .id + .starts_with(&format!("{SYNTHETIC_CALL_ID_PREFIX}_")), + "unexpected synthetic id {}", + calls[0].id + ); + assert!( + calls[0].id.ends_with("_1"), + "slot suffix lost: {}", + calls[0].id + ); assert_eq!(calls[0].arguments, serde_json::json!({"path": "a.txt"})); } @@ -235,7 +248,12 @@ fn prompt_parser_extracts_multiple_calls_and_keeps_prose() { assert_eq!(calls.len(), 2); assert_eq!(calls[0].name, "one"); assert_eq!(calls[1].name, "two"); - assert_eq!(calls[1].id, "call_2"); + assert!( + calls[1].id.ends_with("_2"), + "slot suffix lost: {}", + calls[1].id + ); + assert_ne!(calls[0].id, calls[1].id); } #[test] @@ -614,3 +632,47 @@ fn bare_tool_call_recovery_preserves_a_thinking_block() { "the reasoning must survive while the consumed object does not" ); } + +/// TOOL-2: two turns of the same run must not both mint `call_1`. +/// +/// The recovered id used to be the call's index *within one response*, which +/// resets every turn. A two-turn run therefore produced a transcript with two +/// assistant messages declaring the same tool-call id and two tool messages +/// answering it — a pairing no provider (and no pairing repair) can resolve. +#[test] +fn synthetic_call_ids_are_unique_across_responses() { + let text = r#"{"name":"one","arguments":{}}"#; + let (_, first) = parse_prompt_tool_calls_from_text(text); + let (_, second) = parse_prompt_tool_calls_from_text(text); + + assert_eq!(first.len(), 1); + assert_eq!(second.len(), 1); + assert_ne!( + first[0].id, second[0].id, + "a second turn reused the first turn's synthetic tool-call id" + ); +} + +/// The synthetic scheme must be visibly distinct from real provider ids and +/// from the OpenAI adapter's own positional fallback (`tool-{slot}`), so the +/// two can never collide. +#[test] +fn synthetic_call_ids_do_not_look_like_provider_ids() { + let id = next_synthetic_call_id(1); + assert!(id.starts_with("ptc_"), "{id}"); + assert!(!id.starts_with("call_"), "{id}"); + assert!(!id.starts_with("tool-"), "{id}"); +} + +/// The bare-object recovery path mints ids from the same counter, so a model +/// that alternates between markup and bare objects still cannot collide. +#[test] +fn bare_object_recovery_also_mints_unique_ids() { + let body = r#"{"name":"one","arguments":{}}"#; + let first = apply_prompt_tool_calls(ModelResponse::assistant(body)); + let second = apply_prompt_tool_calls(ModelResponse::assistant(body)); + + let first_id = &first.message.tool_calls[0].id; + let second_id = &second.message.tool_calls[0].id; + assert_ne!(first_id, second_id); +} diff --git a/src/harness/tool/schema_prepare.rs b/src/harness/tool/schema_prepare.rs new file mode 100644 index 0000000..498ceb6 --- /dev/null +++ b/src/harness/tool/schema_prepare.rs @@ -0,0 +1,290 @@ +//! The provider projection seam for tool declarations. +//! +//! # Why this module exists +//! +//! [`SchemaCleanr`][super::SchemaCleanr] has always known how to resolve +//! `$ref`/`$defs`, strip per-provider unsupported keywords, flatten literal +//! unions, and break `$ref` cycles — and nothing in the crate ever called it. +//! It was dead code for a structural reason, not an oversight: the tool layer +//! had no *place* where a provider-specific projection of a +//! [`ToolSchema`] happens. Every adapter took `Tool::schema()` and put it on +//! the wire unchanged, so a schema written with `$defs` (which is what every +//! JSON-Schema generator emits for a nested type) reached Gemini and Anthropic +//! in a shape they reject. +//! +//! LangChain has exactly one such seam and applies it unconditionally: +//! `convert_to_openai_function` dereferences refs and pops `definitions`/`$defs` +//! for **every** conversion. This module is that seam. +//! +//! # Using it +//! +//! A provider adapter converts its declarations through +//! [`prepare_tool_schemas`] instead of reading `Tool::schema()` directly: +//! +//! ``` +//! use tinyagents::harness::tool::{ +//! CleaningStrategy, SchemaPreparation, ToolSchema, prepare_tool_schemas, +//! }; +//! use serde_json::json; +//! +//! let declared = vec![ToolSchema::new( +//! "lookup", +//! "Look a record up", +//! json!({ +//! "type": "object", +//! "$defs": {"Id": {"type": "string"}}, +//! "properties": {"id": {"$ref": "#/$defs/Id"}}, +//! }), +//! )]; +//! +//! let wire = prepare_tool_schemas(&declared, &SchemaPreparation::anthropic()); +//! // The `$ref` is resolved and `$defs` is gone — Anthropic accepts this. +//! assert_eq!(wire[0].parameters["properties"]["id"]["type"], "string"); +//! assert!(wire[0].parameters.get("$defs").is_none()); +//! ``` +//! +//! # Strict mode +//! +//! OpenAI's structured tool calling (`strict: true`) additionally demands that +//! every declared property be listed in `required` and that +//! `additionalProperties` be `false` at every object level. +//! [`SchemaPreparation::strict`] applies both, mirroring +//! `convert_to_openai_function(..., strict=True)`. + +use serde_json::{Map, Value, json}; + +use super::schema::{CleaningStrategy, SchemaCleanr}; +use super::types::ToolSchema; + +/// How a [`ToolSchema`] should be projected for a specific provider. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct SchemaPreparation { + /// Which keyword subset the target provider accepts. + pub strategy: CleaningStrategy, + /// Apply the OpenAI strict-mode sanitizer: force `required` to list every + /// declared property, and set `additionalProperties: false` at every object + /// level (overriding a pre-existing `true`). + /// + /// Off by default. Turning it on **changes the contract the model is given** + /// — a previously optional argument becomes mandatory — so it belongs to the + /// adapter that actually sends `strict: true`, not to the tool author. + pub strict: bool, +} + +impl SchemaPreparation { + /// Projection for Gemini / Google AI / Vertex AI (the most restrictive + /// keyword set). + pub const fn gemini() -> Self { + Self { + strategy: CleaningStrategy::Gemini, + strict: false, + } + } + + /// Projection for Anthropic (local refs must be resolved). + pub const fn anthropic() -> Self { + Self { + strategy: CleaningStrategy::Anthropic, + strict: false, + } + } + + /// Projection for OpenAI (most permissive; refs are still resolved). + pub const fn openai() -> Self { + Self { + strategy: CleaningStrategy::OpenAI, + strict: false, + } + } + + /// The conservative common subset, for an unknown OpenAI-compatible route. + pub const fn conservative() -> Self { + Self { + strategy: CleaningStrategy::Conservative, + strict: false, + } + } + + /// Enables the strict-mode sanitizer. See [`Self::strict`]. + pub const fn with_strict(mut self) -> Self { + self.strict = true; + self + } +} + +impl Default for SchemaPreparation { + fn default() -> Self { + Self::conservative() + } +} + +/// The parameter schema substituted for a missing or non-object declaration. +/// +/// [`ToolSchema::parameters`] is an unconditional [`Value`], so a tool is free +/// to return `Value::Null`. Serialised as-is that becomes `"parameters": null`, +/// which every provider rejects with a `400` — a tool that simply takes no +/// arguments should not be able to break a whole request that way. +fn empty_object_schema() -> Value { + json!({"type": "object", "properties": {}}) +} + +/// Normalises a parameter schema into something a provider will accept as an +/// object schema. +/// +/// `null`, a bare scalar, and an array are all replaced by +/// `{"type":"object","properties":{}}`. An object schema missing `type` gains +/// `type: "object"` when it declares `properties` or `required`, since that is +/// unambiguously what it meant. +pub fn normalize_parameters(parameters: &Value) -> Value { + let Some(object) = parameters.as_object() else { + tracing::debug!( + "[tool::schema] non-object tool parameters ({}) replaced with an empty object schema", + parameters_kind(parameters) + ); + return empty_object_schema(); + }; + + if object.is_empty() { + return empty_object_schema(); + } + + let mut normalized = object.clone(); + if !normalized.contains_key("type") + && (normalized.contains_key("properties") || normalized.contains_key("required")) + { + normalized.insert("type".to_string(), Value::String("object".to_string())); + } + Value::Object(normalized) +} + +fn parameters_kind(value: &Value) -> &'static str { + match value { + Value::Null => "null", + Value::Bool(_) => "boolean", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Array(_) => "array", + Value::Object(_) => "object", + } +} + +/// Forces every declared property to be `required`. +/// +/// OpenAI strict mode has no notion of an optional argument: a property absent +/// from `required` is rejected outright. Port of the `required` half of +/// `convert_to_openai_function(..., strict=True)`. A schema with no properties +/// is left alone. +pub fn require_all_properties(mut schema: Value) -> Value { + let Some(object) = schema.as_object_mut() else { + return schema; + }; + let Some(Value::Object(properties)) = object.get("properties") else { + return schema; + }; + if properties.is_empty() { + return schema; + } + let required: Vec = properties + .keys() + .map(|key| Value::String(key.clone())) + .collect(); + object.insert("required".to_string(), Value::Array(required)); + schema +} + +/// Sets `additionalProperties: false` at every object level that needs it, +/// **overriding** a pre-existing `true`. +/// +/// Port of LangChain's `_recursive_set_additional_properties_false`, including +/// its three trigger conditions: the level declares `required`, it declares an +/// empty `properties`, or it already carries an `additionalProperties` key +/// (schema generators emit `additionalProperties: true` for open dictionaries, +/// and strict mode forbids that). Recurses through `anyOf`, `properties`, and +/// `items`. +pub fn set_additional_properties_false(mut schema: Value) -> Value { + let Some(object) = schema.as_object_mut() else { + return schema; + }; + + let empty_properties = + matches!(object.get("properties"), Some(Value::Object(p)) if p.is_empty()); + if object.contains_key("required") + || empty_properties + || object.contains_key("additionalProperties") + { + object.insert("additionalProperties".to_string(), Value::Bool(false)); + } + + if let Some(Value::Array(variants)) = object.remove("anyOf") { + let cleaned = variants + .into_iter() + .map(set_additional_properties_false) + .collect(); + object.insert("anyOf".to_string(), Value::Array(cleaned)); + } + + if let Some(Value::Object(properties)) = object.remove("properties") { + let cleaned: Map = properties + .into_iter() + .map(|(key, value)| (key, set_additional_properties_false(value))) + .collect(); + object.insert("properties".to_string(), Value::Object(cleaned)); + } + + if let Some(items) = object.remove("items") { + object.insert("items".to_string(), set_additional_properties_false(items)); + } + + schema +} + +/// Projects one parameter schema for a provider: normalise, clean, then (when +/// requested) apply the strict-mode sanitizer. +/// +/// The order matters. Cleaning runs before the strict sanitizer so that +/// `required` is computed from the *resolved* property set — a schema whose +/// properties arrive through a `$ref` would otherwise be marked as having none. +pub fn prepare_parameters(parameters: &Value, preparation: &SchemaPreparation) -> Value { + let normalized = normalize_parameters(parameters); + let cleaned = SchemaCleanr::clean(normalized, preparation.strategy); + if !preparation.strict { + return cleaned; + } + let required = require_all_properties(cleaned); + set_additional_properties_false(required) +} + +/// Projects one [`ToolSchema`] for a provider, leaving name, description, and +/// format untouched. +pub fn prepare_tool_schema(schema: &ToolSchema, preparation: &SchemaPreparation) -> ToolSchema { + let prepared = ToolSchema { + name: schema.name.clone(), + description: schema.description.clone(), + parameters: prepare_parameters(&schema.parameters, preparation), + format: schema.format.clone(), + }; + tracing::trace!( + "[tool::schema] prepared `{}` for {:?} (strict={})", + schema.name, + preparation.strategy, + preparation.strict + ); + prepared +} + +/// Projects a whole declaration set. This is the call a provider adapter makes. +pub fn prepare_tool_schemas( + schemas: &[ToolSchema], + preparation: &SchemaPreparation, +) -> Vec { + tracing::debug!( + "[tool::schema] preparing {} tool declaration(s) for {:?} (strict={})", + schemas.len(), + preparation.strategy, + preparation.strict + ); + schemas + .iter() + .map(|schema| prepare_tool_schema(schema, preparation)) + .collect() +} diff --git a/src/harness/tool/schema_prepare_test.rs b/src/harness/tool/schema_prepare_test.rs new file mode 100644 index 0000000..1d9871b --- /dev/null +++ b/src/harness/tool/schema_prepare_test.rs @@ -0,0 +1,174 @@ +//! Tests for the provider projection seam. +//! +//! `SchemaCleanr` was fully implemented and entirely unreachable: no call site +//! anywhere in the crate. These tests pin the seam that makes it callable, plus +//! the strict-mode sanitizer and the `parameters: null` guard. + +use super::*; +use serde_json::json; + +fn ref_schema() -> ToolSchema { + ToolSchema::new( + "lookup", + "Look a record up", + json!({ + "type": "object", + "$defs": {"Id": {"type": "string", "description": "record id"}}, + "properties": { + "id": {"$ref": "#/$defs/Id"}, + "limit": {"type": "integer"}, + }, + "required": ["id"], + }), + ) +} + +#[test] +fn local_refs_are_resolved_and_defs_dropped_for_anthropic() { + // Anthropic rejects `$ref` / `$defs` outright, and every JSON-Schema + // generator emits them for a nested type. + let prepared = prepare_tool_schema(&ref_schema(), &SchemaPreparation::anthropic()); + + assert_eq!(prepared.parameters["properties"]["id"]["type"], "string"); + assert!(prepared.parameters.get("$defs").is_none()); + assert!( + prepared.parameters["properties"]["id"] + .get("$ref") + .is_none() + ); + // Name/description/format are untouched by the projection. + assert_eq!(prepared.name, "lookup"); + assert_eq!(prepared.description, "Look a record up"); +} + +#[test] +fn gemini_drops_the_keywords_it_rejects() { + let schema = ToolSchema::new( + "search", + "Search", + json!({ + "type": "object", + "additionalProperties": false, + "properties": {"q": {"type": "string", "minLength": 3, "pattern": "^a"}}, + }), + ); + let prepared = prepare_tool_schema(&schema, &SchemaPreparation::gemini()); + + assert!(prepared.parameters.get("additionalProperties").is_none()); + let q = &prepared.parameters["properties"]["q"]; + assert!(q.get("minLength").is_none()); + assert!(q.get("pattern").is_none()); + assert_eq!(q["type"], "string"); +} + +#[test] +fn strict_mode_requires_every_property_and_closes_every_object() { + let prepared = prepare_tool_schema(&ref_schema(), &SchemaPreparation::openai().with_strict()); + + let required: Vec<&str> = prepared.parameters["required"] + .as_array() + .unwrap() + .iter() + .map(|value| value.as_str().unwrap()) + .collect(); + assert!(required.contains(&"id")); + assert!( + required.contains(&"limit"), + "strict mode has no optional arguments" + ); + assert_eq!(prepared.parameters["additionalProperties"], json!(false)); +} + +#[test] +fn strict_mode_overrides_a_pre_existing_additional_properties_true() { + // Schema generators emit `additionalProperties: true` for open dictionaries; + // strict mode forbids it, so it must be overridden rather than preserved. + let schema = ToolSchema::new( + "config", + "Configure", + json!({ + "type": "object", + "properties": { + "options": {"type": "object", "additionalProperties": true}, + }, + "required": ["options"], + }), + ); + let prepared = prepare_tool_schema(&schema, &SchemaPreparation::openai().with_strict()); + + assert_eq!( + prepared.parameters["properties"]["options"]["additionalProperties"], + json!(false) + ); +} + +#[test] +fn strict_mode_recurses_through_items_and_any_of() { + let schema = ToolSchema::new( + "batch", + "Batch", + json!({ + "type": "object", + "properties": { + "rows": { + "type": "array", + "items": { + "type": "object", + "properties": {"a": {"type": "string"}}, + "required": ["a"], + }, + }, + }, + "required": ["rows"], + }), + ); + let prepared = prepare_tool_schema(&schema, &SchemaPreparation::openai().with_strict()); + + assert_eq!( + prepared.parameters["properties"]["rows"]["items"]["additionalProperties"], + json!(false) + ); +} + +#[test] +fn non_strict_preparation_leaves_optionality_alone() { + let prepared = prepare_tool_schema(&ref_schema(), &SchemaPreparation::openai()); + let required = prepared.parameters["required"].as_array().unwrap(); + assert_eq!(required, &vec![json!("id")], "`limit` must stay optional"); +} + +#[test] +fn null_parameters_never_reach_the_wire() { + // `ToolSchema::parameters` is an unconditional `Value`, so a tool taking no + // arguments can return `Value::Null` — which serialises to + // `"parameters": null` and 400s. + for parameters in [json!(null), json!("nonsense"), json!([1, 2]), json!({})] { + let schema = ToolSchema::new("bare", "no args", parameters); + let prepared = prepare_tool_schema(&schema, &SchemaPreparation::default()); + assert_eq!(prepared.parameters["type"], "object"); + assert!(prepared.parameters["properties"].is_object()); + } +} + +#[test] +fn an_object_schema_missing_its_type_gains_one() { + let schema = ToolSchema::new( + "implied", + "implied object", + json!({"properties": {"a": {"type": "string"}}}), + ); + let prepared = prepare_tool_schema(&schema, &SchemaPreparation::default()); + assert_eq!(prepared.parameters["type"], "object"); +} + +#[test] +fn preparing_a_set_preserves_order_and_count() { + let schemas = vec![ + ref_schema(), + ToolSchema::new("other", "Other", json!({"type": "object"})), + ]; + let prepared = prepare_tool_schemas(&schemas, &SchemaPreparation::conservative()); + assert_eq!(prepared.len(), 2); + assert_eq!(prepared[0].name, "lookup"); + assert_eq!(prepared[1].name, "other"); +} diff --git a/src/harness/tool/types.rs b/src/harness/tool/types.rs index a27a32b..34f82db 100644 --- a/src/harness/tool/types.rs +++ b/src/harness/tool/types.rs @@ -20,7 +20,7 @@ use crate::harness::cancel::CancellationToken; use crate::harness::context::RunContext; use crate::harness::events::EventSink; use crate::harness::ids::{RunId, ThreadId}; -use crate::harness::tool::{context_detail_from_args, humanize_tool_name}; +use crate::harness::tool::{ToolErrorPolicy, context_detail_from_args, humanize_tool_name}; /// The model-visible syntax a tool declaration prefers. /// @@ -465,6 +465,50 @@ pub trait Tool: Send + Sync { } } + /// Names of arguments this tool receives from the **host**, never from the + /// model. + /// + /// These keys are projected out of the model-facing declaration by + /// [`ToolRegistry::schemas`] — removed from both `properties` and + /// `required` — so the model neither sees them nor is asked to supply them. + /// The host fills them in at execution time. + /// + /// This is the declarative replacement for reaching around the argument + /// schema into [`ToolExecutionContext`], which is what a tool needing the + /// caller's thread id or recursion depth has to do today. Port of + /// LangChain's `InjectedToolArg` / `InjectedToolCallId`. + /// + /// # Execution-time contract + /// + /// A model-supplied value for an injected key must be **stripped before + /// schema validation**, then replaced by the host's real value — see the + /// ordering rule in [`crate::harness::tool::injected`]. Getting that order + /// wrong lets a model forge a hidden argument. + /// + /// The default is an empty list: a tool opts in. + fn injected_arguments(&self) -> &[&str] { + &[] + } + + /// What the harness should do when this tool returns `Err`. + /// + /// The default is [`ToolErrorPolicy::Fail`], preserving today's behaviour + /// exactly: an `Err` propagates and ends the run. A tool whose failures are + /// routine and recoverable — a lookup that misses, a network call that + /// times out — should return [`ToolErrorPolicy::ReturnToError`] so the + /// model sees the failure and can adapt, or + /// [`ToolErrorPolicy::Message`] to show a fixed message instead of an error + /// text that may not be safe to expose. + /// + /// [`TinyAgentsError::Cancelled`][crate::error::TinyAgentsError::Cancelled] + /// and + /// [`TinyAgentsError::Interrupted`][crate::error::TinyAgentsError::Interrupted] + /// are re-raised regardless of this policy; see + /// [`ToolErrorPolicy::apply`]. + fn error_policy(&self) -> ToolErrorPolicy { + ToolErrorPolicy::default() + } + /// Executes the tool against application state and a validated call. async fn call(&self, state: &State, call: ToolCall) -> Result; diff --git a/src/session/README.md b/src/session/README.md index 4c7f9d2..c045196 100644 --- a/src/session/README.md +++ b/src/session/README.md @@ -96,6 +96,31 @@ and completion read state and then act on it, so they take the write lock up front with `BEGIN IMMEDIATE` — racing claims serialize at `BEGIN` rather than failing at `COMMIT` after one has already decided it won. +**That serialization depends on a busy timeout, which we now set ourselves.** +SQLite's own default is zero — with no busy handler a `BEGIN IMMEDIATE` that +meets a competing writer fails immediately with `SQLITE_BUSY` instead of +waiting. It was never actually zero here: `rusqlite`'s `Connection::open` +installs a 5s timeout unconditionally. `store::BUSY_TIMEOUT` sets the same value +explicitly, so a correctness property the claim/gate/sequence logic relies on is +not silently supplied by a transitive dependency's undocumented default. + +**An upsert reads its own write back inside the same transaction.** Inserting on +an autocommit connection, closing it, then re-opening to `get_*` returns +whatever a concurrent writer left behind rather than what this call wrote. + +**Schema changes go through the versioned migration list** in `migrations.rs`, +not through more `CREATE TABLE IF NOT EXISTS` at the top of an operation. The +list index *is* the version, so it may only be appended to; retire a migration +by replacing its body with `"SELECT 1;"` rather than deleting it. Without a +version marker no column could ever be added to a workspace database that +already existed. + +**Nothing is deleted unless a host asks.** `retention.rs` owns the only delete +paths (sessions, messages, tool calls, run events, telemetry) plus +`reindex_fts`, which rebuilds the search index for rows whose FTS entry was lost +before the row/index pairs became transactional. Retention is a policy decision, +so none of it runs on a schedule of its own. + **A claim is meaningful only while a task is `in_progress`.** An upsert that moves a task off that status clears `claimed_by_member_id` and `claim_token`; leaving them set strands the task, since a new claim sees `AlreadyClaimed`, diff --git a/src/session/migrations.rs b/src/session/migrations.rs new file mode 100644 index 0000000..229f203 --- /dev/null +++ b/src/session/migrations.rs @@ -0,0 +1,370 @@ +//! Versioned, idempotent schema migrations for the session database. +//! +//! # Why a version marker at all +//! +//! The session database used to (re-)execute its full `CREATE TABLE IF NOT +//! EXISTS` DDL on *every* operation. That is idempotent, but it is also a dead +//! end: `CREATE TABLE IF NOT EXISTS` does nothing to a table that already +//! exists, so no column could ever be added to a workspace database that had +//! already been created. There was no way to express "and now also do this". +//! +//! # Shape +//! +//! [`MIGRATIONS`] is an ordered list whose **index is the version**. A +//! `schema_version` table records the highest index applied. On connection open +//! every migration with an index greater than the recorded version runs, in +//! order, each inside its own transaction, and the version is bumped after each +//! one. +//! +//! Two rules keep this sound: +//! +//! 1. **Append only.** Never reorder, insert into the middle of, or delete from +//! the list — the index *is* the version, so any of those silently re-number +//! every later migration. +//! 2. **Retire in place.** A migration that must stop running is replaced by the +//! deliberate no-op `"SELECT 1;"` rather than removed, preserving the +//! numbering of everything after it. (This is the convention LangGraph's +//! Postgres checkpointer uses for the same reason.) +//! +//! # Pre-existing databases +//! +//! A workspace database created before this module existed has the tables but no +//! `schema_version` row, so it reads as version `-1` and every migration runs. +//! Migration 0–2 are exactly the DDL those databases already had, and every +//! statement is `IF NOT EXISTS`, so replaying them is a no-op that ends with the +//! version marker correctly stamped. No special-case detection is needed. + +use rusqlite::{Connection, params}; + +use super::context::StorageContext; +use crate::error::Result; + +/// Grep prefix for migration logging. +const LOG_PREFIX: &str = "[session_db:migrations]"; + +/// The ordered migration list. **Index == schema version.** +/// +/// See the module docs before touching this: append only, and retire a +/// migration by replacing its body with `"SELECT 1;"` rather than deleting it. +pub(super) const MIGRATIONS: &[&str] = &[ + // ---- 0: base session tables ------------------------------------------ + "CREATE TABLE IF NOT EXISTS sessions ( + id TEXT PRIMARY KEY, + agent_definition_id TEXT NOT NULL, + agent_definition_name TEXT NOT NULL, + session_key TEXT NOT NULL, + parent_session_id TEXT, + thread_id TEXT, + source_channel TEXT, + status TEXT NOT NULL DEFAULT 'running', + model TEXT, + turn_count INTEGER NOT NULL DEFAULT 0, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cached_input_tokens INTEGER NOT NULL DEFAULT 0, + cost_usd REAL NOT NULL DEFAULT 0.0, + transcript_path TEXT, + started_at TEXT NOT NULL, + ended_at TEXT, + FOREIGN KEY (parent_session_id) REFERENCES sessions(id) ON DELETE SET NULL + ); + CREATE INDEX IF NOT EXISTS idx_sessions_agent ON sessions(agent_definition_id); + CREATE INDEX IF NOT EXISTS idx_sessions_status ON sessions(status); + CREATE INDEX IF NOT EXISTS idx_sessions_started ON sessions(started_at); + CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id); + CREATE INDEX IF NOT EXISTS idx_sessions_thread ON sessions(thread_id); + CREATE INDEX IF NOT EXISTS idx_sessions_channel ON sessions(source_channel); + CREATE INDEX IF NOT EXISTS idx_sessions_key ON sessions(session_key); + + CREATE TABLE IF NOT EXISTS session_messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + model TEXT, + input_tokens INTEGER, + output_tokens INTEGER, + cost_usd REAL, + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE + ); + CREATE INDEX IF NOT EXISTS idx_messages_session ON session_messages(session_id); + + CREATE TABLE IF NOT EXISTS session_tool_calls ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + message_id INTEGER, + tool_name TEXT NOT NULL, + tool_input TEXT, + tool_output TEXT, + status TEXT NOT NULL DEFAULT 'pending', + duration_ms INTEGER, + created_at TEXT NOT NULL, + FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE, + FOREIGN KEY (message_id) REFERENCES session_messages(id) ON DELETE SET NULL + ); + CREATE INDEX IF NOT EXISTS idx_tool_calls_session ON session_tool_calls(session_id); + CREATE INDEX IF NOT EXISTS idx_tool_calls_name ON session_tool_calls(tool_name);", + // ---- 1: full-text search index --------------------------------------- + "CREATE VIRTUAL TABLE IF NOT EXISTS sessions_fts USING fts5( + session_id, + agent_definition_name, + content, + tool_name + );", + // ---- 2: run ledger --------------------------------------------------- + "CREATE TABLE IF NOT EXISTS agent_runs ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL, + parent_run_id TEXT, + parent_thread_id TEXT, + agent_id TEXT, + status TEXT NOT NULL, + prompt_ref TEXT, + worker_thread_id TEXT, + task_board_id TEXT, + task_card_id TEXT, + checkpoint_path TEXT, + checkpoint_json TEXT, + summary TEXT, + error TEXT, + metadata_json TEXT NOT NULL DEFAULT '{}', + started_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + completed_at TEXT + ); + CREATE INDEX IF NOT EXISTS idx_agent_runs_status ON agent_runs(status); + CREATE INDEX IF NOT EXISTS idx_agent_runs_kind ON agent_runs(kind); + CREATE INDEX IF NOT EXISTS idx_agent_runs_parent ON agent_runs(parent_run_id); + CREATE INDEX IF NOT EXISTS idx_agent_runs_thread ON agent_runs(parent_thread_id); + CREATE INDEX IF NOT EXISTS idx_agent_runs_updated ON agent_runs(updated_at); + CREATE INDEX IF NOT EXISTS idx_agent_runs_worker_thread ON agent_runs(worker_thread_id); + + CREATE TABLE IF NOT EXISTS workflow_runs ( + id TEXT PRIMARY KEY, + definition_id TEXT NOT NULL, + parent_thread_id TEXT, + input_json TEXT NOT NULL DEFAULT '{}', + phase_states_json TEXT NOT NULL DEFAULT '{}', + child_run_ids_json TEXT NOT NULL DEFAULT '[]', + status TEXT NOT NULL, + summary TEXT, + started_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + completed_at TEXT + ); + CREATE INDEX IF NOT EXISTS idx_workflow_runs_definition ON workflow_runs(definition_id); + CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs(status); + CREATE INDEX IF NOT EXISTS idx_workflow_runs_thread ON workflow_runs(parent_thread_id); + + CREATE TABLE IF NOT EXISTS run_events ( + run_id TEXT NOT NULL, + sequence INTEGER NOT NULL, + event_type TEXT NOT NULL, + payload_json TEXT NOT NULL DEFAULT '{}', + timestamp TEXT NOT NULL, + PRIMARY KEY (run_id, sequence) + ); + CREATE INDEX IF NOT EXISTS idx_run_events_timestamp ON run_events(timestamp); + + CREATE TABLE IF NOT EXISTS run_telemetry ( + run_id TEXT PRIMARY KEY, + input_tokens INTEGER NOT NULL DEFAULT 0, + output_tokens INTEGER NOT NULL DEFAULT 0, + cached_input_tokens INTEGER NOT NULL DEFAULT 0, + cost_usd REAL NOT NULL DEFAULT 0.0, + elapsed_ms INTEGER, + tool_count INTEGER NOT NULL DEFAULT 0, + model TEXT, + provider TEXT, + error TEXT, + updated_at TEXT NOT NULL + ); + + CREATE TABLE IF NOT EXISTS agent_teams ( + id TEXT PRIMARY KEY, + parent_thread_id TEXT, + lead_agent_id TEXT NOT NULL, + status TEXT NOT NULL, + summary TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + closed_at TEXT + ); + CREATE INDEX IF NOT EXISTS idx_agent_teams_thread ON agent_teams(parent_thread_id); + CREATE INDEX IF NOT EXISTS idx_agent_teams_status ON agent_teams(status); + + CREATE TABLE IF NOT EXISTS agent_team_members ( + id TEXT PRIMARY KEY, + team_id TEXT NOT NULL, + name TEXT NOT NULL, + agent_id TEXT, + member_status TEXT NOT NULL, + current_task_id TEXT, + worker_thread_id TEXT, + run_id TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + UNIQUE(team_id, name) + ); + CREATE INDEX IF NOT EXISTS idx_agent_team_members_team ON agent_team_members(team_id); + + CREATE TABLE IF NOT EXISTS agent_team_tasks ( + id TEXT PRIMARY KEY, + team_id TEXT NOT NULL, + title TEXT NOT NULL, + objective TEXT, + status TEXT NOT NULL, + owner_member_id TEXT, + claimed_by_member_id TEXT, + claim_token TEXT, + depends_on_json TEXT NOT NULL DEFAULT '[]', + gate_status TEXT NOT NULL DEFAULT 'pending', + gate_reason TEXT, + evidence_json TEXT NOT NULL DEFAULT '[]', + source_run_id TEXT, + order_index INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL + ); + CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_team ON agent_team_tasks(team_id); + CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_status ON agent_team_tasks(status); + CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_claimed ON agent_team_tasks(claimed_by_member_id);", + // ---- 3: index coverage for the ORDER BY / sort columns --------------- + // + // `list_workflow_runs` and `list_agent_teams` both order by `updated_at + // DESC` and `list_agent_team_tasks` by `(order_index, created_at)`, none of + // which had an index — every listing was a full scan plus a sort. The + // `agent_runs` table already had its `updated_at` index; these bring the + // rest up to parity. + "CREATE INDEX IF NOT EXISTS idx_workflow_runs_updated ON workflow_runs(updated_at); + CREATE INDEX IF NOT EXISTS idx_agent_teams_updated ON agent_teams(updated_at); + CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_order + ON agent_team_tasks(team_id, order_index, created_at);", +]; + +/// Applies every migration newer than the database's recorded schema version. +/// +/// Idempotent: running it against an up-to-date database reads one row and +/// returns. Each migration runs inside its own transaction together with the +/// version bump, so a crash mid-migration leaves the version pointing at the +/// last fully applied step rather than at a half-applied one. +pub(super) fn apply(conn: &Connection) -> Result<()> { + conn.execute_batch( + "CREATE TABLE IF NOT EXISTS schema_version ( + id INTEGER PRIMARY KEY CHECK (id = 1), + version INTEGER NOT NULL + );", + ) + .storage_context("failed to create schema_version table")?; + + // `-1` means "nothing applied yet", which is also what a pre-migration + // workspace database reads as. Every migration is `IF NOT EXISTS`, so + // replaying 0..=2 over such a database is a no-op that stamps the marker. + let current: i64 = conn + .query_row( + "SELECT COALESCE((SELECT version FROM schema_version WHERE id = 1), -1)", + [], + |row| row.get(0), + ) + .storage_context("failed to read schema_version")?; + + let latest = MIGRATIONS.len() as i64 - 1; + if current >= latest { + return Ok(()); + } + tracing::debug!("{LOG_PREFIX} applying migrations from version {current} to {latest}"); + + for (version, sql) in MIGRATIONS.iter().enumerate() { + let version = version as i64; + if version <= current { + continue; + } + conn.execute_batch("BEGIN IMMEDIATE") + .storage_context("begin migration transaction")?; + let applied = (|| -> Result<()> { + conn.execute_batch(sql) + .storage_context(&format!("failed to apply session DB migration {version}"))?; + conn.execute( + "INSERT INTO schema_version (id, version) VALUES (1, ?1) + ON CONFLICT(id) DO UPDATE SET version = excluded.version", + params![version], + ) + .storage_context("failed to record schema version")?; + Ok(()) + })(); + match applied { + Ok(()) => { + conn.execute_batch("COMMIT") + .storage_context("commit migration transaction")?; + tracing::debug!("{LOG_PREFIX} applied migration {version}"); + } + Err(err) => { + if let Err(rollback) = conn.execute_batch("ROLLBACK") { + tracing::warn!( + "{LOG_PREFIX} rollback of migration {version} failed: {rollback} (original: {err})" + ); + } + return Err(err); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod test { + use super::*; + + /// The index of a migration **is** its version, so the list may only ever + /// be appended to. This pins the current length: bumping it is the moment + /// to re-read the module docs and confirm nothing was reordered. + #[test] + fn migration_list_is_append_only() { + assert_eq!( + MIGRATIONS.len(), + 4, + "MIGRATIONS is append-only — adding one is fine, reordering or \ + deleting one silently re-numbers every later migration" + ); + } + + #[test] + fn apply_is_idempotent_and_records_the_version() { + let conn = Connection::open_in_memory().expect("open"); + apply(&conn).expect("first apply"); + let version: i64 = conn + .query_row("SELECT version FROM schema_version WHERE id = 1", [], |r| { + r.get(0) + }) + .expect("read version"); + assert_eq!(version, MIGRATIONS.len() as i64 - 1); + // Second run is a no-op and must not error. + apply(&conn).expect("second apply"); + } + + /// A database created by the pre-migration DDL has the tables but no + /// version marker. Applying migrations must bring it forward rather than + /// failing on already-existing objects. + #[test] + fn apply_upgrades_a_pre_migration_database() { + let conn = Connection::open_in_memory().expect("open"); + // Simulate the old world: migrations 0..=2 executed with no marker. + for sql in &MIGRATIONS[..3] { + conn.execute_batch(sql).expect("legacy ddl"); + } + apply(&conn).expect("upgrade"); + // The version-3 index only exists because the migration ran. + let exists: bool = conn + .prepare( + "SELECT 1 FROM sqlite_master WHERE type='index' AND name='idx_agent_teams_updated'", + ) + .expect("prepare") + .exists([]) + .expect("exists"); + assert!( + exists, + "migration 3 added the agent_teams(updated_at) index" + ); + } +} diff --git a/src/session/mod.rs b/src/session/mod.rs index 707275e..744955d 100644 --- a/src/session/mod.rs +++ b/src/session/mod.rs @@ -65,14 +65,21 @@ //! coordination guarantees. mod context; -mod ops; +mod migrations; +pub mod ops; +pub mod retention; pub mod run_ledger; mod store; pub mod types; pub use ops::{ - get_session, list_children, list_messages, list_sessions, list_tool_calls, mark_interrupted, - record_message, record_session_end, record_session_start, record_tool_call, search_sessions, + DEFAULT_FTS_SNIPPET_BYTES, fts_snippet_bytes, get_session, list_children, list_messages, + list_sessions, list_tool_calls, mark_interrupted, record_message, record_session_end, + record_session_start, record_tool_call, search_sessions, set_fts_snippet_bytes, +}; +pub use retention::{ + RetentionReport, apply_retention, prune_run_events_before, prune_run_telemetry_before, + prune_sessions_before, prune_tool_calls_before, reindex_fts, trim_session_messages, }; pub use store::{db_path, with_connection, with_transaction}; pub use types::{ diff --git a/src/session/ops.rs b/src/session/ops.rs index 6299cce..00263ec 100644 --- a/src/session/ops.rs +++ b/src/session/ops.rs @@ -6,7 +6,7 @@ use rusqlite::{Connection, params}; use crate::error::{Result, TinyAgentsError}; use super::context::StorageContext; -use super::store::with_connection; +use super::store::{with_connection, with_transaction}; use super::types::{ SessionMessage, SessionRecord, SessionSearchParams, SessionSearchResult, SessionStatus, SessionToolCall, @@ -39,7 +39,13 @@ pub fn record_session_start( source_channel.unwrap_or("-"), ); - with_connection(workspace_dir, |conn| { + // The row and its FTS entry are ONE unit of work. On an autocommit + // connection they are two independent commits, so a failure between them + // leaves a permanently unsearchable row with no reindex path — and the FTS + // table is external-content-free, so nothing ever notices the gap. The + // transaction makes the pair atomic; `reindex_fts` exists to repair rows + // that predate it. + with_transaction(workspace_dir, |conn| { conn.execute( "INSERT INTO sessions ( id, agent_definition_id, agent_definition_name, session_key, @@ -134,7 +140,13 @@ pub fn record_message( content.len() ); - with_connection(workspace_dir, |conn| { + // The row and its FTS entry are ONE unit of work. On an autocommit + // connection they are two independent commits, so a failure between them + // leaves a permanently unsearchable row with no reindex path — and the FTS + // table is external-content-free, so nothing ever notices the gap. The + // transaction makes the pair atomic; `reindex_fts` exists to repair rows + // that predate it. + with_transaction(workspace_dir, |conn| { conn.execute( "INSERT INTO session_messages ( session_id, role, content, model, @@ -194,7 +206,13 @@ pub fn record_tool_call( } }); - with_connection(workspace_dir, |conn| { + // The row and its FTS entry are ONE unit of work. On an autocommit + // connection they are two independent commits, so a failure between them + // leaves a permanently unsearchable row with no reindex path — and the FTS + // table is external-content-free, so nothing ever notices the gap. The + // transaction makes the pair atomic; `reindex_fts` exists to repair rows + // that predate it. + with_transaction(workspace_dir, |conn| { conn.execute( "INSERT INTO session_tool_calls ( session_id, message_id, tool_name, tool_input, @@ -578,8 +596,35 @@ pub(super) fn fts_match_query(raw: &str) -> String { .join(" AND ") } -/// Longest FTS snippet indexed per message, in bytes. -pub(super) const MAX_FTS_SNIPPET_BYTES: usize = 2000; +/// Default cap on how much of a message body is copied into the search index, +/// in bytes. +/// +/// Indexing whole messages would let the FTS table grow without bound next to +/// the message rows themselves, so only a leading snippet is indexed — which +/// means **a match beyond this offset is not findable**. That is a real +/// behavioural limit, not an implementation detail, so it is public and +/// overridable rather than a private constant nobody could see. +pub const DEFAULT_FTS_SNIPPET_BYTES: usize = 2000; + +/// The live snippet cap. Read through [`fts_snippet_bytes`], set through +/// [`set_fts_snippet_bytes`]. +static FTS_SNIPPET_BYTES: std::sync::atomic::AtomicUsize = + std::sync::atomic::AtomicUsize::new(DEFAULT_FTS_SNIPPET_BYTES); + +/// Returns the current FTS snippet cap in bytes. +pub fn fts_snippet_bytes() -> usize { + FTS_SNIPPET_BYTES.load(std::sync::atomic::Ordering::Relaxed) +} + +/// Overrides the FTS snippet cap. +/// +/// Process-wide and takes effect for subsequently indexed content only — +/// raising it does not retroactively widen what is already indexed. Pair it +/// with [`super::retention::reindex_fts`] to rebuild the index at the new cap. +pub fn set_fts_snippet_bytes(bytes: usize) { + tracing::debug!("[session_db] fts snippet cap set to {bytes} bytes"); + FTS_SNIPPET_BYTES.store(bytes, std::sync::atomic::Ordering::Relaxed); +} pub(super) fn index_fts_content(conn: &Connection, session_id: &str, content: &str) -> Result<()> { // Slice on a character boundary, not a byte offset. `&content[..2000]` @@ -588,8 +633,9 @@ pub(super) fn index_fts_content(conn: &Connection, session_id: &str, content: &s // the message INSERT has already autocommitted by this point, so the row // survives with no FTS entry and is silently unsearchable forever after. // Mirrors the truncation already done in `record_tool_call`. - let snippet = if content.len() > MAX_FTS_SNIPPET_BYTES { - let mut cutoff = MAX_FTS_SNIPPET_BYTES; + let limit = fts_snippet_bytes(); + let snippet = if content.len() > limit { + let mut cutoff = limit; while cutoff > 0 && !content.is_char_boundary(cutoff) { cutoff -= 1; } diff --git a/src/session/retention.rs b/src/session/retention.rs new file mode 100644 index 0000000..b264a49 --- /dev/null +++ b/src/session/retention.rs @@ -0,0 +1,337 @@ +//! Retention, pruning and search-index repair for the session database. +//! +//! # Why this module exists +//! +//! Until it did, the session database had **no delete path at all**: sessions, +//! messages, tool calls, run events and telemetry accumulated forever, and the +//! only `DELETE FROM` anywhere in the crate lived in the graph checkpointer. A +//! long-lived workspace therefore grew without bound, and its full-text index +//! grew with it. +//! +//! It also had no way to *repair* the search index. `record_message` and its +//! siblings used to insert the row and its FTS entry as two separate autocommit +//! statements, so a failure in between left a row that could never be found by +//! search — permanently, because nothing re-derived the index. Those pairs are +//! transactional now, but databases written before that change still carry the +//! gaps, which is what [`reindex_fts`] is for. +//! +//! # Shape +//! +//! Every entry point takes an explicit bound (`older_than` / `keep_last`) and +//! returns the number of rows removed. Nothing here runs on a schedule or on +//! its own: retention is a policy decision, so the host chooses when and how +//! much. All of them run inside one transaction, so a partial prune is never +//! observable. + +use std::path::Path; + +use chrono::{DateTime, Utc}; +use rusqlite::params; + +use super::context::StorageContext; +use super::ops::fts_snippet_bytes; +use super::store::with_transaction; +use crate::error::Result; + +/// Grep prefix for retention logging. +const LOG_PREFIX: &str = "[session_db:retention]"; + +/// What a single retention pass removed. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct RetentionReport { + /// Sessions deleted (their messages and tool calls cascade). + pub sessions: usize, + /// Messages deleted directly (not via a session cascade). + pub messages: usize, + /// Tool calls deleted directly (not via a session cascade). + pub tool_calls: usize, + /// Run-ledger events deleted. + pub run_events: usize, + /// Run-telemetry rows deleted. + pub run_telemetry: usize, +} + +impl RetentionReport { + /// Total rows removed across every table. + pub fn total(&self) -> usize { + self.sessions + self.messages + self.tool_calls + self.run_events + self.run_telemetry + } +} + +/// Deletes every **finished** session that ended before `older_than`, together +/// with its messages, tool calls and search-index entries. +/// +/// Running sessions are never touched, whatever their start time: a long-lived +/// session is not stale data. `session_messages` and `session_tool_calls` +/// cascade through their foreign keys (the connection sets +/// `PRAGMA foreign_keys = ON`), but the FTS entries do not — the virtual table +/// has no foreign keys — so they are removed explicitly. +pub fn prune_sessions_before(workspace_dir: &Path, older_than: DateTime) -> Result { + let cutoff = older_than.to_rfc3339(); + tracing::debug!("{LOG_PREFIX} prune_sessions_before.entry cutoff={cutoff}"); + let removed = with_transaction(workspace_dir, |conn| { + // Collect first so the FTS rows can be removed by session id. + let ids: Vec = { + let mut stmt = conn.prepare( + "SELECT id FROM sessions + WHERE status != 'running' AND ended_at IS NOT NULL AND ended_at < ?1", + )?; + let rows = stmt.query_map(params![cutoff], |row| row.get::<_, String>(0))?; + let mut ids = Vec::new(); + for row in rows { + ids.push(row?); + } + ids + }; + let mut removed = 0usize; + for id in &ids { + conn.execute( + "DELETE FROM sessions_fts WHERE session_id = ?1", + params![id], + ) + .storage_context("delete session FTS rows")?; + removed += conn + .execute("DELETE FROM sessions WHERE id = ?1", params![id]) + .storage_context("delete session")?; + } + Ok(removed) + })?; + tracing::debug!("{LOG_PREFIX} prune_sessions_before.exit removed={removed}"); + Ok(removed) +} + +/// Trims a session's transcript to its most recent `keep_last` messages, +/// returning how many were removed. +/// +/// The oldest messages go first. Their FTS entries are content-addressed by +/// session rather than by message id, so trimming a session's messages +/// necessarily leaves its index entries stale; call [`reindex_fts`] afterwards +/// if search accuracy for trimmed sessions matters. +pub fn trim_session_messages( + workspace_dir: &Path, + session_id: &str, + keep_last: usize, +) -> Result { + tracing::debug!( + "{LOG_PREFIX} trim_session_messages.entry session={session_id} keep_last={keep_last}" + ); + let removed = with_transaction(workspace_dir, |conn| { + let removed = conn + .execute( + "DELETE FROM session_messages + WHERE session_id = ?1 AND id NOT IN ( + SELECT id FROM session_messages WHERE session_id = ?1 + ORDER BY id DESC LIMIT ?2 + )", + params![session_id, keep_last as i64], + ) + .storage_context("trim session messages")?; + Ok(removed) + })?; + tracing::debug!("{LOG_PREFIX} trim_session_messages.exit removed={removed}"); + Ok(removed) +} + +/// Deletes tool-call rows created before `older_than`, returning how many. +pub fn prune_tool_calls_before(workspace_dir: &Path, older_than: DateTime) -> Result { + let cutoff = older_than.to_rfc3339(); + tracing::debug!("{LOG_PREFIX} prune_tool_calls_before.entry cutoff={cutoff}"); + let removed = with_transaction(workspace_dir, |conn| { + conn.execute( + "DELETE FROM session_tool_calls WHERE created_at < ?1", + params![cutoff], + ) + .storage_context("prune tool calls") + })?; + tracing::debug!("{LOG_PREFIX} prune_tool_calls_before.exit removed={removed}"); + Ok(removed) +} + +/// Deletes run-ledger events stamped before `older_than`, returning how many. +/// +/// Event sequences are per run and allocated as `MAX(sequence) + 1`, so pruning +/// the head of a run's log does not renumber or collide with later appends. +pub fn prune_run_events_before(workspace_dir: &Path, older_than: DateTime) -> Result { + let cutoff = older_than.to_rfc3339(); + tracing::debug!("{LOG_PREFIX} prune_run_events_before.entry cutoff={cutoff}"); + let removed = with_transaction(workspace_dir, |conn| { + conn.execute( + "DELETE FROM run_events WHERE timestamp < ?1", + params![cutoff], + ) + .storage_context("prune run events") + })?; + tracing::debug!("{LOG_PREFIX} prune_run_events_before.exit removed={removed}"); + Ok(removed) +} + +/// Deletes run-telemetry rows last updated before `older_than`, returning how +/// many. +pub fn prune_run_telemetry_before( + workspace_dir: &Path, + older_than: DateTime, +) -> Result { + let cutoff = older_than.to_rfc3339(); + tracing::debug!("{LOG_PREFIX} prune_run_telemetry_before.entry cutoff={cutoff}"); + let removed = with_transaction(workspace_dir, |conn| { + conn.execute( + "DELETE FROM run_telemetry WHERE updated_at < ?1", + params![cutoff], + ) + .storage_context("prune run telemetry") + })?; + tracing::debug!("{LOG_PREFIX} prune_run_telemetry_before.exit removed={removed}"); + Ok(removed) +} + +/// Applies one age-based retention pass across every table, returning what it +/// removed. +/// +/// Sessions are pruned first so their messages and tool calls cascade rather +/// than being deleted twice; the remaining passes then catch orphaned rows that +/// outlived their session or belong to the run ledger. +pub fn apply_retention(workspace_dir: &Path, older_than: DateTime) -> Result { + tracing::debug!( + "{LOG_PREFIX} apply_retention.entry cutoff={}", + older_than.to_rfc3339() + ); + let cutoff = older_than.to_rfc3339(); + let report = with_transaction(workspace_dir, |conn| { + let ids: Vec = { + let mut stmt = conn.prepare( + "SELECT id FROM sessions WHERE status != 'running' AND ended_at IS NOT NULL AND ended_at < ?1", + )?; + stmt.query_map(params![cutoff], |row| row.get(0))? + .collect::>()? + }; + for id in &ids { + conn.execute( + "DELETE FROM sessions_fts WHERE session_id = ?1", + params![id], + ) + .storage_context("delete session FTS rows")?; + conn.execute("DELETE FROM sessions WHERE id = ?1", params![id]) + .storage_context("delete session")?; + } + Ok(RetentionReport { + sessions: ids.len(), + messages: 0, + tool_calls: conn + .execute( + "DELETE FROM session_tool_calls WHERE created_at < ?1", + params![cutoff], + ) + .storage_context("prune tool calls")?, + run_events: conn + .execute( + "DELETE FROM run_events WHERE timestamp < ?1", + params![cutoff], + ) + .storage_context("prune run events")?, + run_telemetry: conn + .execute( + "DELETE FROM run_telemetry WHERE updated_at < ?1", + params![cutoff], + ) + .storage_context("prune run telemetry")?, + }) + })?; + tracing::info!( + "{LOG_PREFIX} apply_retention.exit removed total={} sessions={} tool_calls={} \ + run_events={} run_telemetry={}", + report.total(), + report.sessions, + report.tool_calls, + report.run_events, + report.run_telemetry + ); + Ok(report) +} + +/// Rebuilds `sessions_fts` from the authoritative tables, returning the number +/// of index rows written. +/// +/// The repair path the database never had. Any row whose FTS entry was lost — +/// to a failure between the two old autocommit statements, to a retention pass +/// that trimmed messages, or to a change of the +/// [snippet cap](super::ops::set_fts_snippet_bytes) — becomes searchable again. +/// The whole rebuild runs in one transaction, so search is never left with a +/// half-built index. +pub fn reindex_fts(workspace_dir: &Path) -> Result { + tracing::debug!("{LOG_PREFIX} reindex_fts.entry"); + let limit = fts_snippet_bytes(); + let written = with_transaction(workspace_dir, |conn| { + conn.execute("DELETE FROM sessions_fts", []) + .storage_context("clear sessions_fts")?; + let mut written = 0usize; + + // One row per session, carrying the agent name. + { + let mut stmt = + conn.prepare("SELECT id, agent_definition_name FROM sessions ORDER BY rowid ASC")?; + let rows = stmt.query_map([], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) + })?; + let mut insert = conn.prepare( + "INSERT INTO sessions_fts (session_id, agent_definition_name, content, tool_name) + VALUES (?1, ?2, '', '')", + )?; + for row in rows { + let (id, name) = row?; + insert.execute(params![id, name])?; + written += 1; + } + } + + // One row per message, carrying the (capped) content snippet. + { + let mut stmt = + conn.prepare("SELECT session_id, content FROM session_messages ORDER BY id ASC")?; + let rows = stmt.query_map([], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) + })?; + let mut insert = conn.prepare( + "INSERT INTO sessions_fts (session_id, agent_definition_name, content, tool_name) + VALUES (?1, '', ?2, '')", + )?; + for row in rows { + let (session_id, content) = row?; + insert.execute(params![session_id, snippet(&content, limit)])?; + written += 1; + } + } + + // One row per tool call, carrying the tool name. + { + let mut stmt = conn + .prepare("SELECT session_id, tool_name FROM session_tool_calls ORDER BY id ASC")?; + let rows = stmt.query_map([], |row| { + Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?)) + })?; + let mut insert = conn.prepare( + "INSERT INTO sessions_fts (session_id, agent_definition_name, content, tool_name) + VALUES (?1, '', '', ?2)", + )?; + for row in rows { + let (session_id, tool_name) = row?; + insert.execute(params![session_id, tool_name])?; + written += 1; + } + } + Ok(written) + })?; + tracing::info!("{LOG_PREFIX} reindex_fts.exit rows={written}"); + Ok(written) +} + +/// Truncates `content` to at most `limit` bytes on a character boundary. +fn snippet(content: &str, limit: usize) -> &str { + if content.len() <= limit { + return content; + } + let mut cutoff = limit; + while cutoff > 0 && !content.is_char_boundary(cutoff) { + cutoff -= 1; + } + &content[..cutoff] +} diff --git a/src/session/run_ledger/ops.rs b/src/session/run_ledger/ops.rs index 5db3090..518407d 100644 --- a/src/session/run_ledger/ops.rs +++ b/src/session/run_ledger/ops.rs @@ -41,7 +41,11 @@ pub fn upsert_agent_run(workspace_dir: &Path, upsert: AgentRunUpsert) -> Result< upsert.parent_thread_id.as_deref().unwrap_or("-") ); - crate::session::store::with_connection(workspace_dir, |conn| { + // One transaction for the write *and* the read-back. Committing the insert + // on an autocommit connection, closing it, then re-opening to `get_*` hands + // the caller whatever a concurrent writer left behind rather than what this + // call wrote — an upsert that reports someone else's row. + crate::session::store::with_transaction(workspace_dir, |conn| { init_run_ledger_schema(conn)?; conn.execute( "INSERT INTO agent_runs ( @@ -95,10 +99,8 @@ pub fn upsert_agent_run(workspace_dir: &Path, upsert: AgentRunUpsert) -> Result< ], ) .storage_context("upsert agent run")?; - Ok(()) - })?; - - get_agent_run(workspace_dir, &upsert.id)?.storage_context("agent run missing after upsert") + get_agent_run_inner(conn, &upsert.id)?.storage_context("agent run missing after upsert") + }) } pub fn upsert_workflow_run(workspace_dir: &Path, upsert: WorkflowRunUpsert) -> Result { @@ -111,7 +113,11 @@ pub fn upsert_workflow_run(workspace_dir: &Path, upsert: WorkflowRunUpsert) -> R let child_run_ids_json = serde_json::to_string(&upsert.child_run_ids).storage_context("serialize child run ids")?; - crate::session::store::with_connection(workspace_dir, |conn| { + // One transaction for the write *and* the read-back. Committing the insert + // on an autocommit connection, closing it, then re-opening to `get_*` hands + // the caller whatever a concurrent writer left behind rather than what this + // call wrote — an upsert that reports someone else's row. + crate::session::store::with_transaction(workspace_dir, |conn| { init_run_ledger_schema(conn)?; conn.execute( "INSERT INTO workflow_runs ( @@ -143,11 +149,9 @@ pub fn upsert_workflow_run(workspace_dir: &Path, upsert: WorkflowRunUpsert) -> R ], ) .storage_context("upsert workflow run")?; - Ok(()) - })?; - - get_workflow_run(workspace_dir, &upsert.id)? - .storage_context("workflow run missing after upsert") + get_workflow_run_inner(conn, &upsert.id)? + .storage_context("workflow run missing after upsert") + }) } pub fn append_run_event(workspace_dir: &Path, event: RunEventAppend) -> Result { @@ -443,18 +447,24 @@ pub fn list_recent_run_events( }) } +/// Connection-scoped workflow-run lookup, so an upsert can read its own write +/// back inside the same transaction. +fn get_workflow_run_inner(conn: &Connection, id: &str) -> Result> { + let mut stmt = conn.prepare( + "SELECT id, definition_id, parent_thread_id, input_json, phase_states_json, + child_run_ids_json, status, summary, started_at, updated_at, completed_at + FROM workflow_runs WHERE id = ?1", + )?; + Ok(stmt + .query_row(params![id], map_workflow_run_row) + .optional()?) +} + pub fn get_workflow_run(workspace_dir: &Path, id: &str) -> Result> { tracing::debug!("{LOG_PREFIX} get_workflow_run.entry id={id}"); crate::session::store::with_connection(workspace_dir, |conn| { init_run_ledger_schema(conn)?; - let mut stmt = conn.prepare( - "SELECT id, definition_id, parent_thread_id, input_json, phase_states_json, - child_run_ids_json, status, summary, started_at, updated_at, completed_at - FROM workflow_runs WHERE id = ?1", - )?; - let run = stmt - .query_row(params![id], map_workflow_run_row) - .optional()?; + let run = get_workflow_run_inner(conn, id)?; tracing::debug!( "{LOG_PREFIX} get_workflow_run.exit id={id} found={}", run.is_some() @@ -563,7 +573,11 @@ pub fn upsert_agent_team(workspace_dir: &Path, upsert: AgentTeamUpsert) -> Resul upsert.lead_agent_id, upsert.status.as_str() ); - crate::session::store::with_connection(workspace_dir, |conn| { + // One transaction for the write *and* the read-back. Committing the insert + // on an autocommit connection, closing it, then re-opening to `get_*` hands + // the caller whatever a concurrent writer left behind rather than what this + // call wrote — an upsert that reports someone else's row. + let team = crate::session::store::with_transaction(workspace_dir, |conn| { init_run_ledger_schema(conn)?; conn.execute( "INSERT INTO agent_teams ( @@ -589,10 +603,8 @@ pub fn upsert_agent_team(workspace_dir: &Path, upsert: AgentTeamUpsert) -> Resul ], ) .storage_context("upsert agent team")?; - Ok(()) + get_agent_team_inner(conn, &upsert.id)?.storage_context("agent team missing after upsert") })?; - let team = get_agent_team(workspace_dir, &upsert.id)? - .storage_context("agent team missing after upsert")?; tracing::debug!("{LOG_PREFIX} upsert_agent_team.exit id={}", team.id); Ok(team) } @@ -700,7 +712,11 @@ pub fn upsert_agent_team_member( upsert.name, upsert.member_status.as_str() ); - crate::session::store::with_connection(workspace_dir, |conn| { + // One transaction for the write *and* the read-back. Committing the insert + // on an autocommit connection, closing it, then re-opening to `get_*` hands + // the caller whatever a concurrent writer left behind rather than what this + // call wrote — an upsert that reports someone else's row. + let member = crate::session::store::with_transaction(workspace_dir, |conn| { init_run_ledger_schema(conn)?; conn.execute( "INSERT INTO agent_team_members ( @@ -729,10 +745,9 @@ pub fn upsert_agent_team_member( ], ) .storage_context("upsert agent team member")?; - Ok(()) + get_agent_team_member_inner(conn, &upsert.id)? + .storage_context("agent team member missing after upsert") })?; - let member = get_agent_team_member(workspace_dir, &upsert.id)? - .storage_context("agent team member missing after upsert")?; tracing::debug!( "{LOG_PREFIX} upsert_agent_team_member.exit id={}", member.id @@ -800,7 +815,11 @@ pub fn upsert_agent_team_task( upsert.status.as_str(), upsert.depends_on.len() ); - crate::session::store::with_connection(workspace_dir, |conn| { + // One transaction for the write *and* the read-back. Committing the insert + // on an autocommit connection, closing it, then re-opening to `get_*` hands + // the caller whatever a concurrent writer left behind rather than what this + // call wrote — an upsert that reports someone else's row. + let task = crate::session::store::with_transaction(workspace_dir, |conn| { init_run_ledger_schema(conn)?; conn.execute( "INSERT INTO agent_team_tasks ( @@ -854,10 +873,9 @@ pub fn upsert_agent_team_task( ], ) .storage_context("upsert agent team task")?; - Ok(()) + get_agent_team_task_inner(conn, &upsert.id)? + .storage_context("agent team task missing after upsert") })?; - let task = get_agent_team_task(workspace_dir, &upsert.id)? - .storage_context("agent team task missing after upsert")?; tracing::debug!("{LOG_PREFIX} upsert_agent_team_task.exit id={}", task.id); Ok(task) } @@ -961,19 +979,31 @@ pub fn claim_agent_team_task( return Ok(ClaimOutcome::Blocked { unmet }); } - // 3. Compare-and-swap on the unclaimed guard. + // 3. Compare-and-swap on the unclaimed guard **and the status**. + // + // `claimed_by_member_id IS NULL` alone is not a guard: `upsert_agent_team_task` + // deliberately NULLs that column whenever the new status is not + // `in_progress`, so every `done` task also satisfies it. A stale worker + // re-claiming a finished task therefore flipped it straight back to + // `in_progress` — and stranded everything downstream, because the + // completion gate re-checks that each dependency is still `done` and now + // reports it unfinished. A terminal task is not claimable, whatever its + // claim column says. let now = Utc::now(); let rows_affected = conn .execute( "UPDATE agent_team_tasks SET claimed_by_member_id = ?1, claim_token = ?2, status = 'in_progress', updated_at = ?3 - WHERE id = ?4 AND team_id = ?5 AND claimed_by_member_id IS NULL", + WHERE id = ?4 AND team_id = ?5 AND claimed_by_member_id IS NULL + AND status IN ('todo', 'ready', 'blocked')", params![member_id, claim_token, now.to_rfc3339(), task_id, team_id], ) .storage_context("compare-and-swap claim agent team task")?; if rows_affected == 0 { tracing::debug!( - "{LOG_PREFIX} claim_agent_team_task.already_claimed team={team_id} task={task_id}" + "{LOG_PREFIX} claim_agent_team_task.already_claimed team={team_id} task={task_id} \ + status={}", + task.status.as_str() ); return Ok(ClaimOutcome::AlreadyClaimed); } diff --git a/src/session/run_ledger/store.rs b/src/session/run_ledger/store.rs index 9a4498f..04bd0e3 100644 --- a/src/session/run_ledger/store.rs +++ b/src/session/run_ledger/store.rs @@ -1,127 +1,24 @@ +//! Run-ledger schema entry point. +//! +//! The tables this module used to create on every operation now live in the +//! versioned migration list (`crate::session::migrations`, migration 2), which +//! is applied once per database when a connection is opened. Re-running +//! `CREATE TABLE IF NOT EXISTS` per call was not just wasted work: with no +//! version marker there was no way to ever *add* a column to a workspace +//! database that already existed. + use rusqlite::Connection; -use super::super::context::StorageContext; use crate::error::Result; -pub(crate) fn init_run_ledger_schema(conn: &Connection) -> Result<()> { - conn.execute_batch( - "CREATE TABLE IF NOT EXISTS agent_runs ( - id TEXT PRIMARY KEY, - kind TEXT NOT NULL, - parent_run_id TEXT, - parent_thread_id TEXT, - agent_id TEXT, - status TEXT NOT NULL, - prompt_ref TEXT, - worker_thread_id TEXT, - task_board_id TEXT, - task_card_id TEXT, - checkpoint_path TEXT, - checkpoint_json TEXT, - summary TEXT, - error TEXT, - metadata_json TEXT NOT NULL DEFAULT '{}', - started_at TEXT NOT NULL, - updated_at TEXT NOT NULL, - completed_at TEXT - ); - CREATE INDEX IF NOT EXISTS idx_agent_runs_status ON agent_runs(status); - CREATE INDEX IF NOT EXISTS idx_agent_runs_kind ON agent_runs(kind); - CREATE INDEX IF NOT EXISTS idx_agent_runs_parent ON agent_runs(parent_run_id); - CREATE INDEX IF NOT EXISTS idx_agent_runs_thread ON agent_runs(parent_thread_id); - CREATE INDEX IF NOT EXISTS idx_agent_runs_updated ON agent_runs(updated_at); - CREATE INDEX IF NOT EXISTS idx_agent_runs_worker_thread ON agent_runs(worker_thread_id); - - CREATE TABLE IF NOT EXISTS workflow_runs ( - id TEXT PRIMARY KEY, - definition_id TEXT NOT NULL, - parent_thread_id TEXT, - input_json TEXT NOT NULL DEFAULT '{}', - phase_states_json TEXT NOT NULL DEFAULT '{}', - child_run_ids_json TEXT NOT NULL DEFAULT '[]', - status TEXT NOT NULL, - summary TEXT, - started_at TEXT NOT NULL, - updated_at TEXT NOT NULL, - completed_at TEXT - ); - CREATE INDEX IF NOT EXISTS idx_workflow_runs_definition ON workflow_runs(definition_id); - CREATE INDEX IF NOT EXISTS idx_workflow_runs_status ON workflow_runs(status); - CREATE INDEX IF NOT EXISTS idx_workflow_runs_thread ON workflow_runs(parent_thread_id); - - CREATE TABLE IF NOT EXISTS run_events ( - run_id TEXT NOT NULL, - sequence INTEGER NOT NULL, - event_type TEXT NOT NULL, - payload_json TEXT NOT NULL DEFAULT '{}', - timestamp TEXT NOT NULL, - PRIMARY KEY (run_id, sequence) - ); - CREATE INDEX IF NOT EXISTS idx_run_events_timestamp ON run_events(timestamp); - - CREATE TABLE IF NOT EXISTS run_telemetry ( - run_id TEXT PRIMARY KEY, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cached_input_tokens INTEGER NOT NULL DEFAULT 0, - cost_usd REAL NOT NULL DEFAULT 0.0, - elapsed_ms INTEGER, - tool_count INTEGER NOT NULL DEFAULT 0, - model TEXT, - provider TEXT, - error TEXT, - updated_at TEXT NOT NULL - ); - - CREATE TABLE IF NOT EXISTS agent_teams ( - id TEXT PRIMARY KEY, - parent_thread_id TEXT, - lead_agent_id TEXT NOT NULL, - status TEXT NOT NULL, - summary TEXT, - created_at TEXT NOT NULL, - updated_at TEXT NOT NULL, - closed_at TEXT - ); - CREATE INDEX IF NOT EXISTS idx_agent_teams_thread ON agent_teams(parent_thread_id); - CREATE INDEX IF NOT EXISTS idx_agent_teams_status ON agent_teams(status); - - CREATE TABLE IF NOT EXISTS agent_team_members ( - id TEXT PRIMARY KEY, - team_id TEXT NOT NULL, - name TEXT NOT NULL, - agent_id TEXT, - member_status TEXT NOT NULL, - current_task_id TEXT, - worker_thread_id TEXT, - run_id TEXT, - created_at TEXT NOT NULL, - updated_at TEXT NOT NULL, - UNIQUE(team_id, name) - ); - CREATE INDEX IF NOT EXISTS idx_agent_team_members_team ON agent_team_members(team_id); - - CREATE TABLE IF NOT EXISTS agent_team_tasks ( - id TEXT PRIMARY KEY, - team_id TEXT NOT NULL, - title TEXT NOT NULL, - objective TEXT, - status TEXT NOT NULL, - owner_member_id TEXT, - claimed_by_member_id TEXT, - claim_token TEXT, - depends_on_json TEXT NOT NULL DEFAULT '[]', - gate_status TEXT NOT NULL DEFAULT 'pending', - gate_reason TEXT, - evidence_json TEXT NOT NULL DEFAULT '[]', - source_run_id TEXT, - order_index INTEGER NOT NULL DEFAULT 0, - created_at TEXT NOT NULL, - updated_at TEXT NOT NULL - ); - CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_team ON agent_team_tasks(team_id); - CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_status ON agent_team_tasks(status); - CREATE INDEX IF NOT EXISTS idx_agent_team_tasks_claimed ON agent_team_tasks(claimed_by_member_id);", - ) - .storage_context("failed to initialize run ledger schema") +/// No-op retained as the run-ledger's schema entry point. +/// +/// Every `crate::session::store::with_connection` handle is already migrated by +/// the time it reaches a caller, so there is nothing left to do here. The +/// function is kept (rather than deleted along with its ~30 call sites) so the +/// ledger operations still read as "ensure my schema exists", and so a future +/// ledger-only bootstrap has an obvious place to land. +#[inline] +pub(crate) fn init_run_ledger_schema(_conn: &Connection) -> Result<()> { + Ok(()) } diff --git a/src/session/store.rs b/src/session/store.rs index 74a61dc..a214425 100644 --- a/src/session/store.rs +++ b/src/session/store.rs @@ -1,8 +1,10 @@ use std::path::{Path, PathBuf}; +use std::time::Duration; use rusqlite::Connection; use super::context::StorageContext; +use super::migrations; use crate::error::Result; /// Subdirectory of the workspace holding the session database. @@ -10,6 +12,30 @@ const DB_SUBDIR: &str = "session_db"; /// Database filename inside [`DB_SUBDIR`]. const DB_FILE: &str = "sessions.db"; +/// How long a statement waits for a competing writer's lock before giving up +/// with `SQLITE_BUSY`. +/// +/// # This is a guarantee we own, not a bug fix +/// +/// SQLite's own default is zero — a `BEGIN IMMEDIATE` that finds the write lock +/// held would fail instantly rather than wait — and every claim, gate and +/// sequence allocation in this module is written on the assumption that racing +/// writers *serialize* at `BEGIN`. +/// +/// That assumption was, as it happens, already satisfied: `rusqlite`'s +/// `Connection::open` calls `sqlite3_busy_timeout(db, 5000)` unconditionally, +/// so the connections here have never actually had a zero timeout. Setting it +/// explicitly changes no behaviour today. It is worth doing anyway, because the +/// alternative is that a load-bearing correctness property of this module is +/// supplied by an undocumented default of a transitive dependency, invisible at +/// every call site and free to change in a patch release. Stating it here makes +/// the dependency deliberate and greppable. +/// +/// Five seconds is long enough to ride out any transaction this module takes +/// (all of them are a handful of small statements) and short enough to surface +/// a genuine deadlock rather than hang. +const BUSY_TIMEOUT: Duration = Duration::from_secs(5); + /// Resolves the session database path for a workspace root. /// /// Kept public so hosts can locate the file for backup, inspection, or @@ -38,11 +64,31 @@ pub fn with_connection( let conn = Connection::open(&db_path) .storage_context(&format!("failed to open session DB: {}", db_path.display()))?; + prepare_connection(&conn)?; + + // Migrations are idempotent, and checking on every fresh connection also + // handles a database atomically replaced at this same path. + migrations::apply(&conn)?; - init_schema(&conn)?; f(&conn) } +/// Applies the per-connection pragmas every session-DB handle needs. +/// +/// `journal_mode = WAL` is persistent (stored in the file header) but is set +/// here so a freshly created database gets it; `foreign_keys` and +/// `busy_timeout` are **per connection** and must be set on every open. +fn prepare_connection(conn: &Connection) -> Result<()> { + conn.busy_timeout(BUSY_TIMEOUT) + .storage_context("failed to set session DB busy_timeout")?; + conn.execute_batch( + "PRAGMA journal_mode = WAL; + PRAGMA foreign_keys = ON;", + ) + .storage_context("failed to apply session DB pragmas")?; + Ok(()) +} + /// Opens the session database and runs `f` inside a single **immediate** /// write transaction, committing on `Ok` and rolling back on `Err`. /// @@ -88,94 +134,7 @@ pub fn with_transaction( pub fn with_memory_connection(f: impl FnOnce(&Connection) -> Result) -> Result { let conn = Connection::open_in_memory().storage_context("failed to open in-memory session DB")?; - init_schema(&conn)?; + prepare_connection(&conn)?; + migrations::apply(&conn)?; f(&conn) } - -pub(super) fn init_schema(conn: &Connection) -> Result<()> { - conn.execute_batch( - "PRAGMA journal_mode = WAL; - PRAGMA foreign_keys = ON; - - CREATE TABLE IF NOT EXISTS sessions ( - id TEXT PRIMARY KEY, - agent_definition_id TEXT NOT NULL, - agent_definition_name TEXT NOT NULL, - session_key TEXT NOT NULL, - parent_session_id TEXT, - thread_id TEXT, - source_channel TEXT, - status TEXT NOT NULL DEFAULT 'running', - model TEXT, - turn_count INTEGER NOT NULL DEFAULT 0, - input_tokens INTEGER NOT NULL DEFAULT 0, - output_tokens INTEGER NOT NULL DEFAULT 0, - cached_input_tokens INTEGER NOT NULL DEFAULT 0, - cost_usd REAL NOT NULL DEFAULT 0.0, - transcript_path TEXT, - started_at TEXT NOT NULL, - ended_at TEXT, - FOREIGN KEY (parent_session_id) REFERENCES sessions(id) ON DELETE SET NULL - ); - CREATE INDEX IF NOT EXISTS idx_sessions_agent ON sessions(agent_definition_id); - CREATE INDEX IF NOT EXISTS idx_sessions_status ON sessions(status); - CREATE INDEX IF NOT EXISTS idx_sessions_started ON sessions(started_at); - CREATE INDEX IF NOT EXISTS idx_sessions_parent ON sessions(parent_session_id); - CREATE INDEX IF NOT EXISTS idx_sessions_thread ON sessions(thread_id); - CREATE INDEX IF NOT EXISTS idx_sessions_channel ON sessions(source_channel); - CREATE INDEX IF NOT EXISTS idx_sessions_key ON sessions(session_key); - - CREATE TABLE IF NOT EXISTS session_messages ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - session_id TEXT NOT NULL, - role TEXT NOT NULL, - content TEXT NOT NULL, - model TEXT, - input_tokens INTEGER, - output_tokens INTEGER, - cost_usd REAL, - created_at TEXT NOT NULL, - FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE - ); - CREATE INDEX IF NOT EXISTS idx_messages_session ON session_messages(session_id); - - CREATE TABLE IF NOT EXISTS session_tool_calls ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - session_id TEXT NOT NULL, - message_id INTEGER, - tool_name TEXT NOT NULL, - tool_input TEXT, - tool_output TEXT, - status TEXT NOT NULL DEFAULT 'pending', - duration_ms INTEGER, - created_at TEXT NOT NULL, - FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE, - FOREIGN KEY (message_id) REFERENCES session_messages(id) ON DELETE SET NULL - ); - CREATE INDEX IF NOT EXISTS idx_tool_calls_session ON session_tool_calls(session_id); - CREATE INDEX IF NOT EXISTS idx_tool_calls_name ON session_tool_calls(tool_name);", - ) - .storage_context("failed to initialize session_db schema")?; - - init_fts(conn)?; - Ok(()) -} - -fn init_fts(conn: &Connection) -> Result<()> { - let has_fts: bool = conn - .prepare("SELECT 1 FROM sqlite_master WHERE type='table' AND name='sessions_fts'")? - .exists([])?; - - if !has_fts { - conn.execute_batch( - "CREATE VIRTUAL TABLE sessions_fts USING fts5( - session_id, - agent_definition_name, - content, - tool_name - );", - ) - .storage_context("failed to create sessions_fts virtual table")?; - } - Ok(()) -} diff --git a/src/session/test.rs b/src/session/test.rs index ab470b7..a77a326 100644 --- a/src/session/test.rs +++ b/src/session/test.rs @@ -4,8 +4,9 @@ //! than per-file inline `mod tests` blocks. Sections mirror the source files. use super::context::StorageContext; +use super::migrations::apply as init_schema; use super::ops::*; -use super::store::{init_schema, with_memory_connection}; +use super::store::with_memory_connection; use super::types::*; use crate::error::TinyAgentsError; use chrono::Utc; @@ -616,3 +617,53 @@ fn record_tool_call_returns_the_tool_call_row_id() { }) .unwrap(); } + +/// Every session-DB connection waits out a competing writer. +/// +/// SQLite's own default `busy_timeout` is 0: with no handler installed a +/// `BEGIN IMMEDIATE` that meets another writer fails instantly with +/// `SQLITE_BUSY` rather than waiting, which would contradict the +/// serialize-at-BEGIN rationale that `with_transaction`, the task claim CAS and +/// the run-event sequence allocation are all written against. +/// +/// This pins the property, not a fix. It passes against the pre-change code as +/// well, because `rusqlite::Connection::open` installs a 5s busy timeout of its +/// own accord — a fact nothing in this crate stated, and nothing enforced. +/// `store::BUSY_TIMEOUT` now sets it explicitly and this test is what catches +/// it disappearing. +/// +/// The test holds a real write lock from a second connection for a beat, then +/// asserts a `with_transaction` issued concurrently waits and succeeds. +#[test] +fn with_transaction_waits_out_a_competing_writer() { + let dir = tempfile::tempdir().unwrap(); + let workspace = dir.path().to_path_buf(); + + // Create the DB (and run migrations) before contending for it. + super::store::with_connection(&workspace, |_| Ok(())).unwrap(); + + let blocker = Connection::open(super::store::db_path(&workspace)).unwrap(); + blocker.execute_batch("BEGIN IMMEDIATE").unwrap(); + + let releaser = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(300)); + blocker.execute_batch("ROLLBACK").unwrap(); + }); + + let result = super::store::with_transaction(&workspace, |conn| { + conn.execute( + "INSERT INTO sessions ( + id, agent_definition_id, agent_definition_name, session_key, + status, started_at + ) VALUES ('waited', 'a', 'a', 'k', 'running', ?1)", + params![Utc::now().to_rfc3339()], + )?; + Ok(()) + }); + + releaser.join().unwrap(); + assert!( + result.is_ok(), + "BEGIN IMMEDIATE must wait for the competing writer, not fail instantly: {result:?}" + ); +} diff --git a/tests/context_and_schema_compaction.rs b/tests/context_and_schema_compaction.rs new file mode 100644 index 0000000..9a2d4f6 --- /dev/null +++ b/tests/context_and_schema_compaction.rs @@ -0,0 +1,200 @@ +//! End-to-end regression tests for transcript compaction defects. +//! +//! Every test in this file is written **only against APIs that existed before +//! the fixes**, so each one compiles against the unfixed crate and fails there. +//! That is deliberate: a regression test for a structural bug is worth little if +//! it can only be expressed in terms of the fix's own new surface. +//! +//! The three defects covered: +//! +//! - **Orphaned tool pairing.** `SummarizationPolicy::plan` and all three +//! `TrimStrategy` variants cut at a blind index, severing an assistant +//! `tool_calls` turn from the `tool` messages answering it. The rebuilt +//! request then carries a `role:"tool"` with no preceding `tool_calls`, which +//! OpenAI rejects with a `400` and Anthropic rejects as a `tool_result` with +//! no matching `tool_use`. +//! - **Token estimation blind to tool calls.** An assistant turn that only +//! calls tools has empty `content`, so it estimated at zero and no compaction +//! gate ever fired on the runs that needed it most. +//! - **History erased by the default summarizer.** `ConcatSummarizer` built its +//! output from `Message::text()`, which renders tool calls and tool results as +//! the empty string. + +use tinyagents::harness::message::{AssistantMessage, Message}; +use tinyagents::harness::summarization::{ + ConcatSummarizer, SummarizationPolicy, Summarizer, TrimStrategy, trim_messages, +}; +use tinyagents::harness::tool::ToolCall; + +/// The provider invariant, checked locally so this file depends on no new +/// crate surface: every `tool` message must be preceded by an assistant turn +/// declaring its id, and every declared call must be answered. +fn pairing_intact(messages: &[Message]) -> bool { + let mut declared: Vec<&str> = Vec::new(); + let mut answered: Vec<&str> = Vec::new(); + + for message in messages { + match message { + Message::Assistant(assistant) => { + declared.extend(assistant.tool_calls.iter().map(|call| call.id.as_str())); + } + Message::Tool(tool) => { + if !declared.contains(&tool.tool_call_id.as_str()) { + return false; + } + answered.push(tool.tool_call_id.as_str()); + } + _ => {} + } + } + + declared.iter().all(|id| answered.contains(id)) +} + +fn assistant_calling(id: &str) -> Message { + Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ToolCall::new( + id, + "get_weather", + serde_json::json!({"city": "Paris"}), + )], + usage: None, + }) +} + +/// `[system, user, assistant(tool_calls=[c1]), tool(c1), assistant]` — the +/// shape a `keep_last = 2` cut splits down the middle. +fn tool_transcript() -> Vec { + vec![ + Message::system("you are helpful"), + Message::user("weather in Paris?"), + assistant_calling("c1"), + Message::tool("c1", r#"{"temp_c":21}"#), + Message::assistant("It is 21C."), + ] +} + +#[test] +fn summarization_plan_never_orphans_a_tool_result() { + let policy = SummarizationPolicy { + keep_last: 2, + ..Default::default() + }; + let (_to_summarize, to_keep) = policy.plan(&tool_transcript()); + + assert!( + pairing_intact(&to_keep), + "plan kept a tool result whose assistant tool-call turn was summarized away: {to_keep:#?}" + ); +} + +#[test] +fn keep_last_never_orphans_a_tool_result() { + let trimmed = trim_messages(&tool_transcript(), &TrimStrategy::KeepLast(2)); + assert!( + pairing_intact(&trimmed), + "KeepLast produced a slice a provider would reject: {trimmed:#?}" + ); +} + +#[test] +fn keep_first_and_last_never_orphans_either_end() { + let messages = vec![ + Message::user("one"), + assistant_calling("c1"), + Message::tool("c1", "r1"), + Message::user("two"), + assistant_calling("c2"), + Message::tool("c2", "r2"), + Message::assistant("done"), + ]; + let trimmed = trim_messages( + &messages, + &TrimStrategy::KeepFirstAndLast { first: 2, last: 2 }, + ); + assert!( + pairing_intact(&trimmed), + "KeepFirstAndLast left an unanswered call or an unpaired result: {trimmed:#?}" + ); +} + +#[test] +fn max_tokens_never_orphans_a_tool_result() { + // The assistant turn carries visible text as well as its tool call, so the + // budget loop stops *between* the call and its result — the only way a + // front-dropping token trim can orphan one. + let narrating_call = Message::Assistant(AssistantMessage { + id: None, + content: vec![tinyagents::harness::message::ContentBlock::Text( + "Let me look that up. ".repeat(20), + )], + tool_calls: vec![ToolCall::new( + "c1", + "get_weather", + serde_json::json!({"city": "Paris"}), + )], + usage: None, + }); + let messages = vec![ + Message::user("x".repeat(40)), + narrating_call, + Message::tool("c1", "r1"), + Message::assistant("done"), + ]; + let trimmed = trim_messages(&messages, &TrimStrategy::MaxTokens(3)); + assert!( + pairing_intact(&trimmed), + "MaxTokens produced a slice a provider would reject: {trimmed:#?}" + ); +} + +#[test] +fn a_tool_only_assistant_turn_trips_the_compaction_gate() { + // A 50-turn run whose assistant messages are all tool calls with 2 KB + // argument blobs used to estimate at zero tokens, so the window overflowed + // uncompacted. + let heavy = Message::Assistant(AssistantMessage { + id: None, + content: Vec::new(), + tool_calls: vec![ToolCall::new( + "c1", + "search", + serde_json::json!({"query": "q".repeat(2_000)}), + )], + usage: None, + }); + assert_eq!(heavy.text(), "", "precondition: no visible text"); + + let policy = SummarizationPolicy { + trigger_tokens: 100, + ..Default::default() + }; + assert!( + policy.should_summarize(&[heavy]), + "a 2 KB tool-call turn must trip a 100-token trigger" + ); +} + +#[tokio::test] +async fn the_default_summarizer_preserves_tool_history() { + let record = ConcatSummarizer + .summarize(&tool_transcript()) + .await + .expect("summarizing a non-empty transcript must succeed"); + let summary = record.summary.text(); + + assert!( + summary.contains("get_weather"), + "the tool name was erased: {summary}" + ); + assert!( + summary.contains("Paris"), + "the tool arguments were erased: {summary}" + ); + assert!( + summary.contains("temp_c"), + "the tool result was erased: {summary}" + ); +} diff --git a/tests/context_and_schema_tool_surface.rs b/tests/context_and_schema_tool_surface.rs new file mode 100644 index 0000000..ab16d95 --- /dev/null +++ b/tests/context_and_schema_tool_surface.rs @@ -0,0 +1,207 @@ +//! End-to-end tests for the tool-layer surfaces added alongside the compaction +//! fixes: unique synthetic call ids, tool-result artifacts, the provider schema +//! projection seam, injected arguments, and per-tool error policy. +//! +//! Unlike `context_and_schema_compaction.rs`, most of these exercise APIs that +//! did not exist before, so they cannot be run red against the unfixed crate — +//! except [`synthetic_tool_call_ids_are_unique_across_turns`], which uses only +//! the pre-existing `parse_prompt_tool_calls_from_text` and does fail there. + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::json; + +use tinyagents::Result; +use tinyagents::error::TinyAgentsError; +use tinyagents::harness::message::Message; +use tinyagents::harness::tool::{ + SchemaPreparation, Tool, ToolCall, ToolErrorPolicy, ToolRegistry, ToolResult, ToolSchema, + parse_prompt_tool_calls_from_text, prepare_tool_schemas, strip_injected_arguments, +}; + +// --------------------------------------------------------------------------- +// TOOL-2 — synthetic call ids +// --------------------------------------------------------------------------- + +/// Two turns of one run must not both mint `call_1`. When they did, the next +/// request declared the same tool-call id twice and answered it twice, leaving +/// the pairing unresolvable for the provider and for the harness alike. +#[test] +fn synthetic_tool_call_ids_are_unique_across_turns() { + let turn = r#"{"name":"lookup","arguments":{"id":1}}"#; + let (_, first) = parse_prompt_tool_calls_from_text(turn); + let (_, second) = parse_prompt_tool_calls_from_text(turn); + + assert_eq!(first.len(), 1); + assert_eq!(second.len(), 1); + assert_ne!( + first[0].id, second[0].id, + "the second turn reused the first turn's synthetic tool-call id" + ); +} + +// --------------------------------------------------------------------------- +// C7 — content_and_artifact +// --------------------------------------------------------------------------- + +/// A tool returns a short model-facing summary and a large structured payload; +/// the transcript must carry both, with only the summary visible as text. +#[test] +fn a_tool_result_artifact_survives_into_the_transcript() { + let mut result = ToolResult::text("c1", "query_rows", "3 rows matched"); + result.raw = Some(json!({"rows": [{"id": 1}, {"id": 2}, {"id": 3}]})); + + let message = Message::tool_from_result(&result); + + assert_eq!(message.text(), "3 rows matched"); + assert_eq!( + message + .artifact() + .and_then(|a| a["rows"].as_array()) + .map(Vec::len), + Some(3) + ); + // The payload is host-side state and is not charged to the context window. + assert_eq!( + message.estimated_char_weight(), + "3 rows matched".len() + "c1".len() + ); +} + +// --------------------------------------------------------------------------- +// TOOL-6 — the provider projection seam +// --------------------------------------------------------------------------- + +/// A schema written the way every JSON-Schema generator writes one — with +/// `$defs` and a `$ref` — must be projected into a shape Anthropic accepts. +#[test] +fn tool_schemas_are_projected_for_the_target_provider() { + let declared = vec![ToolSchema::new( + "lookup", + "Look a record up", + json!({ + "type": "object", + "$defs": {"Id": {"type": "string"}}, + "properties": {"id": {"$ref": "#/$defs/Id"}}, + "required": ["id"], + }), + )]; + + let wire = prepare_tool_schemas(&declared, &SchemaPreparation::anthropic()); + assert_eq!(wire[0].parameters["properties"]["id"]["type"], "string"); + assert!(wire[0].parameters.get("$defs").is_none()); +} + +/// A tool that declares no arguments must not be able to serialise +/// `"parameters": null` and break the whole request. +#[test] +fn null_tool_parameters_are_normalized_rather_than_sent() { + let declared = vec![ToolSchema::new("ping", "Ping", json!(null))]; + let wire = prepare_tool_schemas(&declared, &SchemaPreparation::openai()); + + assert!(!wire[0].parameters.is_null()); + assert_eq!(wire[0].parameters["type"], "object"); +} + +// --------------------------------------------------------------------------- +// C8 — injected arguments +// --------------------------------------------------------------------------- + +struct ContextualTool; + +#[async_trait] +impl Tool<()> for ContextualTool { + fn name(&self) -> &str { + "contextual" + } + fn description(&self) -> &str { + "Acts within the caller's thread" + } + fn schema(&self) -> ToolSchema { + ToolSchema::new( + "contextual", + "Acts within the caller's thread", + json!({ + "type": "object", + "properties": { + "note": {"type": "string"}, + "thread_id": {"type": "string"}, + }, + "required": ["note", "thread_id"], + }), + ) + } + fn injected_arguments(&self) -> &[&str] { + &["thread_id"] + } + async fn call(&self, _state: &(), call: ToolCall) -> Result { + Ok(ToolResult::text(call.id, call.name, "ok")) + } +} + +#[test] +fn injected_arguments_never_reach_the_model_and_cannot_be_forged() { + let mut registry: ToolRegistry<()> = ToolRegistry::new(); + registry.register(Arc::new(ContextualTool)); + + // Declaration side: hidden from `properties` and from `required`. + let wire = registry.schemas(); + assert!(wire[0].parameters["properties"].get("thread_id").is_none()); + assert_eq!(wire[0].parameters["required"], json!(["note"])); + + // Enforcement primitive: a model-supplied value for the hidden key is + // discarded before anything else looks at the arguments. + let mut forged = json!({"note": "hi", "thread_id": "someone-elses-thread"}); + let removed = strip_injected_arguments(&mut forged, &["thread_id"]); + assert_eq!(removed, vec!["thread_id".to_string()]); + assert_eq!(forged, json!({"note": "hi"})); +} + +// --------------------------------------------------------------------------- +// C9 — per-tool error policy +// --------------------------------------------------------------------------- + +#[test] +fn a_recoverable_tool_failure_can_be_handed_back_to_the_model() { + let call = ToolCall::new("c1", "lookup", json!({})); + let handled = ToolErrorPolicy::ReturnToError + .apply(&call, Err(TinyAgentsError::Tool("no such record".into()))) + .expect("a handled error must not fail the run"); + + assert!(handled.is_error()); + assert!(handled.content.contains("no such record")); +} + +/// The rule that must never be relaxed: a policy may convert tool failures, but +/// never a cancellation or an interrupt. Swallowing one lets the loop keep +/// running after it was told to stop. +#[test] +fn no_error_policy_can_swallow_a_cancellation_or_an_interrupt() { + let call = ToolCall::new("c1", "lookup", json!({})); + + for policy in [ + ToolErrorPolicy::Fail, + ToolErrorPolicy::ReturnToError, + ToolErrorPolicy::Message("masked".into()), + ] { + assert!( + policy + .apply(&call, Err(TinyAgentsError::Cancelled)) + .is_err(), + "{policy:?} swallowed a cancellation" + ); + assert!( + policy + .apply( + &call, + Err(TinyAgentsError::Interrupted { + node: "approval".into(), + message: "waiting".into(), + }) + ) + .is_err(), + "{policy:?} swallowed an interrupt" + ); + } +} diff --git a/tests/e2e_agent_graph.rs b/tests/e2e_agent_graph.rs index 94fa816..de74cbf 100644 --- a/tests/e2e_agent_graph.rs +++ b/tests/e2e_agent_graph.rs @@ -51,6 +51,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -68,6 +69,7 @@ fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_budget.rs b/tests/e2e_budget.rs index 1960f57..29b064b 100644 --- a/tests/e2e_budget.rs +++ b/tests/e2e_budget.rs @@ -59,6 +59,7 @@ fn tool_call_response(id: &str, name: &str, input: u64, output: u64) -> ModelRes raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -76,6 +77,7 @@ fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -293,6 +295,7 @@ async fn cost_pricing_records_and_enforces_money_budget() { source: ModelResolutionSource::RegistryDefault, }), continue_turn: None, + served_from_cache: false, }; stack @@ -519,6 +522,7 @@ async fn cached_input_budget_blocks_next_call() { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, }; stack .run_after_model(&mut ctx, &(), &mut resp) diff --git a/tests/e2e_fuzz_graph_agents.rs b/tests/e2e_fuzz_graph_agents.rs index 64f142e..c424ef3 100644 --- a/tests/e2e_fuzz_graph_agents.rs +++ b/tests/e2e_fuzz_graph_agents.rs @@ -262,6 +262,7 @@ fn tool_call_response(calls: Vec) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -278,6 +279,7 @@ fn text_response(text: impl Into) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_graph_subagent_node.rs b/tests/e2e_graph_subagent_node.rs index 5d11c3f..5f34c25 100644 --- a/tests/e2e_graph_subagent_node.rs +++ b/tests/e2e_graph_subagent_node.rs @@ -32,6 +32,7 @@ fn tool_call_response(id: &str, name: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_graph_todos.rs b/tests/e2e_graph_todos.rs index 320ccda..c3f8dec 100644 --- a/tests/e2e_graph_todos.rs +++ b/tests/e2e_graph_todos.rs @@ -30,6 +30,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -46,6 +47,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_harness_provider_contracts.rs b/tests/e2e_harness_provider_contracts.rs index 48b4b51..8c05371 100644 --- a/tests/e2e_harness_provider_contracts.rs +++ b/tests/e2e_harness_provider_contracts.rs @@ -277,6 +277,7 @@ async fn model_request_response_registry_and_stream_contracts_are_stable() { code: Some("internal".into()), message: "nope".into(), retryable: true, + retry_after_ms: None, raw: Some(json!({ "error": "nope" })), }, )])); @@ -455,6 +456,7 @@ fn structured_output_supports_provider_schema_and_tool_fallbacks() { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, }; let tool_output = StructuredExtractor::new(StructuredStrategy::ToolCall, "score", schema) .extract(&tool_response) diff --git a/tests/e2e_middleware.rs b/tests/e2e_middleware.rs index cad6eda..35a560a 100644 --- a/tests/e2e_middleware.rs +++ b/tests/e2e_middleware.rs @@ -35,7 +35,7 @@ use tinyagents::harness::middleware::{ }; use tinyagents::harness::model::{ChatModel, ModelRequest, ModelResponse}; use tinyagents::harness::retry::RetryPolicy; -use tinyagents::harness::runtime::{AgentHarness, RunPolicy}; +use tinyagents::harness::runtime::{AgentHarness, InvalidArgsPolicy, RunPolicy}; use tinyagents::harness::testkit::{EventRecorder, FakeTool, ScriptedModel, Trajectory}; use tinyagents::harness::tool::{Tool, ToolCall, ToolResult, ToolSchema}; use tinyagents::harness::usage::Usage; @@ -143,6 +143,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -160,6 +161,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -343,6 +345,12 @@ async fn invalid_tool_arguments_are_rejected_before_execution() { .register_tool(Arc::new(StrictLookupTool { calls: Arc::clone(&calls), })); + // `ReturnToolError` is the default now, so the fail-closed schema boundary + // this test pins has to be opted into. + harness.with_policy(RunPolicy { + invalid_args: InvalidArgsPolicy::Fail, + ..RunPolicy::default() + }); let ctx = RunContext::new(RunConfig::new("mw-e2e-tool-schema"), ()).with_events(recorder.sink()); diff --git a/tests/e2e_observability.rs b/tests/e2e_observability.rs index 57e2770..a94ba87 100644 --- a/tests/e2e_observability.rs +++ b/tests/e2e_observability.rs @@ -55,6 +55,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -71,6 +72,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_public_api_contracts.rs b/tests/e2e_public_api_contracts.rs index b7c0533..a4957c6 100644 --- a/tests/e2e_public_api_contracts.rs +++ b/tests/e2e_public_api_contracts.rs @@ -345,11 +345,17 @@ async fn retry_rate_limit_and_summarization_contracts_are_deterministic() { assert!(!retry.should_retry(2)); assert_eq!(retry.backoff_for_attempt(0), Duration::from_millis(50)); assert_eq!(retry.backoff_for_attempt(2), Duration::from_millis(90)); + // Jitter is additive (LOOP-2), not multiplicative: `rand01 = 0.5` is the + // band midpoint and so reproduces the un-jittered value exactly. This used + // to assert 45ms — `base * rand01` — the form that let a hardcoded + // `rand01 = 0.0` on the production path disable backoff altogether. let jittered = retry.clone().with_jitter(true); assert_eq!( jittered.backoff_for_attempt_with(1, 0.5), - Duration::from_millis(45) + Duration::from_millis(90) ); + // And the band never reaches zero. + assert!(jittered.backoff_for_attempt_with(1, 0.0) > Duration::ZERO); assert!(is_retryable(&TinyAgentsError::Model("timeout".into()))); assert!(is_retryable(&TinyAgentsError::Tool("temporary".into()))); diff --git a/tests/e2e_reasoning_and_selection.rs b/tests/e2e_reasoning_and_selection.rs index 99d8f8c..dc8040f 100644 --- a/tests/e2e_reasoning_and_selection.rs +++ b/tests/e2e_reasoning_and_selection.rs @@ -46,6 +46,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_subagents.rs b/tests/e2e_subagents.rs index 74fbe8f..e1a6f6a 100644 --- a/tests/e2e_subagents.rs +++ b/tests/e2e_subagents.rs @@ -47,6 +47,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -64,6 +65,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_tool_policy.rs b/tests/e2e_tool_policy.rs index aaf542d..8490aa9 100644 --- a/tests/e2e_tool_policy.rs +++ b/tests/e2e_tool_policy.rs @@ -49,6 +49,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -65,6 +66,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/e2e_unknown_tool_policy.rs b/tests/e2e_unknown_tool_policy.rs index d26fd05..94f7047 100644 --- a/tests/e2e_unknown_tool_policy.rs +++ b/tests/e2e_unknown_tool_policy.rs @@ -39,6 +39,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -55,6 +56,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -80,7 +82,7 @@ fn single_unknown_tool_event(events: &[AgentEvent]) -> (String, String) { found.expect("an UnknownToolCall event should have been recorded") } -// ── 1. Fail policy (default) ────────────────────────────────────────────────── +// ── 1. Fail policy (opt-in) ─────────────────────────────────────────────────── #[tokio::test] async fn fail_policy_errors_on_unregistered_tool() { @@ -89,7 +91,12 @@ async fn fail_policy_errors_on_unregistered_tool() { "mock", Arc::new(MockModel::with_tool_call("missing", json!({}))), ); - // Default policy is UnknownToolPolicy::Fail; no tool registered. + // `ReturnToolError` is the default now (a hallucinated tool name is a + // routine model mistake), so a hard stop has to be asked for. + harness.with_policy(RunPolicy { + unknown_tool: UnknownToolPolicy::Fail, + ..RunPolicy::default() + }); let err = harness .invoke_default(&(), vec![Message::user("go")]) diff --git a/tests/e2e_workspace_and_registry.rs b/tests/e2e_workspace_and_registry.rs index 4f19b6f..5ee5130 100644 --- a/tests/e2e_workspace_and_registry.rs +++ b/tests/e2e_workspace_and_registry.rs @@ -232,6 +232,7 @@ async fn harness_run_threads_workspace_and_enforces_out_of_root() { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } fn text_response(text: &str) -> ModelResponse { @@ -247,6 +248,7 @@ async fn harness_run_threads_workspace_and_enforces_out_of_root() { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/feature_harness_agent_loop.rs b/tests/feature_harness_agent_loop.rs index 7f49788..9f6c674 100644 --- a/tests/feature_harness_agent_loop.rs +++ b/tests/feature_harness_agent_loop.rs @@ -46,6 +46,7 @@ fn multi_tool_call_response(calls: Vec) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -66,6 +67,7 @@ fn text_response(text: &str) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/feature_harness_structured.rs b/tests/feature_harness_structured.rs index 3715105..1c1a1d8 100644 --- a/tests/feature_harness_structured.rs +++ b/tests/feature_harness_structured.rs @@ -62,6 +62,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -120,8 +121,13 @@ async fn provider_schema_rejects_non_json_text() { } #[tokio::test] -async fn provider_schema_parse_type_mismatch_errors() { - // Valid JSON, but the shape does not match `Answer` (score is a string). +async fn provider_schema_rejects_a_value_that_violates_the_schema() { + // Valid JSON, but the shape does not match the declared schema (`score` is + // a string). This used to extract *successfully* — the extractor stored its + // schema and never read it — and the mismatch only surfaced later, at + // `parse::()`. It is now caught at the boundary, with the failing + // instance path named, so `run.structured` can never hold a value the + // schema rejects. let extractor = StructuredExtractor::new( StructuredStrategy::ProviderSchema, "answer", @@ -129,6 +135,28 @@ async fn provider_schema_parse_type_mismatch_errors() { ); let response = ModelResponse::assistant(r#"{"value":"x","score":"not-a-number"}"#); + let err = extractor + .extract(&response) + .expect_err("a schema violation must fail closed"); + assert!( + matches!(err, TinyAgentsError::StructuredOutput(_)), + "expected a StructuredOutput error, got {err:?}" + ); + assert!( + err.to_string().contains("schema 'answer'.score"), + "the error must name the failing instance path: {err}" + ); +} + +#[tokio::test] +async fn parse_type_mismatch_still_errors_for_a_schema_free_extractor() { + // The `parse::()` mismatch path is still reachable: an extractor with no + // declared schema validates nothing, so the deserialisation boundary is + // where the mismatch is caught. + let extractor = + StructuredExtractor::new(StructuredStrategy::ProviderSchema, "answer", json!({})); + let response = ModelResponse::assistant(r#"{"value":"x","score":"not-a-number"}"#); + let output = extractor.extract(&response).expect("valid JSON extracts"); let err = output .parse::() @@ -163,6 +191,7 @@ async fn tool_call_strategy_reads_named_tool_arguments() { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, }; let output = extractor @@ -337,6 +366,7 @@ async fn provider_schema_reads_text_content_blocks() { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, }; let parsed: Answer = extractor diff --git a/tests/feature_infra_resilience.rs b/tests/feature_infra_resilience.rs index 3e7cd7c..ca74956 100644 --- a/tests/feature_infra_resilience.rs +++ b/tests/feature_infra_resilience.rs @@ -37,20 +37,33 @@ fn backoff_grows_exponentially_then_caps() { } #[test] -fn jitter_scales_backoff_by_supplied_random_value() { +fn jitter_spreads_backoff_additively_around_the_base() { + // Updated for LOOP-2. This test used to assert the *multiplicative* form + // (`base * rand01`), which is the defect: `rand01 = 0.0` collapsed the + // window to `Duration::ZERO`, and the production path passed a hardcoded + // `0.0`, so turning jitter on disabled backoff entirely. Jitter is now + // additive (`base * (1 ± JITTER_FRACTION)`), matching both LangGraph + // (`interval + uniform(0, 1)`) and LangChain (`delay ± 25%`), and can only + // ever widen the delay band. let policy = RetryPolicy::default() .with_initial_backoff_ms(1000) .with_multiplier(1.0) + .with_max_backoff_ms(u64::MAX) .with_jitter(true); - // rand01 = 0.0 collapses the window to zero; 1.0 keeps the full base. - assert_eq!(policy.backoff_for_attempt_with(0, 0.0), Duration::ZERO); + + // The band is [750, 1250] at the default ±25%. + assert_eq!( + policy.backoff_for_attempt_with(0, 0.0), + Duration::from_millis(750) + ); assert_eq!( - policy.backoff_for_attempt_with(0, 0.25), - Duration::from_millis(250) + policy.backoff_for_attempt_with(0, 0.5), + Duration::from_millis(1000), + "the midpoint must reproduce the un-jittered value" ); assert_eq!( policy.backoff_for_attempt_with(0, 1.0), - Duration::from_millis(1000) + Duration::from_millis(1250) ); } @@ -107,6 +120,7 @@ fn provider_error_retryability_follows_its_flag() { code: None, message: "service unavailable".into(), retryable: true, + retry_after_ms: None, raw: None, }; assert!(is_retryable(&TinyAgentsError::Provider(Box::new( @@ -116,6 +130,7 @@ fn provider_error_retryability_follows_its_flag() { let terminal = ProviderError { retryable: false, + retry_after_ms: None, status: Some(401), message: "invalid api key".into(), ..retryable.clone() @@ -179,6 +194,7 @@ fn classify_provider_error_and_reason_labels_are_stable() { code: None, message: "internal server error".into(), retryable: true, + retry_after_ms: None, raw: None, }; let class = classify_provider_error(&err); diff --git a/tests/harness_agent_loop.rs b/tests/harness_agent_loop.rs index d2349c9..ca90e1d 100644 --- a/tests/harness_agent_loop.rs +++ b/tests/harness_agent_loop.rs @@ -67,6 +67,7 @@ fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> Mod raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } @@ -83,6 +84,7 @@ fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { raw: None, resolved_model: None, continue_turn: None, + served_from_cache: false, } } diff --git a/tests/live_local_models.rs b/tests/live_local_models.rs index e11d68e..e4d0720 100644 --- a/tests/live_local_models.rs +++ b/tests/live_local_models.rs @@ -79,11 +79,11 @@ const TIMEOUT_MS: u64 = 180_000; /// that never reached the answer. /// /// The crate already treats this as a first-class failure mode — see -/// [`RunPolicy::empty_response_retries`], whose documentation names `qwen3` via +/// [`RunPolicy::truncated_empty_retries`], whose documentation names `qwen3` via /// Ollama specifically — so the tests must not reintroduce it by being frugal. /// A budget this size is what a host talking to local models should use. /// -/// [`RunPolicy::empty_response_retries`]: tinyagents::harness::runtime::RunPolicy::empty_response_retries +/// [`RunPolicy::truncated_empty_retries`]: tinyagents::harness::runtime::RunPolicy::truncated_empty_retries const MAX_TOKENS: u32 = 1024; /// The weather the fake tool always reports. Distinctive enough that finding it @@ -268,19 +268,26 @@ async fn reachable_runtimes() -> Vec { /// The [`RunPolicy`] a host should use to drive a small local model. /// -/// The crate default is [`InvalidArgsPolicy::Fail`], which aborts the entire -/// run the first time a model calls a registered tool with schema-invalid -/// arguments. That is a reasonable default for a frontier model, where the case -/// is nearly always a genuine bug — but a 3B quantised model omits a required -/// argument often enough that `Fail` makes the loop unusably brittle: one bad -/// call and the run dies rather than the model getting a chance to correct -/// itself. [`InvalidArgsPolicy::NormalizeThenReturnToolError`] repairs the -/// common provider-shape defects and otherwise hands the validation error back -/// to the model as a tool result, and the recovery still consumes a tool-call -/// budget slot so the loop stays bounded. +/// The crate default is now [`InvalidArgsPolicy::ReturnToolError`], which hands +/// the validation error back to the model as a tool result instead of aborting +/// the run. That change removed the original reason this helper existed: the +/// default used to be [`InvalidArgsPolicy::Fail`], which killed the whole run +/// the first time a model called a registered tool with schema-invalid +/// arguments — reasonable for a frontier model, where that is nearly always a +/// genuine bug, but unusably brittle for a 3B quantised model that omits a +/// required argument often enough to end most runs on the first tool call. /// -/// This is observed behaviour, not a hypothetical: with the default policy, -/// `llama3.2:3b` fails this file's tool-loop test with +/// The helper still earns its keep, for a narrower reason. +/// [`InvalidArgsPolicy::NormalizeThenReturnToolError`] additionally repairs the +/// common provider-shape defects — JSON emitted as a string, a scalar sent +/// where an array is declared — *before* deciding the call is invalid. Those +/// defects are characteristic of small local models specifically, so the extra +/// normalization pass is worth requesting locally and not worth paying for by +/// default. Either way the recovery consumes a tool-call budget slot, so the +/// correction loop stays bounded. +/// +/// This is observed behaviour, not a hypothetical: under the old default, +/// `llama3.2:3b` failed this file's tool-loop test with /// `tool `get_weather` arguments.city is required`. fn local_run_policy() -> RunPolicy { RunPolicy { @@ -790,18 +797,24 @@ fn local_presets_need_no_credential_and_default_to_their_own_ports() { ); } -/// Pins the default that makes local tool loops brittle, and the opt-in that -/// fixes them. +/// Pins the recovering crate default, and the stronger normalizing policy a +/// local host still opts into on top of it. /// -/// If the crate default ever changes to a recovering policy, this test fails -/// and [`local_run_policy`]'s rationale (and the docs pointing hosts at it) -/// should be revisited rather than the assertion simply flipped. +/// This test previously asserted the default was [`InvalidArgsPolicy::Fail`] +/// and said that if the default ever became a recovering policy, this +/// assertion should not simply be flipped — [`local_run_policy`]'s rationale +/// had to be revisited first. The default did change, and it was: the helper's +/// doc comment no longer justifies itself by "the default aborts the run", but +/// by the provider-shape normalization pass that `ReturnToolError` alone does +/// not perform. The assertion below is updated only because that reasoning was +/// re-derived, not to make a red test green. #[test] -fn invalid_tool_arguments_abort_the_run_unless_recovery_is_opted_into() { +fn invalid_tool_arguments_recover_by_default_and_local_hosts_add_normalization() { assert_eq!( RunPolicy::default().invalid_args, - InvalidArgsPolicy::Fail, - "the crate default aborts a run on schema-invalid tool arguments" + InvalidArgsPolicy::ReturnToolError, + "the crate default hands the validation error back to the model \ + rather than aborting the run" ); assert_eq!( local_run_policy().invalid_args, diff --git a/tests/persistence_conformance.rs b/tests/persistence_conformance.rs new file mode 100644 index 0000000..7be81d7 --- /dev/null +++ b/tests/persistence_conformance.rs @@ -0,0 +1,48 @@ +//! Checkpointer conformance for the lineage and pending-writes contracts. +//! +//! The pre-existing `tests/conformance.rs` covers put/get/list/delete_thread/ +//! prune. These two suites cover what it did not — `get_tuple`, +//! `state_history`, `copy_thread`, `delete_checkpoints` and the whole +//! pending-writes protocol — against all three bundled backends, so a defect in +//! any one of them cannot hide behind the other two. + +use tinyagents::graph::checkpoint::{FileCheckpointer, InMemoryCheckpointer}; +use tinyagents::graph::testkit::conformance::{ + checkpointer_lineage_contract, checkpointer_writes_contract, +}; + +#[tokio::test] +async fn in_memory_checkpointer_satisfies_the_writes_contract() { + checkpointer_writes_contract(InMemoryCheckpointer::::new()).await; +} + +#[tokio::test] +async fn in_memory_checkpointer_satisfies_the_lineage_contract() { + checkpointer_lineage_contract(InMemoryCheckpointer::::new()).await; +} + +#[tokio::test] +async fn file_checkpointer_satisfies_the_writes_contract() { + let dir = tempfile::tempdir().unwrap(); + checkpointer_writes_contract(FileCheckpointer::::new(dir.path())).await; +} + +#[tokio::test] +async fn file_checkpointer_satisfies_the_lineage_contract() { + let dir = tempfile::tempdir().unwrap(); + checkpointer_lineage_contract(FileCheckpointer::::new(dir.path())).await; +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_checkpointer_satisfies_the_writes_contract() { + use tinyagents::graph::checkpoint::SqliteCheckpointer; + checkpointer_writes_contract(SqliteCheckpointer::::in_memory().unwrap()).await; +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_checkpointer_satisfies_the_lineage_contract() { + use tinyagents::graph::checkpoint::SqliteCheckpointer; + checkpointer_lineage_contract(SqliteCheckpointer::::in_memory().unwrap()).await; +} diff --git a/tests/persistence_session.rs b/tests/persistence_session.rs new file mode 100644 index 0000000..39b78d6 --- /dev/null +++ b/tests/persistence_session.rs @@ -0,0 +1,293 @@ +//! Regression tests for the session database's durability and correctness +//! fixes: schema migrations, the busy handler, the task-claim guard, atomic +//! FTS indexing, and the retention/reindex entry points. +//! +//! These live as integration tests rather than module-local ones so they +//! exercise the same public surface a host uses. +#![cfg(feature = "sqlite")] + +use chrono::{Duration, Utc}; +use tinyagents::session; +use tinyagents::session::run_ledger::{ + self, + types::{ + AgentTeamStatus, AgentTeamTaskStatus, AgentTeamTaskUpsert, AgentTeamUpsert, ClaimOutcome, + }, +}; + +fn team(id: &str) -> AgentTeamUpsert { + AgentTeamUpsert { + id: id.into(), + parent_thread_id: None, + lead_agent_id: "lead".into(), + status: AgentTeamStatus::Active, + summary: None, + created_at: None, + closed_at: None, + } +} + +fn task(id: &str, status: AgentTeamTaskStatus) -> AgentTeamTaskUpsert { + AgentTeamTaskUpsert { + id: id.into(), + team_id: "team".into(), + title: format!("work {id}"), + objective: None, + status, + owner_member_id: None, + depends_on: Vec::new(), + gate_status: None, + gate_reason: None, + evidence: Vec::new(), + source_run_id: None, + order_index: 0, + created_at: None, + } +} + +fn workspace() -> tempfile::TempDir { + tempfile::tempdir().unwrap() +} + +/// A competing writer must be waited out, not failed on. +/// +/// This pins a property rather than a fix: it passes against the pre-change +/// code too, because `rusqlite` installs a 5s busy timeout on every +/// `Connection::open` of its own accord. The explicit `store::BUSY_TIMEOUT` is +/// about owning that guarantee instead of inheriting it from an undocumented +/// dependency default — and this test is what would catch it going away. +#[test] +fn a_competing_writer_is_waited_out_rather_than_failed_on() { + let dir = workspace(); + session::with_connection(dir.path(), |_| Ok(())).unwrap(); + + let blocker = rusqlite::Connection::open(session::db_path(dir.path())).unwrap(); + blocker.execute_batch("BEGIN IMMEDIATE").unwrap(); + let releaser = std::thread::spawn(move || { + std::thread::sleep(std::time::Duration::from_millis(300)); + blocker.execute_batch("ROLLBACK").unwrap(); + }); + + let result = session::record_session_start( + dir.path(), + "s-waited", + "agent", + "Agent", + "key", + None, + None, + None, + None, + None, + ); + releaser.join().unwrap(); + assert!( + result.is_ok(), + "the write must wait for the lock, not fail instantly: {result:?}" + ); +} + +/// C1: the schema carries a version marker, so a column could actually be +/// added to an existing workspace database. +#[test] +fn the_schema_records_a_version() { + let dir = workspace(); + let version: i64 = session::with_connection(dir.path(), |conn| { + Ok(conn + .query_row("SELECT version FROM schema_version WHERE id = 1", [], |r| { + r.get(0) + }) + .unwrap()) + }) + .unwrap(); + assert!( + version >= 0, + "a freshly created database records the schema version it was built at" + ); +} + +/// SESS-11: the listing sort columns are indexed. +#[test] +fn listing_sort_columns_are_indexed() { + let dir = workspace(); + session::with_connection(dir.path(), |conn| { + for name in [ + "idx_workflow_runs_updated", + "idx_agent_teams_updated", + "idx_agent_team_tasks_order", + ] { + let exists: bool = conn + .prepare("SELECT 1 FROM sqlite_master WHERE type='index' AND name=?1") + .unwrap() + .exists([name]) + .unwrap(); + assert!(exists, "missing index {name}"); + } + Ok(()) + }) + .unwrap(); +} + +/// SESS-2: a finished task must not be re-claimable. +/// +/// `upsert_agent_team_task` NULLs `claimed_by_member_id` whenever the status is +/// not `in_progress`, so the old `claimed_by_member_id IS NULL` guard was +/// satisfied by every `done` task. A stale worker could flip a completed task +/// back to `in_progress` and strand everything that depended on it. +#[test] +fn a_done_task_cannot_be_reclaimed() { + let dir = workspace(); + run_ledger::upsert_agent_team(dir.path(), team("team")).unwrap(); + run_ledger::upsert_agent_team_task(dir.path(), task("task", AgentTeamTaskStatus::Done)) + .unwrap(); + + let outcome = + run_ledger::claim_agent_team_task(dir.path(), "team", "task", "stale-worker", "tok") + .unwrap(); + assert!( + matches!(outcome, ClaimOutcome::AlreadyClaimed), + "a terminal task is not claimable, whatever its claim column says: {outcome:?}" + ); + + let task = run_ledger::get_agent_team_task(dir.path(), "task") + .unwrap() + .unwrap(); + assert_eq!( + task.status, + AgentTeamTaskStatus::Done, + "the rejected claim must not have flipped the task back to in_progress" + ); +} + +/// A claimable task still claims — the guard must not be so tight that it +/// breaks the happy path. +#[test] +fn a_todo_task_is_still_claimable() { + let dir = workspace(); + run_ledger::upsert_agent_team(dir.path(), team("team")).unwrap(); + run_ledger::upsert_agent_team_task(dir.path(), task("task", AgentTeamTaskStatus::Todo)) + .unwrap(); + let outcome = + run_ledger::claim_agent_team_task(dir.path(), "team", "task", "worker", "tok").unwrap(); + assert!(matches!(outcome, ClaimOutcome::Claimed(_)), "{outcome:?}"); +} + +/// SESS-6: retention actually deletes, and reindexing restores searchability. +#[test] +fn retention_prunes_finished_sessions_and_reindex_restores_search() { + let dir = workspace(); + session::record_session_start( + dir.path(), + "old", + "agent", + "Agent", + "key", + None, + None, + None, + None, + None, + ) + .unwrap(); + session::record_message( + dir.path(), + "old", + "user", + "findable haystack", + None, + None, + None, + None, + ) + .unwrap(); + session::record_session_end( + dir.path(), + "old", + session::SessionStatus::Completed, + 1, + 0, + 0, + 0, + 0.0, + ) + .unwrap(); + + // Searchable to begin with. + let found = session::search_sessions( + dir.path(), + &session::SessionSearchParams { + query: Some("haystack".into()), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(found.total, 1); + + // Simulate the pre-fix damage: an FTS entry lost between the two old + // autocommit statements. Before `reindex_fts` existed there was no way back. + session::with_connection(dir.path(), |conn| { + conn.execute("DELETE FROM sessions_fts", []).unwrap(); + Ok(()) + }) + .unwrap(); + let lost = session::search_sessions( + dir.path(), + &session::SessionSearchParams { + query: Some("haystack".into()), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(lost.total, 0, "the index really is gone"); + + let rows = session::reindex_fts(dir.path()).unwrap(); + assert!(rows >= 2, "reindex writes a row per session and message"); + let recovered = session::search_sessions( + dir.path(), + &session::SessionSearchParams { + query: Some("haystack".into()), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!( + recovered.total, 1, + "reindexing makes the row findable again" + ); + + // Retention removes the finished session and everything hanging off it. + let report = session::apply_retention(dir.path(), Utc::now() + Duration::seconds(1)).unwrap(); + assert_eq!(report.sessions, 1, "the finished session was pruned"); + assert!(session::get_session(dir.path(), "old").is_err()); + let messages = session::with_connection(dir.path(), |conn| { + Ok(conn + .query_row("SELECT COUNT(*) FROM session_messages", [], |r| { + r.get::<_, i64>(0) + }) + .unwrap()) + }) + .unwrap(); + assert_eq!(messages, 0, "messages cascade with their session"); +} + +/// A running session is never pruned, however old. +#[test] +fn retention_never_prunes_a_running_session() { + let dir = workspace(); + session::record_session_start( + dir.path(), + "live", + "agent", + "Agent", + "key", + None, + None, + None, + None, + None, + ) + .unwrap(); + let report = session::apply_retention(dir.path(), Utc::now() + Duration::days(365)).unwrap(); + assert_eq!(report.sessions, 0); + assert!(session::get_session(dir.path(), "live").is_ok()); +} diff --git a/tests/persistence_store.rs b/tests/persistence_store.rs new file mode 100644 index 0000000..1dbf17a --- /dev/null +++ b/tests/persistence_store.rs @@ -0,0 +1,228 @@ +//! Regression tests for the harness store/memory backends and the file +//! checkpointer's on-disk durability. + +use std::sync::Arc; + +use tinyagents::graph::checkpoint::{Checkpoint, Checkpointer, FileCheckpointer}; +use tinyagents::harness::ids::NodeId; +use tinyagents::harness::memory::{ChatHistory, StoreChatHistory}; +use tinyagents::harness::message::Message; +use tinyagents::harness::store::{AppendStore, FileStore, JsonlAppendStore}; + +fn checkpoint(thread: &str, id: &str) -> Checkpoint { + Checkpoint { + thread_id: thread.to_string(), + checkpoint_id: id.to_string(), + run_id: None, + parent_checkpoint_id: None, + namespace: vec![], + state: 1, + next_nodes: vec![NodeId::from("n")], + completed_tasks: vec![], + pending_writes: vec![], + interrupts: vec![], + pending_activations: None, + barrier_arrivals: vec![], + metadata: serde_json::json!({ "source": "loop", "step": 1 }), + } +} + +// ── SESS-5: thread-id escaping ─────────────────────────────────────────────── + +/// Thread ids that differ only by case must not share a file. On APFS/NTFS +/// `Alice.jsonl` and `alice.jsonl` are the same file, so the old +/// `[A-Za-z0-9._-]` safe set silently merged two unrelated lineages. +#[tokio::test] +async fn thread_ids_differing_only_by_case_do_not_share_a_file() { + let dir = tempfile::tempdir().unwrap(); + let cp = FileCheckpointer::::new(dir.path()); + cp.put(checkpoint("Alice", "upper")).await.unwrap(); + cp.put(checkpoint("alice", "lower")).await.unwrap(); + + let upper = cp.list("Alice").await.unwrap(); + let lower = cp.list("alice").await.unwrap(); + assert_eq!(upper.len(), 1, "`Alice` holds only its own checkpoint"); + assert_eq!(lower.len(), 1, "`alice` holds only its own checkpoint"); + assert_eq!(upper[0].checkpoint_id, "upper"); + assert_eq!(lower[0].checkpoint_id, "lower"); + + // And the filenames themselves differ case-insensitively. + let mut names: Vec = std::fs::read_dir(dir.path()) + .unwrap() + .map(|e| e.unwrap().file_name().to_string_lossy().to_lowercase()) + .collect(); + names.sort(); + names.dedup(); + assert_eq!( + names.len(), + 2, + "two distinct case-folded filenames: {names:?}" + ); +} + +/// The empty thread id escapes to `""`, so its file is the dotfile `.jsonl`, +/// whose `Path::extension()` is `None` — it used to be invisible to listing. +#[tokio::test] +async fn the_empty_thread_id_is_visible_to_list_threads() { + let dir = tempfile::tempdir().unwrap(); + let cp = FileCheckpointer::::new(dir.path()); + cp.put(checkpoint("", "c1")).await.unwrap(); + let threads = cp.list_threads().await.unwrap(); + assert!( + threads.iter().any(|t| t.is_empty()), + "list_threads reports the empty thread id: {threads:?}" + ); +} + +// ── SESS-4: torn writes ────────────────────────────────────────────────────── + +/// A crash mid-append leaves a partial final line. It used to make the whole +/// thread permanently unreadable; now the torn tail is discarded and every +/// intact record in front of it still loads. +#[tokio::test] +async fn a_torn_trailing_line_does_not_destroy_the_thread() { + let dir = tempfile::tempdir().unwrap(); + let cp = FileCheckpointer::::new(dir.path()); + cp.put(checkpoint("t", "c1")).await.unwrap(); + cp.put(checkpoint("t", "c2")).await.unwrap(); + + let path = std::fs::read_dir(dir.path()) + .unwrap() + .map(|e| e.unwrap().path()) + .find(|p| p.to_string_lossy().ends_with("t.jsonl")) + .expect("thread file"); + let mut text = std::fs::read_to_string(&path).unwrap(); + text.push_str("{\"thread_id\":\"t\",\"checkpoint"); + std::fs::write(&path, text).unwrap(); + + let listed = cp.list("t").await.expect("a torn tail is survivable"); + assert_eq!(listed.len(), 2, "both intact records still load"); +} + +/// One poisoned file must not break listing for every other thread. +#[tokio::test] +async fn one_unreadable_thread_file_does_not_break_list_threads() { + let dir = tempfile::tempdir().unwrap(); + let cp = FileCheckpointer::::new(dir.path()); + cp.put(checkpoint("good", "c1")).await.unwrap(); + std::fs::write(dir.path().join("poison.jsonl"), "not json at all\n").unwrap(); + + let threads = cp + .list_threads() + .await + .expect("listing survives one bad file"); + assert!(threads.iter().any(|t| t == "good")); +} + +// ── SESS-9: StoreChatHistory::append ───────────────────────────────────────── + +/// `append` is a read-modify-write over the store. Concurrent appends used to +/// drop messages, giving the two `ChatHistory` backends different guarantees +/// for one trait method. +/// +/// The runtime flavour is load-bearing. `#[tokio::test]` defaults to a +/// **current-thread** runtime, and `FileStore`'s methods are `async fn`s with +/// no interior `.await`, so each read-modify-write runs to completion before +/// the next task is polled — the race cannot occur and the test passes even +/// against the unsynchronised version. Only a multi-threaded runtime actually +/// interleaves them. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn concurrent_store_appends_do_not_lose_messages() { + let dir = tempfile::tempdir().unwrap(); + let history = Arc::new(StoreChatHistory::new(FileStore::new(dir.path()))); + + const N: usize = 24; + let mut handles = Vec::new(); + for i in 0..N { + let history = history.clone(); + handles.push(tokio::spawn(async move { + history + .append("thread", Message::user(format!("m{i}"))) + .await + .unwrap(); + })); + } + for h in handles { + h.await.unwrap(); + } + let messages = history.messages("thread").await.unwrap(); + assert_eq!( + messages.len(), + N, + "every concurrent append must survive; got {} of {N}", + messages.len() + ); +} + +// ── SESS-10: JsonlAppendStore offsets ──────────────────────────────────────── + +/// Two store instances over one directory must keep handing out fresh offsets. +/// The old per-instance counter learned the length once, so a second instance +/// re-issued offsets the first had already used. +#[tokio::test] +async fn offsets_stay_unique_across_store_instances() { + let dir = tempfile::tempdir().unwrap(); + let first = JsonlAppendStore::new(dir.path()); + let second = JsonlAppendStore::new(dir.path()); + + let mut offsets = Vec::new(); + offsets.push(first.append("s", serde_json::json!(0)).await.unwrap()); + offsets.push(second.append("s", serde_json::json!(1)).await.unwrap()); + offsets.push(first.append("s", serde_json::json!(2)).await.unwrap()); + offsets.push(second.append("s", serde_json::json!(3)).await.unwrap()); + + assert_eq!( + offsets, + vec![0, 1, 2, 3], + "a second instance continues the stream instead of restarting it" + ); + assert_eq!(second.len("s").await.unwrap(), 4); +} + +/// `read_from` resolves by offset, not by position — the documented contract, +/// and what the in-memory backend already did. +#[tokio::test] +async fn read_from_resolves_by_offset_not_position() { + let dir = tempfile::tempdir().unwrap(); + let store = JsonlAppendStore::new(dir.path()); + for i in 0..4 { + store.append("s", serde_json::json!(i)).await.unwrap(); + } + let window = store.read_from("s", 2).await.unwrap(); + assert_eq!( + window.iter().map(|(o, _)| *o).collect::>(), + vec![2, 3], + "the window and the offsets labelling it must agree" + ); + assert_eq!(window[0].1, serde_json::json!(2)); +} + +/// A stream whose offsets are sparse (a hand-written or partially pruned log) +/// must still window correctly. Positional `.skip` gets this wrong. +#[tokio::test] +async fn read_from_handles_a_sparse_offset_sequence() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("sparse.jsonl"); + std::fs::create_dir_all(dir.path()).unwrap(); + std::fs::write( + &path, + "{\"offset\":10,\"value\":\"a\",\"created_at_ms\":0}\n\ + {\"offset\":11,\"value\":\"b\",\"created_at_ms\":0}\n", + ) + .unwrap(); + let store = JsonlAppendStore::new(dir.path()); + let window = store.read_from("sparse", 11).await.unwrap(); + assert_eq!( + window.iter().map(|(o, _)| *o).collect::>(), + vec![11], + "offset 11 selects the entry labelled 11, not the entry at index 11" + ); + // A new append continues from the file's own numbering. + assert_eq!( + store + .append("sparse", serde_json::json!("c")) + .await + .unwrap(), + 12 + ); +} diff --git a/tests/provider_local_wire.rs b/tests/provider_local_wire.rs new file mode 100644 index 0000000..17a6e14 --- /dev/null +++ b/tests/provider_local_wire.rs @@ -0,0 +1,984 @@ +//! Regression tests for the local-runtime provider fixes, asserted **at the +//! wire**. +//! +//! # Why these go through a real socket +//! +//! The pre-existing local tests assert that a request body *contains* a field — +//! for example that `{"options": {"num_ctx": 8192}}` appears in the JSON. That +//! shape of assertion is what let LOCAL-2 survive: `num_ctx` was present in the +//! body of a request sent to `POST /chat/completions`, an endpoint that does not +//! read it. The field was there; the behaviour was not. +//! +//! So these tests assert the thing that actually matters — **which URL the bytes +//! went to, and what the adapter did with the answer** — by standing up a +//! throwaway HTTP server on a loopback port and inspecting what arrives. +//! +//! The server is hand-rolled on `std::net` rather than a mock crate because the +//! dev-dependency `tokio` is built without the `net` feature, and adding an HTTP +//! mocking dependency to exercise four endpoints is not a trade worth making. + +use std::io::{BufRead, BufReader, Read, Write}; +use std::net::{TcpListener, TcpStream}; +use std::sync::{Arc, Mutex}; + +use serde_json::{Value, json}; + +use tinyagents::harness::message::Message; +use tinyagents::harness::model::{ + ChatModel, ModelRequest, ReasoningEffort, ResponseFormat, ToolChoice, +}; +use tinyagents::harness::providers::openai::OpenAiModel; +use tinyagents::harness::tool::ToolSchema; + +// --------------------------------------------------------------------------- +// Minimal recording HTTP server +// --------------------------------------------------------------------------- + +/// One request the server received. +#[derive(Clone, Debug)] +struct Recorded { + method: String, + path: String, + body: Value, +} + +/// A canned reply. +#[derive(Clone)] +struct Canned { + status: u16, + body: String, +} + +impl Canned { + fn ok(body: Value) -> Self { + Self { + status: 200, + body: body.to_string(), + } + } + + fn error(status: u16, body: Value) -> Self { + Self { + status, + body: body.to_string(), + } + } +} + +/// A loopback HTTP server that replies from a scripted queue and records every +/// request it saw. +struct MockServer { + base_url: String, + seen: Arc>>, +} + +impl MockServer { + /// Starts a server that answers each successive request with the next + /// scripted reply, repeating the last one once the script is exhausted. + fn start(script: Vec) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").expect("bind a loopback port"); + let port = listener.local_addr().expect("local addr").port(); + let seen = Arc::new(Mutex::new(Vec::new())); + let recorder = Arc::clone(&seen); + + std::thread::spawn(move || { + for (index, stream) in listener.incoming().enumerate() { + let Ok(stream) = stream else { break }; + let reply = script + .get(index) + .or_else(|| script.last()) + .cloned() + .unwrap_or_else(|| Canned::ok(json!({}))); + serve_one(stream, &reply, &recorder); + } + }); + + Self { + base_url: format!("http://127.0.0.1:{port}"), + seen, + } + } + + fn requests(&self) -> Vec { + self.seen.lock().expect("recorder lock").clone() + } + + /// The single request sent to `path`, panicking with the full request log + /// when there is not exactly one — a far more useful failure than + /// `Option::unwrap`. + fn request_to(&self, path: &str) -> Recorded { + let all = self.requests(); + let matched: Vec<&Recorded> = all.iter().filter(|r| r.path.starts_with(path)).collect(); + assert_eq!( + matched.len(), + 1, + "expected exactly one request to {path}; saw {:?}", + all.iter() + .map(|r| format!("{} {}", r.method, r.path)) + .collect::>() + ); + matched[0].clone() + } + + fn paths(&self) -> Vec { + self.requests().into_iter().map(|r| r.path).collect() + } +} + +/// Reads one HTTP/1.1 request off `stream`, records it, then writes `reply`. +/// +/// **Recording happens before the reply is written**, and that ordering is +/// load-bearing rather than incidental. Recording afterwards is a race the +/// client always wins: `invoke` can receive the full response body and return to +/// the test while the server thread has not yet reached its `push`, so an +/// assertion made immediately after the call sees an empty log. That is exactly +/// how two of these tests came out flaky under a loaded parallel test run and +/// green when the file ran alone. +fn serve_one(mut stream: TcpStream, reply: &Canned, recorder: &Arc>>) { + let Some(record) = read_request(&mut stream) else { + return; + }; + recorder.lock().expect("recorder lock").push(record); + + let response = format!( + "HTTP/1.1 {} OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}", + reply.status, + reply.body.len(), + reply.body + ); + let _ = stream.write_all(response.as_bytes()); + let _ = stream.flush(); +} + +/// Reads one HTTP/1.1 request off `stream`. Returns `None` for a malformed one. +fn read_request(stream: &mut TcpStream) -> Option { + let mut reader = BufReader::new(stream.try_clone().ok()?); + + let mut request_line = String::new(); + reader.read_line(&mut request_line).ok()?; + let mut parts = request_line.split_whitespace(); + let method = parts.next()?.to_string(); + let path = parts.next()?.to_string(); + + let mut content_length = 0usize; + loop { + let mut header = String::new(); + if reader.read_line(&mut header).ok()? == 0 { + break; + } + if header.trim().is_empty() { + break; + } + if let Some(value) = header + .to_ascii_lowercase() + .strip_prefix("content-length:") + .map(str::trim) + .and_then(|v| v.parse::().ok()) + { + content_length = value; + } + } + + let mut raw = vec![0u8; content_length]; + if content_length > 0 { + reader.read_exact(&mut raw).ok()?; + } + let body = serde_json::from_slice::(&raw).unwrap_or(Value::Null); + + Some(Recorded { method, path, body }) +} + +/// A minimal, valid Chat Completions reply. +fn chat_reply(text: &str) -> Value { + json!({ + "id": "chatcmpl-test", + "choices": [{ + "message": { "role": "assistant", "content": text }, + "finish_reason": "stop" + }], + "usage": { "prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7 } + }) +} + +fn user(text: &str) -> ModelRequest { + ModelRequest::new(vec![Message::user(text)]) +} + +// --------------------------------------------------------------------------- +// LOCAL-2 — `num_ctx` must reach an endpoint that reads it +// --------------------------------------------------------------------------- + +/// Before the fix this was unreachable: `num_ctx` was flattened onto the +/// `/chat/completions` body, where Ollama's compatibility layer ignores it, and +/// the only tests asserted the field's *presence in that body*. The behaviour +/// under test is that the value goes to `/api/chat`, the endpoint that reads it. +#[tokio::test] +async fn num_ctx_and_keep_alive_reach_the_native_ollama_endpoint() { + let server = MockServer::start(vec![Canned::ok(json!({ "model": "llama3.2" }))]); + + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2") + .expect("valid local URL") + .with_local_num_ctx(8192) + .with_keep_alive("30m"); + + model.warm_up().await.expect("warm-up succeeds"); + + let request = server.request_to("/api/chat"); + assert_eq!(request.method, "POST"); + assert_eq!( + request.path, "/api/chat", + "num_ctx is an /api/chat field; sending it to /chat/completions is how it went nowhere" + ); + assert_eq!(request.body["options"]["num_ctx"], json!(8192)); + assert_eq!(request.body["keep_alive"], json!("30m")); + // An empty `messages` array is Ollama's documented load request. + assert_eq!(request.body["messages"], json!([])); +} + +/// `with_local_num_ctx` must also move the *advertised* window. A window you +/// requested and a window you advertise being different numbers is precisely the +/// LOCAL-1 failure: compaction is gated on the advertised one. +#[test] +fn requesting_a_context_window_also_advertises_it() { + let model = OpenAiModel::ollama().with_local_num_ctx(8192); + let profile = >::profile(&model).expect("local profile"); + assert_eq!(profile.max_input_tokens, Some(8192)); +} + +/// A runtime with no native API must not attempt one. +#[tokio::test] +async fn warm_up_is_a_no_op_for_a_runtime_without_a_native_api() { + let server = MockServer::start(vec![Canned::ok(json!({}))]); + let model = OpenAiModel::llama_cpp(&server.base_url, "local-model").expect("valid URL"); + model.warm_up().await.expect("warm-up is a no-op"); + assert!( + server.requests().is_empty(), + "llama.cpp-server has no /api/chat; inventing one would 404 every startup" + ); +} + +// --------------------------------------------------------------------------- +// LOCAL-1 / C10 — the real context window comes from the server +// --------------------------------------------------------------------------- + +/// `llama3.2` matches the generic hint table's `("llama3", Substring, 128_000)` +/// entry. Ollama's real default `num_ctx` is 2048. The probe must replace the +/// guess with what the server reports, and the un-probed default must be `None` +/// rather than the guess. +#[tokio::test] +async fn probing_replaces_the_model_id_guess_with_the_servers_own_window() { + let server = MockServer::start(vec![Canned::ok(json!({ + "model_info": { "general.architecture": "llama", "llama.context_length": 8192 }, + "capabilities": ["completion", "tools"] + }))]); + + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + // Un-probed: unknown, not 128 000. + assert_eq!( + >::profile(&model) + .expect("profile") + .max_input_tokens, + None, + "an invented window is worse than None: compaction is gated on it" + ); + + let model = model.probed().await.expect("probe succeeds"); + let profile = >::profile(&model).expect("profile"); + assert_eq!(profile.max_input_tokens, Some(8192)); + assert!( + profile.tool_calling, + "the server reported the `tools` capability" + ); + assert!( + !profile.modalities.image_in, + "no `vision` capability reported" + ); + + assert_eq!( + server.request_to("/api/show").body["model"], + json!("llama3.2") + ); +} + +/// A server without the probe endpoint must degrade to "learned nothing", not +/// fail the caller's startup. +#[tokio::test] +async fn a_probe_against_an_older_server_is_not_an_error() { + let server = MockServer::start(vec![Canned::error(404, json!({ "error": "not found" }))]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + let probe = model + .probe_local_profile() + .await + .expect("a 404 is not fatal"); + assert!(probe.is_empty()); +} + +/// Probing a hosted endpoint is a programming error, and says so. +#[tokio::test] +async fn probing_a_hosted_endpoint_is_rejected() { + let error = OpenAiModel::new("k") + .probe_local_profile() + .await + .expect_err("hosted OpenAI is not a local runtime"); + assert!(error.to_string().contains("local runtime"), "{error}"); +} + +// --------------------------------------------------------------------------- +// LOCAL-3 / C11 — native tools are sent, then degraded only on evidence +// --------------------------------------------------------------------------- + +/// Native tools used to be hard-disabled for every local runtime, so every call +/// took the prompt-guided branch and injected the protocol plus each tool's JSON +/// Schema into the system prompt. They must now go on the wire. +#[tokio::test] +async fn a_local_runtime_sends_native_tools() { + let server = MockServer::start(vec![Canned::ok(chat_reply("done"))]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + let request = user("hi").with_tools(vec![ToolSchema::new( + "get_weather", + "look up weather", + json!({"type": "object", "properties": {"city": {"type": "string"}}}), + )]); + let _ = ChatModel::<()>::invoke(&model, &(), request).await; + + let sent = server.request_to("/v1/chat/completions"); + let tools = sent.body["tools"] + .as_array() + .expect("native tools on the wire"); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0]["function"]["name"], json!("get_weather")); + + // The prompt-guided protocol block must be absent — it is what the tiny real + // context window could not afford. + let system = sent.body["messages"][0]["content"].as_str().unwrap_or(""); + assert!( + !system.contains(""), + "prompt-guided protocol leaked into a native-tools request: {system}" + ); +} + +/// The auto-degrade half: a 400 that implicates `tools` must retry +/// prompt-guided **and latch**, so the rejection is paid once per process rather +/// than being assumed forever up front. +#[tokio::test] +async fn a_tools_rejection_degrades_once_and_then_latches() { + let server = MockServer::start(vec![ + Canned::error( + 400, + json!({ "error": { "message": "registry.ollama.ai/library/gemma3 does not support tools" } }), + ), + Canned::ok(chat_reply("done")), + Canned::ok(chat_reply("again")), + ]); + let model = OpenAiModel::ollama_at(&server.base_url, "gemma3").expect("valid local URL"); + let tools = vec![ToolSchema::new("t", "d", json!({"type": "object"}))]; + + let first = user("hi").with_tools(tools.clone()); + ChatModel::<()>::invoke(&model, &(), first) + .await + .expect("the degraded retry succeeds"); + + let seen = server.requests(); + assert_eq!( + seen.len(), + 2, + "one rejected attempt, then one degraded retry" + ); + assert!( + seen[0].body["tools"] + .as_array() + .is_some_and(|t| !t.is_empty()), + "the first attempt must try native tools" + ); + assert!( + seen[1] + .body + .get("tools") + .is_none_or(|t| t.as_array().is_none_or(Vec::is_empty)), + "the retry must drop native tools" + ); + + // Latched: the next call goes straight to the degraded shape. + let second = user("again").with_tools(tools); + ChatModel::<()>::invoke(&model, &(), second) + .await + .expect("second call succeeds"); + let seen = server.requests(); + assert_eq!( + seen.len(), + 3, + "the latch must skip the doomed baseline attempt" + ); + assert!( + seen[2] + .body + .get("tools") + .is_none_or(|t| t.as_array().is_none_or(Vec::is_empty)), + "the latch was not applied to the following call" + ); +} + +// --------------------------------------------------------------------------- +// LOCAL-6 — self-hosted servers other than Ollama/LM Studio get local treatment +// --------------------------------------------------------------------------- + +#[test] +fn llama_cpp_and_vllm_are_treated_as_local_runtimes() { + for model in [ + OpenAiModel::llama_cpp("127.0.0.1:8080", "local-model").expect("valid URL"), + OpenAiModel::vllm("127.0.0.1:8000", "", "meta-llama/Llama-3.3-70B").expect("valid URL"), + ] { + // `/v1` normalisation, which the hosted `Compatible` path did not do. + assert!( + model.base_url().ends_with("/v1"), + "unexpected base url: {}", + model.base_url() + ); + assert!(model.local_runtime_kind().is_some()); + + let profile = >::profile(&model).expect("profile"); + // No invented window even for `Llama-3.3-70B`, which the hint table + // matches through `("llama-3", Substring, 128_000)`. + assert_eq!(profile.max_input_tokens, None); + } +} + +/// The degrade knobs the local presets exist for must be pre-set, so the first +/// call does not pay a guaranteed 400 to rediscover a documented rejection. +#[tokio::test] +async fn local_presets_pre_set_the_shapes_local_servers_reject() { + let server = MockServer::start(vec![Canned::ok(chat_reply("{}"))]); + let model = OpenAiModel::llama_cpp(&server.base_url, "local-model").expect("valid URL"); + + let request = user("hi") + .with_tools(vec![ToolSchema::new("t", "d", json!({"type": "object"}))]) + .with_tool_choice(ToolChoice::Tool("t".to_string())) + .with_response_format(ResponseFormat::JsonObject); + let _ = ChatModel::<()>::invoke(&model, &(), request).await; + + let sent = server.request_to("/v1/chat/completions"); + assert_eq!( + sent.body["tool_choice"], + json!("required"), + "the named tool_choice object must already be degraded on the first call" + ); + assert_eq!( + sent.body["response_format"]["type"], + json!("json_schema"), + "json_object must already be degraded on the first call" + ); +} + +// --------------------------------------------------------------------------- +// C12 — an actionable missing-model error +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn a_missing_local_model_names_the_fix() { + let server = MockServer::start(vec![Canned::error( + 404, + json!({ "error": { "message": "model 'llama3.2' not found" } }), + )]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + let error = ChatModel::<()>::invoke(&model, &(), user("hi")) + .await + .expect_err("a 404 fails the call"); + let rendered = error.to_string(); + assert!( + rendered.contains("ollama pull llama3.2"), + "an opaque 404 tells the operator nothing: {rendered}" + ); +} + +#[tokio::test] +async fn validate_model_lists_what_the_server_actually_serves() { + let server = MockServer::start(vec![Canned::ok(json!({ + "object": "list", + "data": [{ "id": "qwen3:8b" }, { "id": "bge-m3" }] + }))]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + let error = model + .validate_model() + .await + .expect_err("llama3.2 is not served"); + let rendered = error.to_string(); + assert!(rendered.contains("ollama pull llama3.2"), "{rendered}"); + assert!(rendered.contains("qwen3:8b"), "{rendered}"); +} + +// --------------------------------------------------------------------------- +// REASON-4 — `strict` is a knob, and the schema is sanitized when it is on +// --------------------------------------------------------------------------- + +/// `strict: true` was hardcoded and paired with the caller's raw schema, so a +/// valid JSON Schema 400d. Local runtimes must default it off entirely. +#[tokio::test] +async fn a_local_runtime_does_not_send_strict_structured_output() { + let server = MockServer::start(vec![Canned::ok(chat_reply("{}"))]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + let request = user("hi").with_response_format(ResponseFormat::json_schema( + "answer", + json!({"type": "object", "properties": {"a": {"type": "string"}}}), + )); + let _ = ChatModel::<()>::invoke(&model, &(), request).await; + + let sent = server.request_to("/v1/chat/completions"); + assert_eq!( + sent.body["response_format"]["json_schema"]["strict"], + json!(false) + ); +} + +/// On hosted OpenAI strict stays on, but the schema is now sanitized to satisfy +/// it: every property in `required`, `additionalProperties: false`. +#[tokio::test] +async fn hosted_strict_mode_sanitizes_the_schema_it_sends() { + let server = MockServer::start(vec![Canned::ok(chat_reply("{}"))]); + let model = OpenAiModel::new("k").with_base_url(format!("{}/v1", server.base_url)); + + let request = user("hi").with_response_format(ResponseFormat::json_schema( + "answer", + json!({ + "type": "object", + "properties": { "a": {"type": "string"}, "b": {"type": "number"} } + }), + )); + let _ = ChatModel::<()>::invoke(&model, &(), request).await; + + let schema = &server.request_to("/v1/chat/completions").body["response_format"]["json_schema"]; + assert_eq!(schema["strict"], json!(true)); + assert_eq!( + schema["schema"]["additionalProperties"], + json!(false), + "strict mode rejects an object without this" + ); + let required = schema["schema"]["required"] + .as_array() + .expect("strict mode requires every property to be listed") + .iter() + .filter_map(Value::as_str) + .collect::>(); + assert!( + required.contains(&"a") && required.contains(&"b"), + "{required:?}" + ); +} + +/// A `JsonSchema` request had no degradation path at all — only `JsonObject` was +/// considered — so a 400 on it was terminal. +#[tokio::test] +async fn a_strict_schema_rejection_retries_without_strict() { + let server = MockServer::start(vec![ + Canned::error( + 400, + json!({ "error": { "message": "response_format: 'strict' is not supported" } }), + ), + Canned::ok(chat_reply("{}")), + ]); + let model = OpenAiModel::new("k").with_base_url(format!("{}/v1", server.base_url)); + + let request = user("hi").with_response_format(ResponseFormat::json_schema( + "answer", + json!({"type": "object", "properties": {"a": {"type": "string"}}}), + )); + ChatModel::<()>::invoke(&model, &(), request) + .await + .expect("the degraded retry succeeds"); + + let seen = server.requests(); + assert_eq!(seen.len(), 2, "a JsonSchema 400 must have a retry path"); + assert_eq!( + seen[0].body["response_format"]["json_schema"]["strict"], + json!(true) + ); + assert_eq!( + seen[1].body["response_format"]["json_schema"]["strict"], + json!(false) + ); +} + +// --------------------------------------------------------------------------- +// REASON-6 — the gpt-5 family +// --------------------------------------------------------------------------- + +/// gpt-5 was routed to `max_tokens`, which OpenAI rejects outright. +#[tokio::test] +async fn gpt5_sends_max_completion_tokens_not_max_tokens() { + let server = MockServer::start(vec![Canned::ok(chat_reply("hi"))]); + let model = OpenAiModel::new("k") + .with_base_url(format!("{}/v1", server.base_url)) + .with_model("gpt-5-mini"); + + let _ = ChatModel::<()>::invoke(&model, &(), user("hi").with_max_tokens(64)).await; + + let sent = server.request_to("/v1/chat/completions"); + assert_eq!(sent.body["max_completion_tokens"], json!(64)); + assert!( + sent.body.get("max_tokens").is_none(), + "OpenAI rejects max_tokens for gpt-5: {}", + sent.body + ); +} + +#[test] +fn gpt5_profiles_as_a_reasoning_model_with_native_structured_output() { + let model = OpenAiModel::new("k").with_model("gpt-5"); + let profile = >::profile(&model).expect("profile"); + assert!( + profile.reasoning, + "CapabilitySet{{reasoning:true}} used to reject gpt-5" + ); + assert!(profile.native_structured_output); + assert!(profile.reasoning_effort); + assert_eq!(profile.max_input_tokens, Some(400_000)); +} + +// --------------------------------------------------------------------------- +// C13 — provider-neutral reasoning config +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn reasoning_effort_lowers_onto_the_chat_completions_field() { + let server = MockServer::start(vec![Canned::ok(chat_reply("hi"))]); + let model = OpenAiModel::new("k") + .with_base_url(format!("{}/v1", server.base_url)) + .with_model("gpt-5"); + + let _ = ChatModel::<()>::invoke( + &model, + &(), + user("hi").with_reasoning_effort(ReasoningEffort::High), + ) + .await; + + assert_eq!( + server.request_to("/v1/chat/completions").body["reasoning_effort"], + json!("high") + ); +} + +/// `provider_options` stays the escape hatch and wins, rather than the key +/// landing on the wire twice. +#[tokio::test] +async fn provider_options_win_over_the_typed_reasoning_field() { + let server = MockServer::start(vec![Canned::ok(chat_reply("hi"))]); + let model = OpenAiModel::new("k") + .with_base_url(format!("{}/v1", server.base_url)) + .with_model("gpt-5"); + + let request = user("hi") + .with_reasoning_effort(ReasoningEffort::High) + .with_provider_option("reasoning_effort", json!("minimal")); + let _ = ChatModel::<()>::invoke(&model, &(), request).await; + + assert_eq!( + server.request_to("/v1/chat/completions").body["reasoning_effort"], + json!("minimal") + ); +} + +// --------------------------------------------------------------------------- +// CACHE-6b — cache tokens are accounted for +// --------------------------------------------------------------------------- + +/// `Usage::cache_creation_tokens` is summed and priced, and no provider ever set +/// it, so cache writes were billed as ordinary input everywhere. +#[tokio::test] +async fn cache_write_tokens_are_recorded() { + let server = MockServer::start(vec![Canned::ok(json!({ + "id": "c1", + "choices": [{ "message": { "role": "assistant", "content": "hi" }, "finish_reason": "stop" }], + "usage": { + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "prompt_tokens_details": { "cached_tokens": 40, "cache_write_tokens": 25 } + } + }))]); + let model = OpenAiModel::new("k").with_base_url(format!("{}/v1", server.base_url)); + + let response = ChatModel::<()>::invoke(&model, &(), user("hi")) + .await + .expect("call succeeds"); + let usage = response.usage.expect("usage reported"); + assert_eq!(usage.cache_read_tokens, 40); + assert_eq!(usage.cache_creation_tokens, 25); + assert_eq!( + usage.input_tokens, 100, + "OpenAI includes cache tokens in the input total" + ); +} + +// --------------------------------------------------------------------------- +// TOOL-2b — synthetic tool-call ids are unique across a run +// --------------------------------------------------------------------------- + +/// A build that omits `id` used to emit `tool-0` on every turn, so one +/// transcript held several assistant messages declaring the same id. +#[tokio::test] +async fn id_less_tool_calls_get_distinct_ids_on_successive_turns() { + let reply = json!({ + "id": "c1", + "choices": [{ + "message": { + "role": "assistant", + "tool_calls": [{ "function": { "name": "ping", "arguments": "{}" } }] + }, + "finish_reason": "tool_calls" + }] + }); + let server = MockServer::start(vec![Canned::ok(reply.clone()), Canned::ok(reply)]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + let mut ids = Vec::new(); + for _ in 0..2 { + let response = ChatModel::<()>::invoke(&model, &(), user("hi")) + .await + .expect("call succeeds"); + ids.push(response.tool_calls()[0].id.clone()); + } + + assert_ne!( + ids[0], ids[1], + "two turns sharing a synthetic id is an unresolvable pairing" + ); + for id in &ids { + assert!(id.starts_with("tacall-"), "{id}"); + // The prompt-guided protocol mints `ptc_{seq}_{slot}`; the schemes must + // be unmistakably disjoint. + assert!(!id.starts_with("ptc_"), "{id}"); + } +} + +/// Gateways emit ids other providers reject on the way back +/// (`functions.write_todos:0`). Normalizing at the provider boundary keeps the +/// id and its paired result consistent, and must be deterministic. +#[tokio::test] +async fn non_conforming_provider_ids_are_normalized_deterministically() { + let reply = json!({ + "id": "c1", + "choices": [{ + "message": { + "role": "assistant", + "tool_calls": [{ + "id": "functions.write_todos:0", + "function": { "name": "ping", "arguments": "{}" } + }] + }, + "finish_reason": "tool_calls" + }] + }); + let server = MockServer::start(vec![Canned::ok(reply.clone()), Canned::ok(reply)]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + let mut ids = Vec::new(); + for _ in 0..2 { + let response = ChatModel::<()>::invoke(&model, &(), user("hi")) + .await + .expect("call succeeds"); + ids.push(response.tool_calls()[0].id.clone()); + } + + assert_eq!(ids[0], ids[1], "normalization must be deterministic"); + assert!( + ids[0] + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-'), + "unexpected id: {}", + ids[0] + ); +} + +// --------------------------------------------------------------------------- +// C15 — context overflow is classified +// --------------------------------------------------------------------------- + +#[tokio::test] +async fn a_context_overflow_is_classified_with_a_stable_code() { + let server = MockServer::start(vec![Canned::error( + 400, + json!({ "error": { "message": "This model's maximum context length is 4096 tokens" } }), + )]); + let model = OpenAiModel::ollama_at(&server.base_url, "llama3.2").expect("valid local URL"); + + let error = ChatModel::<()>::invoke(&model, &(), user("hi")) + .await + .expect_err("a 400 fails the call"); + let tinyagents::TinyAgentsError::ContextOverflow { + provider, + model, + message, + } = error + else { + panic!("expected a typed context-overflow error"); + }; + assert_eq!(provider, "ollama"); + assert_eq!(model.as_deref(), Some("llama3.2")); + assert!(message.contains("maximum context length")); +} + +// --------------------------------------------------------------------------- +// REASON-7 — the Responses path carries the whole request +// --------------------------------------------------------------------------- + +/// The body used to be `{model, input, instructions, stream, store, +/// max_output_tokens}` and everything else was silently dropped. +#[tokio::test] +async fn the_responses_path_no_longer_drops_the_request() { + let server = MockServer::start(vec![Canned::ok(json!({ "output_text": "hi" }))]); + let model = OpenAiModel::new("k") + .with_base_url(format!("{}/v1", server.base_url)) + .with_model("gpt-5") + .with_responses_api_primary(); + + let request = user("hi") + .with_tools(vec![ToolSchema::new("t", "d", json!({"type": "object"}))]) + .with_response_format(ResponseFormat::json_schema( + "answer", + json!({"type": "object"}), + )) + .with_temperature(0.3) + .with_top_p(0.9) + .with_seed(7) + .with_stop_sequences(["STOP"]) + .with_continuation_id("resp_123") + .with_reasoning_effort(ReasoningEffort::Low); + let _ = ChatModel::<()>::invoke(&model, &(), request).await; + + let body = server.request_to("/v1/responses").body; + assert_eq!(body["tools"][0]["name"], json!("t")); + assert!(body.get("tool_choice").is_some()); + assert_eq!(body["text"]["format"]["type"], json!("json_schema")); + assert_eq!(body["temperature"], json!(0.3)); + assert_eq!(body["top_p"], json!(0.9)); + assert_eq!(body["seed"], json!(7)); + assert_eq!(body["stop"], json!(["STOP"])); + // `continuation_id` had a builder and no reader anywhere in the crate. + assert_eq!(body["previous_response_id"], json!("resp_123")); + assert_eq!(body["reasoning"]["effort"], json!("low")); + // With `store: false`, reasoning is droppable unless it carries + // `encrypted_content`, which only arrives when explicitly requested. + assert_eq!(body["store"], json!(false)); + assert_eq!(body["include"], json!(["reasoning.encrypted_content"])); +} + +/// Reasoning items, the encrypted payload, and the usage breakdowns were all +/// unread on this path. +#[tokio::test] +async fn the_responses_path_reads_reasoning_and_cache_usage() { + let server = MockServer::start(vec![Canned::ok(json!({ + "output": [ + { + "type": "reasoning", + "summary": [{ "type": "summary_text", "text": "thinking about it" }], + "encrypted_content": "enc-abc" + }, + { + "type": "message", + "content": [{ "type": "output_text", "text": "the answer" }] + } + ], + "usage": { + "input_tokens": 100, + "output_tokens": 20, + "input_tokens_details": { "cached_tokens": 30, "cache_write_tokens": 10 }, + "output_tokens_details": { "reasoning_tokens": 12 } + } + }))]); + let model = OpenAiModel::new("k") + .with_base_url(format!("{}/v1", server.base_url)) + .with_responses_api_primary(); + + let response = ChatModel::<()>::invoke(&model, &(), user("hi")) + .await + .expect("call succeeds"); + + assert_eq!( + response.text(), + "the answer", + "reasoning must not leak into the answer" + ); + + let usage = response.usage.expect("usage reported"); + assert_eq!( + usage.cache_read_tokens, 30, + "every cached token was billed at full rate" + ); + assert_eq!(usage.cache_creation_tokens, 10); + assert_eq!(usage.reasoning_tokens, 12); + + let thinking = response + .message + .content + .iter() + .find_map(|block| match block { + tinyagents::harness::message::ContentBlock::Thinking { text, signature } => { + Some((text.clone(), signature.clone())) + } + _ => None, + }) + .expect("reasoning surfaces as a Thinking block"); + assert_eq!(thinking.0, "thinking about it"); + assert_eq!( + thinking.1.as_deref(), + Some("enc-abc"), + "the encrypted payload is what makes reasoning replayable under store:false" + ); +} + +/// Tool results folded into anonymous assistant turns, erasing which call they +/// answered. +#[tokio::test] +async fn tool_results_keep_their_call_identity_on_the_responses_path() { + let server = MockServer::start(vec![Canned::ok(json!({ "output_text": "ok" }))]); + let model = OpenAiModel::new("k") + .with_base_url(format!("{}/v1", server.base_url)) + .with_responses_api_primary(); + + let request = ModelRequest::new(vec![ + Message::user("what is the weather"), + Message::tool("call_abc", "sunny"), + ]); + let _ = ChatModel::<()>::invoke(&model, &(), request).await; + + let input = server.request_to("/v1/responses").body["input"].clone(); + let rendered = input.to_string(); + assert!( + rendered.contains("call_abc"), + "the tool call id must survive into the input: {rendered}" + ); + let tool_item = input + .as_array() + .expect("input items") + .last() + .expect("last item"); + assert_eq!( + tool_item["role"], + json!("user"), + "a tool result is not the assistant asserting a fact" + ); +} + +// --------------------------------------------------------------------------- +// Cross-cutting: probing never happens on its own +// --------------------------------------------------------------------------- + +/// Construction must stay free of network I/O — a constructor that blocks on a +/// round trip is unusable where this crate is embedded. +#[test] +fn construction_performs_no_network_io() { + let server = MockServer::start(vec![Canned::ok(json!({}))]); + let _ = OpenAiModel::ollama_at(&server.base_url, "llama3.2") + .expect("valid local URL") + .with_local_num_ctx(8192) + .with_keep_alive("30m"); + assert!( + server.paths().is_empty(), + "probing and warm-up must be opt-in, not a side effect of construction" + ); +} diff --git a/tests/runtime_primitives_resilience.rs b/tests/runtime_primitives_resilience.rs new file mode 100644 index 0000000..d255a97 --- /dev/null +++ b/tests/runtime_primitives_resilience.rs @@ -0,0 +1,371 @@ +//! Public-API regression tests for the harness runtime primitives. +//! +//! Each test corresponds to a specific defect and is written so it **fails +//! against the pre-fix code**: +//! +//! - LOOP-2 — enabling jitter zeroed the backoff, so nothing slept. +//! - LOCAL-4 — `backoff_sleep` defaulted off, so retries fired back-to-back. +//! - LOOP-5 — a `Retry-After` was parsed but never honored. +//! - LOOP-5b — `is_retryable` retried every `Model(_)` and had no extension point. +//! - LOOP-1 — reconciling two limit sources could *widen* a caller's cap. +//! - LOOP-9b — a reached cap was always a hard error. +//! - LOOP-6 — `Started` events had no terminal partner on error paths. +//! - LOOP-8 — steering applied part of a rejected batch, and a pause was unresumable. +//! - LOOP-9 — a panicking listener wedged the event sink forever. +//! - C6 — `StreamChunk` had no producer. + +use std::sync::Arc; +use std::time::Duration; + +use serde_json::json; + +use tinyagents::error::TinyAgentsError; +use tinyagents::harness::context::{RunConfig, RunContext}; +use tinyagents::harness::events::{ + AgentEvent, EventListener, EventRecord, EventSink, RecordingListener, +}; +use tinyagents::harness::ids::{CallId, RunId}; +use tinyagents::harness::limits::{LimitBehavior, LimitOutcome, LimitTracker, RunLimits}; +use tinyagents::harness::message::{Message, MessageDelta}; +use tinyagents::harness::no_progress::fingerprint_arguments; +use tinyagents::harness::retry::{RetryPolicy, is_retryable, retry_after_hint}; +use tinyagents::harness::steering::{ + SteeringCommand, SteeringCommandKind, SteeringHandle, SteeringOutcome, SteeringPolicy, + apply_pending_steering, +}; +use tinyagents::harness::stream::{StreamChunk, StreamMode, StreamSink, project_event}; + +// ── LOOP-2: jitter must never collapse the backoff ─────────────────────────── + +#[test] +fn loop2_enabling_jitter_does_not_disable_backoff() { + let policy = RetryPolicy::default().with_jitter(true); + for attempt in 0..6 { + for _ in 0..50 { + assert!( + policy.backoff_for_attempt(attempt) > Duration::ZERO, + "jitter zeroed the backoff at attempt {attempt}" + ); + } + } +} + +#[tokio::test(start_paused = true)] +async fn loop2_the_production_hardened_config_actually_sleeps() { + let policy = RetryPolicy::default() + .with_backoff_sleep(true) + .with_jitter(true); + let start = tokio::time::Instant::now(); + policy.sleep_backoff(1).await; + assert!( + start.elapsed() > Duration::ZERO, + "with_backoff_sleep(true).with_jitter(true) never slept" + ); +} + +// ── LOCAL-4: backoff sleeps by default ─────────────────────────────────────── + +#[test] +fn local4_backoff_sleep_is_on_by_default_and_opt_out() { + assert!(RetryPolicy::default().backoff_sleep); + assert!( + !RetryPolicy::default() + .with_backoff_sleep(false) + .backoff_sleep + ); +} + +// ── LOOP-5: Retry-After is honored ─────────────────────────────────────────── + +#[test] +fn loop5_retry_after_is_read_and_lengthens_the_wait() { + let policy = RetryPolicy::default(); + let rate_limited = TinyAgentsError::Model("429 Too Many Requests, Retry-After: 30".into()); + + assert_eq!( + retry_after_hint(&rate_limited), + Some(Duration::from_secs(30)) + ); + assert_eq!( + policy.backoff_for_error(0, &rate_limited), + Duration::from_secs(30), + "a 429 saying Retry-After: 30 was still retried after 200ms" + ); +} + +#[test] +fn loop5_a_retry_after_can_only_lengthen_never_shorten() { + let policy = RetryPolicy::default(); + let zero_hint = TinyAgentsError::Model("429 rate limited, Retry-After: 0".into()); + assert_eq!( + policy.backoff_for_error(2, &zero_hint), + policy.backoff_for_attempt(2) + ); +} + +#[test] +fn loop5_a_hostile_retry_after_is_clamped() { + let policy = RetryPolicy::default().with_max_retry_after_ms(1_000); + let absurd = TinyAgentsError::Model("429, Retry-After: 999999".into()); + assert_eq!( + policy.backoff_for_error(0, &absurd), + Duration::from_millis(1_000) + ); +} + +// ── LOOP-5b: classification and the retry_on extension point ───────────────── + +#[test] +fn loop5b_model_errors_consult_the_provider_failure_class() { + assert!(!is_retryable(&TinyAgentsError::Model( + "401 Unauthorized: invalid api key".into() + ))); + assert!(is_retryable(&TinyAgentsError::Model( + "503 Service Unavailable".into() + ))); +} + +#[test] +fn loop5b_retry_on_predicate_is_honored_by_should_retry_error() { + let policy = RetryPolicy::default() + .with_max_attempts(3) + .with_retry_on(Arc::new(|err: &TinyAgentsError| { + !matches!(err, TinyAgentsError::Tool(_)) + })); + + assert!(!policy.should_retry_error(0, &TinyAgentsError::Tool("no".into()))); + assert!(policy.should_retry_error(0, &TinyAgentsError::Model("yes".into()))); +} + +// ── LOOP-1: reconciling two limit sources is fail-closed ───────────────────── + +#[test] +fn loop1_an_explicit_cap_is_never_widened_by_a_second_source() { + let mut tracker = LimitTracker::new(RunLimits::default().with_max_model_calls(2)); + tracker.tighten_call_limits(25, 50); + assert_eq!(tracker.limits().max_model_calls, 2); + + tracker.record_model_call().unwrap(); + tracker.record_model_call().unwrap(); + assert!( + tracker.record_model_call().is_err(), + "the caller's cap of 2 ran to 25" + ); +} + +// ── LOOP-9b: exhaustion can stop cleanly instead of discarding the run ─────── + +#[test] +fn loop9b_stop_with_partial_returns_an_outcome_not_an_error() { + let mut tracker = LimitTracker::new( + RunLimits::default() + .with_max_model_calls(1) + .with_behavior(LimitBehavior::StopWithPartial), + ); + assert_eq!( + tracker.try_record_model_call().unwrap(), + LimitOutcome::Proceed + ); + assert!(matches!( + tracker.try_record_model_call().unwrap(), + LimitOutcome::Stop(_) + )); +} + +#[test] +fn loop9b_error_remains_the_default_behavior() { + let mut tracker = LimitTracker::new(RunLimits::default().with_max_model_calls(1)); + tracker.try_record_model_call().unwrap(); + assert!(tracker.try_record_model_call().is_err()); +} + +// ── LOOP-6: every Started has a terminal partner on the error path ─────────── + +#[test] +fn loop6_failure_variants_exist_for_tool_model_and_subagent() { + let events = [ + AgentEvent::ToolFailed { + call_id: CallId::new("c1"), + tool_name: "search".into(), + started_at_ms: Some(1), + duration_ms: Some(2), + error: "boom".into(), + }, + AgentEvent::ModelFailed { + call_id: CallId::new("c2"), + model: "gpt-4o".into(), + started_at_ms: Some(1), + attempts: Some(4), + error: "boom".into(), + }, + AgentEvent::SubAgentFailed { + name: "child".into(), + depth: 1, + error: "boom".into(), + }, + ]; + let kinds: Vec<&str> = events.iter().map(AgentEvent::kind).collect(); + assert_eq!( + kinds, + vec!["tool.failed", "model.failed", "subagent.failed"] + ); + + // They must survive a durable journal round trip. + for event in events { + let json = serde_json::to_value(&event).unwrap(); + let back: AgentEvent = serde_json::from_value(json).unwrap(); + assert_eq!(back, event); + } +} + +// ── LOOP-9: a panicking listener must not wedge the sink ───────────────────── + +struct OneShotBomb; + +impl EventListener for OneShotBomb { + fn on_event(&self, _record: &EventRecord) { + panic!("listener exploded"); + } +} + +#[test] +fn loop9_a_panicking_listener_does_not_stop_later_delivery() { + let sink = EventSink::new(); + let recorder = Arc::new(RecordingListener::new()); + sink.subscribe(Arc::new(OneShotBomb)); + sink.subscribe(recorder.clone()); + + let clone = sink.clone(); + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + clone.emit(AgentEvent::StateUpdate) + })); + + // Every later emit is swallowed by the panicking listener, but the *sink* + // must still dispatch — before the fix nothing was delivered ever again. + for _ in 0..3 { + let clone = sink.clone(); + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + clone.emit(AgentEvent::StateUpdate) + })); + } + assert_eq!( + recorder.len(), + 0, + "the bomb listener is first, so it aborts each dispatch" + ); + + // With the bomb removed, delivery resumes immediately — proof the sink was + // never latched into a permanently non-dispatching state. + let clean = EventSink::new(); + let recorder2 = Arc::new(RecordingListener::new()); + clean.subscribe(recorder2.clone()); + clean.subscribe(Arc::new(OneShotBomb)); + for _ in 0..3 { + let clone = clean.clone(); + let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + clone.emit(AgentEvent::StateUpdate) + })); + } + assert_eq!( + recorder2.len(), + 3, + "sink stayed wedged after a listener panic" + ); +} + +// ── LOOP-8: steering batches are atomic and pauses are resumable ───────────── + +#[test] +fn loop8_a_rejected_batch_applies_nothing() { + let handle = + SteeringHandle::new(SteeringPolicy::new().allow(SteeringCommandKind::InjectMessage)); + handle.send(SteeringCommand::InjectMessage(Message::user("first"))); + handle.send(SteeringCommand::Cancel); // not allowed + + let mut ctx: RunContext = RunContext::new(RunConfig::new("r"), ()).with_steering(handle); + let mut messages: Vec = Vec::new(); + + assert!(apply_pending_steering(&mut ctx, &mut messages).is_err()); + assert!( + messages.is_empty(), + "a command before the rejected one was still applied" + ); +} + +#[test] +fn loop8_a_pause_is_latched_and_resumable_across_checkpoints() { + let handle = SteeringHandle::allow_all(); + let mut ctx: RunContext = + RunContext::new(RunConfig::new("r"), ()).with_steering(handle.clone()); + let mut messages: Vec = Vec::new(); + + handle.send(SteeringCommand::PauseWith { + reason: "human review".into(), + }); + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Pause + ); + // Still paused at the next checkpoint, with an empty queue. + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Pause + ); + + // The state that makes a pause distinguishable from an empty answer. + let state = handle.pause_state().expect("a pause carries state"); + assert_eq!(state.reason.as_deref(), Some("human review")); + + // Resumable from a later batch — impossible before. + handle.send(SteeringCommand::Resume); + assert_eq!( + apply_pending_steering(&mut ctx, &mut messages).unwrap(), + SteeringOutcome::Continue + ); + assert!(handle.pause_state().is_none()); +} + +// ── C6: the stream projection ──────────────────────────────────────────────── + +#[test] +fn c6_events_project_onto_stream_chunks() { + let delta = AgentEvent::ModelDelta { + run_id: RunId::new("r1"), + call_id: CallId::new("c1"), + delta: MessageDelta::text("hi"), + }; + assert!(matches!( + project_event(&delta), + Some(StreamChunk::Message(_)) + )); + + // The Interrupt variant finally has a producer. + let interrupt = AgentEvent::ControlApplied { + control: "interrupt".into(), + detail: "needs approval".into(), + }; + assert!(matches!( + project_event(&interrupt), + Some(StreamChunk::Interrupt(_)) + )); + + // Mode filtering happens producer-side. + let sink = StreamSink::new([StreamMode::Interrupts]); + assert!(!sink.push_event(&delta)); + assert!(sink.push_event(&interrupt)); + assert_eq!(sink.drain().len(), 1); +} + +// ── C5: the no-progress tracker is drivable from a hook ────────────────────── + +#[test] +fn c5_argument_fingerprints_are_canonical() { + assert_eq!( + fingerprint_arguments(&json!({"a": 1, "b": 2})), + fingerprint_arguments(&json!({"b": 2, "a": 1})) + ); + assert_ne!( + fingerprint_arguments(&json!({"a": 1})), + fingerprint_arguments(&json!({"a": 2})) + ); +} diff --git a/tests/wave2_cache_key_scope.rs b/tests/wave2_cache_key_scope.rs new file mode 100644 index 0000000..aa2c465 --- /dev/null +++ b/tests/wave2_cache_key_scope.rs @@ -0,0 +1,306 @@ +//! Wave 2 — response-cache **key** regressions. +//! +//! Covers CACHE-1 (the key carried no provider/model identity), CACHE-4's key +//! half (`streaming` was not in the key), CACHE-7 (the envelope was +//! over-inclusive) and CACHE-8 (a non-array `messages`/`tools` was dropped +//! without being hashed). + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use async_trait::async_trait; + +use tinyagents::Result; +use tinyagents::harness::cache::{ + CachePolicy, InMemoryResponseCache, cache_key, credential_fingerprint, model_cache_identity, + scoped_cache_key, +}; +use tinyagents::harness::message::Message; +use tinyagents::harness::model::{ChatModel, ModelRequest, ModelResponse}; +use tinyagents::harness::runtime::AgentHarness; + +// ── Fixtures ───────────────────────────────────────────────────────────────── + +/// A model that always answers `answer` and reports a distinct cache identity. +struct IdentifiedModel { + identity: String, + answer: String, + calls: Arc, +} + +#[async_trait] +impl ChatModel<()> for IdentifiedModel { + fn cache_identity(&self) -> Option { + Some(self.identity.clone()) + } + + async fn invoke(&self, _state: &(), _request: ModelRequest) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(ModelResponse::assistant(self.answer.clone())) + } +} + +fn request(prompt: &str) -> ModelRequest { + ModelRequest::new(vec![Message::user(prompt)]) +} + +// ── CACHE-1 ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn shared_cache_does_not_cross_serve_between_two_providers() { + // Two harnesses — think "hosted" and "local" — sharing one cache, asked the + // same question. Before the identity was folded into the key, the second + // harness was served the first harness's answer: `request.model` is never + // set by the loop, and the endpoint/credentials live inside the + // `Arc`, so nothing in the key distinguished them. + let cache = Arc::new(InMemoryResponseCache::new()); + let hosted_calls = Arc::new(AtomicUsize::new(0)); + let local_calls = Arc::new(AtomicUsize::new(0)); + + let mut hosted: AgentHarness<()> = AgentHarness::new(); + hosted.register_model( + "chat", + Arc::new(IdentifiedModel { + identity: "openai|gpt-5|https://api.openai.com/v1||abc".to_string(), + answer: "hosted answer".to_string(), + calls: hosted_calls.clone(), + }), + ); + hosted.with_response_cache(cache.clone()); + + let mut local: AgentHarness<()> = AgentHarness::new(); + local.register_model( + "chat", + Arc::new(IdentifiedModel { + identity: "ollama|llama3.2|http://localhost:11434/v1||no-credential".to_string(), + answer: "local answer".to_string(), + calls: local_calls.clone(), + }), + ); + local.with_response_cache(cache.clone()); + + let hosted_run = hosted + .invoke_default(&(), vec![Message::user("what is 2+2?")]) + .await + .expect("hosted run"); + let local_run = local + .invoke_default(&(), vec![Message::user("what is 2+2?")]) + .await + .expect("local run"); + + assert_eq!(hosted_run.text().as_deref(), Some("hosted answer")); + assert_eq!( + local_run.text().as_deref(), + Some("local answer"), + "the local harness must not be served the hosted harness's cached answer" + ); + assert_eq!(local_calls.load(Ordering::SeqCst), 1, "local really ran"); +} + +#[tokio::test] +async fn identical_identity_still_shares_the_cache() { + // The fix must not disable caching: two harnesses on the *same* model + // identity still share entries. + let cache = Arc::new(InMemoryResponseCache::new()); + let calls = Arc::new(AtomicUsize::new(0)); + let identity = "openai|gpt-5|https://api.openai.com/v1||abc"; + + let build = |calls: Arc| { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "chat", + Arc::new(IdentifiedModel { + identity: identity.to_string(), + answer: "same answer".to_string(), + calls, + }), + ); + harness.with_response_cache(cache.clone()); + harness + }; + + let first = build(calls.clone()); + let second = build(calls.clone()); + + first + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("first run"); + second + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("second run"); + + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "the second harness must be served from cache" + ); +} + +#[test] +fn scoped_key_separates_identity_streaming_and_namespace() { + let base = cache_key(&request("hello")); + + let a = scoped_cache_key(&base, Some("provider-a"), false, None); + let b = scoped_cache_key(&base, Some("provider-b"), false, None); + let anon = scoped_cache_key(&base, None, false, None); + let streamed = scoped_cache_key(&base, Some("provider-a"), true, None); + let namespaced = scoped_cache_key(&base, Some("provider-a"), false, Some("tenant-7")); + + assert_ne!(a, b, "two identities must not share a key"); + assert_ne!( + a, anon, + "an anonymous model must not collide with a named one" + ); + assert_ne!( + a, streamed, + "streaming is a call parameter and must be keyed" + ); + assert_ne!(a, namespaced, "the policy namespace must be keyed"); + assert_eq!( + a, + scoped_cache_key(&base, Some("provider-a"), false, None), + "the composition must stay deterministic" + ); + assert_eq!(a.len(), 64); +} + +#[test] +fn model_identity_never_carries_the_raw_credential() { + // The identity ends up folded into keys that reach logs, events, and + // durable cache files, so a raw key must never survive into it. + let secret = "sk-super-secret-value"; + let identity = + model_cache_identity("openai", "gpt-5", "https://api.openai.com/v1", None, secret); + assert!( + !identity.contains(secret), + "the raw credential leaked into the cache identity: {identity}" + ); + assert!(identity.contains(&credential_fingerprint(secret))); + // Two different credentials must still be distinguishable. + assert_ne!( + credential_fingerprint("key-one"), + credential_fingerprint("key-two") + ); + assert_eq!(credential_fingerprint(""), "no-credential"); +} + +// ── CACHE-7 ────────────────────────────────────────────────────────────────── + +#[test] +fn key_ignores_fields_that_cannot_change_the_answer() { + let base = request("hello"); + let key = cache_key(&base); + + // `metadata` is free-form and its natural use is a run id. Folding it in + // gave such a caller a permanent 0% hit rate with no diagnostic. + let mut with_metadata = base.clone(); + with_metadata.metadata = serde_json::json!({ "run_id": "run-1234" }); + assert_eq!(cache_key(&with_metadata), key, "metadata must not be keyed"); + + // Documented as "propagated to events and traces". + let mut with_tags = base.clone(); + with_tags.tags = vec!["trace-a".to_string()]; + assert_eq!(cache_key(&with_tags), key, "tags must not be keyed"); + + // A transport deadline cannot change what the model says. + let mut with_timeout = base.clone(); + with_timeout.timeout_ms = Some(30_000); + assert_eq!( + cache_key(&with_timeout), + key, + "timeout_ms must not be keyed" + ); + + // The policy selects *whether* to cache. Folding it in meant flipping the + // (previously dead) `protect_prompt_prefix` flag invalidated every entry. + let mut with_policy = base.clone(); + with_policy.cache_policy = Some(CachePolicy { + response_cache_enabled: true, + protect_prompt_prefix: true, + ..CachePolicy::default() + }); + assert_eq!( + cache_key(&with_policy), + key, + "cache_policy must not be keyed" + ); + + // Derived from the messages that are already folded. + let mut with_fingerprint = base.clone(); + with_fingerprint.prompt_fingerprint = Some("deadbeef".to_string()); + assert_eq!( + cache_key(&with_fingerprint), + key, + "prompt_fingerprint is derived and must not be keyed" + ); +} + +#[test] +fn key_still_reflects_every_behaviour_affecting_field() { + let base = request("hello"); + let key = cache_key(&base); + + let mut hotter = base.clone(); + hotter.temperature = Some(0.9); + assert_ne!(cache_key(&hotter), key, "temperature changes the answer"); + + let mut capped = base.clone(); + capped.max_tokens = Some(16); + assert_ne!(cache_key(&capped), key, "max_tokens changes the answer"); + + let mut seeded = base.clone(); + seeded.seed = Some(7); + assert_ne!(cache_key(&seeded), key, "seed changes the answer"); + + let mut opted = base.clone(); + opted.provider_options = serde_json::json!({ "hotness": 3 }); + assert_ne!(cache_key(&opted), key, "provider_options change the answer"); + + let mut stopped = base.clone(); + stopped.stop_sequences = vec!["STOP".to_string()]; + assert_ne!(cache_key(&stopped), key, "stop_sequences change the answer"); + + let mut continued = base.clone(); + continued.continuation_id = Some("resp_123".to_string()); + assert_ne!( + cache_key(&continued), + key, + "continuation_id changes provider state" + ); + + let mut named = base.clone(); + named.model = Some("model-b".to_string()); + assert_ne!( + cache_key(&named), + key, + "an explicit model override is keyed" + ); + + let mut longer = base.clone(); + longer.messages.push(Message::user("and one more")); + assert_ne!(cache_key(&longer), key, "messages are keyed"); +} + +// ── CACHE-8 ────────────────────────────────────────────────────────────────── + +#[test] +fn every_message_and_tool_participates_in_the_key() { + // The old envelope removed `messages`/`tools` from a serialized `Value` + // with `map.remove(..)` *inside* an `if let Some(Value::Array(..))`: the + // removal ran unconditionally and the value was dropped unhashed when the + // pattern did not match. Hashing the typed fields directly removes the + // shape assumption; assert the guarantee the docstring promised. + let one = ModelRequest::new(vec![Message::user("a")]); + let two = ModelRequest::new(vec![Message::user("a"), Message::user("b")]); + let swapped = ModelRequest::new(vec![Message::user("b"), Message::user("a")]); + + assert_ne!(cache_key(&one), cache_key(&two)); + assert_ne!(cache_key(&two), cache_key(&swapped), "order is significant"); + + // An empty transcript and empty tools still produce a well-defined key. + let empty = ModelRequest::new(vec![]); + assert_eq!(cache_key(&empty).len(), 64); + assert_eq!(cache_key(&empty), cache_key(&ModelRequest::new(vec![]))); +} diff --git a/tests/wave2_cache_layout.rs b/tests/wave2_cache_layout.rs new file mode 100644 index 0000000..fdd2ab0 --- /dev/null +++ b/tests/wave2_cache_layout.rs @@ -0,0 +1,308 @@ +//! Wave 2 — prompt-cache layout protection and provider breakpoints. +//! +//! Covers CACHE-6 (layout protection was inert and compared ids only) and +//! C-BREAKPOINT (the tooling only ever *observed* a prefix; it now injects a +//! provider `prompt_cache_key`). + +use tinyagents::harness::cache::{ + CacheLayoutEvent, CachePolicy, PROMPT_CACHE_KEY_OPTION, PromptCacheLayout, + apply_prompt_cache_breakpoints, prompt_cache_key, +}; +use tinyagents::harness::message::Message; +use tinyagents::harness::model::{ModelRequest, PromptSegment, SegmentRole}; +use tinyagents::harness::prompt::PromptBuilder; +use tinyagents::harness::tool::{ToolFormat, ToolSchema}; + +fn segment(id: &str, role: SegmentRole, cacheable: bool) -> PromptSegment { + PromptSegment { + id: id.to_string(), + role, + cacheable, + } +} + +/// Builds a request through [`PromptBuilder`] so `prompt_fingerprint` — the +/// content digest the layout now consults — is actually populated. +fn built_with_system(system: &str, question: &str) -> ModelRequest { + let mut builder = PromptBuilder::new(); + builder.push_system("sys", vec![Message::system(system)]); + builder.build(vec![Message::user(question)]) +} + +fn tool(name: &str, description: &str) -> ToolSchema { + ToolSchema { + name: name.to_string(), + description: description.to_string(), + parameters: serde_json::json!({ "type": "object", "properties": {} }), + format: ToolFormat::Json, + } +} + +// ── CACHE-6 ────────────────────────────────────────────────────────────────── + +#[test] +fn editing_a_stable_segments_text_is_reported_as_a_prefix_change() { + // This is the exact failure the module exists to catch: the segment *ids* + // are unchanged, so an id-only comparison reported "prefix stable" while + // the provider's KV prefix was already destroyed. The content digest comes + // from `PromptBuilder::fingerprint`, which hashes the segments' messages + // and was previously ignored by the layout entirely. + let before_request = built_with_system("You are a careful assistant.", "q"); + let after_request = built_with_system("You are a RECKLESS assistant.", "q"); + + let before = PromptCacheLayout::from_request(&before_request); + let after = PromptCacheLayout::from_request(&after_request); + + assert_eq!( + before.prefix_ids(), + after.prefix_ids(), + "the ids are deliberately identical — that is the trap" + ); + assert!( + !before.is_prefix_stable_against(&after), + "rewriting a stable segment's TEXT must not report a stable prefix" + ); + assert!( + before.is_content_only_change(&after), + "the change is content-only: same ids, different bytes" + ); + assert_ne!(before.fingerprint(), after.fingerprint()); + assert_eq!(before.fingerprint().len(), 16); +} + +#[test] +fn editing_a_tool_schema_invalidates_the_prefix() { + // Tool declarations sit inside the stable prefix on every provider that + // caches prompts. + let base = |description: &str| { + ModelRequest::new(vec![Message::user("q")]) + .with_cache_segments(vec![segment("sys", SegmentRole::System, true)]) + .with_tools(vec![tool("search", description)]) + }; + let before = PromptCacheLayout::from_request(&base("Search the web.")); + let after = PromptCacheLayout::from_request(&base("Search the web, carefully.")); + + assert_eq!(before.prefix_ids(), after.prefix_ids()); + assert!( + !before.is_prefix_stable_against(&after), + "a tool-schema edit must invalidate the prefix" + ); +} + +#[test] +fn appending_to_the_tail_keeps_the_prefix_stable() { + // A provider prompt cache is a byte-prefix cache: appending is the one edit + // it tolerates, and the common multi-turn case must not be flagged. + let segments = vec![ + segment("sys", SegmentRole::System, true), + segment("turn", SegmentRole::Volatile, false), + ]; + let before = PromptCacheLayout::from_request( + &ModelRequest::new(vec![Message::user("q1")]).with_cache_segments(segments.clone()), + ); + let after = PromptCacheLayout::from_request( + &ModelRequest::new(vec![ + Message::user("q1"), + Message::assistant("a1"), + Message::user("q2"), + ]) + .with_cache_segments(segments), + ); + + assert!( + before.is_prefix_stable_against(&after), + "appending turns must not be reported as a prefix invalidation" + ); + assert!(!CacheLayoutEvent::new(&before, &after).changed_prefix); +} + +#[test] +fn rewriting_history_mid_stream_invalidates_the_prefix() { + // A summarizer that compacts history rewrites bytes the provider already + // cached. Ids are unchanged, so only content awareness catches it. + let segments = vec![segment("sys", SegmentRole::System, true)]; + let before = PromptCacheLayout::from_request( + &ModelRequest::new(vec![Message::user("q1"), Message::assistant("a1")]) + .with_cache_segments(segments.clone()), + ); + let after = PromptCacheLayout::from_request( + &ModelRequest::new(vec![ + Message::user("[summary of earlier turns]"), + Message::assistant("a1"), + ]) + .with_cache_segments(segments), + ); + assert!(!before.is_prefix_stable_against(&after)); +} + +#[test] +fn protect_prompt_prefix_is_load_bearing_for_the_layout_event() { + // The flag had no reader anywhere in the crate: only struct literals and + // one assertion. It now decides whether a detected invalidation counts as a + // policy violation. + let before = PromptCacheLayout::from_request( + &ModelRequest::new(vec![Message::user("q")]).with_cache_segments(vec![segment( + "sys", + SegmentRole::System, + true, + )]), + ); + let after = PromptCacheLayout::from_request( + &ModelRequest::new(vec![Message::user("q")]).with_cache_segments(vec![segment( + "turn", + SegmentRole::Volatile, + false, + )]), + ); + + let unprotected = CacheLayoutEvent::under_policy(&CachePolicy::default(), &before, &after) + .expect("the prefix did change"); + assert!(unprotected.changed_prefix); + assert!( + !unprotected.violates_policy, + "without protection a change is reported but is not a violation" + ); + + let protected = CachePolicy { + protect_prompt_prefix: true, + ..CachePolicy::default() + }; + let violation = + CacheLayoutEvent::under_policy(&protected, &before, &after).expect("the prefix did change"); + assert!(violation.violates_policy, "the flag must be load-bearing"); + assert!(violation.volatile_only); + + // No change, no event, under either policy. + assert!(CacheLayoutEvent::under_policy(&protected, &before, &before).is_none()); +} + +// ── C-BREAKPOINT ───────────────────────────────────────────────────────────── + +#[test] +fn a_prompt_cache_key_is_derived_from_the_stable_prefix() { + let request = |question: &str| built_with_system("You are a careful assistant.", question); + + let first = prompt_cache_key(&request("q1")).expect("a stable prefix exists"); + let second = prompt_cache_key(&request("q2")).expect("a stable prefix exists"); + assert_eq!( + first, second, + "every turn of one logical thread must route to the same provider cache shard" + ); + + let other = built_with_system("You are a different assistant.", "q1"); + assert_ne!( + prompt_cache_key(&other).expect("a stable prefix exists"), + first, + "a different stable prefix must route to a different shard" + ); + + // No declared prefix, nothing to route. + assert!(prompt_cache_key(&ModelRequest::new(vec![Message::user("q")])).is_none()); +} + +#[test] +fn breakpoints_are_injected_only_under_the_protection_policy() { + let build = |policy: Option| { + let mut request = ModelRequest::new(vec![Message::user("q")]) + .with_cache_segments(vec![segment("sys", SegmentRole::System, true)]); + request.cache_policy = policy; + request + }; + + // Off by default: no policy, no injection. + let mut none = build(None); + assert!(!apply_prompt_cache_breakpoints(&mut none)); + assert!(none.provider_options.get(PROMPT_CACHE_KEY_OPTION).is_none()); + + // On: a routing key is written into `provider_options`. + let protected = CachePolicy { + protect_prompt_prefix: true, + ..CachePolicy::default() + }; + let mut on = build(Some(protected.clone())); + assert!(apply_prompt_cache_breakpoints(&mut on)); + let injected = on + .provider_options + .get(PROMPT_CACHE_KEY_OPTION) + .and_then(|v| v.as_str()) + .expect("a prompt_cache_key was injected"); + assert!(injected.starts_with("tap-")); + + // A caller who already set the key wins, matching the rest of the crate's + // provider-options precedence. + let mut explicit = build(Some(protected)); + explicit.provider_options = serde_json::json!({ PROMPT_CACHE_KEY_OPTION: "mine" }); + assert!(!apply_prompt_cache_breakpoints(&mut explicit)); + assert_eq!( + explicit.provider_options[PROMPT_CACHE_KEY_OPTION], + serde_json::json!("mine") + ); +} + +// ── Guard baseline is scoped to one run ─────────────────────────────────────── + +/// A `PromptCacheGuardMiddleware` must not compare across run boundaries. +/// +/// A KV-cache prefix is only meaningful *within* one conversation. A single +/// guard instance is routinely shared across runs — a sub-agent's middleware +/// stack is built once and its agent invoked many times — so without run +/// scoping the guard compares the last request of one run against the first +/// request of the next, two unrelated transcripts, and reports an invalidation +/// that never happened. +/// +/// This went unnoticed while stability was compared by segment id alone: any +/// two requests carrying the same segment ids compared equal regardless of +/// content, so the cross-run comparison was vacuously stable. Making the +/// comparison content-aware (CACHE-6) turned that latent bug into a false +/// positive on every multi-run sub-agent, which is how it surfaced. +#[tokio::test] +async fn the_guard_does_not_compare_layouts_across_runs() { + use tinyagents::harness::context::{RunConfig, RunContext}; + use tinyagents::harness::middleware::{Middleware, PromptCacheGuardMiddleware}; + + let guard = PromptCacheGuardMiddleware::new(); + let segments = vec![ + segment("system", SegmentRole::System, true), + segment("turn", SegmentRole::Volatile, false), + ]; + + // Two independent runs that share a stable prefix but ask different + // questions — exactly what two invocations of one sub-agent look like. + for (run, question) in [("run-a", "investigate topic"), ("run-b", "ask for more")] { + let mut ctx: RunContext<()> = RunContext::new(RunConfig::new(run), ()); + let mut request = + ModelRequest::new(vec![Message::user(question)]).with_cache_segments(segments.clone()); + Middleware::<(), ()>::before_model(&guard, &mut ctx, &(), &mut request) + .await + .expect("guard pass succeeds"); + } + + assert!( + guard.layout_events().is_empty(), + "a fresh run must not be diffed against the previous run's last request" + ); +} + +/// The run scoping must not blunt the detection it was added around: within a +/// single run, rewriting a stable segment's content is still reported. +#[tokio::test] +async fn the_guard_still_reports_an_invalidation_inside_one_run() { + use tinyagents::harness::context::{RunConfig, RunContext}; + use tinyagents::harness::middleware::{Middleware, PromptCacheGuardMiddleware}; + + let guard = PromptCacheGuardMiddleware::new(); + let mut ctx: RunContext<()> = RunContext::new(RunConfig::new("one-run"), ()); + + for system in ["you are a helpful assistant", "you are a terse assistant"] { + let mut request = built_with_system(system, "q"); + Middleware::<(), ()>::before_model(&guard, &mut ctx, &(), &mut request) + .await + .expect("guard pass succeeds"); + } + + assert_eq!( + guard.layout_events().len(), + 1, + "rewriting a stable segment's text inside one run is still an invalidation" + ); +} diff --git a/tests/wave2_cache_loop.rs b/tests/wave2_cache_loop.rs new file mode 100644 index 0000000..de99bb2 --- /dev/null +++ b/tests/wave2_cache_loop.rs @@ -0,0 +1,423 @@ +//! Wave 2 — agent-loop regressions around the response cache and the model +//! wrap onion. +//! +//! Covers CACHE-2 (fallback answers cached under the primary's key), CACHE-3 +//! (cache hits re-billing tokens), CACHE-4 (a streaming cache hit emitted zero +//! deltas), CACHE-5 (a cache read/write failure killing the run) and LOOP-3 +//! (`ModelFallbackMiddleware` could not actually switch models). + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use async_trait::async_trait; + +use tinyagents::harness::cache::{InMemoryResponseCache, ResponseCache}; +use tinyagents::harness::context::{RunConfig, RunContext}; +use tinyagents::harness::message::Message; +use tinyagents::harness::middleware::ModelFallbackMiddleware; +use tinyagents::harness::model::{ChatModel, ModelRequest, ModelResponse}; +use tinyagents::harness::retry::{FallbackPolicy, RetryPolicy}; +use tinyagents::harness::runtime::{AgentHarness, RunPolicy}; +use tinyagents::harness::testkit::EventRecorder; +use tinyagents::harness::usage::Usage; +use tinyagents::{Result, TinyAgentsError}; + +// ── Fixtures ───────────────────────────────────────────────────────────────── + +/// A model that answers with a fixed text plus a fixed usage, counting calls. +struct FixedModel { + identity: &'static str, + answer: &'static str, + usage: Option, + calls: Arc, +} + +impl FixedModel { + fn new(identity: &'static str, answer: &'static str, calls: Arc) -> Self { + Self { + identity, + answer, + usage: Some(Usage::new(100, 50)), + calls, + } + } +} + +#[async_trait] +impl ChatModel<()> for FixedModel { + fn cache_identity(&self) -> Option { + Some(self.identity.to_string()) + } + + async fn invoke(&self, _state: &(), _request: ModelRequest) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + let mut response = ModelResponse::assistant(self.answer); + response.usage = self.usage; + Ok(response) + } +} + +/// A model that always fails with a retryable provider-shaped error. +struct AlwaysFailing { + identity: &'static str, + calls: Arc, +} + +#[async_trait] +impl ChatModel<()> for AlwaysFailing { + fn cache_identity(&self) -> Option { + Some(self.identity.to_string()) + } + + async fn invoke(&self, _state: &(), _request: ModelRequest) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(TinyAgentsError::Model( + "openai returned HTTP 503: service unavailable".to_string(), + )) + } +} + +/// A [`ResponseCache`] whose every operation fails, standing in for a poisoned +/// mutex or an unavailable third-party backend. +struct BrokenCache { + gets: Arc, + puts: Arc, +} + +#[async_trait] +impl ResponseCache for BrokenCache { + async fn get(&self, _key: &str) -> Result> { + self.gets.fetch_add(1, Ordering::SeqCst); + Err(TinyAgentsError::Validation( + "cache lock poisoned".to_string(), + )) + } + + async fn put(&self, _key: &str, _value: ModelResponse) -> Result<()> { + self.puts.fetch_add(1, Ordering::SeqCst); + Err(TinyAgentsError::Validation( + "cache lock poisoned".to_string(), + )) + } +} + +// ── CACHE-2 ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn a_fallback_answer_is_never_cached_under_the_primary_key() { + // The primary always fails, so the harness-level fallback chain answers. + // Writing that answer under the primary's key poisons it permanently (no + // TTL), so every later run of the primary silently gets the fallback's + // answer while `ModelStarted` announced the primary. + let primary_calls = Arc::new(AtomicUsize::new(0)); + let backup_calls = Arc::new(AtomicUsize::new(0)); + let cache = Arc::new(InMemoryResponseCache::new()); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "primary", + Arc::new(AlwaysFailing { + identity: "primary-identity", + calls: primary_calls.clone(), + }), + ); + harness.register_model( + "backup", + Arc::new(FixedModel::new( + "backup-identity", + "backup answer", + backup_calls.clone(), + )), + ); + harness.set_default_model("primary"); + harness.with_response_cache(cache.clone()); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default().with_max_attempts(1), + fallback: Some(FallbackPolicy { + models: vec!["primary".to_string(), "backup".to_string()], + }), + ..RunPolicy::default() + }); + + let run = harness + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("run falls back to the backup"); + assert_eq!(run.text().as_deref(), Some("backup answer")); + + // The backup's answer must not be sitting under the primary's key: a second + // run has to try the primary again (and fall back again). + let before = primary_calls.load(Ordering::SeqCst); + let run2 = harness + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("second run"); + assert_eq!(run2.text().as_deref(), Some("backup answer")); + assert!( + primary_calls.load(Ordering::SeqCst) > before, + "the primary must be retried; a fallback answer must not be cached under its key" + ); + assert_eq!( + cache.stats().writes, + 0, + "no fallback response may be written under the primary's key" + ); +} + +// ── CACHE-3 ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn a_cache_hit_is_marked_served_from_cache() { + // Accounting needs to tell a replay from a real call: the cached response + // retains the provider's `usage`, and re-billing it prices spend that never + // happened (which a cost budget can abort a run over). + let calls = Arc::new(AtomicUsize::new(0)); + let cache = Arc::new(InMemoryResponseCache::new()); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "chat", + Arc::new(FixedModel::new("id", "answer", calls.clone())), + ); + harness.with_response_cache(cache); + + let cold = harness + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("cold run"); + let warm = harness + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("warm run"); + + assert_eq!(calls.load(Ordering::SeqCst), 1, "warm run hit the cache"); + let cold_response = cold.final_response.expect("cold response"); + let warm_response = warm.final_response.expect("warm response"); + assert!( + !cold_response.served_from_cache, + "a real provider call is not served from cache" + ); + assert!( + warm_response.served_from_cache, + "a cache hit must be marked so accounting does not re-bill its tokens" + ); + // The usage itself is preserved so a caller can still inspect it; only the + // accounting sites are expected to skip it. + assert!(warm_response.usage.is_some()); +} + +// ── CACHE-4 ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn a_streaming_cache_hit_still_emits_deltas() { + // A hit returned before the streaming path was reached, so a warm streaming + // run emitted zero `ModelDelta` events — a UI concatenating deltas rendered + // nothing at all, contradicting the documented streaming contract. + let calls = Arc::new(AtomicUsize::new(0)); + let cache = Arc::new(InMemoryResponseCache::new()); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "chat", + Arc::new(FixedModel::new("id", "streamed answer", calls.clone())), + ); + harness.with_response_cache(cache); + + let cold_events = EventRecorder::new(); + let cold_ctx = + RunContext::new(RunConfig::new("stream-cold"), ()).with_events(cold_events.sink()); + harness + .invoke_streaming_in_context(&(), cold_ctx, vec![Message::user("q")]) + .await + .expect("cold streaming run"); + + let warm_events = EventRecorder::new(); + let warm_ctx = + RunContext::new(RunConfig::new("stream-warm"), ()).with_events(warm_events.sink()); + let warm = harness + .invoke_streaming_in_context(&(), warm_ctx, vec![Message::user("q")]) + .await + .expect("warm streaming run"); + + assert_eq!(calls.load(Ordering::SeqCst), 1, "warm run hit the cache"); + assert_eq!(warm.text().as_deref(), Some("streamed answer")); + + let cold_deltas = cold_events + .kinds() + .iter() + .filter(|k| *k == "model.delta") + .count(); + let warm_deltas = warm_events + .kinds() + .iter() + .filter(|k| *k == "model.delta") + .count(); + assert!(cold_deltas > 0, "cold streaming run emits deltas"); + assert!( + warm_deltas > 0, + "a warm streaming run must replay the cached response as deltas so warm and \ + cold runs are observationally identical" + ); + assert!( + warm_events.kinds().iter().any(|k| k == "cache.hit"), + "the warm run really was served from cache" + ); +} + +#[tokio::test] +async fn a_streaming_run_does_not_reuse_a_unary_cache_entry() { + // `streaming` is a parameter of the call, not a field of the request, so it + // never reached the key. Deliberate sharing is a caller decision, not a + // silent default. + let calls = Arc::new(AtomicUsize::new(0)); + let cache = Arc::new(InMemoryResponseCache::new()); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "chat", + Arc::new(FixedModel::new("id", "answer", calls.clone())), + ); + harness.with_response_cache(cache); + + harness + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("unary run"); + harness + .invoke_streaming_default(&(), vec![Message::user("q")]) + .await + .expect("streaming run"); + + assert_eq!( + calls.load(Ordering::SeqCst), + 2, + "a streaming call must not be served an entry written by a unary call" + ); +} + +// ── CACHE-5 ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn a_failing_cache_never_fails_the_run() { + // Both the read and the write used `?`. The write case is the worse one: + // the provider call already succeeded and was paid for, and its answer was + // discarded because the cache was unavailable. + let gets = Arc::new(AtomicUsize::new(0)); + let puts = Arc::new(AtomicUsize::new(0)); + let calls = Arc::new(AtomicUsize::new(0)); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "chat", + Arc::new(FixedModel::new("id", "answer", calls.clone())), + ); + harness.with_response_cache(Arc::new(BrokenCache { + gets: gets.clone(), + puts: puts.clone(), + })); + + let run = harness + .invoke_default(&(), vec![Message::user("q")]) + .await + .expect("a broken cache must not fail the run"); + assert_eq!(run.text().as_deref(), Some("answer")); + assert_eq!(gets.load(Ordering::SeqCst), 1, "the read was attempted"); + assert_eq!(puts.load(Ordering::SeqCst), 1, "the write was attempted"); +} + +// ── LOOP-3 ─────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn model_fallback_middleware_actually_switches_models() { + // This is deliberately driven through the REAL `ModelCallBase` (the + // innermost base of the model-wrap onion) rather than a `FakeModelBase` + // that dispatches on `req.model`. Every pre-existing test used such a fake, + // which is exactly why the bug shipped: the real base rebuilt its binding + // from fields captured *before* the wrap onion ran and never re-resolved + // `request.model`, so the "fallback" re-invoked the same failing model once + // per configured fallback name and returned the same error. + let primary_calls = Arc::new(AtomicUsize::new(0)); + let backup_calls = Arc::new(AtomicUsize::new(0)); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "primary", + Arc::new(AlwaysFailing { + identity: "primary-identity", + calls: primary_calls.clone(), + }), + ); + harness.register_model( + "backup", + Arc::new(FixedModel::new( + "backup-identity", + "backup answer", + backup_calls.clone(), + )), + ); + harness.set_default_model("primary"); + // No harness-level `FallbackPolicy` — the switch must come from the wrap + // middleware alone, which steers by mutating `request.model`. + harness.with_policy(RunPolicy { + retry: RetryPolicy::default() + .with_max_attempts(1) + .with_backoff_sleep(false), + ..RunPolicy::default() + }); + harness.push_model_middleware(Arc::new(ModelFallbackMiddleware::new(["backup"]))); + + let events = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("fallback"), ()).with_events(events.sink()); + let run = harness + .invoke_in_context(&(), ctx, vec![Message::user("q")]) + .await + .expect("the wrap-layer fallback must reach the backup model"); + + assert_eq!(run.text().as_deref(), Some("backup answer")); + assert!( + backup_calls.load(Ordering::SeqCst) >= 1, + "the backup model must actually be invoked, not merely announced" + ); + assert!( + events + .kinds() + .iter() + .any(|k| k == "model.fallback_selected"), + "the fallback event is still emitted" + ); +} + +#[tokio::test] +async fn an_unresolvable_wrap_override_keeps_the_resolved_binding() { + // Fail-closed: naming a model the registry cannot resolve must not silently + // substitute a different one — it keeps the resolved binding and makes the + // skip observable, matching what `run_loop` already does for a pre-wrap + // override. + let primary_calls = Arc::new(AtomicUsize::new(0)); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "primary", + Arc::new(AlwaysFailing { + identity: "primary-identity", + calls: primary_calls.clone(), + }), + ); + harness.set_default_model("primary"); + harness.with_policy(RunPolicy { + retry: RetryPolicy::default() + .with_max_attempts(1) + .with_backoff_sleep(false), + ..RunPolicy::default() + }); + harness.push_model_middleware(Arc::new(ModelFallbackMiddleware::new(["nonexistent"]))); + + let events = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("bad-override"), ()).with_events(events.sink()); + let outcome = harness + .invoke_in_context(&(), ctx, vec![Message::user("q")]) + .await; + + assert!(outcome.is_err(), "no model could answer"); + assert!( + events.kinds().iter().any(|k| k == "model.override_skipped"), + "an unresolvable wrap-layer override must be observable" + ); +} diff --git a/tests/wave2_cache_retry_after.rs b/tests/wave2_cache_retry_after.rs new file mode 100644 index 0000000..a47a365 --- /dev/null +++ b/tests/wave2_cache_retry_after.rs @@ -0,0 +1,85 @@ +//! Wave 2 — LOOP-5b: the HTTP `Retry-After` header is read. +//! +//! Wave 1 landed the retry side (`retry_after_hint`, `backoff_for_error`, the +//! `max_retry_after_ms` clamp) but it read only the error's **message text**. +//! A provider that sends the header without echoing it into the JSON body was +//! simply not honored. This pins the structured path: the transport parses the +//! header into [`ProviderError::retry_after_ms`], and `retry_after_hint` reads +//! that field ahead of the text fallback. + +use std::time::Duration; + +use tinyagents::TinyAgentsError; +use tinyagents::harness::model::ProviderError; +use tinyagents::harness::retry::{RetryPolicy, retry_after_hint}; + +fn provider_error(message: &str, retry_after_ms: Option) -> TinyAgentsError { + TinyAgentsError::Provider(Box::new(ProviderError { + provider: "openai".to_string(), + model: Some("gpt-5".to_string()), + status: Some(429), + code: Some("rate_limit_exceeded".to_string()), + message: message.to_string(), + retryable: true, + retry_after_ms, + raw: None, + })) +} + +#[test] +fn the_structured_field_is_honored_when_the_body_says_nothing() { + // The header-only case: the provider sent `Retry-After: 30` but the JSON + // body carries no wait at all. Before the field existed there was nothing + // to read and the hint was `None`. + let error = provider_error("Rate limit reached for gpt-5.", Some(30_000)); + assert_eq!( + retry_after_hint(&error), + Some(Duration::from_millis(30_000)), + "the header value must be honored even when the body is silent" + ); +} + +#[test] +fn the_structured_field_wins_over_the_message_text() { + // Both present and disagreeing: the header is the contract, the text is a + // string-matching fallback. + let error = provider_error("Rate limited. retry-after: 2", Some(45_000)); + assert_eq!( + retry_after_hint(&error), + Some(Duration::from_millis(45_000)) + ); +} + +#[test] +fn the_message_text_fallback_still_works() { + // Adapters that do not yet populate the field keep the wave-1 behaviour. + let error = provider_error("Rate limited. retry-after: 2", None); + assert_eq!(retry_after_hint(&error), Some(Duration::from_secs(2))); +} + +#[test] +fn a_server_supplied_wait_drives_the_backoff_and_stays_clamped() { + let policy = RetryPolicy::default(); + let honored = policy.backoff_for_error(0, &provider_error("rate limited", Some(30_000))); + assert!( + honored >= Duration::from_millis(30_000), + "the computed backoff must be at least the server-supplied wait, got {honored:?}" + ); + + // The wave-1 clamp still applies to an absurd server value. + let clamped = policy.backoff_for_error(0, &provider_error("rate limited", Some(86_400_000))); + assert!( + clamped <= Duration::from_millis(policy.max_retry_after_ms), + "an absurd Retry-After must stay clamped, got {clamped:?}" + ); +} + +#[test] +fn provider_error_round_trips_without_the_new_field() { + // `#[serde(default)]` keeps payloads written before the field existed + // decodable — important for anything that persists provider errors. + let legacy: ProviderError = + serde_json::from_str(r#"{"provider":"openai","message":"boom","retryable":true}"#) + .expect("legacy payload decodes"); + assert_eq!(legacy.retry_after_ms, None); +} diff --git a/tests/wave2_cache_store.rs b/tests/wave2_cache_store.rs new file mode 100644 index 0000000..c69eabe --- /dev/null +++ b/tests/wave2_cache_store.rs @@ -0,0 +1,377 @@ +//! Wave 2 — [`ResponseCache`] storage capabilities. +//! +//! Covers C-TTL (expiry, `clear`, namespacing), C-BYTES (a byte bound as well +//! as an entry count), C-STATS (counters), C-SINGLEFLIGHT (stampede +//! protection), C-SQLITE-CACHE (durability) and CACHE-9 (LRU recency without a +//! linear scan). + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::Duration; + +use tinyagents::harness::cache::{CachePolicy, InMemoryResponseCache, ResponseCache, SingleFlight}; +use tinyagents::harness::model::ModelResponse; + +// ── C-TTL ──────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn an_expired_entry_is_a_miss_and_is_dropped() { + // Without a TTL, a poisoned entry is permanent in a cache that is never + // cleared — which is what made CACHE-2's cross-model poisoning forever. + let cache = InMemoryResponseCache::new(); + cache + .put_with_ttl( + "k", + ModelResponse::assistant("stale"), + Some(Duration::from_millis(20)), + ) + .await + .unwrap(); + assert!( + cache.get("k").await.unwrap().is_some(), + "live before expiry" + ); + + tokio::time::sleep(Duration::from_millis(40)).await; + assert!( + cache.get("k").await.unwrap().is_none(), + "an expired entry must read as a miss" + ); + assert_eq!( + cache.stats().entries, + 0, + "the expired entry must be dropped on the way past, not merely hidden" + ); + assert_eq!(cache.stats().expirations, 1); +} + +#[tokio::test] +async fn put_without_a_ttl_never_expires() { + let cache = InMemoryResponseCache::new(); + cache.put("k", ModelResponse::assistant("v")).await.unwrap(); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(cache.get("k").await.unwrap().is_some()); +} + +#[tokio::test] +async fn clear_drops_every_entry() { + let cache = InMemoryResponseCache::new(); + cache.put("a", ModelResponse::assistant("a")).await.unwrap(); + cache.put("b", ModelResponse::assistant("b")).await.unwrap(); + cache.clear().await.unwrap(); + assert!(cache.get("a").await.unwrap().is_none()); + assert!(cache.get("b").await.unwrap().is_none()); + assert_eq!(cache.stats().entries, 0); +} + +#[test] +fn cache_policy_carries_ttl_and_namespace() { + let policy = CachePolicy::enabled() + .with_ttl(Duration::from_secs(90)) + .with_namespace("tenant-7"); + assert!(policy.response_cache_enabled); + assert_eq!(policy.ttl(), Some(Duration::from_secs(90))); + assert_eq!(policy.namespace.as_deref(), Some("tenant-7")); + // Round-trips, so a policy can be persisted with a saved agent config. + let json = serde_json::to_string(&policy).unwrap(); + let back: CachePolicy = serde_json::from_str(&json).unwrap(); + assert_eq!(back, policy); + // And a policy serialized before these fields existed still deserializes. + let legacy: CachePolicy = + serde_json::from_str(r#"{"response_cache_enabled":true,"protect_prompt_prefix":false}"#) + .unwrap(); + assert!(legacy.ttl().is_none()); +} + +// ── C-BYTES ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn the_cache_is_bounded_by_bytes_as_well_as_entries() { + // 1024 long-context responses carrying large tool payloads is hundreds of + // megabytes; an entry count alone does not bound memory. + let cache = InMemoryResponseCache::with_bounds(1024, 4_000); + let big = || ModelResponse::assistant("x".repeat(1_500)); + + for i in 0..10 { + cache.put(&format!("k{i}"), big()).await.unwrap(); + } + let stats = cache.stats(); + assert!( + stats.bytes <= 4_000, + "the byte budget must bound the cache: {} bytes retained", + stats.bytes + ); + assert!( + stats.entries < 10, + "entries must have been evicted to stay under the byte budget" + ); + assert!(stats.evictions > 0); + assert!( + cache.get("k9").await.unwrap().is_some(), + "the most recent write always survives" + ); +} + +// ── C-STATS ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn stats_report_hits_misses_and_writes() { + let cache = InMemoryResponseCache::new(); + assert!(cache.get("nope").await.unwrap().is_none()); + cache.put("k", ModelResponse::assistant("v")).await.unwrap(); + assert!(cache.get("k").await.unwrap().is_some()); + assert!(cache.get("k").await.unwrap().is_some()); + + let stats = cache.stats(); + assert_eq!(stats.hits, 2); + assert_eq!(stats.misses, 1); + assert_eq!(stats.writes, 1); + assert_eq!(stats.entries, 1); + assert!(stats.bytes > 0); +} + +// ── CACHE-9 ────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn recency_tracking_stays_correct_at_scale() { + // The recency index moved from a linear `VecDeque` scan (up to `capacity` + // string comparisons plus a memmove on every hit) to an ordered map. This + // asserts the *behaviour* the rewrite had to preserve. + let cache = InMemoryResponseCache::with_capacity(64); + for i in 0..64 { + cache + .put(&format!("k{i}"), ModelResponse::assistant("v")) + .await + .unwrap(); + } + // Keep the oldest key hot. + for _ in 0..5 { + assert!(cache.get("k0").await.unwrap().is_some()); + } + // Overflow by one: the victim must be `k1`, not the freshly-touched `k0`. + cache + .put("overflow", ModelResponse::assistant("v")) + .await + .unwrap(); + assert!(cache.get("k0").await.unwrap().is_some(), "k0 was hot"); + assert!(cache.get("k1").await.unwrap().is_none(), "k1 was the LRU"); + assert_eq!(cache.stats().entries, 64); +} + +// ── C-SINGLEFLIGHT ─────────────────────────────────────────────────────────── + +#[tokio::test] +async fn concurrent_identical_calls_collapse_into_one() { + // N concurrent identical requests all missed and all called the provider; + // N-1 of those calls were paid for and thrown away. + let flight = SingleFlight::new(); + let calls = Arc::new(AtomicUsize::new(0)); + + let mut handles = Vec::new(); + for _ in 0..8 { + let flight = flight.clone(); + let calls = calls.clone(); + handles.push(tokio::spawn(async move { + flight + .run("same-key", || async move { + calls.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(50)).await; + Ok(ModelResponse::assistant("one answer")) + }) + .await + })); + } + + let mut followers = 0; + for handle in handles { + let (response, was_follower) = handle.await.unwrap().expect("call succeeds"); + assert_eq!(response.text(), "one answer"); + if was_follower { + followers += 1; + } + } + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "eight identical concurrent calls must reach the provider once" + ); + assert_eq!( + followers, 7, + "seven callers rode along on the leader's call" + ); + assert_eq!(flight.inflight_len(), 0, "the key is retired when done"); +} + +#[tokio::test] +async fn a_follower_runs_its_own_call_when_the_leader_fails() { + // An error is not a value worth sharing: one caller's transient 503 must + // not become every concurrent caller's failure. + let flight = SingleFlight::new(); + let calls = Arc::new(AtomicUsize::new(0)); + + let leader = { + let flight = flight.clone(); + let calls = calls.clone(); + tokio::spawn(async move { + flight + .run("k", || async move { + calls.fetch_add(1, Ordering::SeqCst); + tokio::time::sleep(Duration::from_millis(50)).await; + Err(tinyagents::TinyAgentsError::Model("boom".to_string())) + }) + .await + }) + }; + tokio::time::sleep(Duration::from_millis(10)).await; + let follower = { + let flight = flight.clone(); + let calls = calls.clone(); + tokio::spawn(async move { + flight + .run("k", || async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(ModelResponse::assistant("recovered")) + }) + .await + }) + }; + + assert!(leader.await.unwrap().is_err()); + let (response, was_follower) = follower.await.unwrap().expect("follower recovers"); + assert_eq!(response.text(), "recovered"); + assert!(!was_follower, "the follower ran its own call"); + assert_eq!(calls.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn distinct_keys_do_not_block_each_other() { + let flight = SingleFlight::new(); + let calls = Arc::new(AtomicUsize::new(0)); + for key in ["a", "b", "c"] { + let calls = calls.clone(); + let (_, was_follower) = flight + .run(key, || async move { + calls.fetch_add(1, Ordering::SeqCst); + Ok(ModelResponse::assistant(key)) + }) + .await + .unwrap(); + assert!(!was_follower); + } + assert_eq!(calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn cancelling_a_leader_releases_followers() { + let flight = SingleFlight::new(); + let (started, started_rx) = tokio::sync::oneshot::channel(); + let leader_flight = flight.clone(); + let leader = tokio::spawn(async move { + leader_flight + .run("cancelled", || async move { + let _ = started.send(()); + std::future::pending::>().await + }) + .await + }); + started_rx.await.unwrap(); + leader.abort(); + let recovered = tokio::time::timeout( + Duration::from_secs(1), + flight.run("cancelled", || async { + Ok(ModelResponse::assistant("recovered")) + }), + ) + .await + .expect("follower must not wait behind a cancelled leader") + .unwrap(); + assert_eq!(recovered.0.text(), "recovered"); + assert!(!recovered.1); + assert_eq!(flight.inflight_len(), 0); +} + +// ── C-SQLITE-CACHE ─────────────────────────────────────────────────────────── + +#[cfg(feature = "sqlite")] +mod sqlite_backend { + use super::*; + use tinyagents::Result; + use tinyagents::harness::cache::SqliteResponseCache; + + #[tokio::test] + async fn sqlite_cache_round_trips_and_expires() -> Result<()> { + let cache = SqliteResponseCache::in_memory()?; + assert!(cache.get("k").await?.is_none()); + + cache.put("k", ModelResponse::assistant("durable")).await?; + let hit = cache.get("k").await?.expect("stored"); + assert_eq!(hit.text(), "durable"); + assert_eq!(cache.stats().entries, 1); + + cache + .put_with_ttl( + "short", + ModelResponse::assistant("stale"), + Some(Duration::from_millis(20)), + ) + .await?; + tokio::time::sleep(Duration::from_millis(40)).await; + assert!( + cache.get("short").await?.is_none(), + "an expired row must read as a miss" + ); + assert_eq!( + cache.stats().entries, + 1, + "the expired row must be purged lazily on read" + ); + Ok(()) + } + + #[tokio::test] + async fn sqlite_cache_survives_a_dropped_handle() -> Result<()> { + // Durability is the whole point: `InMemoryResponseCache` loses + // everything when the value is dropped, so every restart pays the + // provider bill again. + let dir = std::env::temp_dir().join(format!("tinyagents-cache-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("responses.sqlite3"); + let _ = std::fs::remove_file(&path); + + { + let cache = SqliteResponseCache::open(&path)?; + cache + .put("k", ModelResponse::assistant("persisted")) + .await?; + } + let reopened = SqliteResponseCache::open(&path)?; + let hit = reopened.get("k").await?.expect("survives a reopen"); + assert_eq!(hit.text(), "persisted"); + + let _ = std::fs::remove_file(&path); + Ok(()) + } + + #[tokio::test] + async fn sqlite_namespaces_do_not_cross_serve_and_clear_independently() -> Result<()> { + let base = SqliteResponseCache::in_memory()?; + let tenant_a = base.with_namespace("tenant-a"); + let tenant_b = base.with_namespace("tenant-b"); + + tenant_a.put("k", ModelResponse::assistant("a")).await?; + assert!( + tenant_b.get("k").await?.is_none(), + "one tenant must never be served another's entry" + ); + + tenant_b.put("k", ModelResponse::assistant("b")).await?; + tenant_a.clear().await?; + assert!(tenant_a.get("k").await?.is_none()); + assert_eq!( + tenant_b.get("k").await?.expect("untouched").text(), + "b", + "clearing one namespace must not touch another" + ); + Ok(()) + } +} diff --git a/tests/wave2_loop_cache_accounting.rs b/tests/wave2_loop_cache_accounting.rs new file mode 100644 index 0000000..cc3329b --- /dev/null +++ b/tests/wave2_loop_cache_accounting.rs @@ -0,0 +1,96 @@ +//! Regression coverage for CACHE-3: a cache-served response must not be billed. +//! +//! A replay makes no provider call and consumes no tokens, but the loop folded +//! its (replayed) `usage` into the run totals and the budget tracker anyway. On +//! a cache-heavy run that is phantom spend, and enough of it can abort a run +//! through `BudgetMiddleware` on money that was never spent. + +use std::sync::Arc; + +use async_trait::async_trait; + +use tinyagents::harness::message::Message; +use tinyagents::harness::middleware::{BudgetLimits, BudgetMiddleware}; +use tinyagents::harness::model::{ChatModel, ModelRequest, ModelResponse}; +use tinyagents::harness::runtime::AgentHarness; +use tinyagents::harness::usage::Usage; + +/// A model that answers with a fixed response, optionally flagged as a replay. +struct FlaggedModel { + served_from_cache: bool, +} + +#[async_trait] +impl ChatModel<()> for FlaggedModel { + async fn invoke( + &self, + _state: &(), + _request: ModelRequest, + ) -> tinyagents::Result { + let mut response = ModelResponse::assistant("done").with_usage(Usage::new(100, 50)); + response.served_from_cache = self.served_from_cache; + Ok(response) + } +} + +async fn run_usage(served_from_cache: bool) -> tinyagents::harness::usage::UsageTotals { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("m", Arc::new(FlaggedModel { served_from_cache })); + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds") + .usage +} + +/// The control: a real provider call is billed exactly as before. +#[tokio::test] +async fn a_real_model_call_is_still_billed() { + let usage = run_usage(false).await; + assert_eq!(usage.usage.input_tokens, 100); + assert_eq!(usage.usage.output_tokens, 50); +} + +/// The fix: a replay is not. +#[tokio::test] +async fn a_cache_served_response_is_not_billed_to_the_run() { + let usage = run_usage(true).await; + assert_eq!( + usage.usage.input_tokens, 0, + "a cache replay consumed no provider tokens" + ); + assert_eq!(usage.usage.output_tokens, 0); +} + +/// The same rule in `BudgetMiddleware`: a cache hit must not consume budget, so +/// a budget that only fits one real call still admits any number of replays. +#[tokio::test] +async fn a_cache_served_response_does_not_consume_the_budget() { + let budget = BudgetMiddleware::new(BudgetLimits { + max_total_tokens: Some(120), + ..BudgetLimits::default() + }); + let tracker = budget.tracker(); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "m", + Arc::new(FlaggedModel { + served_from_cache: true, + }), + ); + harness.push_middleware(Arc::new(budget)); + + for _ in 0..5 { + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("cache replays must never exhaust a budget they did not spend"); + } + + let snapshot = tracker.snapshot(); + assert_eq!( + snapshot.usage.usage.input_tokens, 0, + "phantom spend was recorded for cache replays: {snapshot:?}" + ); +} diff --git a/tests/wave2_loop_control.rs b/tests/wave2_loop_control.rs new file mode 100644 index 0000000..0402644 --- /dev/null +++ b/tests/wave2_loop_control.rs @@ -0,0 +1,154 @@ +//! Regression coverage for loop control-flow checkpoints. +//! +//! - **LOOP-7**: `MiddlewareControl` was drained only *after* a model call, so a +//! control requested from `after_tool` was honored one full model call late — +//! an extra billable provider round trip after a guardrail (or a human gate) +//! had already said stop. +//! - **LOOP-8**: a steering pause fell through to the success epilogue, so a +//! paused run was indistinguishable from one whose model returned an empty +//! final answer. + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::json; + +use tinyagents::TinyAgentsError; +use tinyagents::harness::context::{MiddlewareControl, RunConfig, RunContext}; +use tinyagents::harness::ids::ExecutionStatus; +use tinyagents::harness::message::Message; +use tinyagents::harness::middleware::Middleware; +use tinyagents::harness::providers::MockModel; +use tinyagents::harness::runtime::AgentHarness; +use tinyagents::harness::steering::{SteeringCommand, SteeringHandle, SteeringPolicy}; +use tinyagents::harness::testkit::FakeTool; +use tinyagents::harness::tool::ToolResult; + +/// Requests a control outcome from `after_tool` — the natural place for a +/// post-hoc guardrail or a budget stop that only knows once the result is in. +struct StopAfterToolMiddleware { + control: MiddlewareControl, +} + +#[async_trait] +impl Middleware<(), ()> for StopAfterToolMiddleware { + fn name(&self) -> &str { + "stop-after-tool" + } + + async fn after_tool( + &self, + ctx: &mut RunContext<()>, + _state: &(), + _result: &mut ToolResult, + ) -> tinyagents::Result<()> { + ctx.request_control(self.control.clone()); + Ok(()) + } +} + +fn spinning_harness() -> (AgentHarness<()>, Arc) { + let mut harness: AgentHarness<()> = AgentHarness::new(); + let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); + harness.register_model("mock", model.clone()); + harness.register_tool(Arc::new(FakeTool::returning("spin", "again"))); + (harness, model) +} + +/// LOOP-7: `StopWithFinal` raised from `after_tool` stops the loop **on that +/// turn**, not after another model call. +#[tokio::test] +async fn stop_with_final_from_after_tool_is_honored_before_the_next_model_call() { + let (mut harness, model) = spinning_harness(); + harness.push_middleware(Arc::new(StopAfterToolMiddleware { + control: MiddlewareControl::StopWithFinal("stopped".to_string()), + })); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run stops cleanly"); + + assert_eq!(run.text(), Some("stopped".to_string())); + assert_eq!( + model.call_count(), + 1, + "a control raised from after_tool must not cost another model call" + ); +} + +/// LOOP-7: the same for `Interrupt`, where the wasted call is the more +/// expensive mistake — a human gate already said stop. +#[tokio::test] +async fn interrupt_from_after_tool_is_honored_before_the_next_model_call() { + let (mut harness, model) = spinning_harness(); + harness.push_middleware(Arc::new(StopAfterToolMiddleware { + control: MiddlewareControl::Interrupt { + node: "approval".to_string(), + message: "needs a human".to_string(), + }, + })); + + let err = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("the interrupt surfaces"); + + assert!( + matches!(err, TinyAgentsError::Interrupted { .. }), + "got {err:?}" + ); + assert_eq!( + model.call_count(), + 1, + "an interrupt raised from after_tool must not cost another billable model call" + ); +} + +/// LOOP-8: a steering pause is a distinct, resumable outcome — reported as +/// interrupted, carrying the pause reason, and keeping the transcript. +#[tokio::test] +async fn a_steering_pause_is_distinguishable_from_a_clean_finish() { + let (harness, _model) = spinning_harness(); + + let steering = SteeringHandle::new(SteeringPolicy::allow_all()); + steering.send(SteeringCommand::PauseWith { + reason: "waiting for a human".to_string(), + }); + + let ctx: RunContext<()> = RunContext::new(RunConfig::new("paused"), ()).with_steering(steering); + + let result = harness + .invoke_in_context_with_status(&(), ctx, vec![Message::user("go")]) + .await + .expect("a pause is not a failure"); + + let pause = result + .run + .paused + .as_ref() + .expect("a paused run must say so, not look like an empty final answer"); + assert_eq!(pause.reason.as_deref(), Some("waiting for a human")); + assert!(result.run.final_response.is_none()); + assert_eq!( + result.status.status, + ExecutionStatus::Interrupted, + "a paused run must not be reported as completed" + ); +} + +/// A run with no steering still completes normally — the pause path must not +/// leak into the ordinary case. +#[tokio::test] +async fn an_unsteered_run_still_reports_completed() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("mock", Arc::new(MockModel::constant("done"))); + + let result = harness + .invoke_with_status(&(), (), RunConfig::new("plain"), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert!(result.run.paused.is_none()); + assert_eq!(result.status.status, ExecutionStatus::Completed); +} diff --git a/tests/wave2_loop_estimators.rs b/tests/wave2_loop_estimators.rs new file mode 100644 index 0000000..4f0404f --- /dev/null +++ b/tests/wave2_loop_estimators.rs @@ -0,0 +1,130 @@ +//! Regression coverage for the two text-only token estimators and for +//! micro-compaction's `trusted_verbatim` violation. +//! +//! Both estimators summed `estimate_tokens(&m.text())`, and `text()` returns +//! only *textual* content blocks — so a transcript dominated by large JSON tool +//! results or image blocks estimated to nearly nothing and sailed past the very +//! gate that exists to catch it. + +use serde_json::json; + +use tinyagents::harness::context::{RunConfig, RunContext}; +use tinyagents::harness::message::{ContentBlock, Message, ToolMessage}; +use tinyagents::harness::middleware::{MicrocompactMiddleware, Middleware}; +use tinyagents::harness::model::ModelRequest; + +const PLACEHOLDER: &str = "[elided]"; + +/// A tool result whose payload lives in a JSON content block, which +/// `Message::text()` does not see at all — the shape a tool returning +/// structured data produces. +fn json_tool_message(id: &str, trusted_verbatim: bool) -> Message { + let payload = json!({ + "id": id, + "rows": (0..80).map(|i| json!({"n": i, "label": format!("row-{i}-{id}")})) + .collect::>(), + }); + Message::Tool(ToolMessage { + tool_call_id: id.to_string(), + content: vec![ContentBlock::Json(payload)], + trusted_verbatim, + artifact: None, + }) +} + +/// Runs the middleware's `before_model` hook against a throwaway context. +async fn compact(middleware: &MicrocompactMiddleware, request: &mut ModelRequest) { + let mut ctx: RunContext<()> = RunContext::new(RunConfig::new("estimator"), ()); + Middleware::<(), ()>::before_model(middleware, &mut ctx, &(), request) + .await + .expect("before_model succeeds"); +} + +/// REASON-3: the shared estimator charges non-textual content; the old +/// text-only sum scored the identical transcript at zero. +#[test] +fn the_shared_estimator_charges_non_textual_payloads() { + let messages: Vec = (0..6) + .map(|i| json_tool_message(&format!("c{i}"), false)) + .collect(); + + let text_only: u64 = messages + .iter() + .map(|m| tinyagents::harness::summarization::estimate_tokens(&m.text())) + .sum(); + let counted = tinyagents::harness::message::count_tokens_approximately(&messages); + + assert_eq!( + text_only, 0, + "precondition: the old text-only estimator scores a JSON transcript at zero" + ); + assert!( + counted > 1_000, + "the shared estimator must charge JSON tool results, got {counted}" + ); +} + +/// REASON-3, micro-compaction: with a budget far below the transcript's real +/// weight but above its text-only estimate of zero, the gate must fire. +/// +/// Before the switch `total_message_tokens` returned 0 here, so the gate +/// concluded the transcript still fit and micro-compaction never ran. +#[tokio::test] +async fn microcompaction_budget_gate_trips_on_a_non_textual_transcript() { + let messages: Vec = (0..6) + .map(|i| json_tool_message(&format!("c{i}"), false)) + .collect(); + + let middleware = MicrocompactMiddleware::new(2, PLACEHOLDER).with_token_budget(100); + let mut request = ModelRequest::new(messages.clone()); + compact(&middleware, &mut request).await; + + assert_ne!( + request.messages, messages, + "the budget gate must trip on a transcript whose weight is non-textual" + ); + assert!( + request.messages.iter().any(|m| m.text() == PLACEHOLDER), + "older tool bodies should have been blanked" + ); +} + +/// MICROCOMPACT: a tool result flagged `trusted_verbatim` asked to reach the +/// model byte-for-byte. Blanking it is exactly the rewrite the flag's contract +/// forbids — it produces content that reads fine and is wrong. +#[tokio::test] +async fn microcompaction_leaves_trusted_verbatim_tool_results_alone() { + let messages = vec![ + // Oldest — the first one micro-compaction would blank. + Message::Tool(ToolMessage { + tool_call_id: "c0".to_string(), + content: vec![ContentBlock::Text( + "argument schema for `write`".to_string(), + )], + trusted_verbatim: true, + artifact: None, + }), + Message::Tool(ToolMessage { + tool_call_id: "c1".to_string(), + content: vec![ContentBlock::Text("ordinary tool output".to_string())], + trusted_verbatim: false, + artifact: None, + }), + Message::tool("c2", "kept recent"), + ]; + + let middleware = MicrocompactMiddleware::new(1, PLACEHOLDER); + let mut request = ModelRequest::new(messages); + compact(&middleware, &mut request).await; + + assert_eq!( + request.messages[0].text(), + "argument schema for `write`", + "a trusted_verbatim tool result must survive compaction verbatim" + ); + assert_eq!( + request.messages[1].text(), + PLACEHOLDER, + "untrusted tool results should still be blanked" + ); +} diff --git a/tests/wave2_loop_limits.rs b/tests/wave2_loop_limits.rs new file mode 100644 index 0000000..051df4b --- /dev/null +++ b/tests/wave2_loop_limits.rs @@ -0,0 +1,181 @@ +//! Regression coverage for run-scoped call limits in the agent loop. +//! +//! Three defects are pinned here: +//! +//! - **LOOP-1**: an explicitly-set `RunConfig` call cap was silently widened by +//! the harness `RunPolicy` default, so `with_max_model_calls(2)` ran 25. +//! - **LOOP-9b**: cap exhaustion was only ever a hard error, discarding the +//! whole run even under `LimitBehavior::StopWithPartial`. +//! - **LOOP-12**: the working transcript was dropped on every failure path. + +use std::sync::Arc; + +use serde_json::json; + +use tinyagents::TinyAgentsError; +use tinyagents::harness::context::RunConfig; +use tinyagents::harness::limits::{LimitBehavior, RunLimits}; +use tinyagents::harness::message::Message; +use tinyagents::harness::providers::MockModel; +use tinyagents::harness::runtime::{AgentHarness, RunPolicy}; +use tinyagents::harness::testkit::FakeTool; + +/// A harness whose model always asks for the same tool, so the loop only ever +/// stops because a cap stops it. +fn spinning_harness() -> (AgentHarness<()>, Arc) { + let mut harness: AgentHarness<()> = AgentHarness::new(); + let model = Arc::new(MockModel::with_tool_call("spin", json!({}))); + harness.register_model("mock", model.clone()); + harness.register_tool(Arc::new(FakeTool::returning("spin", "again"))); + (harness, model) +} + +/// LOOP-1: an explicitly-set `RunConfig` model-call cap is a **ceiling**. The +/// permissive `RunPolicy` default (25) must not widen it. +/// +/// Before the fix the loop overwrote the tracker's caps with the policy's by +/// plain assignment, so this ran 25 model calls instead of 2 — an embedder +/// budgeting a cheap sub-agent at 2 calls silently burned 25. +#[tokio::test] +async fn explicit_run_config_model_cap_is_not_widened_by_the_policy_default() { + let (harness, model) = spinning_harness(); + + let err = harness + .invoke( + &(), + (), + RunConfig::new("capped").with_max_model_calls(2), + vec![Message::user("go")], + ) + .await + .expect_err("the explicit cap should trip"); + + assert!( + matches!(err, TinyAgentsError::LimitExceeded(_)), + "got {err:?}" + ); + assert_eq!( + model.call_count(), + 2, + "the explicitly-set cap of 2 must bound the run, not the policy default of 25" + ); +} + +/// LOOP-1, tool axis: the same ceiling rule applies to `max_tool_calls`. +#[tokio::test] +async fn explicit_run_config_tool_cap_is_not_widened_by_the_policy_default() { + let (harness, model) = spinning_harness(); + + let err = harness + .invoke( + &(), + (), + RunConfig::new("capped").with_max_tool_calls(1), + vec![Message::user("go")], + ) + .await + .expect_err("the explicit tool cap should trip"); + + assert!( + matches!(err, TinyAgentsError::LimitExceeded(_)), + "got {err:?}" + ); + // Turn 1 spends the single tool call; turn 2 requests a second and trips. + assert_eq!( + model.call_count(), + 2, + "the explicit tool cap of 1 must bound the run, not the policy default of 50" + ); +} + +/// LOOP-1, the legitimate loosening case that must keep working: an **unset** +/// `RunConfig` cap merely defaulted, so a policy configuring a higher cap is +/// the only real source of truth and wins. +#[tokio::test] +async fn unset_run_config_cap_lets_the_policy_raise_the_limit() { + let (mut harness, model) = spinning_harness(); + harness.with_policy(RunPolicy { + limits: RunLimits::default() + .with_max_model_calls(30) + .with_max_tool_calls(1000), + ..RunPolicy::default() + }); + + let err = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("the policy cap should trip"); + + assert!( + matches!(err, TinyAgentsError::LimitExceeded(_)), + "got {err:?}" + ); + assert!( + err.to_string().contains("30"), + "expected the policy limit (30) to be reported, got: {err}" + ); + assert_eq!( + model.call_count(), + 30, + "an unset config cap must not clamp the policy's higher cap back to the default" + ); +} + +/// LOOP-9b: under `LimitBehavior::StopWithPartial` the model-call cap ends the +/// run **cleanly**, keeping the transcript, counters, and usage rather than +/// discarding all of it behind a `LimitExceeded`. +#[tokio::test] +async fn model_cap_stops_with_the_partial_run_under_stop_with_partial() { + let (mut harness, _model) = spinning_harness(); + harness.with_policy(RunPolicy { + limits: RunLimits::default() + .with_max_model_calls(2) + .with_behavior(LimitBehavior::StopWithPartial), + ..RunPolicy::default() + }); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("StopWithPartial must not fail the run"); + + assert_eq!(run.model_calls, 2); + assert!( + run.messages.len() > 1, + "the partial transcript must survive, got {:?}", + run.messages.len() + ); + assert!( + run.tool_calls >= 1, + "work completed before the cap must be reported" + ); +} + +/// LOOP-12: even on the hard-error path the accumulated transcript is worth +/// keeping. `invoke_collecting_partial` hands back the partial run alongside +/// the error instead of dropping it. +#[tokio::test] +async fn a_failed_run_still_returns_its_partial_transcript() { + let (harness, _model) = spinning_harness(); + + let outcome = harness + .invoke_collecting_partial( + &(), + (), + RunConfig::new("partial").with_max_model_calls(2), + vec![Message::user("go")], + ) + .await; + + let error = outcome.error.expect("the cap should trip"); + assert!( + matches!(error, TinyAgentsError::LimitExceeded(_)), + "got {error:?}" + ); + assert!( + outcome.run.messages.len() > 1, + "the working transcript must be preserved on the failure path, got {:?}", + outcome.run.messages + ); + assert_eq!(outcome.run.model_calls, 2); +} diff --git a/tests/wave2_loop_recovery.rs b/tests/wave2_loop_recovery.rs new file mode 100644 index 0000000..65553e5 --- /dev/null +++ b/tests/wave2_loop_recovery.rs @@ -0,0 +1,180 @@ +//! Regression coverage for TOOL-4: an unknown tool name and invalid tool +//! arguments used to abort the whole run by default. +//! +//! That made the crate inconsistent with itself — an *unparseable* arguments +//! blob has always recovered unconditionally, so `{city:` survived while +//! `{"city": 5}` killed the run — and diverged from LangGraph, which never +//! fails on either. + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use serde_json::json; + +use tinyagents::harness::message::Message; +use tinyagents::harness::model::{ChatModel, ModelRequest, ModelResponse}; +use tinyagents::harness::runtime::{AgentHarness, InvalidArgsPolicy, RunPolicy, UnknownToolPolicy}; +use tinyagents::harness::tool::{Tool, ToolCall, ToolResult, ToolSchema}; + +/// The defaults are the whole point of this file, so pin them directly too. +#[test] +fn recoverable_variants_are_the_defaults() { + assert_eq!( + UnknownToolPolicy::default(), + UnknownToolPolicy::ReturnToolError + ); + assert_eq!( + InvalidArgsPolicy::default(), + InvalidArgsPolicy::ReturnToolError + ); + assert_eq!( + RunPolicy::default().unknown_tool, + UnknownToolPolicy::ReturnToolError + ); + assert_eq!( + RunPolicy::default().invalid_args, + InvalidArgsPolicy::ReturnToolError + ); +} + +/// A model that emits one scripted tool call and then a plain final answer. +struct TwoTurnModel { + call: Mutex>, +} + +impl TwoTurnModel { + fn new(call: ToolCall) -> Self { + Self { + call: Mutex::new(Some(call)), + } + } +} + +#[async_trait] +impl ChatModel<()> for TwoTurnModel { + async fn invoke( + &self, + _state: &(), + _request: ModelRequest, + ) -> tinyagents::Result { + match self.call.lock().expect("poisoned").take() { + Some(call) => { + let mut response = ModelResponse::assistant(""); + response.message.content.clear(); + response.message.tool_calls = vec![call]; + response.finish_reason = Some("tool_calls".to_string()); + Ok(response) + } + None => Ok(ModelResponse::assistant("recovered")), + } + } +} + +/// A tool whose schema genuinely requires a string `city`. +struct WeatherTool; + +#[async_trait] +impl Tool<()> for WeatherTool { + fn name(&self) -> &str { + "weather" + } + + fn description(&self) -> &str { + "Looks up the weather for a city." + } + + fn schema(&self) -> ToolSchema { + ToolSchema::new( + "weather".to_string(), + "Looks up the weather for a city.".to_string(), + json!({ + "type": "object", + "properties": { "city": { "type": "string" } }, + "required": ["city"], + }), + ) + } + + async fn call(&self, _state: &(), call: ToolCall) -> tinyagents::Result { + Ok(ToolResult::text(call.id, call.name, "sunny")) + } +} + +/// TOOL-4a: a hallucinated tool name is answered with a tool-error message the +/// model can act on, and the run continues. +#[tokio::test] +async fn an_unknown_tool_name_recovers_by_default() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(TwoTurnModel::new(ToolCall::new("c1", "nope", json!({})))), + ); + harness.register_tool(Arc::new(WeatherTool)); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("an unknown tool name must not abort the run"); + + assert_eq!(run.text(), Some("recovered".to_string())); + let tool_text = run + .messages + .iter() + .find(|m| matches!(m, Message::Tool(_))) + .map(|m| m.text()) + .expect("a tool-error message must be injected"); + assert!( + tool_text.contains("nope"), + "the recovery message should name the unknown tool, got: {tool_text}" + ); +} + +/// TOOL-4b: arguments that fail schema validation are answered with the +/// validation detail rather than aborting the run. +#[tokio::test] +async fn invalid_tool_arguments_recover_by_default() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(TwoTurnModel::new(ToolCall::new( + "c1", + "weather", + json!({ "city": 5 }), + ))), + ); + harness.register_tool(Arc::new(WeatherTool)); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("invalid arguments must not abort the run"); + + assert_eq!(run.text(), Some("recovered".to_string())); + assert!( + run.messages + .iter() + .any(|m| matches!(m, Message::Tool(_)) && !m.text().is_empty()), + "a validation-error tool message must be injected, got {:?}", + run.messages + ); +} + +/// `Fail` is still available for callers that genuinely want a hard stop. +#[tokio::test] +async fn the_fail_variants_are_still_available() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(TwoTurnModel::new(ToolCall::new("c1", "nope", json!({})))), + ); + harness.register_tool(Arc::new(WeatherTool)); + harness.with_policy(RunPolicy { + unknown_tool: UnknownToolPolicy::Fail, + ..RunPolicy::default() + }); + + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("an opted-in Fail policy still aborts"); +} diff --git a/tests/wave2_loop_structured.rs b/tests/wave2_loop_structured.rs new file mode 100644 index 0000000..dc88a2b --- /dev/null +++ b/tests/wave2_loop_structured.rs @@ -0,0 +1,250 @@ +//! Regression coverage for structured output inside a *tool-using* loop. +//! +//! - **TOOL-7**: under the tool-call structured strategy the loop forced +//! `tool_choice` onto the artificial schema tool on every turn, so the model +//! had to emit the structured call immediately and no registered tool could +//! ever run. +//! - **TOOL-8**: a structured hit terminated the turn even when the model asked +//! for real tools alongside it, and the schema name was never checked against +//! the registered tool names. + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use serde_json::json; + +use tinyagents::TinyAgentsError; +use tinyagents::harness::message::Message; +use tinyagents::harness::model::{ + ChatModel, ModelProfile, ModelRequest, ModelResponse, ToolChoice, +}; +use tinyagents::harness::runtime::{AgentHarness, RunPolicy}; +use tinyagents::harness::testkit::FakeTool; +use tinyagents::harness::tool::ToolCall; + +/// The schema the tests ask the model to fill in. +fn schema() -> serde_json::Value { + json!({ + "type": "object", + "properties": { "answer": { "type": "string" } }, + "required": ["answer"], + }) +} + +/// A model that records every request it receives and replays a script. +/// +/// A profile with `native_structured_output = false` is the whole point: it is +/// what selects `StructuredStrategy::ToolCall`, the path the bug lived on, and +/// it is what `ModelProfile::default()` yields for that field — so this is the +/// ordinary case for a tool-calling provider without native constrained JSON, +/// not an exotic corner. +struct RecordingModel { + profile: ModelProfile, + script: Mutex>, + seen: Mutex>, +} + +impl RecordingModel { + fn new(script: Vec) -> Self { + Self { + profile: ModelProfile { + tool_calling: true, + parallel_tool_calls: true, + native_structured_output: false, + json_schema: false, + ..ModelProfile::default() + }, + script: Mutex::new(script), + seen: Mutex::new(Vec::new()), + } + } + + fn requests(&self) -> Vec { + self.seen.lock().expect("poisoned").clone() + } +} + +#[async_trait] +impl ChatModel<()> for RecordingModel { + fn profile(&self) -> Option<&ModelProfile> { + Some(&self.profile) + } + + async fn invoke( + &self, + _state: &(), + request: ModelRequest, + ) -> tinyagents::Result { + self.seen.lock().expect("poisoned").push(request); + let mut script = self.script.lock().expect("poisoned"); + if script.len() > 1 { + Ok(script.remove(0)) + } else { + Ok(script[0].clone()) + } + } +} + +/// Builds a response carrying exactly the supplied tool calls. +fn tool_call_response(calls: Vec) -> ModelResponse { + let mut response = ModelResponse::assistant(""); + response.message.content.clear(); + response.message.tool_calls = calls; + response.finish_reason = Some("tool_calls".to_string()); + response +} + +/// TOOL-7: with real tools registered, the artificial schema tool is *offered* +/// but never forced, so the model is free to call a registered tool first. +/// +/// Before the fix the loop set `tool_choice = Tool("result")` on every request, +/// which compels the provider to emit the structured call on turn 1 and ends +/// the loop before any registered tool can run. The symptom — "my agent never +/// uses its tools" — points nowhere near the structured-output code. +#[tokio::test] +async fn structured_tool_strategy_does_not_force_the_schema_tool_when_real_tools_exist() { + let model = Arc::new(RecordingModel::new(vec![tool_call_response(vec![ + ToolCall::new("c1", "result", json!({"answer": "42"})), + ])])); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("rec", model.clone()); + harness.register_tool(Arc::new(FakeTool::returning("search", "hits"))); + harness.with_policy(RunPolicy { + default_response_format: Some(tinyagents::harness::model::ResponseFormat::auto( + "result", + schema(), + )), + ..RunPolicy::default() + }); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!(run.structured, Some(json!({"answer": "42"}))); + + let requests = model.requests(); + assert_eq!(requests.len(), 1); + assert!( + !matches!(requests[0].tool_choice, ToolChoice::Tool(_)), + "the schema tool must not be forced while real tools are registered, got {:?}", + requests[0].tool_choice + ); + assert!( + requests[0].tools.iter().any(|t| t.name == "search"), + "the registered tool must still be offered" + ); + assert!( + requests[0].tools.iter().any(|t| t.name == "result"), + "the schema tool must still be offered" + ); +} + +/// The terminal case is unchanged: with **no** registered tools the schema tool +/// is the only thing the model can call, so forcing it is correct. +#[tokio::test] +async fn structured_tool_strategy_still_forces_the_schema_tool_with_no_real_tools() { + let model = Arc::new(RecordingModel::new(vec![tool_call_response(vec![ + ToolCall::new("c1", "result", json!({"answer": "42"})), + ])])); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("rec", model.clone()); + harness.with_policy(RunPolicy { + default_response_format: Some(tinyagents::harness::model::ResponseFormat::auto( + "result", + schema(), + )), + ..RunPolicy::default() + }); + + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + let requests = model.requests(); + assert!( + matches!(&requests[0].tool_choice, ToolChoice::Tool(name) if name == "result"), + "with no registered tools the schema tool should be forced, got {:?}", + requests[0].tool_choice + ); +} + +/// TOOL-8a: a turn that returns the schema call **alongside** real tool calls +/// must still execute the real ones. +/// +/// Before the fix `structured_tool_hit` was true if *any* returned call matched +/// the schema name, and the loop broke out regardless — `search` never ran and +/// no event said so. +#[tokio::test] +async fn a_structured_hit_does_not_discard_sibling_real_tool_calls() { + let model = Arc::new(RecordingModel::new(vec![ + tool_call_response(vec![ + ToolCall::new("c1", "search", json!({})), + ToolCall::new("c2", "result", json!({"answer": "42"})), + ]), + tool_call_response(vec![ToolCall::new( + "c3", + "result", + json!({"answer": "final"}), + )]), + ])); + + let search = Arc::new(FakeTool::returning("search", "hits")); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("rec", model.clone()); + harness.register_tool(search.clone()); + harness.with_policy(RunPolicy { + default_response_format: Some(tinyagents::harness::model::ResponseFormat::auto( + "result", + schema(), + )), + ..RunPolicy::default() + }); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run succeeds"); + + assert_eq!( + search.calls().len(), + 1, + "the real tool requested in the same turn must still be executed" + ); + assert_eq!(run.structured, Some(json!({"answer": "final"}))); +} + +/// TOOL-8b: a schema name that collides with a registered tool would put two +/// identically-named `function` entries in one request — which OpenAI rejects — +/// and makes every returned call ambiguous. Fail closed, up front. +#[tokio::test] +async fn a_schema_name_colliding_with_a_registered_tool_fails_closed() { + let model = Arc::new(RecordingModel::new(vec![ModelResponse::assistant("hi")])); + + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model("rec", model); + harness.register_tool(Arc::new(FakeTool::returning("result", "hits"))); + harness.with_policy(RunPolicy { + default_response_format: Some(tinyagents::harness::model::ResponseFormat::auto( + "result", + schema(), + )), + ..RunPolicy::default() + }); + + let err = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("the collision must be rejected"); + + assert!(matches!(err, TinyAgentsError::Validation(_)), "got {err:?}"); + assert!( + err.to_string().contains("collides"), + "the error should name the collision, got: {err}" + ); +} diff --git a/tests/wave2_tools_execution.rs b/tests/wave2_tools_execution.rs new file mode 100644 index 0000000..5d3bcba --- /dev/null +++ b/tests/wave2_tools_execution.rs @@ -0,0 +1,586 @@ +//! Regression coverage for the wave-2 tool-execution defects in +//! `harness::agent_loop::tools`. +//! +//! Each test pins one defect: +//! +//! | Test | Defect | +//! |------|--------| +//! | `tool_result_call_id_is_overwritten_with_the_admitted_id` | TOOL-1: the transcript took the tool's `call_id`, not the admitted one | +//! | `concurrent_admission_failure_emits_no_tool_started` | TOOL-3: `ToolStarted` was emitted for calls that never ran | +//! | `serial_tool_error_becomes_a_model_visible_result` | TOOL-5 (serial): an `Err` killed the run | +//! | `concurrent_tool_error_becomes_a_model_visible_result` | TOOL-5 (concurrent) | +//! | `fatal_tool_error_emits_tool_failed_and_clears_active_calls` | TOOL-6 | +//! | `duplicate_call_ids_do_not_clear_each_others_active_entry` | TOOL-10 | +//! | `unknown_tool_recovery_emits_started_and_completed` | TOOL-11 | +//! | `before_tool_rejection_does_not_consume_a_tool_call_slot` | TOOL-12 | +//! | `injected_arguments_are_stripped_before_validation` | injected-argument enforcement | + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use tinyagents::TinyAgentsError; +use tinyagents::harness::context::{RunConfig, RunContext}; +use tinyagents::harness::events::AgentEvent; +use tinyagents::harness::limits::RunLimits; +use tinyagents::harness::message::Message; +use tinyagents::harness::middleware::Middleware; +use tinyagents::harness::model::ModelResponse; +use tinyagents::harness::providers::MockModel; +use tinyagents::harness::runtime::{AgentHarness, RunPolicy, UnknownToolPolicy}; +use tinyagents::harness::testkit::EventRecorder; +use tinyagents::harness::tool::{Tool, ToolCall, ToolErrorPolicy, ToolResult, ToolSchema}; + +// ── Scripted model helpers ──────────────────────────────────────────────────── + +/// Builds an assistant turn that requests `calls`. +/// +/// Deliberately built from [`ModelResponse::assistant`] rather than a struct +/// literal so the fixture survives new fields being added to `ModelResponse`. +fn tool_calls_response(calls: Vec) -> ModelResponse { + let mut response = ModelResponse::assistant(""); + response.message.content = Vec::new(); + response.message.tool_calls = calls; + response.finish_reason = Some("tool_calls".into()); + response +} + +fn text_response(text: &str) -> ModelResponse { + let mut response = ModelResponse::assistant(text); + response.finish_reason = Some("stop".into()); + response +} + +fn empty_object_schema(name: &str) -> ToolSchema { + ToolSchema::new(name, "test tool", json!({ "type": "object" })) +} + +// ── Test tools ──────────────────────────────────────────────────────────────── + +/// A third-party tool that stamps its own (wrong) `call_id` on the result. +struct WrongCallIdTool; + +#[async_trait] +impl Tool<()> for WrongCallIdTool { + fn name(&self) -> &str { + "wrong_id" + } + fn description(&self) -> &str { + "returns a result carrying a hard-coded call id" + } + fn schema(&self) -> ToolSchema { + empty_object_schema("wrong_id") + } + async fn call(&self, _state: &(), _call: ToolCall) -> tinyagents::Result { + // The classic third-party mistake: `call_id` is the first positional + // argument, so a hard-coded or empty string slips in unnoticed. + Ok(ToolResult::text("", "wrong_id", "ok")) + } +} + +/// A tool that always returns `Err`, with a configurable error policy. +struct ErringTool { + name: String, + policy: ToolErrorPolicy, + calls: Arc, +} + +#[async_trait] +impl Tool<()> for ErringTool { + fn name(&self) -> &str { + &self.name + } + fn description(&self) -> &str { + "always fails" + } + fn schema(&self) -> ToolSchema { + empty_object_schema(&self.name) + } + fn error_policy(&self) -> ToolErrorPolicy { + self.policy.clone() + } + async fn call(&self, _state: &(), _call: ToolCall) -> tinyagents::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(TinyAgentsError::Tool("transient 503".into())) + } +} + +/// Records the arguments it was invoked with, and declares one injected key. +struct InjectedArgTool { + seen: Arc>>, +} + +#[async_trait] +impl Tool<()> for InjectedArgTool { + fn name(&self) -> &str { + "injected" + } + fn description(&self) -> &str { + "declares a host-injected argument" + } + fn schema(&self) -> ToolSchema { + ToolSchema::new( + "injected", + "declares a host-injected argument", + json!({ + "type": "object", + "properties": { + "query": { "type": "string" }, + "thread_id": { "type": "string" } + }, + "required": ["query", "thread_id"], + "additionalProperties": false + }), + ) + } + fn injected_arguments(&self) -> &[&str] { + &["thread_id"] + } + async fn call(&self, _state: &(), call: ToolCall) -> tinyagents::Result { + self.seen.lock().unwrap().push(call.arguments.clone()); + Ok(ToolResult::text(call.id, "injected", "ok")) + } +} + +/// A plain echo tool used to fill multi-call turns. +struct EchoTool { + name: String, + calls: Arc, +} + +#[async_trait] +impl Tool<()> for EchoTool { + fn name(&self) -> &str { + &self.name + } + fn description(&self) -> &str { + "echoes" + } + fn schema(&self) -> ToolSchema { + empty_object_schema(&self.name) + } + async fn call(&self, _state: &(), call: ToolCall) -> tinyagents::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(ToolResult::text(call.id, self.name.clone(), "echo")) + } +} + +/// Middleware that rejects every tool call from `before_tool`. +struct RejectingMiddleware; + +#[async_trait] +impl Middleware<(), ()> for RejectingMiddleware { + fn name(&self) -> &str { + "rejecting" + } + async fn before_tool( + &self, + _ctx: &mut RunContext<()>, + _state: &(), + _call: &mut ToolCall, + ) -> tinyagents::Result<()> { + Err(TinyAgentsError::Middleware("approval denied".into())) + } +} + +// ── Event helpers ───────────────────────────────────────────────────────────── + +fn started_call_ids(events: &[AgentEvent]) -> Vec { + events + .iter() + .filter_map(|event| match event { + AgentEvent::ToolStarted { call_id, .. } => Some(call_id.as_str().to_string()), + _ => None, + }) + .collect() +} + +fn completed_call_ids(events: &[AgentEvent]) -> Vec { + events + .iter() + .filter_map(|event| match event { + AgentEvent::ToolCompleted { call_id, .. } => Some(call_id.as_str().to_string()), + _ => None, + }) + .collect() +} + +fn tool_message_ids(messages: &[Message]) -> Vec { + messages + .iter() + .filter_map(|message| match message { + Message::Tool(tool) => Some(tool.tool_call_id.clone()), + _ => None, + }) + .collect() +} + +// ── TOOL-1 ──────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn tool_result_call_id_is_overwritten_with_the_admitted_id() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new("call_abc", "wrong_id", json!({}))]), + text_response("done"), + ])), + ); + harness.register_tool(Arc::new(WrongCallIdTool)); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("run should complete"); + + assert_eq!( + tool_message_ids(&run.messages), + vec!["call_abc".to_string()], + "the transcript must answer the admitted tool_call_id, not the id the tool stamped" + ); +} + +// ── TOOL-3 ──────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn concurrent_admission_failure_emits_no_tool_started() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ + ToolCall::new("c1", "echo_a", json!({})), + ToolCall::new("c2", "echo_b", json!({})), + ToolCall::new("c3", "echo_c", json!({})), + ])])), + ); + for name in ["echo_a", "echo_b", "echo_c"] { + harness.register_tool(Arc::new(EchoTool { + name: name.to_string(), + calls: calls.clone(), + })); + } + + // The harness policy is the enforced source of truth for the cap (the run + // config's value is overwritten by `sync_call_limits` at run start). + harness.with_policy(RunPolicy { + limits: RunLimits::default().with_max_tool_calls(2), + ..RunPolicy::default() + }); + + let recorder = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("tool3"), ()).with_events(recorder.sink()); + + let err = harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect_err("the third call must trip the tool-call cap"); + assert!(matches!(err, TinyAgentsError::LimitExceeded(_)), "{err:?}"); + + let events = recorder.events(); + assert!( + started_call_ids(&events).is_empty(), + "no ToolStarted may be emitted for calls that never run: {:?}", + started_call_ids(&events) + ); + assert_eq!( + calls.load(Ordering::SeqCst), + 0, + "no tool may execute when admission fails" + ); +} + +// ── TOOL-5 ──────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn serial_tool_error_becomes_a_model_visible_result() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new("c1", "flaky", json!({}))]), + text_response("recovered"), + ])), + ); + harness.register_tool(Arc::new(ErringTool { + name: "flaky".into(), + policy: ToolErrorPolicy::ReturnToError, + calls: calls.clone(), + })); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("a ReturnToError tool failure must not kill the run"); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + let tool_text = run + .messages + .iter() + .find_map(|m| match m { + Message::Tool(t) => Some(t.content.clone()), + _ => None, + }) + .expect("a tool message should have been appended"); + let rendered = format!("{tool_text:?}"); + assert!( + rendered.contains("transient 503"), + "the model should see the tool error: {rendered}" + ); +} + +#[tokio::test] +async fn concurrent_tool_error_becomes_a_model_visible_result() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ + ToolCall::new("c1", "flaky", json!({})), + ToolCall::new("c2", "echo_a", json!({})), + ]), + text_response("recovered"), + ])), + ); + harness.register_tool(Arc::new(ErringTool { + name: "flaky".into(), + policy: ToolErrorPolicy::ReturnToError, + calls: calls.clone(), + })); + harness.register_tool(Arc::new(EchoTool { + name: "echo_a".into(), + calls: calls.clone(), + })); + + let run = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("a ReturnToError tool failure must not kill a concurrent turn"); + + assert_eq!( + tool_message_ids(&run.messages), + vec!["c1".to_string(), "c2".to_string()], + "both calls must be answered, in original order" + ); +} + +// ── TOOL-6 ──────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn fatal_tool_error_emits_tool_failed_and_clears_active_calls() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ + ToolCall::new("c1", "fatal", json!({})), + ])])), + ); + harness.register_tool(Arc::new(ErringTool { + name: "fatal".into(), + policy: ToolErrorPolicy::Fail, + calls: calls.clone(), + })); + + let recorder = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("tool6"), ()).with_events(recorder.sink()); + + harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect_err("a Fail-policy tool error must abort the run"); + + let events = recorder.events(); + let failed: Vec<_> = events + .iter() + .filter(|event| matches!(event, AgentEvent::ToolFailed { .. })) + .collect(); + assert_eq!( + failed.len(), + 1, + "every ToolStarted needs a terminal partner; got kinds {:?}", + events.iter().map(AgentEvent::kind).collect::>() + ); + match failed[0] { + AgentEvent::ToolFailed { + call_id, + tool_name, + error, + .. + } => { + assert_eq!(call_id.as_str(), "c1"); + assert_eq!(tool_name, "fatal"); + assert!(error.contains("transient 503"), "{error}"); + } + other => panic!("expected ToolFailed, got {other:?}"), + } +} + +// ── TOOL-10 ─────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn duplicate_call_ids_do_not_clear_each_others_active_entry() { + let calls = Arc::new(AtomicUsize::new(0)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + // A provider that reuses one id across two calls in the same turn. + tool_calls_response(vec![ + ToolCall::new("dup", "echo_a", json!({})), + ToolCall::new("dup", "flaky", json!({})), + ]), + text_response("done"), + ])), + ); + harness.register_tool(Arc::new(EchoTool { + name: "echo_a".into(), + calls: calls.clone(), + })); + harness.register_tool(Arc::new(ErringTool { + name: "flaky".into(), + policy: ToolErrorPolicy::Fail, + calls: calls.clone(), + })); + + let recorder = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("tool10"), ()).with_events(recorder.sink()); + + harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect_err("the second (failing) call aborts the run"); + + // `status.active_tool_calls` is internal to the loop, so the observable + // proxy is the started/terminal pairing: two spans open, one closes with + // `ToolCompleted` and one with `ToolFailed`. Positional release is what + // keeps the two duplicate entries independent; a `retain` would have + // cleared both on the first completion. + let events = recorder.events(); + assert_eq!( + started_call_ids(&events), + vec!["dup".to_string(), "dup".to_string()] + ); + assert_eq!(completed_call_ids(&events), vec!["dup".to_string()]); + assert_eq!( + events + .iter() + .filter(|event| matches!(event, AgentEvent::ToolFailed { .. })) + .count(), + 1, + "the failing duplicate needs its own terminal event" + ); +} + +// ── TOOL-11 ─────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn unknown_tool_recovery_emits_started_and_completed() { + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new("c1", "missing", json!({}))]), + text_response("recovered"), + ])), + ); + harness.with_policy(RunPolicy { + unknown_tool: UnknownToolPolicy::ReturnToolError, + ..RunPolicy::default() + }); + + let recorder = EventRecorder::new(); + let ctx = RunContext::new(RunConfig::new("tool11"), ()).with_events(recorder.sink()); + + harness + .invoke_in_context(&(), ctx, vec![Message::user("go")]) + .await + .expect("ReturnToolError recovers"); + + let events = recorder.events(); + assert_eq!( + started_call_ids(&events), + vec!["c1".to_string()], + "a recovery must still open a tool span" + ); + assert_eq!( + completed_call_ids(&events), + vec!["c1".to_string()], + "a recovery must close its tool span" + ); +} + +// ── TOOL-12 ─────────────────────────────────────────────────────────────────── + +#[tokio::test] +async fn before_tool_rejection_surfaces_as_a_middleware_error() { + // The tool-call slot released when `before_tool` refuses a call is not + // externally observable (a refusal aborts the run, so no later admission + // ever reads the counter). What is observable, and what this pins, is that + // the cap is still checked *before* the hook runs: the refusal, not a limit + // error, is what surfaces while budget remains. + let calls = Arc::new(AtomicUsize::new(0)); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![tool_calls_response(vec![ + ToolCall::new("c1", "echo_a", json!({})), + ])])), + ); + harness.register_tool(Arc::new(EchoTool { + name: "echo_a".into(), + calls: calls.clone(), + })); + harness.push_middleware(Arc::new(RejectingMiddleware)); + + let err = harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect_err("the rejecting middleware aborts the run"); + assert!(matches!(err, TinyAgentsError::Middleware(_)), "{err:?}"); + assert_eq!( + calls.load(Ordering::SeqCst), + 0, + "a refused call must never reach the tool" + ); +} + +// ── Injected-argument enforcement ───────────────────────────────────────────── + +#[tokio::test] +async fn injected_arguments_are_stripped_before_validation() { + let seen = Arc::new(std::sync::Mutex::new(Vec::new())); + let mut harness: AgentHarness<()> = AgentHarness::new(); + harness.register_model( + "mock", + Arc::new(MockModel::with_responses(vec![ + tool_calls_response(vec![ToolCall::new( + "c1", + "injected", + // The model forges the hidden key it was never shown. + json!({ "query": "hi", "thread_id": "forged" }), + )]), + text_response("done"), + ])), + ); + harness.register_tool(Arc::new(InjectedArgTool { seen: seen.clone() })); + + harness + .invoke_default(&(), vec![Message::user("go")]) + .await + .expect("the call validates against the model-facing schema"); + + let observed = seen.lock().unwrap().clone(); + assert_eq!(observed.len(), 1); + assert!( + observed[0].get("thread_id").is_none(), + "a forged injected argument must be stripped: {:?}", + observed[0] + ); + assert_eq!(observed[0]["query"], "hi"); +} diff --git a/tests/wave2_tools_structured.rs b/tests/wave2_tools_structured.rs new file mode 100644 index 0000000..825dcae --- /dev/null +++ b/tests/wave2_tools_structured.rs @@ -0,0 +1,222 @@ +//! Regression coverage for the wave-2 structured-output defects (REASON-5, +//! REASON-9) and the typed context-overflow error (C15). +//! +//! | Test | Defect | +//! |------|--------| +//! | `extraction_validates_against_the_stored_schema` | REASON-5(a): the schema was stored and never read | +//! | `validation_error_names_the_failing_instance_path` | REASON-5(a) | +//! | `extraction_repairs_a_fenced_response` | REASON-5(b) | +//! | `extraction_repairs_a_truncated_response` | REASON-5(b) | +//! | `extraction_repairs_relaxed_json` | REASON-5(b) | +//! | `extract_outcome_records_a_failure_instead_of_raising` | REASON-5(c) | +//! | `non_tool_calling_profile_does_not_select_the_tool_call_strategy` | REASON-9 | +//! | `context_overflow_is_a_typed_variant` | C15 | + +use serde_json::json; + +use tinyagents::TinyAgentsError; +use tinyagents::harness::model::{ModelProfile, ModelResponse}; +use tinyagents::harness::structured::{StructuredExtractor, StructuredStrategy}; + +fn score_schema() -> serde_json::Value { + json!({ + "type": "object", + "properties": { "score": { "type": "integer" } }, + "required": ["score"], + "additionalProperties": false + }) +} + +fn extractor() -> StructuredExtractor { + StructuredExtractor::new(StructuredStrategy::ProviderSchema, "score", score_schema()) +} + +// ── REASON-5(a): the schema is now enforced ─────────────────────────────────── + +#[test] +fn extraction_validates_against_the_stored_schema() { + // Well-formed JSON of entirely the wrong shape used to succeed, leaving + // `run.structured` holding something no caller asked for. + let err = extractor() + .extract(&ModelResponse::assistant(r#"{"wrong_key": 1}"#)) + .expect_err("a value that does not conform must not be reported as success"); + assert!( + matches!(err, TinyAgentsError::StructuredOutput(_)), + "{err:?}" + ); +} + +#[test] +fn validation_error_names_the_failing_instance_path() { + let err = extractor() + .extract(&ModelResponse::assistant(r#"{"score": "four"}"#)) + .expect_err("a string is not an integer"); + let message = err.to_string(); + assert!(message.contains("schema 'score'.score"), "{message}"); + assert!(message.contains("integer"), "{message}"); +} + +#[test] +fn a_conforming_value_still_extracts() { + let output = extractor() + .extract(&ModelResponse::assistant(r#"{"score": 4}"#)) + .expect("a conforming value extracts"); + assert_eq!(output.value["score"], 4); +} + +// ── REASON-5(b): the repair ladder ──────────────────────────────────────────── + +#[test] +fn extraction_repairs_a_fenced_response() { + let output = extractor() + .extract(&ModelResponse::assistant("```json\n{\"score\": 4}\n```")) + .expect("a fenced value must not end the run"); + assert_eq!(output.value["score"], 4); +} + +#[test] +fn extraction_repairs_a_truncated_response() { + let schema = json!({ + "type": "object", + "properties": { "summary": { "type": "string" } }, + "required": ["summary"] + }); + let extractor = StructuredExtractor::new(StructuredStrategy::ProviderSchema, "review", schema); + let output = extractor + .extract(&ModelResponse::assistant( + r#"{"summary": "cut off mid-sente"#, + )) + .expect("a truncated value is closed rather than discarding the run"); + assert_eq!(output.value["summary"], "cut off mid-sente"); +} + +#[test] +fn extraction_repairs_relaxed_json() { + let output = extractor() + .extract(&ModelResponse::assistant("{score: 4}")) + .expect("unquoted keys are repaired, as they already are for tool arguments"); + assert_eq!(output.value["score"], 4); +} + +#[test] +fn extraction_still_fails_on_text_that_is_not_json() { + let err = extractor() + .extract(&ModelResponse::assistant("I could not answer that.")) + .expect_err("the repair ladder must not launder prose into a value"); + assert!(err.to_string().contains("no conservative repair"), "{err}"); +} + +// ── REASON-5(c): non-fatal extraction ───────────────────────────────────────── + +#[test] +fn extract_outcome_records_a_failure_instead_of_raising() { + let response = ModelResponse::assistant("I could not answer that."); + let outcome = extractor().extract_outcome(&response); + + assert!(!outcome.is_success()); + assert!(outcome.value.is_none()); + assert!( + outcome + .error + .as_deref() + .is_some_and(|e| e.contains("score")), + "the recorded error must be usable as a repair prompt: {:?}", + outcome.error + ); + assert_eq!( + outcome.raw.text(), + "I could not answer that.", + "the raw response must survive so a caller can re-ask or return it" + ); +} + +#[test] +fn extract_outcome_carries_the_value_on_success() { + let outcome = extractor().extract_outcome(&ModelResponse::assistant(r#"{"score": 4}"#)); + assert!(outcome.is_success()); + assert_eq!(outcome.into_result().unwrap()["score"], 4); +} + +// ── REASON-9: strategy selection respects `tool_calling` ────────────────────── + +#[test] +fn non_tool_calling_profile_does_not_select_the_tool_call_strategy() { + // A profile that can do neither native structured output nor tool calling + // used to route into `ToolCall`, whose forced `ToolChoice::Tool` is dropped + // in prompt-guided mode — an extractor waiting for a call that can never + // arrive. + let profile = ModelProfile { + tool_calling: false, + json_schema: false, + native_structured_output: false, + ..ModelProfile::default() + }; + assert_ne!( + StructuredStrategy::for_profile(Some(&profile)), + StructuredStrategy::ToolCall + ); + assert_eq!( + StructuredStrategy::for_profile(Some(&profile)), + StructuredStrategy::ProviderSchema + ); +} + +#[test] +fn a_tool_calling_profile_still_selects_the_tool_call_strategy() { + let profile = ModelProfile { + tool_calling: true, + ..ModelProfile::default() + }; + assert_eq!( + StructuredStrategy::for_profile(Some(&profile)), + StructuredStrategy::ToolCall + ); +} + +// ── C15: typed context overflow ─────────────────────────────────────────────── + +#[test] +fn context_overflow_is_a_typed_variant() { + let error = TinyAgentsError::ContextOverflow { + provider: "ollama".to_string(), + model: Some("qwen3:8b".to_string()), + message: "this model's maximum context length is 8192 tokens".to_string(), + }; + assert!(error.is_context_overflow()); + assert!(error.to_string().contains("context overflow"), "{error}"); +} + +#[test] +fn a_provider_error_carrying_the_overflow_code_classifies_the_same() { + // Providers construct `TinyAgentsError::Provider` directly today, so + // classification must recognise the code as well as the typed variant — + // otherwise the two paths disagree about the same failure. + use tinyagents::harness::model::ProviderError; + use tinyagents::harness::providers::openai::CONTEXT_OVERFLOW_CODE; + + let provider_error = ProviderError { + provider: "openai".to_string(), + status: Some(400), + code: Some(CONTEXT_OVERFLOW_CODE.to_string()), + message: "context length exceeded".to_string(), + ..ProviderError::default() + }; + let error = TinyAgentsError::from_provider_error(provider_error); + + assert!(error.is_context_overflow(), "{error:?}"); + assert!(matches!(error, TinyAgentsError::ContextOverflow { .. })); +} + +#[test] +fn an_unrelated_provider_error_is_not_a_context_overflow() { + use tinyagents::harness::model::ProviderError; + + let error = TinyAgentsError::from_provider_error(ProviderError { + provider: "openai".to_string(), + status: Some(429), + message: "rate limited".to_string(), + ..ProviderError::default() + }); + assert!(!error.is_context_overflow()); + assert!(matches!(error, TinyAgentsError::Provider(_))); +}