diff --git a/Cargo.lock b/Cargo.lock index 1105a7d39..94f71a533 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1509,6 +1509,7 @@ dependencies = [ "tempfile", "thiserror", "tinyagents-definition", + "tinyinference-core", "tinyinference-embeddings", "tinyinference-image", "tinyinference-llm", diff --git a/crates/tinyagents-harness/Cargo.toml b/crates/tinyagents-harness/Cargo.toml index 3a5b7e936..6e83cb4ca 100644 --- a/crates/tinyagents-harness/Cargo.toml +++ b/crates/tinyagents-harness/Cargo.toml @@ -31,6 +31,7 @@ tracing = { workspace = true } tinyagents-definition = { path = "../tinyagents-definition", version = "2.1.2" } tinytools-agent = { path = "../../vendor/tinytools/crates/tinytools-agent", version = "0.4.1", default-features = false } tinyinference-llm = { path = "../../vendor/tinyinference/crates/tinyinference-llm", version = "0.3.0" } +tinyinference-core = { path = "../../vendor/tinyinference/crates/tinyinference-core", version = "0.3.0" } tinyinference-embeddings = { path = "../../vendor/tinyinference/crates/tinyinference-embeddings", version = "0.3.0" } tinyinference-image = { path = "../../vendor/tinyinference/crates/tinyinference-image", version = "0.3.0", optional = true } tinyinference-video = { path = "../../vendor/tinyinference/crates/tinyinference-video", version = "0.3.0", optional = true } diff --git a/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs b/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs new file mode 100644 index 000000000..9321d8aad --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs @@ -0,0 +1,180 @@ +//! [`ArgRecoveryMiddleware`]: repair malformed tool-call arguments before the +//! harness runs its fatal schema gate. Generic over the run-context payload. + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::context::RunContext; +use crate::error::Result as TaResult; +use crate::middleware::Middleware; +use tinyinference_llm::tool::{ToolCall as TaToolCall, ToolSchema}; + +use tinytools::Tool; + +/// `before_tool`: repair a tool call's arguments *before* the harness runs its +/// fatal pre-execution schema gate (issues #4249 / #4451). A model can emit +/// arguments the model adapter parses to a non-object `Value` — invalid JSON +/// decodes to `Value::Null`, and some providers emit the whole arguments blob as +/// a JSON-encoded *string* (optionally wrapped in a ```json markdown fence). Left +/// alone the harness rejects those against an object schema and aborts the whole +/// turn. +/// +/// Recovery, in order: +/// 1. Already a JSON object → leave it (the common, valid case). +/// 2. A JSON-encoded string (optionally fenced) that decodes to an object → +/// decode and use it. +/// 3. Otherwise a non-object whose tool schema declares **no** required fields → +/// coerce to `{}` (legacy-engine parity: the tool runs and produces its own +/// recoverable error). +/// 4. Otherwise (non-object + schema has required fields) → leave the arguments +/// untouched so the crate's `InvalidArgsPolicy::ReturnToolError` path reports +/// the original validation failure. Coercing to `{}` would discard the raw +/// malformed value without improving recovery. +pub struct ArgRecoveryMiddleware { + /// The same `Arc`-shared tool sets the runner registers, used to resolve a + /// call's schema so we can tell whether coercing to `{}` is safe. + tool_sets: Vec>>>, +} + +impl ArgRecoveryMiddleware { + /// Build the middleware over the runner's shared tool sets. + pub fn new(tool_sets: Vec>>>) -> Self { + Self { tool_sets } + } + + fn schema_for(&self, name: &str) -> Option { + schema_for_tool(&self.tool_sets, name) + } +} + +#[async_trait] +impl Middleware<(), C> for ArgRecoveryMiddleware { + fn name(&self) -> &str { + "arg_recovery" + } + + async fn before_tool( + &self, + _ctx: &mut RunContext, + _state: &(), + call: &mut TaToolCall, + ) -> TaResult<()> { + // (1) Already valid object shape — nothing to do. + if call.arguments.is_object() { + return Ok(()); + } + + // (2) JSON-encoded-string arguments (optionally markdown-fenced): decode + // and adopt the inner object. + if let Some(raw) = call.arguments.as_str() + && let Some(obj) = recover_object_from_json_string(raw) + { + tracing::debug!( + tool = call.name.as_str(), + "[tinyagents::mw] arg_recovery: decoded JSON-encoded-string tool arguments to object" + ); + call.arguments = obj; + return Ok(()); + } + + // (3) Non-object with a permissive schema (no required fields): coerce to + // `{}` so the tool runs and produces its own recoverable error — engine + // parity for tools that predate the schema gate. + let has_required = self + .schema_for(&call.name) + .map(|schema| schema_has_required_fields(&schema.parameters)) + .unwrap_or(false); + if !has_required { + tracing::debug!( + tool = call.name.as_str(), + "[tinyagents::mw] arg_recovery: coercing non-object tool arguments to {{}} (schema declares no required fields)" + ); + call.arguments = serde_json::json!({}); + return Ok(()); + } + + // (4) Non-object + schema has required fields: leave untouched. The + // crate admission policy surfaces a descriptive, recoverable tool + // result without executing the real tool. + tracing::debug!( + tool = call.name.as_str(), + args_kind = json_value_kind(&call.arguments), + "[tinyagents::mw] arg_recovery: leaving non-object tool arguments for crate invalid-args recovery" + ); + Ok(()) + } +} + +/// Resolves the harness [`ToolSchema`] for `name` across the runner's shared +/// tool sets. +/// +/// It constructs the canonical TinyInference schema directly from the +/// TinyTools declaration, so the `parameters` we validate against are +/// byte-identical to the ones the crate's fatal `validate_call` gate checks — +/// otherwise our pre-validation could disagree with the crate and either miss +/// a fatal case or stub a call the crate still rejects. +fn schema_for_tool(tool_sets: &[Arc>>], name: &str) -> Option { + tool_sets + .iter() + .flat_map(|set| set.iter()) + .find(|tool| tool.name() == name) + .map(|tool| { + let spec = tool.spec(); + ToolSchema::new(spec.name, spec.description, spec.parameters) + }) +} + +/// Whether a tool's JSON-schema `parameters` declares any `required` field. +fn schema_has_required_fields(parameters: &serde_json::Value) -> bool { + parameters + .get("required") + .and_then(serde_json::Value::as_array) + .map(|required| required.iter().any(serde_json::Value::is_string)) + .unwrap_or(false) +} + +/// Attempts to recover a JSON **object** from a string-encoded arguments payload: +/// providers sometimes emit the whole arguments blob as a JSON string, optionally +/// wrapped in a ```json markdown fence. Returns `None` when the string does not +/// decode to a JSON object. +fn recover_object_from_json_string(raw: &str) -> Option { + let candidate = strip_code_fence(raw); + serde_json::from_str::(candidate) + .ok() + .filter(serde_json::Value::is_object) +} + +/// Strips a surrounding markdown code fence (```` ```json … ``` ````) and its +/// optional language tag, returning the inner text. A string with no fence is +/// returned trimmed and unchanged. +fn strip_code_fence(raw: &str) -> &str { + let trimmed = raw.trim(); + let Some(after_open) = trimmed.strip_prefix("```") else { + return trimmed; + }; + // Drop an optional language tag on the opening fence line (e.g. `json`). + let body = match after_open.find('\n') { + Some(newline) + if after_open[..newline] + .chars() + .all(|c| c.is_ascii_alphanumeric()) => + { + &after_open[newline + 1..] + } + _ => after_open, + }; + body.trim().strip_suffix("```").unwrap_or(body).trim() +} + +/// A short, human-readable kind label for a JSON value, for debug logging. +fn json_value_kind(value: &serde_json::Value) -> &'static str { + match value { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "boolean", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} diff --git a/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs b/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs new file mode 100644 index 000000000..0d8df8306 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs @@ -0,0 +1,99 @@ +//! Tests for [`ArgRecoveryMiddleware`]. + +use std::sync::Arc; + +use serde_json::json; + +use super::*; +use crate::context::{RunConfig, RunContext}; +use crate::middleware::Middleware; +use tinyinference_llm::tool::ToolCall; +use tinytools::{Tool, ToolResult}; + +struct Stub { + name: &'static str, + required: bool, +} + +#[async_trait::async_trait] +impl Tool for Stub { + fn name(&self) -> &str { + self.name + } + + fn description(&self) -> &str { + "stub" + } + + fn parameters_schema(&self) -> serde_json::Value { + if self.required { + json!({"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}) + } else { + json!({"type": "object", "properties": {}}) + } + } + + async fn execute(&self, _args: serde_json::Value) -> anyhow::Result { + Ok(ToolResult::success("ok")) + } +} + +fn mw() -> ArgRecoveryMiddleware { + let set: Vec> = vec![ + Box::new(Stub { + name: "loose", + required: false, + }), + Box::new(Stub { + name: "strict", + required: true, + }), + ]; + ArgRecoveryMiddleware::new(vec![Arc::new(set)]) +} + +async fn recover(name: &str, arguments: serde_json::Value) -> serde_json::Value { + let mut ctx: RunContext = RunContext::new(RunConfig::new("mw-test"), ()); + let mut call = ToolCall::new("c1", name, arguments); + mw().before_tool(&mut ctx, &(), &mut call).await.unwrap(); + call.arguments +} + +#[tokio::test] +async fn an_object_is_left_alone() { + assert_eq!( + recover("strict", json!({"q": "x"})).await, + json!({"q": "x"}) + ); +} + +#[tokio::test] +async fn a_json_encoded_string_is_decoded() { + assert_eq!( + recover("strict", json!("{\"q\": \"x\"}")).await, + json!({"q": "x"}) + ); +} + +#[tokio::test] +async fn a_fenced_json_string_is_decoded() { + assert_eq!( + recover("strict", json!("```json\n{\"q\": \"x\"}\n```")).await, + json!({"q": "x"}) + ); +} + +#[tokio::test] +async fn a_non_object_for_a_schema_without_required_fields_becomes_an_empty_object() { + assert_eq!(recover("loose", json!(null)).await, json!({})); +} + +#[tokio::test] +async fn a_non_object_for_a_schema_with_required_fields_is_left_for_the_crate_gate() { + assert_eq!(recover("strict", json!(null)).await, json!(null)); +} + +#[tokio::test] +async fn an_unknown_tool_is_treated_as_permissive() { + assert_eq!(recover("nope", json!(3)).await, json!({})); +} diff --git a/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs new file mode 100644 index 000000000..8621ca731 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs @@ -0,0 +1,307 @@ +//! [`ArtifactIndexTocMiddleware`]: render the run's persisted-artifact index +//! into every request as a bounded contents list (issue #6014). + +use async_trait::async_trait; + +use crate::context::RunContext; +use crate::error::Result as TaResult; +use crate::middleware::Middleware; +use tinyinference_llm::message::Message as TaMessage; +use tinyinference_llm::model::ModelRequest; + +use super::image_trim::estimate_text_tokens; + +// ── ArtifactIndexTocMiddleware (issue #6014) ───────────────────────────────── + +/// The namespace a host's tool-output middleware conventionally writes each +/// persisted artifact under; the default for [`ArtifactIndexTocMiddleware::new`]. +pub const ARTIFACT_INDEX_NAMESPACE: &str = "tool_results"; + +/// The input allowance the reductions divide up when no context window is +/// advertised. +/// +/// A window is what every other bound here is derived from, so without one +/// there is nothing to take a share of — and no `ImageAwareMessageTrimMiddleware` +/// installed either, which is precisely why neither share may be left unbounded +/// on that path (CodeRabbit on #6068): these are the two additions nothing +/// downstream can shrink, on the one configuration where nothing downstream is +/// watching. Sized so the contents list's tenth lands on the 512 tokens a +/// realistic handful of artifacts needs. +pub const NO_WINDOW_ALLOWANCE: u64 = 5_120; + +/// The floor under the contents list's proportional share, so a small window +/// still yields a cap rather than collapsing to `0` — which both middlewares +/// read as "unbounded". +const TOC_ALLOWANCE_MIN: u64 = 64; + +/// Reserved for the omitted-count line, which cannot be measured before the +/// row cap decides how many rows were dropped. One short sentence. +pub const FOOTER_ALLOWANCE: u64 = 48; + +/// Split a turn's input allowance between the two things that add to the +/// request: the artifact contents list and the wrap-up's result restoration. +/// +/// Returns `(toc, restore)`. **Neither is ever `0`** — that is the sentinel both +/// middlewares read as "no cap", so handing it to either is the bug this +/// function exists to make unrepresentable. It bit twice on #6068: first when +/// the two bounds were computed independently and the pair went unbounded, then +/// when the no-window fallback was given to the contents list alone and +/// `saturating_sub` handed restoration the sentinel it had just been rescued +/// from. +/// +/// The contents list takes a tenth — generous for the realistic case, firm +/// about the pathological one — floored so a small window cannot round it away, +/// and capped at half so restoration keeps a real share of a small allowance. +pub fn split_input_allowance(trim_allowance: u64) -> (u64, u64) { + let total = if trim_allowance == 0 { + NO_WINDOW_ALLOWANCE + } else { + trim_allowance + }; + let toc = (total / 10) + .max(TOC_ALLOWANCE_MIN) + .min(total.div_ceil(2)) + .max(1); + (toc, total.saturating_sub(toc).max(1)) +} + +/// Renders the run's persisted-artifact index into the request as a short +/// contents list, so the model can always see what this turn has gathered and +/// where the full copy lives. +/// +/// # The reference used to live somewhere it could be destroyed +/// +/// A tool result over the per-result budget is written to disk and replaced by +/// a preview plus an `artifact_path` pointer (`apply_per_result_persistence`). +/// That pointer's only home was the tool-result message itself — and tool-result +/// messages are exactly what the two reduction steps act on. Microcompact +/// replaces a body with `CLEARED_PLACEHOLDER`; compression folds the older slice +/// into a summary that *should* carry paths forward ("prefer concrete facts — +/// paths, names, values") but is a model call, not a guarantee. +/// +/// Either way the outcome is the same and it is silent: the data sits intact on +/// disk with nothing in context saying it exists. Every component did its job — +/// the result was persisted, the transcript was reduced — and the turn quietly +/// lost the ability to reach its own findings. No log fires, because nothing +/// failed. +/// +/// The index has been maintained all along (`ToolResultArtifactIndexStore`, +/// registered on `RunContext.stores`, written on every persist). Nothing read +/// it. This reads it. +/// +/// # Why re-rendered rather than injected once +/// +/// The contents list is built fresh into each request and never enters the +/// loop's own transcript, which is what makes it un-destroyable rather than +/// merely durable: there is no stored copy for a later pass to blank, fold or +/// evict, and it cannot accumulate into a stack of stale lists. It is also +/// always current — an artifact persisted on the previous round appears on the +/// next call with no bookkeeping. +/// +/// Emitted as a **system** message for the same reason: compression keeps every +/// system message verbatim and the trim never drops one, so the one thing that +/// says where the data went is the one thing the ladder may not take. It is +/// appended at the tail rather than the head so the cacheable prompt prefix is +/// untouched. +pub struct ArtifactIndexTocMiddleware { + /// This middleware's share of the turn's input allowance (a tenth, split at + /// the install site so restoration and this list cannot each claim the + /// whole). `0` disables the cap — no advertised window, and no trim + /// installed either. + input_budget: u64, + /// Name of the [`StoreRegistry`](crate::store::StoreRegistry) entry the + /// host records persisted artifacts in. + store_name: String, + /// Namespace within that store the entries live under. + namespace: String, +} + +impl ArtifactIndexTocMiddleware { + /// Build the middleware over the registry store named `store_name`, reading + /// entries from [`ARTIFACT_INDEX_NAMESPACE`]. + pub fn new(input_budget: u64, store_name: impl Into) -> Self { + Self { + input_budget, + store_name: store_name.into(), + namespace: ARTIFACT_INDEX_NAMESPACE.to_string(), + } + } + + /// Read entries from `namespace` instead of [`ARTIFACT_INDEX_NAMESPACE`]. + pub fn with_namespace(mut self, namespace: impl Into) -> Self { + self.namespace = namespace.into(); + self + } +} + +#[async_trait] +impl Middleware<(), C> for ArtifactIndexTocMiddleware { + fn name(&self) -> &str { + "artifact_index_toc" + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &(), + request: &mut ModelRequest, + ) -> TaResult<()> { + let Some(store) = ctx.stores.get(&self.store_name) else { + return Ok(()); + }; + // A read failure and an empty index produce the same contents list — + // none — but they are not the same event: the second is the ordinary + // case, the first means every persisted result is unreachable for this + // call with nothing said about it. Degrading is still right (a contents + // list is not worth failing a turn over), so the difference has to show + // up in the log or it shows up nowhere (CodeRabbit on #6068). + let keys = match store.list(&self.namespace).await { + Ok(keys) => keys, + Err(err) => { + tracing::warn!( + error = %err, + "[tinyagents::mw] could not read the persisted-artifact index; this call gets \ + no contents list and cannot see what was offloaded" + ); + return Ok(()); + } + }; + if keys.is_empty() { + // No result has been offloaded, so there is nothing to point at and + // no reason to spend context saying so. + return Ok(()); + } + + let mut rows: Vec = Vec::new(); + let mut failed_reads = 0usize; + for key in &keys { + let entry = match store.get(&self.namespace, key).await { + Ok(Some(entry)) => entry, + Ok(None) => continue, + Err(err) => { + failed_reads += 1; + tracing::warn!( + key = %key, + error = %err, + "[tinyagents::mw] could not read a persisted-artifact index entry" + ); + continue; + } + }; + let tool = entry.get("tool").and_then(|v| v.as_str()).unwrap_or("tool"); + let Some(path) = entry.get("artifact_path").and_then(|v| v.as_str()) else { + continue; + }; + let bytes = entry + .get("original_bytes") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + rows.push(format!("- `{tool}` → `{path}` ({bytes} bytes)")); + } + if failed_reads > 0 && rows.is_empty() { + tracing::warn!( + failed_reads, + "[tinyagents::mw] no persisted-artifact index entries could be read" + ); + } + if rows.is_empty() { + return Ok(()); + } + // Sorted so the list is stable between calls when nothing was added — + // the index is a `HashMap`, whose iteration order is not. An unstable + // ordering would rewrite this message on every call for no reason, + // costing cache and making the diff unreadable in a trace. + rows.sort(); + let total = rows.len(); + // Cap the list against the input allowance (CodeRabbit on #6068). + // + // This message is a system message precisely so the ladder cannot take + // it — compression keeps system messages verbatim and the trim never + // evicts one — which means nothing downstream can shrink it either. A + // turn that persisted enough oversized results would otherwise grow an + // un-evictable message without limit, and the guarantee that made the + // pointers safe would be the thing that broke the request. + // + // A tenth of the input allowance: generous for the realistic case (a + // handful of artifacts is a few hundred bytes) and firm about the + // pathological one. The omitted count is reported rather than the list + // silently ending, so the model knows more exist — the same disclosure + // rule the rest of this ladder follows, and the reason the artifacts are + // findable at all. + let header = format!( + "## Stored results from this turn\n\n\ + {total} tool result(s) were too large to keep inline and were written to disk. The \ + text you saw for them is a preview; the full content is at the path below and can be \ + read with the file-reading tool when you need detail the preview does not carry.\n\n" + ); + // One line that still names the count, for a share too small to hold + // the header at all (CodeRabbit on #6068). Truncating to zero rows was + // half the fix: the fixed text is ~118 tokens and the floor a small + // window gets is 64, so the message still cleared its share with no + // rows in it. + let compact = format!( + "_{total} tool result(s) were written to disk — ask by tool name to locate one._" + ); + if self.input_budget > 0 { + let cap = self.input_budget.max(1); + let fixed = estimate_text_tokens(&header).saturating_add(FOOTER_ALLOWANCE); + if fixed > cap { + // Even the compact line has to fit. Saying nothing loses the + // pointer, which is bad; pushing an unshrinkable system message + // over the bound makes the trim evict transcript instead, which + // is worse. + if estimate_text_tokens(&compact) <= cap { + request.messages.push(TaMessage::system(compact)); + } + return Ok(()); + } + // Already this middleware's share of the turn's allowance — the + // split happens once, at the install site, so the two things that + // add to the request cannot each spend the whole of it. + // Seeded with the fixed text, so the cap bounds the whole message + // rather than the rows alone (CodeRabbit on #6068). The header is a + // paragraph; against the 64-token floor a small window gets, it is + // comparable to the rows it introduces. `FOOTER_ALLOWANCE` stands in + // for the omitted-count line, which cannot be measured before the + // count is known — it is one short sentence, and reserving it when + // nothing is omitted only spends a row. + let mut used: u64 = estimate_text_tokens(&header).saturating_add(FOOTER_ALLOWANCE); + let mut kept = 0usize; + for row in &rows { + used = used.saturating_add(estimate_text_tokens(row)); + if used > cap { + break; + } + kept += 1; + } + // No `max(1)`: once `used` is seeded with the fixed text, a share + // smaller than that text leaves `kept == 0`, and forcing a row then + // puts the whole message over a bound nothing downstream can shrink + // — this is a system message, so compression keeps it and the trim + // never evicts it, and the overshoot is charged against the + // transcript instead (CodeRabbit on #6068). The header and the + // omitted count still say the results exist and how to ask for one. + rows.truncate(kept); + } + let shown = rows.len(); + let omitted = total - shown; + let count = total; + tracing::debug!( + artifacts = count, + "[tinyagents::mw] rendering the persisted-artifact contents list" + ); + request.messages.push(TaMessage::system(format!( + "{header}{}{}", + rows.join("\n"), + if omitted > 0 { + format!( + "\n\n_({omitted} more stored result(s) not listed here — ask for the one you \ + need by tool name and it can be located.)_" + ) + } else { + String::new() + } + ))); + Ok(()) + } +} diff --git a/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs new file mode 100644 index 000000000..973867808 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs @@ -0,0 +1,403 @@ +//! Tests for [`ArtifactIndexTocMiddleware`] and its allowance split. + +use std::sync::Arc; + +use super::*; +use crate::context::{RunConfig, RunContext}; +use crate::middleware::Middleware; +use tinyinference_llm::message::Message as TaMessage; +use tinyinference_llm::model::ModelRequest; + +/// A fixed set of captured outcomes, in the shape the wrap-up reads them. +struct Sink(Vec<(String, String)>); + +impl CapturedOutcomes for Sink { + fn content_for( + &self, + call_id: &str, + ) -> std::result::Result, OutcomesUnavailable> { + Ok(self + .0 + .iter() + .find(|(id, _)| id == call_id) + .map(|(_, content)| content.clone())) + } +} + +fn sink_with(entries: &[(&str, &str)]) -> Arc { + Arc::new(Sink( + entries + .iter() + .map(|(id, content)| ((*id).to_string(), (*content).to_string())) + .collect(), + )) +} + +const STORE: &str = "tool_result_artifact_index"; + +async fn ctx_with_artifacts(entries: &[(&str, &str, &str, u64)]) -> RunContext { + ctx_with_artifacts_and_config(entries, RunConfig::new("mw-test")).await +} + +async fn ctx_with_artifacts_and_config( + entries: &[(&str, &str, &str, u64)], + config: RunConfig, +) -> RunContext { + use crate::store::{InMemoryStore, Store, StoreRegistry}; + let index = Arc::new(InMemoryStore::new()); + for (call_id, tool, path, bytes) in entries { + let mut fields = serde_json::Map::new(); + fields.insert("tool".to_string(), (*tool).into()); + fields.insert("call_id".to_string(), (*call_id).into()); + fields.insert("artifact_path".to_string(), (*path).into()); + fields.insert( + "original_bytes".to_string(), + serde_json::Value::from(*bytes), + ); + Store::put(index.as_ref(), "tool_results", call_id, fields.into()) + .await + .unwrap(); + } + let mut registry = StoreRegistry::new(); + registry.register(STORE, index); + RunContext::new(config, ()).with_stores(registry) +} + +#[tokio::test] +async fn toc_is_absent_when_no_result_was_offloaded() { + let mw = ArtifactIndexTocMiddleware::new(0, STORE); + let mut ctx = ctx_with_artifacts(&[]).await; + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert_eq!(request.messages.len(), 1, "no artifacts, no contents list"); +} + +/// The paths reach the model even though nothing in the transcript carries +/// them — which is the whole point: the pointer's old home (the tool body) is +/// what the reduction steps act on. +#[tokio::test] +async fn toc_lists_every_persisted_artifact_as_a_system_message() { + let mw = ArtifactIndexTocMiddleware::new(0, STORE); + let mut ctx = ctx_with_artifacts(&[ + ("call-1", "fetch_issues", "outputs/issues-p1.json", 240_000), + ("call-2", "web_search", "outputs/search-2.json", 91_000), + ]) + .await; + // A transcript that has already lost both pointers: one body blanked by + // microcompact, the other never present. + let mut request = ModelRequest { + messages: vec![ + TaMessage::user("what did you find?"), + TaMessage::tool("call-1", DEFAULT_CLEARED_PLACEHOLDER), + ], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let last = request.messages.last().expect("a contents list"); + assert!( + matches!(last, TaMessage::System(_)), + "must be a system message — compression keeps those verbatim and the trim never drops \ + one, so the pointer cannot be reduced away: {last:?}" + ); + let text = last.text(); + assert!( + text.contains("outputs/issues-p1.json"), + "missing path: {text}" + ); + assert!( + text.contains("outputs/search-2.json"), + "missing path: {text}" + ); + assert!(text.contains("fetch_issues"), "missing tool name: {text}"); +} + +/// Rendered fresh per request, so repeated passes cannot stack stale lists. +#[tokio::test] +async fn toc_does_not_accumulate_across_calls() { + let mw = ArtifactIndexTocMiddleware::new(0, STORE); + let mut ctx = ctx_with_artifacts(&[("call-1", "fetch", "outputs/a.json", 100)]).await; + + let mut first = ModelRequest { + messages: vec![TaMessage::user("hi")], + ..Default::default() + }; + mw.before_model(&mut ctx, &(), &mut first).await.unwrap(); + + // The loop rebuilds the request from its own transcript each iteration, so + // the second call starts from the same base — not from the mutated first. + let mut second = ModelRequest { + messages: vec![TaMessage::user("hi")], + ..Default::default() + }; + mw.before_model(&mut ctx, &(), &mut second).await.unwrap(); + + assert_eq!(first.messages.len(), 2); + assert_eq!( + second.messages.len(), + 2, + "exactly one contents list per request" + ); +} + +/// The contents list is capped, and says how many it left out. +/// +/// It is a system message so the ladder cannot take it — which also means +/// nothing downstream can shrink it, so it has to bound itself. +#[tokio::test] +async fn toc_is_capped_and_reports_what_it_omitted() { + let entries: Vec<(String, String, String, u64)> = (0..200) + .map(|i| { + ( + format!("call-{i}"), + format!("tool_with_a_long_name_{i}"), + format!("outputs/a-fairly-long-artifact-path-{i}.json"), + 1_000, + ) + }) + .collect(); + let borrowed: Vec<(&str, &str, &str, u64)> = entries + .iter() + .map(|(a, b, c, d)| (a.as_str(), b.as_str(), c.as_str(), *d)) + .collect(); + let mut ctx = ctx_with_artifacts(&borrowed).await; + // This middleware's share of the allowance, already split at the install + // site — not the whole turn allowance. + let mw = ArtifactIndexTocMiddleware::new(200, STORE); + let mut request = ModelRequest { + messages: vec![TaMessage::user("what did you find?")], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let text = request.messages.last().expect("a contents list").text(); + assert!( + text.contains("not listed here"), + "an omitted count must be disclosed, not silently dropped: {text}" + ); + // The nominal share, not a loose multiple of it: the cap now seeds `used` + // with the header and the reserved footer, so it bounds the whole message + // rather than the rows alone (CodeRabbit on #6068). + assert!( + estimate_text_tokens(&text) <= 200, + "the list must stay inside the share it was given: {}", + estimate_text_tokens(&text) + ); +} + +/// Neither share may ever be `0`, because `0` is the sentinel both middlewares +/// read as "unbounded" — so the allowances that round down into it are the ones +/// to pin (CodeRabbit on #6068, twice). A model advertising no window is the +/// sharp case: no window means no `ImageAwareMessageTrimMiddleware` either, so +/// these two additions are the ones nothing downstream would shrink. +#[test] +fn neither_share_is_ever_the_unbounded_sentinel() { + for trim_allowance in [0_u64, 1, 5, 9, 99, 100, 2_000, 100_000] { + let (toc, restore) = split_input_allowance(trim_allowance); + assert!(toc > 0, "allowance {trim_allowance} disabled the TOC cap"); + assert!( + restore > 0, + "allowance {trim_allowance} left restoration unbounded" + ); + // The same effective total the function computes — not + // `trim_allowance.max(NO_WINDOW_ALLOWANCE)`, which accepted any + // over-split below 5_120 (and whose `.min(u64::MAX)` was a no-op). + let total = if trim_allowance == 0 { + NO_WINDOW_ALLOWANCE + } else { + trim_allowance + }; + // `1` stays excluded: the function deliberately keeps both shares + // positive, which costs one token at that degenerate size. + if trim_allowance > 1 { + assert!( + toc + restore <= total, + "allowance {trim_allowance} split into more than it had: {toc} + {restore}" + ); + } + } + // With no window the split is taken from the fallback total, not from zero. + let (toc, restore) = split_input_allowance(0); + assert_eq!(toc, 512); + assert_eq!(toc + restore, NO_WINDOW_ALLOWANCE); + // Proportional sizing is untouched where there is something to divide. + assert_eq!(split_input_allowance(2_000), (200, 1_800)); + assert_eq!(split_input_allowance(100_000), (10_000, 90_000)); +} + +/// The end of the same argument: with no window, a run holding far more +/// artifacts than any list should carry still produces a bounded message. +#[tokio::test] +async fn the_contents_list_stays_bounded_when_no_window_is_advertised() { + let entries: Vec<(String, String, String, u64)> = (0..400) + .map(|i| { + ( + format!("call-{i}"), + "file_read".to_string(), + format!("outputs/a-fairly-long-artifact-path-{i}.json"), + 1_000, + ) + }) + .collect(); + let borrowed: Vec<(&str, &str, &str, u64)> = entries + .iter() + .map(|(a, b, c, d)| (a.as_str(), b.as_str(), c.as_str(), *d)) + .collect(); + let mut ctx = ctx_with_artifacts(&borrowed).await; + let mw = ArtifactIndexTocMiddleware::new(split_input_allowance(0).0, STORE); + let mut request = ModelRequest { + messages: vec![TaMessage::user("what did you find?")], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let text = request.messages.last().expect("a contents list").text(); + assert!( + text.contains("not listed here"), + "an omitted count must be disclosed, not silently dropped: {text}" + ); + assert!( + estimate_text_tokens(&text) <= split_input_allowance(0).0, + "the no-window fallback must bound the whole list message: {}", + estimate_text_tokens(&text) + ); +} + +/// The two consumers share one allowance, so the bound that matters is the one +/// on the request they *both* wrote to (CodeRabbit on #6068). Runs them in +/// registration order — wrap-up, then contents list — on the no-window path, +/// where nothing downstream would trim what they overshoot by. +#[tokio::test] +async fn both_middlewares_together_stay_inside_the_no_window_allowance() { + let big = "x".repeat(8_000); + let outcomes: Vec<(String, String)> = (0..40) + .map(|i| (format!("call-{i}"), big.clone())) + .collect(); + let borrowed: Vec<(&str, &str)> = outcomes + .iter() + .map(|(a, b)| (a.as_str(), b.as_str())) + .collect(); + let artifacts: Vec<(String, String, String, u64)> = (0..400) + .map(|i| { + ( + format!("call-{i}"), + "file_read".to_string(), + format!("outputs/a-fairly-long-artifact-path-{i}.json"), + 1_000, + ) + }) + .collect(); + let artifact_refs: Vec<(&str, &str, &str, u64)> = artifacts + .iter() + .map(|(a, b, c, d)| (a.as_str(), b.as_str(), c.as_str(), *d)) + .collect(); + + let mut ctx = ctx_with_artifacts_and_config( + &artifact_refs, + RunConfig::new("mw-test").with_max_model_calls(2), + ) + .await; + ctx.limits.record_model_call().unwrap(); + ctx.limits.record_model_call().unwrap(); + + let (toc_allowance, restore_allowance) = split_input_allowance(0); + let mut request = ModelRequest { + messages: borrowed + .iter() + .map(|(id, _)| TaMessage::tool(*id, DEFAULT_CLEARED_PLACEHOLDER)) + .collect(), + ..Default::default() + }; + + // Registration order at the install site: wrap-up first, contents list after. + FinalCallWrapUpMiddleware::new( + "CONCLUDE NOW", + "WRITE NOW", + sink_with(&borrowed), + restore_allowance, + ) + .before_model(&mut ctx, &(), &mut request) + .await + .unwrap(); + ArtifactIndexTocMiddleware::new(toc_allowance, STORE) + .before_model(&mut ctx, &(), &mut request) + .await + .unwrap(); + + let used: u64 = request.messages.iter().map(estimate_message_tokens).sum(); + assert!( + used <= NO_WINDOW_ALLOWANCE, + "the combined request must stay inside the shared allowance: {used} > {NO_WINDOW_ALLOWANCE}" + ); + // Both actually contributed — otherwise the bound is met vacuously. + assert!( + request.messages.iter().any(|m| m.text().starts_with("xxx")), + "the wrap-up restored nothing, so the bound proves nothing" + ); + assert!( + request + .messages + .iter() + .any(|m| m.text().contains("Stored results from this turn")), + "the contents list is absent, so the bound proves nothing" + ); +} + +/// A share smaller than the fixed text renders no rows at all. Forcing one +/// would push a message nothing downstream can shrink over its bound, and the +/// header plus the omitted count already disclose that the results exist +/// (CodeRabbit on #6068). +#[tokio::test] +async fn a_share_smaller_than_the_header_renders_no_rows() { + let entries: Vec<(String, String, String, u64)> = (0..10) + .map(|i| { + ( + format!("call-{i}"), + "file_read".to_string(), + format!("outputs/artifact-{i}.json"), + 1_000, + ) + }) + .collect(); + let borrowed: Vec<(&str, &str, &str, u64)> = entries + .iter() + .map(|(a, b, c, d)| (a.as_str(), b.as_str(), c.as_str(), *d)) + .collect(); + let mut ctx = ctx_with_artifacts(&borrowed).await; + // The floor a small window gets — below the header's own cost. + let mw = ArtifactIndexTocMiddleware::new(64, STORE); + let mut request = ModelRequest { + messages: vec![TaMessage::user("what did you find?")], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let text = request.messages.last().expect("a contents list").text(); + assert!( + !text.contains("file_read"), + "no row fits the share, so none may be forced: {text}" + ); + assert!( + text.contains("10 tool result(s) were written to disk"), + "the count must still be named so the results stay findable: {text}" + ); + // The bound that matters: the whole rendered message inside the share. + // Truncating to zero rows was not enough — the header and footer alone are + // ~118 tokens against this 64-token floor, and the earlier form of this + // assertion (`< 64 + FOOTER_ALLOWANCE + 64`) was loose enough to accept + // exactly that overshoot (CodeRabbit on #6068). + assert!( + estimate_text_tokens(&text) <= 64, + "the message must fit the share it was given: {}", + estimate_text_tokens(&text) + ); +} diff --git a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs new file mode 100644 index 000000000..93c0e8e68 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs @@ -0,0 +1,170 @@ +//! [`CredentialScrubMiddleware`]: scrub credential-shaped secrets out of every +//! tool result before it leaves the tool boundary. + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::context::RunContext; +use crate::error::Result as TaResult; +use crate::middleware::{MiddlewareToolOutcome, ToolHandler, ToolMiddleware}; +use tinyinference_llm::tool::ToolCall as TaToolCall; + +/// The placeholder `scrub_credentials` emits once per redacted value. Counted +/// to tell the model *how many* values went, which is the difference between +/// "something was withheld" and a number it can relay. +pub const REDACTION_PLACEHOLDER: &str = "*[REDACTED]"; + +/// Appended to a scrubbed tool result so the **model** learns what the log +/// already knew. +/// +/// Without it the model receives a result that silently differs from what the +/// tool returned: it cannot find the content it was asked for, re-runs the same +/// call, gets an identically scrubbed result, and never converges — until the +/// successful-repeat tracker halts the run and the user is told "Incomplete" +/// with no reason (#6416). The redaction itself is correct and unchanged; only +/// its silence was the defect. +/// +/// The "do not retry" clause is load-bearing: a retry is guaranteed to be +/// scrubbed identically, so it is the one action that cannot help. +/// +/// Deliberately contains no `: ` shape, so it cannot match +/// `SENSITIVE_KV_REGEX` and scrub itself on a second pass — pinned by +/// `the_notice_does_not_scrub_itself`. +pub fn redaction_notice(count: usize) -> String { + format!( + "[credential_scrub] {count} value(s) in this result were redacted as credentials. \ + Re-running this tool returns the same redaction, so do not retry — tell the user \ + which values were withheld and that they can view them directly in the source app." + ) +} + +/// Scrub `content`, returning the replacement text **and** how many values +/// went — or `None` when nothing was credential-shaped. +/// +/// Split out of `wrap_tool` so the decision and the composed result are +/// directly testable. Exercising this is the difference between proving the +/// notice text is well-formed and proving a scrubbed result actually carries +/// it; only `replace_tool_result_text` plumbing stays untested. +pub fn scrub_with_notice(content: &str) -> Option<(String, usize)> { + let scrubbed = tinyinference_core::sanitize::scrub_credentials(content); + if scrubbed == content { + return None; + } + // Count what this pass removed, not what the text already carried — a + // result may legitimately contain the placeholder already. + let redactions = scrubbed + .matches(REDACTION_PLACEHOLDER) + .count() + .saturating_sub(content.matches(REDACTION_PLACEHOLDER).count()); + Some(( + format!("{scrubbed}\n\n{}", redaction_notice(redactions)), + redactions, + )) +} + +/// Decides, per tool, how a result's text is scrubbed: `(tool_name, content)` +/// to the replacement text and redaction count, or `None` when nothing needs +/// to change. The default is [`scrub_with_notice`] for every tool; a host that +/// must protect a host-minted field in one tool's output supplies its own and +/// falls back to [`scrub_with_notice`] for everything else. +pub type ToolScrubber = Arc Option<(String, usize)> + Send + Sync>; + +/// `wrap_tool`: scrub credential-shaped secrets out of every tool result before +/// it leaves the tool boundary. Secrets in tool output (env dumps, config +/// reads, API responses, shell output) would otherwise reach model context, +/// on-disk transcripts, worker-thread mirrors, and any tool-outcome capture +/// sink. +/// +/// Install it as the **innermost** tool wrap (pushed last), so it observes the +/// RAW tool result first and scrubs it before any outer wrap, the `after_tool` +/// chain (summarization/caps), the transcript push, or an outcome-capture sink +/// can see the unredacted content. Generic over the run-context payload. +pub struct CredentialScrubMiddleware { + scrubber: ToolScrubber, +} + +impl Default for CredentialScrubMiddleware { + fn default() -> Self { + Self::new() + } +} + +impl CredentialScrubMiddleware { + /// Scrub every tool's result with [`scrub_with_notice`]. + pub fn new() -> Self { + Self { + scrubber: Arc::new(|_tool, content| scrub_with_notice(content)), + } + } + + /// Scrub with a host-supplied per-tool scrubber (see [`ToolScrubber`]). + pub fn with_scrubber(scrubber: ToolScrubber) -> Self { + Self { scrubber } + } +} + +#[async_trait] +impl ToolMiddleware<(), C> for CredentialScrubMiddleware { + fn name(&self) -> &str { + "credential_scrub" + } + + async fn wrap_tool( + &self, + ctx: &mut RunContext, + state: &(), + call: TaToolCall, + next: ToolHandler<'_, (), C>, + ) -> TaResult { + let tool_name = call.name.clone(); + let outcome = next.run(ctx, state, call).await?; + // `MiddlewareToolOutcome` is `#[non_exhaustive]`; today it only carries a + // `Result`, but match rather than irrefutable-let so a future variant + // fails loud instead of silently bypassing scrubbing. + let mut result = match outcome { + MiddlewareToolOutcome::Result(result) => result, + other => return Ok(other), + }; + + let mut redactions = 0usize; + for block in &mut result.content { + let (text, is_json) = match block { + tinytools::ToolContent::Text { text } => (text.clone(), false), + tinytools::ToolContent::Json { data } => (data.to_string(), true), + tinytools::ToolContent::Image { .. } | tinytools::ToolContent::File { .. } => { + continue; + } + }; + if let Some((scrubbed, count)) = (self.scrubber)(&tool_name, &text) { + if is_json { + if let Ok(value) = serde_json::from_str(&scrubbed) { + if let tinytools::ToolContent::Json { data } = block { + *data = value; + } + } else if let tinytools::ToolContent::Json { data } = block { + *data = serde_json::Value::String(scrubbed); + } + } else if let tinytools::ToolContent::Text { text } = block { + *text = scrubbed; + } + redactions = redactions.saturating_add(count); + } + } + if redactions > 0 { + tracing::warn!( + tool = %tool_name, + redactions, + "[tinyagents::mw] credential_scrub redacted secret(s) from tool result content" + ); + // The notice goes to the model, in the result itself. The warning + // above goes to the log, where no model will ever read it. + result.content.push(tinytools::ToolContent::Text { + text: redaction_notice(redactions), + }); + result.markdown_formatted = None; + } + + Ok(MiddlewareToolOutcome::Result(result)) + } +} diff --git a/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs new file mode 100644 index 000000000..4378f6fb2 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs @@ -0,0 +1,209 @@ +//! #6416: the redaction must reach the model, not only the log. +//! +//! The scrubber's behaviour is deliberately NOT relaxed here — a signed +//! magic-link token in an email is a credential, and letting it through would +//! re-open #4453 for the commonest credential shape in mail. What changes is +//! that a scrubbed result now says so, so the agent stops re-fetching content +//! that will be redacted identically every time. +//! +//! The first test is the one that matters: it proves the redaction still +//! happens. A notice that appeared while the scrubbing silently weakened would +//! be the worst outcome this change could have. + +use super::*; +use crate::context::{RunConfig, RunContext}; +use std::sync::Arc; + +use crate::middleware::{BoxToolFuture, MiddlewareStack, ToolBaseCall}; +use tinyinference_core::sanitize::scrub_credentials; +use tinyinference_llm::tool::ToolCall; +use tinytools::ToolResult; + +/// A magic-link token — exactly the shape QA hit in #6416 — must still be +/// removed. This is the security half of the contract and comes first. +#[test] +fn a_link_token_in_email_content_is_still_redacted() { + let email = "Click https://acme.example/verify?token=aB3dE5fG7hJ9kL1mN3pQ to sign in."; + let scrubbed = scrub_credentials(email); + + assert!( + !scrubbed.contains("aB3dE5fG7hJ9kL1mN3pQ"), + "the token body must not survive scrubbing: {scrubbed}" + ); + assert!( + scrubbed.contains(REDACTION_PLACEHOLDER), + "a redaction must leave its placeholder: {scrubbed}" + ); +} + +/// The notice must not match the scrubber's own patterns, or a second pass +/// would redact the explanation and the model would be told even less than +/// before. Verified rather than reasoned about: `credential` is followed by +/// `]` and a space here, never by `["']?\s*[:=]`. +#[test] +fn the_notice_does_not_scrub_itself() { + let notice = redaction_notice(2); + assert_eq!( + scrub_credentials(¬ice), + notice, + "the notice must survive the scrubber unchanged, or it would redact its own explanation" + ); +} + +/// Scrubbing already-scrubbed content must be a no-op, which is what makes the +/// middleware idempotent: the `scrubbed != content` guard is false on a second +/// pass, so no second notice is appended and the count is not restated. +#[test] +fn scrubbing_is_idempotent_so_a_second_pass_adds_no_second_notice() { + let email = "Click https://acme.example/verify?token=aB3dE5fG7hJ9kL1mN3pQ to sign in."; + let once = scrub_credentials(email); + let twice = scrub_credentials(&once); + assert_eq!(twice, once, "scrub must be stable on its own output"); + + // The value the middleware would hand on: scrubbed body plus the notice. + let annotated = format!("{once}\n\n{}", redaction_notice(1)); + assert_eq!( + scrub_credentials(&annotated), + annotated, + "a result already carrying the notice must pass through untouched" + ); +} + +/// The count is what the model relays to the user, so it must describe this +/// pass rather than every placeholder in the text. +#[test] +fn the_notice_counts_only_what_this_pass_removed() { + let two = "api_key=aB3dE5fG7hJ9kL1mN3pQ and token=zX9yW8vU7tS6rQ5pO4nM"; + let scrubbed = scrub_credentials(two); + let removed = scrubbed.matches(REDACTION_PLACEHOLDER).count() + - two.matches(REDACTION_PLACEHOLDER).count(); + assert_eq!(removed, 2, "both values were redacted: {scrubbed}"); + + assert!(redaction_notice(removed).contains("2 value(s)")); + assert!( + redaction_notice(removed).contains("do not retry"), + "the notice must tell the model retrying cannot help — that loop is the defect" + ); +} + +/// Ordinary email content with nothing credential-shaped must be untouched, so +/// the notice never appears on a result that lost nothing. +#[test] +fn unremarkable_content_is_left_alone_and_gets_no_notice() { + let email = "Lunch at 12:30 tomorrow? The meeting room is booked until 2pm."; + assert_eq!( + scrub_credentials(email), + email, + "nothing here is credential-shaped; a notice would be a lie" + ); +} + +/// The behaviour this change exists for: a scrubbed result **carries** the +/// notice. The tests above prove the notice text is well-formed and that the +/// redaction still happens; this proves the two are actually joined, which is +/// the part a caller sees. +#[test] +fn a_scrubbed_result_carries_the_notice_and_the_count() { + let email = "Click https://acme.example/verify?token=aB3dE5fG7hJ9kL1mN3pQ to sign in."; + let (annotated, redactions) = + scrub_with_notice(email).expect("credential-shaped content must be scrubbed"); + + assert_eq!(redactions, 1); + assert!( + !annotated.contains("aB3dE5fG7hJ9kL1mN3pQ"), + "the token must not survive: {annotated}" + ); + assert!( + annotated.contains("[credential_scrub]"), + "the model must be told the result was altered: {annotated}" + ); + assert!( + annotated.contains("do not retry"), + "the model must be told retrying cannot help: {annotated}" + ); +} + +/// And a result that lost nothing must be left exactly alone — no notice, and +/// no rewrite of the tool's own output. +#[test] +fn an_unscrubbed_result_is_not_annotated_at_all() { + assert!( + scrub_with_notice("Lunch at 12:30 tomorrow? Room booked until 2pm.").is_none(), + "nothing was redacted, so the result must pass through untouched" + ); +} + +struct FixedBase(&'static str); + +impl ToolBaseCall<(), ()> for FixedBase { + fn call<'a>( + &'a self, + _ctx: &'a mut RunContext, + _state: &'a (), + _call: ToolCall, + ) -> BoxToolFuture<'a> { + Box::pin(async move { Ok(ToolResult::success(self.0)) }) + } +} + +async fn run(mw: CredentialScrubMiddleware, tool: &str, body: &'static str) -> String { + let mut stack: MiddlewareStack<()> = MiddlewareStack::new(); + stack.push_tool_middleware(Arc::new(mw)); + let mut ctx: RunContext = RunContext::new(RunConfig::new("mw-test"), ()); + stack + .run_wrapped_tool( + &mut ctx, + &(), + ToolCall::new("c1", tool, serde_json::json!({})), + &FixedBase(body), + ) + .await + .unwrap() + .into_result() + .output() +} + +/// The model-visible result carries the notice, and nothing else changes. +#[tokio::test] +async fn the_middleware_annotates_a_scrubbed_result() { + let out = run( + CredentialScrubMiddleware::new(), + "fetch", + "Click https://acme.example/verify?token=aB3dE5fG7hJ9kL1mN3pQ to sign in.", + ) + .await; + assert!(!out.contains("aB3dE5fG7hJ9kL1mN3pQ"), "{out}"); + assert!(out.contains("[credential_scrub] 1 value(s)"), "{out}"); +} + +#[tokio::test] +async fn the_middleware_leaves_a_clean_result_exactly_alone() { + let out = run( + CredentialScrubMiddleware::new(), + "fetch", + "Lunch at 12:30 tomorrow?", + ) + .await; + assert_eq!(out, "Lunch at 12:30 tomorrow?"); +} + +/// A host scrubber is handed the tool name and its verdict is what lands. +#[tokio::test] +async fn a_host_scrubber_sees_the_tool_name() { + let scrubber: ToolScrubber = + Arc::new(|tool, content| (tool == "special").then(|| (format!("{content} [handled]"), 0))); + let handled = run( + CredentialScrubMiddleware::with_scrubber(scrubber.clone()), + "special", + "body", + ) + .await; + assert_eq!(handled, "body [handled]"); + let untouched = run( + CredentialScrubMiddleware::with_scrubber(scrubber), + "other", + "body", + ) + .await; + assert_eq!(untouched, "body"); +} diff --git a/crates/tinyagents-harness/src/middleware/library/image_trim.rs b/crates/tinyagents-harness/src/middleware/library/image_trim.rs new file mode 100644 index 000000000..fc5083a58 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/image_trim.rs @@ -0,0 +1,260 @@ +//! [`ImageAwareMessageTrimMiddleware`] and the image-aware token estimation it +//! (and the wrap-up / contents-list middlewares) budget with. +//! +//! Generic over the run-context payload `C`: nothing here reads it. + +use async_trait::async_trait; + +use crate::context::RunContext; +use crate::error::Result as TaResult; +use crate::middleware::Middleware; +use tinyinference_llm::message::{ContentBlock, Message as TaMessage}; +use tinyinference_llm::model::ModelRequest; + +// ── ImageAwareMessageTrimMiddleware ─────────────────────────────────────────── + +/// Flat token cost charged per image — an inline `[IMAGE:…]` marker or a native +/// [`ContentBlock::Image`] block — instead of counting the base64 payload as +/// text. Restores the legacy `harness/token_budget.rs` semantics (issue #4462): +/// the crate `estimate_tokens` prices text at chars/4, so a single large base64 +/// image reads as ~2M tokens and the trim believes the context is massively over +/// budget, evicting the whole transcript (system messages included). Providers +/// bill an image at ≈85–1100 tokens by detail; 1200 is a conservative upper +/// bound that keeps the budget realistic without the base64 payload inflating it. +pub const IMAGE_MARKER_TOKEN_COST: u64 = 1_200; + +/// Inline image-marker prefix produced by the multimodal composer +/// (`agent/multimodal.rs`, `compose_multimodal_message`). Priced at +/// [`IMAGE_MARKER_TOKEN_COST`] rather than by its base64 length. +const IMAGE_MARKER_PREFIX: &str = "[IMAGE:"; + +/// Minimum reply/output reserve — mirrors the legacy `MIN_OUTPUT_RESERVE_TOKENS`. +const MIN_OUTPUT_RESERVE_TOKENS: u64 = 512; + +/// Upper anchor for the reply/output reserve — mirrors the legacy +/// `DEFAULT_OUTPUT_RESERVE_TOKENS`. +const DEFAULT_OUTPUT_RESERVE_TOKENS: u64 = 8_192; + +/// Rough token estimate (~4 characters per token) with inline `[IMAGE:…]` +/// markers charged a flat [`IMAGE_MARKER_TOKEN_COST`] instead of their base64 +/// length. Mirrors the deleted `token_budget::estimate_tokens` (issue #4462). +/// Markerless text takes the fast char/4 path. +pub fn estimate_text_tokens(text: &str) -> u64 { + if !text.contains(IMAGE_MARKER_PREFIX) { + return (text.len() as u64).saturating_add(3) / 4; + } + let mut text_bytes: u64 = 0; + let mut images: u64 = 0; + let mut cursor = 0usize; + while let Some(rel) = text[cursor..].find(IMAGE_MARKER_PREFIX) { + let start = cursor + rel; + text_bytes = text_bytes.saturating_add((start - cursor) as u64); // preceding text + let after = start + IMAGE_MARKER_PREFIX.len(); + match text[after..].find(']') { + Some(rel_end) => { + images += 1; + cursor = after + rel_end + 1; // skip the whole marker payload + } + None => { + // Unterminated marker — count the remainder as text and stop. + text_bytes = text_bytes.saturating_add((text.len() - start) as u64); + cursor = text.len(); + break; + } + } + } + text_bytes = text_bytes.saturating_add((text.len() - cursor) as u64); // trailing text + (text_bytes.saturating_add(3) / 4) + .saturating_add(images.saturating_mul(IMAGE_MARKER_TOKEN_COST)) +} + +/// Count native [`ContentBlock::Image`] blocks on a message. `Message::text()` +/// concatenates only text blocks, so a native multimodal image would otherwise +/// contribute zero tokens; we charge each one [`IMAGE_MARKER_TOKEN_COST`]. +fn count_native_image_blocks(msg: &TaMessage) -> u64 { + let content = match msg { + TaMessage::System(m) => &m.content, + TaMessage::User(m) => &m.content, + TaMessage::Assistant(m) => &m.content, + TaMessage::Tool(m) => &m.content, + // Out-of-band host record; carries no content blocks. + TaMessage::Custom(_) => return 0, + }; + content + .iter() + .filter(|b| matches!(b, ContentBlock::Image(_))) + .count() as u64 +} + +/// Estimate the tokens of a crate [`TaMessage`]: image-aware text tokens, a flat +/// [`IMAGE_MARKER_TOKEN_COST`] per native image block, and the assistant's +/// tool-call name/arguments (which `Message::text()` drops). Mirrors the legacy +/// `estimate_conversation_message_tokens` (issue #4462). +pub fn estimate_message_tokens(msg: &TaMessage) -> u64 { + let mut total = estimate_text_tokens(&msg.text()); + total = total + .saturating_add(count_native_image_blocks(msg).saturating_mul(IMAGE_MARKER_TOKEN_COST)); + let content = match msg { + TaMessage::System(m) => &m.content, + TaMessage::User(m) => &m.content, + TaMessage::Assistant(m) => &m.content, + TaMessage::Tool(m) => &m.content, + TaMessage::Custom(_) => return total, + }; + // `Message::text()` omits structured JSON and provider-visible thinking / + // extension blocks. Charge their serialized representation so large + // structured tool results cannot bypass the prompt budget. + for block in content { + match block { + ContentBlock::Json(value) => { + total = total.saturating_add(estimate_text_tokens(&value.to_string())); + } + ContentBlock::Thinking { text, .. } => { + total = total.saturating_add(estimate_text_tokens(text)); + } + ContentBlock::RedactedThinking { data } => { + total = total.saturating_add(estimate_text_tokens(data)); + } + ContentBlock::ProviderExtension(value) => { + total = total.saturating_add(estimate_text_tokens(&value.to_string())); + } + ContentBlock::Image(_) + | ContentBlock::Audio(_) + | ContentBlock::Video(_) + | ContentBlock::Document(_) => { + total = total.saturating_add(IMAGE_MARKER_TOKEN_COST); + } + ContentBlock::Text(_) => {} + } + } + if let TaMessage::Assistant(m) = msg { + for call in &m.tool_calls { + total = total.saturating_add(estimate_text_tokens(&call.name)); + total = total.saturating_add(estimate_text_tokens(&call.arguments.to_string())); + } + } + total +} + +/// Reply/output reserve, mirroring the legacy proportional clamp +/// `clamp(window/10, ≥512, ≤max(8192, window/4))`. Restores the small-window +/// budget the fixed `window − AGENT_TURN_MAX_OUTPUT_TOKENS` regressed (issue +/// #4462): an 8k model reserves ~819 tokens (input budget ~7373), not +/// 16384 → floored 1024. +fn legacy_output_reserve_tokens(window: u64) -> u64 { + let pct = window / 10; + pct.max(MIN_OUTPUT_RESERVE_TOKENS) + .min(DEFAULT_OUTPUT_RESERVE_TOKENS.max(window / 4)) +} + +/// Input-prompt token budget after reserving room for the reply. Public to the +/// seam so the install site (and tests) can assert the legacy proportional +/// formula (issue #4462). +pub fn legacy_max_input_tokens(window: u64) -> u64 { + window.saturating_sub(legacy_output_reserve_tokens(window)) +} + +/// Deterministic history trim that replaces the crate `MessageTrimMiddleware` +/// (issue #4462), restoring three regression guards the crate trim lost: +/// +/// 1. **Image-aware token estimate** — inline `[IMAGE:…]` markers and native +/// image blocks are each charged a flat [`IMAGE_MARKER_TOKEN_COST`] instead +/// of their base64 length, so one large image can no longer read as ~2M +/// tokens and evict the whole transcript. +/// 2. **System messages never dropped** — only non-system history is evictable; +/// the crate trim reorders system messages to the front and drops them as a +/// last resort. +/// 3. **Order preserved + observable** — retained messages keep their original +/// relative order, leading orphaned tool results are snapped past (so no +/// provider 400), and any eviction logs a grep-able `warn` carrying +/// (messages dropped, messages/tokens before-and-after). +pub struct ImageAwareMessageTrimMiddleware { + /// Input-prompt token budget (already net of the proportional reply reserve). + budget: u64, +} + +impl ImageAwareMessageTrimMiddleware { + /// Build a trim middleware whose budget is the legacy proportional + /// [`legacy_max_input_tokens`] for `window` (issue #4462) — NOT the crate's + /// fixed `window − 16384`. Floored at 1 so the budget is always positive. + pub fn for_context_window(window: u64) -> Self { + Self { + budget: legacy_max_input_tokens(window).max(1), + } + } +} + +#[async_trait] +impl Middleware<(), C> for ImageAwareMessageTrimMiddleware { + fn name(&self) -> &str { + "image_aware_message_trim" + } + + async fn before_model( + &self, + _ctx: &mut RunContext, + _state: &(), + request: &mut ModelRequest, + ) -> TaResult<()> { + let messages = &mut request.messages; + let original_tokens: u64 = messages.iter().map(estimate_message_tokens).sum(); + if original_tokens <= self.budget { + return Ok(()); + } + let original_len = messages.len(); + + // Evict oldest non-system messages first, preserving the relative order + // of every retained message (rebuilding as `system ++ other` would + // reorder history when a system message appears after non-system ones — + // exactly the crate-trim regression). System messages are NEVER dropped. + let mut removable_positions: Vec = messages + .iter() + .enumerate() + .filter_map(|(idx, m)| (!matches!(m, TaMessage::System(_))).then_some(idx)) + .collect(); + + let mut removed = 0usize; + while !removable_positions.is_empty() { + let total: u64 = messages.iter().map(estimate_message_tokens).sum(); + if total <= self.budget { + break; + } + let absolute_idx = removable_positions.remove(0); + // Subsequent positions shift left by one for every prior removal. + let remove_at = absolute_idx - removed; + messages.remove(remove_at); + removed += 1; + } + + // Snap the window forward past any leading orphaned tool results: dropping + // an `assistant(tool_calls)` while keeping its `tool` answer leaves the + // transcript opening on a tool message with no preceding tool-call, which + // native providers reject with a 400. Drop leading tool results until the + // first non-system message is a clean turn boundary. + while let Some(first_non_system) = messages + .iter() + .position(|m| !matches!(m, TaMessage::System(_))) + { + if matches!(messages[first_non_system], TaMessage::Tool(_)) { + messages.remove(first_non_system); + removed += 1; + } else { + break; + } + } + + if removed > 0 { + let final_tokens: u64 = messages.iter().map(estimate_message_tokens).sum(); + tracing::warn!( + messages_dropped = removed, + messages_before = original_len, + messages_after = messages.len(), + tokens_before = original_tokens, + tokens_after = final_tokens, + budget = self.budget, + "[tinyagents::mw] message_trim evicted oldest history to fit the token budget" + ); + } + Ok(()) + } +} diff --git a/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs new file mode 100644 index 000000000..2f36cfaed --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs @@ -0,0 +1,138 @@ +//! Tests for the image-aware token estimate and trim. + +use super::*; +use crate::context::{RunConfig, RunContext}; +use crate::middleware::Middleware; +use tinyinference_llm::message::Message as TaMessage; +use tinyinference_llm::model::ModelRequest; + +// Image-aware token estimation. A base64 image marker must be priced at the +// flat IMAGE_MARKER_TOKEN_COST, not chars/4 of its payload — otherwise one +// image reads as millions of tokens and the trimmer evicts everything, +// including the system prompt. + +#[test] +fn estimate_text_tokens_markerless_is_chars_over_four() { + assert_eq!(estimate_text_tokens(&"a".repeat(40)), 40_u64.div_ceil(4)); + assert_eq!(estimate_text_tokens(""), 0); +} + +#[test] +fn estimate_text_tokens_prices_image_marker_flat_not_by_length() { + let huge = "x".repeat(40_000); + let text = format!("[IMAGE:{huge}]"); + let tokens = estimate_text_tokens(&text); + // chars/4 of the payload would be ~10_000; the flat price is 1_200. + assert!( + tokens >= IMAGE_MARKER_TOKEN_COST, + "at least the flat image cost: {tokens}" + ); + assert!( + tokens < 2_000, + "image priced flat, not by base64 length: {tokens}" + ); +} + +#[test] +fn estimate_text_tokens_charges_each_image_marker_once() { + let tokens = estimate_text_tokens("[IMAGE:aaaa] and [IMAGE:bbbb]"); + assert!( + tokens >= 2 * IMAGE_MARKER_TOKEN_COST, + "two images each priced: {tokens}" + ); + assert!( + tokens < 2 * IMAGE_MARKER_TOKEN_COST + 100, + "no runaway from the surrounding text: {tokens}" + ); +} + +#[test] +fn an_unterminated_marker_is_counted_as_text() { + let text = format!("[IMAGE:{}", "y".repeat(400)); + assert_eq!(estimate_text_tokens(&text), (text.len() as u64).div_ceil(4)); +} + +#[test] +fn the_input_budget_reserves_a_proportional_reply() { + // window/10, floored at 512 and capped at max(8192, window/4). + assert_eq!(legacy_max_input_tokens(8_192), 8_192 - 819); + assert_eq!(legacy_max_input_tokens(1_000), 1_000 - 512); + assert_eq!(legacy_max_input_tokens(200_000), 200_000 - 20_000); + assert_eq!(legacy_max_input_tokens(2_000_000), 2_000_000 - 200_000); +} + +async fn trim(window: u64, messages: Vec) -> Vec { + let mw = ImageAwareMessageTrimMiddleware::for_context_window(window); + let mut ctx: RunContext = RunContext::new(RunConfig::new("mw-test"), ()); + let mut request = ModelRequest { + messages, + ..Default::default() + }; + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + request.messages +} + +#[tokio::test] +async fn a_request_inside_the_budget_is_untouched() { + let messages = vec![TaMessage::system("sys"), TaMessage::user("hi")]; + let kept = trim(8_192, messages).await; + assert_eq!(kept.len(), 2); +} + +#[tokio::test] +async fn one_large_image_does_not_evict_the_transcript() { + let image = format!("[IMAGE:{}]", "x".repeat(400_000)); + let messages = vec![ + TaMessage::system("sys"), + TaMessage::user("earlier"), + TaMessage::user(image), + ]; + let kept = trim(32_000, messages).await; + assert_eq!( + kept.len(), + 3, + "the image is priced flat, so nothing is evicted" + ); +} + +#[tokio::test] +async fn eviction_keeps_system_messages_and_order_and_drops_oldest_first() { + let big = "z".repeat(4_000); + let messages = vec![ + TaMessage::system("sys-1"), + TaMessage::user(format!("old {big}")), + TaMessage::system("sys-2"), + TaMessage::user(format!("mid {big}")), + TaMessage::user("newest"), + ]; + // Budget ~1_000 input tokens: two big bodies cannot both stay. + let kept = trim(1_600, messages).await; + let texts: Vec = kept.iter().map(|m| m.text()).collect(); + assert!(texts.contains(&"sys-1".to_string()), "{texts:?}"); + assert!(texts.contains(&"sys-2".to_string()), "{texts:?}"); + assert!(texts.contains(&"newest".to_string()), "{texts:?}"); + assert!(!texts.iter().any(|t| t.starts_with("old ")), "{texts:?}"); + let sys1 = texts.iter().position(|t| t == "sys-1").unwrap(); + let sys2 = texts.iter().position(|t| t == "sys-2").unwrap(); + assert!(sys1 < sys2, "relative order is preserved: {texts:?}"); +} + +#[tokio::test] +async fn eviction_never_leaves_an_orphaned_leading_tool_result() { + let big = "z".repeat(4_000); + let messages = vec![ + TaMessage::system("sys"), + TaMessage::user(format!("ask {big}")), + TaMessage::tool("call-1", "answer"), + TaMessage::user("newest"), + ]; + let kept = trim(1_100, messages).await; + let first_non_system = kept + .iter() + .find(|m| !matches!(m, TaMessage::System(_))) + .expect("something remains"); + assert!( + !matches!(first_non_system, TaMessage::Tool(_)), + "a leading tool result without its call is a provider 400: {kept:?}" + ); +} diff --git a/crates/tinyagents-harness/src/middleware/library/mod.rs b/crates/tinyagents-harness/src/middleware/library/mod.rs index bc8524b22..1f5fc857f 100644 --- a/crates/tinyagents-harness/src/middleware/library/mod.rs +++ b/crates/tinyagents-harness/src/middleware/library/mod.rs @@ -13,6 +13,14 @@ //! [`DynamicPromptMiddleware`], [`RedactionMiddleware`], and //! [`TracingMiddleware`] implement the lifecycle [`Middleware`] trait. //! +//! - **Context ladder and loop guards** — [`ImageAwareMessageTrimMiddleware`], +//! [`ArtifactIndexTocMiddleware`], [`FinalCallWrapUpMiddleware`], +//! [`RepeatProgressMiddleware`] / [`RepeatEvictionObserver`], +//! [`ArgRecoveryMiddleware`] and [`CredentialScrubMiddleware`]. All are +//! generic over the run-context payload; host policy (tool names, the +//! cleared-result placeholder, exemptions, per-tool scrubbing) is injected +//! through constructor arguments. +//! //! Type definitions live in `types`; this file holds the constructors and //! trait impls. Tests live in `test.rs`. //! @@ -50,11 +58,49 @@ use tinyinference_llm::model::{ModelDelta, ModelRequest, ModelResponse, Response use tinyinference_llm::tool::{ToolCall, ToolDelta, ToolSchema}; use tinytools::ToolResult; +mod arg_recovery; +mod artifact_toc; mod budget; mod context; +mod credential_scrub; +mod image_trim; mod observe; +mod repeat_progress; mod resilience; mod tool_policy; +mod wrap_up; + +pub use arg_recovery::ArgRecoveryMiddleware; +pub use artifact_toc::{ + ARTIFACT_INDEX_NAMESPACE, ArtifactIndexTocMiddleware, FOOTER_ALLOWANCE, NO_WINDOW_ALLOWANCE, + split_input_allowance, +}; +pub use credential_scrub::{ + CredentialScrubMiddleware, REDACTION_PLACEHOLDER, ToolScrubber, redaction_notice, + scrub_with_notice, +}; +pub use image_trim::{ + IMAGE_MARKER_TOKEN_COST, ImageAwareMessageTrimMiddleware, estimate_message_tokens, + estimate_text_tokens, legacy_max_input_tokens, +}; +pub use repeat_progress::{ + HaltSummarySlot, RepeatEvictionObserver, RepeatExemption, RepeatProgressMiddleware, +}; +pub use wrap_up::{ + CapturedOutcomes, DEFAULT_CLEARED_PLACEHOLDER, FinalCallWrapUpMiddleware, OutcomesUnavailable, +}; +#[cfg(test)] +mod arg_recovery_test; +#[cfg(test)] +mod artifact_toc_test; +#[cfg(test)] +mod credential_scrub_test; +#[cfg(test)] +mod image_trim_test; +#[cfg(test)] +mod repeat_progress_test; #[cfg(test)] mod test; +#[cfg(test)] +mod wrap_up_test; diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs new file mode 100644 index 000000000..02dd57817 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs @@ -0,0 +1,468 @@ +//! [`RepeatProgressMiddleware`]: host adapter for the crate successful-repeat +//! tracker — halts identical-output / identical-call loops that succeed but +//! make no progress (#4088 / #4095), including loops whose repeats are not +//! back to back (#6275). + +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; + +use crate::context::RunContext; +use crate::error::Result as TaResult; +use crate::middleware::{Middleware, ToolInvocationIdentity}; +use crate::no_progress::{SuccessfulRepeat, SuccessfulRepeatTracker, fingerprint_arguments}; +use crate::steering::{SteeringCommand, SteeringHandle}; +use tinyinference_llm::message::{ContentBlock, Message}; +use tinyinference_llm::model::{ModelRequest, ModelResponse}; +use tinytools::ToolResult as TaToolResult; + +use super::wrap_up::DEFAULT_CLEARED_PLACEHOLDER; + +/// Shared slot a guard writes its root-cause halt summary into when it trips, so +/// the turn can surface the cause instead of an empty or last-model reply. +pub type HaltSummarySlot = Arc>>; + +/// Whether a tool is contractually re-invoked with identical arguments (a +/// polling/wait tool), so an identical repeat is progress rather than a loop. +pub type RepeatExemption = Arc bool + Send + Sync>; + +/// Extract the assistant's visible text (concatenated [`ContentBlock::Text`] +/// blocks) from a model response message, for the repeat-output signature. +fn assistant_visible_text(message: &tinyinference_llm::message::AssistantMessage) -> String { + let mut out = String::new(); + for block in &message.content { + if let ContentBlock::Text(t) = block { + if !out.is_empty() { + out.push('\n'); + } + out.push_str(t); + } + } + out +} + +/// Per-batch state the repeat-CALL guard needs but can only fully evaluate once +/// every tool result in the assistant's batch has come back: the canonical +/// `(tool, args)` signature captured at `after_model`, plus the running +/// success/remaining accounting folded in at each `after_tool`. +#[derive(Default)] +struct PendingCallBatch { + /// Canonical `(tool, args)` signature of the batch, from `after_model`. + call_sig: String, + /// Tool results still outstanding for this batch. + remaining: usize, + /// `true` while every result so far in the batch has succeeded. + all_ok: bool, + /// `true` when every call in the batch is a polling/wait exemption. + exempt: bool, + /// `call_id` → per-call `(tool, argument fingerprint)` signature for the + /// recurrence ledger. Polling/wait calls are left out. + call_sigs: HashMap, + /// `true` once a result in this batch has already halted the run, so the + /// batch does not pause it a second time. + halted: bool, +} + +/// Tracker state shared between [`RepeatProgressMiddleware`] and its +/// [`RepeatEvictionObserver`]. +struct RepeatState { + tracker: Mutex>, + /// The body a cleared tool result carries. + cleared_placeholder: String, + /// `call_id`s of the results fed to the recurrence ledger since its last reset. + recorded: Mutex>>, + /// Recorded results still verbatim in the current request before any + /// reduction step ran; compared against the final request by the observer. + visible_before_reduction: Mutex>>, +} + +impl RepeatState { + fn new(placeholder: impl Into) -> Self { + Self { + tracker: Mutex::default(), + cleared_placeholder: placeholder.into(), + recorded: Mutex::default(), + visible_before_reduction: Mutex::default(), + } + } +} + +/// The `ids` whose tool result is still in `request` with its body intact. +fn visible_tool_results( + request: &ModelRequest, + ids: &HashSet, + placeholder: &str, +) -> HashSet { + request + .messages + .iter() + .filter_map(|message| match message { + Message::Tool(tool) + if ids.contains(&tool.tool_call_id) && message.text() != placeholder => + { + Some(tool.tool_call_id.clone()) + } + _ => None, + }) + .collect() +} + +/// Host adapter for the crate's successful-repeat tracker (#4088 / #4095). +/// [`SuccessfulRepeatTracker`] owns the generic streak accounting; this adapter +/// builds canonical tool signatures, applies the host's polling-tool +/// exemption ([`RepeatExemption`]), and maps a crate halt verdict into the shared halt summary and +/// steering pause: +/// +/// - **Repeat-output** (`after_model`, checked before the tools run): halts when +/// the assistant's visible text + tool-call `(name, args)` batch is byte +/// identical [`DEFAULT_REPEAT_OUTPUT_THRESHOLD`] iterations in a row. +/// - **Repeat-call** (evaluated once the batch's tool results are all back, gated +/// on every call succeeding): halts when the `(tool, args)` batch alone repeats +/// [`DEFAULT_REPEAT_CALL_THRESHOLD`] times — catching successful no-op loops +/// that vary only their narration. +/// - **Recurrence** (each successful `after_tool`, #6275): halts when one call +/// returns the identical result [`DEFAULT_REPEAT_CALL_THRESHOLD`] times in the +/// run, adjacent or not — catching a model cycling A, B, A, B through steps +/// whose results it already has. The result is the content after the earlier +/// `after_tool` layers (caps, summarizer), i.e. what the model saw, so a +/// re-read whose output changed does not count. A result that compaction +/// later evicts stops counting; see [`RepeatEvictionObserver`]. +/// +/// Polling/wait tools (per the [`RepeatExemption`]) are exempt from all three: +/// their contract is to be re-invoked identically, so an all-poll batch resets +/// the streaks instead of recording. On a trip it writes the legacy root-cause +/// summary into the shared [`HaltSummarySlot`] and pauses +/// the run through the shared steering handle — the same halt mechanism as the +/// repeated-failure breaker. +/// +/// [`DEFAULT_REPEAT_OUTPUT_THRESHOLD`]: crate::no_progress::DEFAULT_REPEAT_OUTPUT_THRESHOLD +/// [`DEFAULT_REPEAT_CALL_THRESHOLD`]: crate::no_progress::DEFAULT_REPEAT_CALL_THRESHOLD +pub struct RepeatProgressMiddleware { + handle: SteeringHandle, + halt_summary: HaltSummarySlot, + exempt: RepeatExemption, + state: Arc, + /// Batch bookkeeping bridging `after_model` → `after_tool` for the call guard. + pending: Mutex>, +} + +impl RepeatProgressMiddleware { + /// Build the guard. `exempt` names the polling/wait tools that are exempt + /// from all three checks; the cleared-result placeholder defaults to + /// [`DEFAULT_CLEARED_PLACEHOLDER`]. + pub fn new( + handle: SteeringHandle, + halt_summary: HaltSummarySlot, + exempt: RepeatExemption, + ) -> Self { + Self { + handle, + halt_summary, + exempt, + state: Arc::new(RepeatState::new(DEFAULT_CLEARED_PLACEHOLDER)), + pending: Mutex::default(), + } + } + + /// Override the placeholder body treated as an evicted tool result. Must be + /// called before [`eviction_observer`](Self::eviction_observer). + pub fn with_cleared_placeholder(mut self, placeholder: impl Into) -> Self { + self.state = Arc::new(RepeatState::new(placeholder)); + self + } + + /// The companion that resets the recurrence ledger when a recorded result is + /// evicted from context. It must be registered after every reduction + /// middleware, and this guard before them. + pub fn eviction_observer(&self) -> RepeatEvictionObserver { + RepeatEvictionObserver { + state: Arc::clone(&self.state), + } + } + + /// Latch a root-cause halt: record the summary the turn surfaces instead of an + /// empty/last-model reply, and pause at the top of the next iteration (before + /// the next model call), matching the repeated-failure breaker's halt path. + fn halt(&self, summary: String) { + if let Ok(mut slot) = self.halt_summary.lock() { + *slot = Some(summary); + } + self.handle.send(SteeringCommand::Pause); + } +} + +#[async_trait] +impl Middleware<(), C> for RepeatProgressMiddleware { + fn name(&self) -> &str { + "repeat_progress" + } + + async fn after_agent( + &self, + ctx: &mut RunContext, + _state: &(), + _run: &mut crate::middleware::AgentRun, + ) -> TaResult<()> { + let run_id = ctx.instance_id(); + if let Ok(mut trackers) = self.state.tracker.lock() { + trackers.remove(&run_id); + } + if let Ok(mut pending) = self.pending.lock() { + pending.remove(&run_id); + } + if let Ok(mut recorded) = self.state.recorded.lock() { + recorded.remove(&run_id); + } + if let Ok(mut visible) = self.state.visible_before_reduction.lock() { + visible.remove(&run_id); + } + Ok(()) + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &(), + request: &mut ModelRequest, + ) -> TaResult<()> { + // Registered ahead of the reduction steps, so this is the request as the + // loop built it. The observer compares the same ids after they ran. + let run_id = ctx.instance_id(); + let visible = match self.state.recorded.lock() { + Ok(recorded) if recorded.get(&run_id).is_some_and(|ids| !ids.is_empty()) => { + visible_tool_results( + request, + recorded.get(&run_id).expect("checked above"), + &self.state.cleared_placeholder, + ) + } + _ => HashSet::new(), + }; + if let Ok(mut slot) = self.state.visible_before_reduction.lock() { + slot.insert(run_id, visible); + } + Ok(()) + } + + async fn after_model( + &self, + ctx: &mut RunContext, + _state: &(), + response: &mut ModelResponse, + ) -> TaResult<()> { + let tool_calls = &response.message.tool_calls; + if tool_calls.is_empty() { + // A final answer (no tool calls) ends the loop; nothing to guard, and + // there is no batch to track for the call guard. + if let Ok(mut pending) = self.pending.lock() { + pending.remove(&ctx.instance_id()); + } + return Ok(()); + } + + // Polling/wait tools are contractually re-invoked with identical args + + // narration each timeout while the work is still running, so an all-poll + // batch is legitimate progress, not a no-progress repeat. + let all_exempt = tool_calls.iter().all(|c| (self.exempt)(&c.name)); + + // Canonical `(tool, args)` batch signature (call guard) and the broader + // narration+call signature (output guard). Both fold each call in order + // with a `\u{1}` separator, matching the legacy signatures. + let mut call_sig = String::new(); + for call in tool_calls { + call_sig.push('\u{1}'); + call_sig.push_str(&call.name); + call_sig.push('\u{1}'); + call_sig.push_str(&call.arguments.to_string()); + } + let output_sig = format!( + "{}{}", + assistant_visible_text(&response.message).trim(), + call_sig + ); + // Per-call signatures for the recurrence ledger. The invocation + // identity joins each post-tool result to the provider call id. + let call_sigs = tool_calls + .iter() + .filter(|call| !(self.exempt)(&call.name)) + .map(|call| { + ( + call.id.clone(), + format!( + "{}\u{1}{}", + call.name, + fingerprint_arguments(&call.arguments) + ), + ) + }) + .collect(); + + // Stage output with the crate tracker. Its halt verdict is intentionally + // deferred until the matching tool batch is confirmed successful. + if let Ok(mut trackers) = self.state.tracker.lock() { + let _ = trackers + .entry(ctx.instance_id()) + .or_default() + .record_output(&output_sig, all_exempt); + } + + // Stage the batch for the repeat-CALL guard, evaluated once every result + // is back (gated on success) in `after_tool`. + if let Ok(mut pending) = self.pending.lock() { + pending.insert( + ctx.instance_id(), + PendingCallBatch { + call_sig, + remaining: tool_calls.len(), + all_ok: true, + exempt: all_exempt, + call_sigs, + halted: false, + }, + ); + } + Ok(()) + } + + async fn after_tool( + &self, + ctx: &mut RunContext, + _state: &(), + invocation: &ToolInvocationIdentity, + result: &mut TaToolResult, + ) -> TaResult<()> { + let tool_name = invocation.tool_name(); + let call_id = invocation.call_id().to_string(); + // Fold this result into the pending batch; the call guard only acts once + // the batch is complete so it sees whole-batch success. + let (already_halted, recurrence, completed) = { + let Ok(mut pending) = self.pending.lock() else { + return Ok(()); + }; + let Some(batch) = pending.get_mut(&ctx.instance_id()) else { + return Ok(()); + }; + let already_halted = batch.halted; + let mut recurrence = SuccessfulRepeat::Continue; + if result.is_error { + batch.all_ok = false; + } else if let Some(sig) = batch.call_sigs.get(&call_id) { + if let Ok(mut trackers) = self.state.tracker.lock() { + recurrence = trackers + .entry(ctx.instance_id()) + .or_default() + .record_call_outcome(sig, &result.output()); + } + if let Ok(mut recorded) = self.state.recorded.lock() { + recorded + .entry(ctx.instance_id()) + .or_default() + .insert(call_id); + } + } + if matches!(recurrence, SuccessfulRepeat::Halt(_)) { + batch.halted = true; + } + batch.remaining = batch.remaining.saturating_sub(1); + let completed = if batch.remaining == 0 { + pending.remove(&ctx.instance_id()) + } else { + None + }; + (already_halted, recurrence, completed) + }; + if already_halted { + // An earlier result in this batch paused the run; keep the streak + // accounting current without pausing again. + if let Some(batch) = completed + && let Ok(mut trackers) = self.state.tracker.lock() + { + let _ = trackers + .entry(ctx.instance_id()) + .or_default() + .record_call_batch(&batch.call_sig, batch.all_ok, batch.exempt); + } + return Ok(()); + } + + let batch_verdict = completed.and_then(|batch| { + self.state.tracker.lock().ok().map(|mut trackers| { + trackers + .entry(ctx.instance_id()) + .or_default() + .record_call_batch(&batch.call_sig, batch.all_ok, batch.exempt) + }) + }); + // When both fire on the same result, the batch summary wins: it is the + // more specific description of an adjacent repeat. + let summary = match (batch_verdict, recurrence) { + (Some(SuccessfulRepeat::Halt(summary)), _) => summary, + (_, SuccessfulRepeat::Halt(summary)) => summary, + _ => return Ok(()), + }; + tracing::warn!( + tool = tool_name, + "[tinyagents::mw] crate successful-repeat tracker halted the run" + ); + self.halt(summary); + Ok(()) + } +} + +/// Resets [`RepeatProgressMiddleware`]'s recurrence ledger when context +/// reduction evicts a recorded tool result (#6275). +/// +/// Compression, microcompact and trim rewrite only the outgoing request, never +/// the loop's transcript, so the eviction is visible only inside the +/// `before_model` chain. The guard snapshots which recorded results are intact +/// before those steps; this observer, registered after them, checks the same +/// ids in the final request. A result that was blanked or dropped is no longer +/// in front of the model, so re-reading it is not a repeat the model can see, +/// and the tracker restarts. Keyed on what disappeared rather than on which +/// middleware removed it, so any reduction step is covered, and a dialect that +/// never carries tool results as tool messages is never snapshotted. +pub struct RepeatEvictionObserver { + state: Arc, +} + +#[async_trait] +impl Middleware<(), C> for RepeatEvictionObserver { + fn name(&self) -> &str { + "repeat_progress_eviction" + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &(), + request: &mut ModelRequest, + ) -> TaResult<()> { + let run_id = ctx.instance_id(); + let before = match self.state.visible_before_reduction.lock() { + Ok(mut slot) => slot.remove(&run_id).unwrap_or_default(), + Err(_) => return Ok(()), + }; + if before.is_empty() { + return Ok(()); + } + let evicted = before.len() + - visible_tool_results(request, &before, &self.state.cleared_placeholder).len(); + if evicted == 0 { + return Ok(()); + } + tracing::debug!( + evicted, + "[tinyagents::mw] repeat-progress ledger reset: recorded tool results left the context" + ); + if let Ok(mut trackers) = self.state.tracker.lock() + && let Some(tracker) = trackers.get_mut(&run_id) + { + tracker.reset(); + } + if let Ok(mut recorded) = self.state.recorded.lock() { + recorded.remove(&run_id); + } + Ok(()) + } +} diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs new file mode 100644 index 000000000..1765c365b --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs @@ -0,0 +1,265 @@ +//! Tests for the repeat-progress guard and its eviction observer. + +use std::sync::Arc; + +use serde_json::json; + +use super::wrap_up::DEFAULT_CLEARED_PLACEHOLDER; +use super::*; +use crate::context::{RunConfig, RunContext}; +use crate::middleware::{Middleware, ToolInvocationIdentity}; +use crate::no_progress::{DEFAULT_REPEAT_CALL_THRESHOLD, DEFAULT_REPEAT_OUTPUT_THRESHOLD}; +use crate::steering::{SteeringCommand, SteeringHandle}; +use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message as TaMessage}; +use tinyinference_llm::model::{ModelRequest, ModelResponse}; +use tinyinference_llm::tool::ToolCall as TaToolCall; +use tinytools::ToolResult as TaToolResult; + +fn ctx() -> RunContext { + let mut ctx = RunContext::new(RunConfig::new("mw-test"), ()); + // Hook-level fixtures create a fresh wrapper for each phase; model those + // phases as the same logical invocation, as the agent loop does in life. + ctx.instance_id = 1; + ctx +} + +/// The polling tool the host exempts, as the OpenHuman host does. +fn exempt() -> RepeatExemption { + Arc::new(|tool| tool == "wait_subagent") +} + +fn new_mw(handle: SteeringHandle, summary: HaltSummarySlot) -> RepeatProgressMiddleware { + RepeatProgressMiddleware::new(handle, summary, exempt()) +} + +fn drain_pause_count(handle: &SteeringHandle) -> usize { + handle + .drain() + .into_iter() + .filter(|c| matches!(c, SteeringCommand::Pause)) + .count() +} + +fn repeated_success_response(tool: &str, args: serde_json::Value) -> ModelResponse { + let mut response = ModelResponse::assistant("working"); + response.message = AssistantMessage { + id: None, + content: vec![ContentBlock::Text("working".to_string())], + tool_calls: vec![TaToolCall::new("repeat-1", tool, args)], + usage: None, + origin: None, + }; + response.finish_reason = Some("tool_calls".to_string()); + response +} + +async fn run_successful_repeat_cycle( + mw: &RepeatProgressMiddleware, + tool: &str, + args: serde_json::Value, + output: &str, + error: Option<&str>, +) { + let mut response = repeated_success_response(tool, args); + mw.after_model(&mut ctx(), &(), &mut response) + .await + .unwrap(); + let mut result = match error { + Some(error) => TaToolResult::error(error), + None => TaToolResult::success(output), + }; + let invocation = ToolInvocationIdentity::new("repeat-1", tool); + mw.after_tool(&mut ctx(), &(), &invocation, &mut result) + .await + .unwrap(); +} + +// ── #6275: repeats that are not back to back ──────────────────────────── + +/// One A, B round of two different successful calls, each returning the given +/// output. The adjacent-batch streak restarts on every step of such a cycle. +async fn run_alternating_round(mw: &RepeatProgressMiddleware, a_output: &str, b_output: &str) { + run_successful_repeat_cycle(mw, "use_skill", json!({"skill": "skills"}), a_output, None).await; + run_successful_repeat_cycle( + mw, + "use_skill", + json!({"skill": "skills", "tool": "skill_search"}), + b_output, + None, + ) + .await; +} + +#[tokio::test] +async fn alternating_identical_successful_calls_halt_on_recurrence() { + let handle = SteeringHandle::allow_all(); + let summary = Arc::new(std::sync::Mutex::new(None)); + let mw = new_mw(handle.clone(), summary.clone()); + + for _ in 0..DEFAULT_REPEAT_CALL_THRESHOLD - 1 { + run_alternating_round(&mw, "doc", "hits").await; + } + assert_eq!(drain_pause_count(&handle), 0); + + run_successful_repeat_cycle(&mw, "use_skill", json!({"skill": "skills"}), "doc", None).await; + assert_eq!( + drain_pause_count(&handle), + 1, + "an alternating call returning the identical result for the third time must halt the run" + ); + assert!( + summary + .lock() + .unwrap() + .as_deref() + .is_some_and(|text| text.contains("identical result")), + "the recurrence halt summary should reach the host turn result" + ); +} + +#[tokio::test] +async fn alternating_calls_whose_output_changed_do_not_halt() { + let handle = SteeringHandle::allow_all(); + let mw = new_mw(handle.clone(), Arc::new(std::sync::Mutex::new(None))); + for i in 0..DEFAULT_REPEAT_CALL_THRESHOLD * 3 { + run_alternating_round(&mw, &format!("doc-{i}"), &format!("hits-{i}")).await; + } + assert_eq!( + drain_pause_count(&handle), + 0, + "a re-read with identical arguments whose output changed is progress, not a repeat" + ); +} + +#[tokio::test] +async fn alternating_exempt_polling_calls_do_not_halt() { + let handle = SteeringHandle::allow_all(); + let mw = new_mw(handle.clone(), Arc::new(std::sync::Mutex::new(None))); + for i in 0..DEFAULT_REPEAT_CALL_THRESHOLD * 3 { + run_successful_repeat_cycle( + &mw, + "wait_subagent", + json!({"task_id": "t"}), + "still running", + None, + ) + .await; + let output = format!("status-{i}"); + run_successful_repeat_cycle(&mw, "lookup", json!({"id": 1}), &output, None).await; + } + assert_eq!( + drain_pause_count(&handle), + 0, + "identical polling results must not feed the recurrence ledger" + ); +} + +#[tokio::test] +async fn evicting_a_recorded_result_resets_the_recurrence_ledger() { + // What the last `before_model` sees for the recorded result `repeat-1`: + // untouched, blanked by microcompact, or dropped by compression / trim. + let cases: [(Vec, usize, &str); 3] = [ + ( + vec![TaMessage::tool("repeat-1", "doc")], + 1, + "a result still in context keeps counting", + ), + ( + vec![TaMessage::tool("repeat-1", DEFAULT_CLEARED_PLACEHOLDER)], + 0, + "a result blanked by compaction must not count toward a repeat", + ), + ( + vec![], + 0, + "a result dropped by compaction must not count toward a repeat", + ), + ]; + for (compacted, expected_pauses, why) in cases { + let handle = SteeringHandle::allow_all(); + let mw = new_mw(handle.clone(), Arc::new(std::sync::Mutex::new(None))); + let observer = mw.eviction_observer(); + for _ in 0..DEFAULT_REPEAT_CALL_THRESHOLD - 1 { + run_alternating_round(&mw, "doc", "hits").await; + } + + // The next model call: the guard sees the request before the reduction + // steps run, the observer after them. + let mut request = ModelRequest::new(vec![TaMessage::tool("repeat-1", "doc")]); + mw.before_model(&mut ctx(), &(), &mut request) + .await + .unwrap(); + let mut request = ModelRequest::new(compacted); + observer + .before_model(&mut ctx(), &(), &mut request) + .await + .unwrap(); + + run_successful_repeat_cycle(&mw, "use_skill", json!({"skill": "skills"}), "doc", None) + .await; + assert_eq!(drain_pause_count(&handle), expected_pauses, "{why}"); + } +} + +#[tokio::test] +async fn successful_repeat_tracker_halt_maps_to_summary_and_pause() { + let handle = SteeringHandle::allow_all(); + let summary = Arc::new(std::sync::Mutex::new(None)); + let mw = new_mw(handle.clone(), summary.clone()); + + for _ in 0..DEFAULT_REPEAT_CALL_THRESHOLD - 1 { + run_successful_repeat_cycle(&mw, "lookup", json!({"id": 1}), "ok", None).await; + assert_eq!(drain_pause_count(&handle), 0); + } + run_successful_repeat_cycle(&mw, "lookup", json!({"id": 1}), "ok", None).await; + + assert_eq!(drain_pause_count(&handle), 1); + assert!( + summary + .lock() + .unwrap() + .as_deref() + .is_some_and(|text| text.contains("successful tool-call batch")), + "crate halt summary should be preserved for the host turn result" + ); +} + +#[tokio::test] +async fn successful_repeat_tracker_resets_failed_and_exempt_batches() { + let handle = SteeringHandle::allow_all(); + let mw = new_mw(handle.clone(), Arc::new(std::sync::Mutex::new(None))); + + // Distinct outputs keep the run-wide recurrence ledger out of this test: it + // pins the adjacent-batch streak, which a failure resets. + for i in 0..DEFAULT_REPEAT_CALL_THRESHOLD - 1 { + let output = format!("before-{i}"); + run_successful_repeat_cycle(&mw, "lookup", json!({"id": 1}), &output, None).await; + } + run_successful_repeat_cycle( + &mw, + "lookup", + json!({"id": 1}), + "ok", + Some("temporary failure"), + ) + .await; + for i in 0..DEFAULT_REPEAT_CALL_THRESHOLD - 1 { + let output = format!("after-{i}"); + run_successful_repeat_cycle(&mw, "lookup", json!({"id": 1}), &output, None).await; + } + assert_eq!( + drain_pause_count(&handle), + 0, + "a failed batch resets the successful-repeat streak" + ); + + for _ in 0..DEFAULT_REPEAT_OUTPUT_THRESHOLD + 1 { + run_successful_repeat_cycle(&mw, "wait_subagent", json!({"task_id": "t"}), "ok", None) + .await; + } + assert_eq!( + drain_pause_count(&handle), + 0, + "polling tools remain exempt from successful-repeat halts" + ); +} diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs new file mode 100644 index 000000000..c1bb93416 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs @@ -0,0 +1,430 @@ +//! [`FinalCallWrapUpMiddleware`]: turn the last permitted model call of a +//! capped turn into the turn's conclusion, inside the loop (issue #6014). + +use std::sync::Arc; + +use async_trait::async_trait; + +use crate::context::RunContext; +use crate::error::Result as TaResult; +use crate::middleware::Middleware; +use tinyinference_llm::message::{ContentBlock, Message as TaMessage}; +use tinyinference_llm::model::ModelRequest; + +use super::image_trim::{estimate_message_tokens, estimate_text_tokens}; + +/// The body microcompact swaps in for a cleared tool result, and therefore the +/// only body this middleware treats as restorable. The default for +/// [`FinalCallWrapUpMiddleware::with_cleared_placeholder`]; it must match the +/// placeholder the run's `MicrocompactMiddleware` was built with. +pub const DEFAULT_CLEARED_PLACEHOLDER: &str = "[Old tool result content cleared]"; + +/// The captured outcome store could not be read (a poisoned lock, say). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct OutcomesUnavailable; + +/// Where [`FinalCallWrapUpMiddleware`] finds each tool call's captured result, +/// so the concluding call can be given back what microcompact blanked. +/// +/// The host owns the capture (it records each call as the result enters the +/// transcript, after any per-result byte cap); this is the read side. +pub trait CapturedOutcomes: Send + Sync { + /// The captured result text for `call_id`, `Ok(None)` when nothing was + /// captured for it, or `Err` when the store cannot be read at all (the + /// middleware then concludes without restoring anything). + fn content_for(&self, call_id: &str) -> Result, OutcomesUnavailable>; +} + +/// Turns the **last permitted model call of a capped turn** into the turn's +/// conclusion, in the loop, instead of leaving the answer to an extra call +/// made after the loop has already exited. +/// +/// # Why the loop and not afterwards +/// +/// A capped turn used to end like this: the loop exits on the model-call cap, +/// and `OpenHumanSessionHost::summarize_turn_wrapup` then dispatches a second, out-of-band +/// request straight at the `ChatModel` asking for a checkpoint. Being outside +/// the harness, that request ran with **none** of the loop's context +/// management — no microcompact, no compression, no trim — while being built +/// from the largest transcript the turn would ever hold (every tool result a +/// full iteration budget produced). It was therefore the likeliest call of the +/// whole turn to overflow the window, and each of its failure paths returns +/// `("", None)` silently, so the answer degraded to a deterministic digest of +/// tool names exactly when the turn had the most to report. It also bypassed +/// usage accounting (folded back by hand at the call site) and the progress +/// bridge (re-implemented there as buffer-then-validate-then-forward). +/// +/// Doing it here removes all of that rather than compensating for it: the +/// concluding call is an ordinary loop iteration, so it inherits the entire +/// middleware stack, its usage rides `UsageCarryMiddleware` like any other +/// call, and its text streams through the normal event bridge. It is also one +/// provider call cheaper — the wrap-up was an extra call *past* the cap the +/// operator configured. +/// +/// # What "last permitted call" means, exactly +/// +/// The loop records the model call **before** it builds the request +/// (`agent_loop::run_loop`), so by the time `before_model` runs for the Nth +/// call of an N-call budget, `remaining_model_calls()` is already `0`. That is +/// the trigger, and it needs no new plumbing or counter of its own. +/// +/// The trade this makes is explicit: a 25-call budget becomes 24 tool rounds +/// plus a conclusion, rather than 25 tool rounds plus a 26th call nobody +/// budgeted for. +/// +/// # Why the tools are cleared rather than merely discouraged +/// +/// The instruction alone is a request the model may ignore — the out-of-band +/// wrap-up had to re-parse its own response through the dispatcher to catch a +/// model that emitted a tool call anyway. Removing the schemas from the +/// request makes it structural instead: there is nothing to call. `tool_choice` +/// is reset alongside them because a `Required` choice with an empty tool array +/// is a provider 400. +/// The middleware itself. Generic over the run-context payload: nothing here +/// reads it. Configure the deliverable-tool names with +/// [`with_deliverable_tools`](Self::with_deliverable_tools) — the tools left on +/// the belt for the **penultimate** call (see `reserve_final_write`). +/// +/// The membership rule is "can only emit, never gather": a tool that writes a +/// file the caller already knows the contents of cannot be spent discovering +/// something the turn then has no room to report, which is what makes reserving +/// the call for them a safe trade rather than a gamble. A shell is deliberately +/// not one — it can equally run a crawler, so keeping it would leave the belt +/// effectively unnarrowed. With none configured the penultimate call is left +/// ordinary. +pub struct FinalCallWrapUpMiddleware { + /// The synthetic user turn appended on the final call. + instruction: String, + /// The synthetic user turn appended on the call before it, when the belt is + /// narrowed to the deliverable tools instead of cleared. + final_write_instruction: String, + /// Names of the tools the penultimate call keeps. + deliverable_tools: Vec, + /// The body a cleared tool result carries (see + /// [`DEFAULT_CLEARED_PLACEHOLDER`]). + cleared_placeholder: String, + /// Every tool call's captured outcome, so the concluding call can be given + /// back the results microcompact blanked (see `before_model`). + outcomes: Arc, + /// The input-token allowance the trim downstream enforces, so restoration + /// can stay under it rather than provoking an eviction. `0` disables the + /// bound (a model advertising no context window). + input_budget: u64, + /// Set when the injection fires, so the caller can report the turn as + /// capped. Necessary because this turn now ends *naturally* — the model + /// returns text and requests no tools, which is the loop's ordinary + /// terminal condition — so the old `final_response.is_none()` tell no + /// longer distinguishes a capped turn from a finished one. + fired: std::sync::Mutex>, +} + +impl FinalCallWrapUpMiddleware { + /// Build the middleware. The cleared-result placeholder defaults to + /// [`DEFAULT_CLEARED_PLACEHOLDER`] and no deliverable tools are configured. + pub fn new( + instruction: impl Into, + final_write_instruction: impl Into, + outcomes: Arc, + input_budget: u64, + ) -> Self { + Self { + instruction: instruction.into(), + final_write_instruction: final_write_instruction.into(), + deliverable_tools: Vec::new(), + cleared_placeholder: DEFAULT_CLEARED_PLACEHOLDER.to_string(), + outcomes, + input_budget, + fired: std::sync::Mutex::default(), + } + } + + /// Tools the penultimate call keeps when the belt is narrowed. + pub fn with_deliverable_tools(mut self, tools: I) -> Self + where + I: IntoIterator, + S: Into, + { + self.deliverable_tools = tools.into_iter().map(Into::into).collect(); + self + } + + /// Override the placeholder body treated as a cleared tool result. + pub fn with_cleared_placeholder(mut self, placeholder: impl Into) -> Self { + self.cleared_placeholder = placeholder.into(); + self + } + + /// Whether a tool is one the penultimate call keeps. + fn is_deliverable_tool(&self, name: &str) -> bool { + self.deliverable_tools.iter().any(|tool| tool == name) + } + + /// Whether the wrap-up injection fired for this run context. + pub fn fired(&self, ctx: &RunContext) -> bool { + self.fired + .lock() + .is_ok_and(|fired| fired.contains(&ctx.instance_id())) + } + + /// Give a concluding-or-persisting call back the tool results microcompact + /// blanked, newest-first and only while the request still fits. + /// + /// Shared by both of this middleware's calls, because both need the same + /// thing for the same reason: the content the turn gathered. The final call + /// needs it to *report* findings, and the penultimate one needs it to + /// *write them into a file* — and a turn asked to produce an artifact from + /// nineteen rounds of "[Old tool result content cleared]" produces the same + /// empty-handed result from either direction. + /// + /// `instruction` is the text the caller will append afterwards; it is + /// seeded into the token accounting here rather than counted later, because + /// restoration fills the budget to its boundary and an unaccounted fixed + /// addition after it is exactly the overshoot the budget exists to prevent + /// (CodeRabbit on #6068). + fn restore_cleared_outcomes(&self, request: &mut ModelRequest, instruction: &str) -> usize { + let budget = self.input_budget; + // Seeded with the instruction this middleware appends unconditionally + // below, not just with what the request already holds (CodeRabbit on + // #6068). Restoration fills the budget to its boundary, so an + // unaccounted fixed addition after it is exactly the overshoot the + // budget exists to prevent. + let mut used: u64 = request + .messages + .iter() + .map(estimate_message_tokens) + .sum::() + .saturating_add(estimate_text_tokens(instruction)); + let restored = { + let mut restored = 0usize; + let mut skipped = 0usize; + let mut unavailable = false; + for message in request.messages.iter_mut().rev() { + let TaMessage::Tool(tool) = message else { + continue; + }; + if tool + .content + .iter() + .any(|block| !matches!(block, ContentBlock::Text(_))) + { + continue; + } + let body: String = tool + .content + .iter() + .filter_map(|block| match block { + ContentBlock::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect(); + if body.trim() != self.cleared_placeholder { + continue; + } + let outcome = match self.outcomes.content_for(&tool.tool_call_id) { + Ok(Some(outcome)) => outcome, + Ok(None) => continue, + Err(OutcomesUnavailable) => { + unavailable = true; + break; + } + }; + if outcome.trim().is_empty() { + continue; + } + // What restoring this body would add, against what the + // placeholder already costs. + let added = estimate_text_tokens(&outcome) + .saturating_sub(estimate_text_tokens(&self.cleared_placeholder)); + if budget > 0 && used.saturating_add(added) > budget { + // Everything older is at least as likely to overflow, but + // keep counting so the log reports the true shortfall + // rather than stopping at the first one that did not fit. + skipped += 1; + continue; + } + used = used.saturating_add(added); + tool.content = vec![ContentBlock::Text(outcome)]; + restored += 1; + } + if unavailable { + tracing::warn!( + "[tinyagents::mw] tool-outcome sink poisoned; concluding without restoring \ + cleared tool results" + ); + return 0; + } + if skipped > 0 { + tracing::info!( + skipped, + restored, + budget, + used, + "[tinyagents::mw] left some cleared tool results cleared: restoring them \ + would have pushed the concluding call past its input budget, and an \ + eviction there costs whole messages rather than one body" + ); + } + restored + }; + if restored > 0 { + tracing::info!( + restored, + "[tinyagents::mw] restored cleared tool results for the concluding call" + ); + } + restored + } +} + +impl FinalCallWrapUpMiddleware { + /// The call *before* the conclusion: narrow the belt to + /// the deliverable tools so a turn that owes a file can still write it. + /// + /// Clearing the belt one call later makes the conclusion structural, which + /// is right — but for a turn whose product is an artifact rather than + /// prose it makes *failure* structural too. The host's final-write + /// instruction documents the case that motivated this and the trade it + /// accepts. + /// + /// Returns `true` when the narrowing fired, so the caller can skip the + /// instruction otherwise. + fn reserve_final_write(&self, ctx: &RunContext, request: &mut ModelRequest) -> bool { + // Below three, reserving would eat the turn rather than shape its end: + // a two-call budget would be one write-only call plus the conclusion, + // leaving no round in which anything could be gathered to write. + if ctx.limits.limits().max_model_calls <= 2 { + return false; + } + // Nothing to reserve the call *for*. A read-only or delegating agent + // has no writer on its belt, and telling it "the only tools left are + // the ones that write files" would be false — so leave the call as an + // ordinary one and let the conclusion handle the cap. + if !request + .tools + .iter() + .any(|t| self.is_deliverable_tool(&t.name)) + { + return false; + } + let before = request.tools.len(); + request.tools.retain(|t| self.is_deliverable_tool(&t.name)); + // `Auto`, never `Required`: a turn that has already written its file, + // or was only ever asked for an answer, must be free to spend this call + // on text instead. Forcing a call here would make it invent a write. + request.tool_choice = tinyinference_llm::model::ToolChoice::Auto; + tracing::info!( + model_calls = ctx.limits.model_calls(), + max_model_calls = ctx.limits.limits().max_model_calls, + tools_withdrawn = before.saturating_sub(request.tools.len()), + tools_kept = request.tools.len(), + "[tinyagents::mw] penultimate model call — narrowing the belt to the tools that can \ + persist a deliverable" + ); + // The same restoration the conclusion gets, and for a sharper reason: + // this call is being asked to write the findings into a file, so it + // needs to be able to read them. + self.restore_cleared_outcomes(request, &self.final_write_instruction); + request + .messages + .push(TaMessage::user(self.final_write_instruction.clone())); + true + } +} + +#[async_trait] +impl Middleware<(), C> for FinalCallWrapUpMiddleware { + fn name(&self) -> &str { + "final_call_wrap_up" + } + + async fn before_model( + &self, + ctx: &mut RunContext, + _state: &(), + request: &mut ModelRequest, + ) -> TaResult<()> { + let remaining = ctx.limits.remaining_model_calls(); + if remaining > 1 { + return Ok(()); + } + // A budget of one call would make the very first call the concluding + // one, so the turn could never run a tool at all. That is a + // misconfiguration rather than a cap being reached, and silently + // answering it with a "you have run out of tool calls" instruction + // would misreport it — leave such a run alone. + if ctx.limits.limits().max_model_calls <= 1 { + return Ok(()); + } + // One call before the conclusion: keep the writers rather than clear + // the belt, so a turn whose deliverable is a file can still produce it. + if remaining == 1 { + self.reserve_final_write(ctx, request); + return Ok(()); + } + tracing::info!( + model_calls = ctx.limits.model_calls(), + max_model_calls = ctx.limits.limits().max_model_calls, + tools_withdrawn = request.tools.len(), + "[tinyagents::mw] final permitted model call — withdrawing tools and asking for the \ + turn's conclusion" + ); + request.tools.clear(); + request.tool_choice = tinyinference_llm::model::ToolChoice::None; + // Give the concluding call back the results microcompact blanked. + // + // `MicrocompactMiddleware` replaces every tool-result body past the + // most recent `keep_recent` (5, by default) with `CLEARED_PLACEHOLDER`, + // and — constructed without a token budget — it does so on every call, + // not only under context pressure. That is right for an intermediate + // call, which needs recent context to choose the next tool and nothing + // more. It is exactly wrong for this one: a turn that spent 24 rounds + // gathering would be asked to report its findings with 19 rounds of + // them replaced by "[Old tool result content cleared]", which is the + // same empty-handed answer this whole mechanism exists to prevent, + // arrived at from the other direction. + // + // Restored from the captured outcomes rather than by exempting the + // turn from microcompact, because the blanking has already happened by + // the time this runs (registration order: microcompact is installed by + // `context_mw.install`, this middleware immediately after it) and + // because the sink is the honest source — it holds each result as it + // entered the transcript, after the per-result byte cap. + // + // Only a body that IS the placeholder is replaced, so a result the + // model legitimately saw in full is never rewritten. + // + // This deliberately runs BEFORE the compression and trim middlewares, + // which are installed after it: restoring can make the request large, + // and those two are what bound it. The resulting degradation ladder is + // the one this call wants — everything when it fits, an LLM summary of + // the older slice when it does not, and oldest-first eviction only in + // extremis. What it never does again is silently blank the middle. + // Restore newest-first, and only while the request still fits. + // + // CodeRabbit on #6068: restoration runs AFTER `ContextCompressionMiddleware` + // and before `ImageAwareMessageTrimMiddleware`, so an unbounded restore can + // push the request over the window and the trim then evicts whole + // messages — which is strictly more destructive than the blanking being + // undone, and can discard the very results just restored. + // + // The review suggested a compression phase after restoration. That is a + // second summarizer model call on the one call already about to produce + // the conclusion, and it is avoidable: the overflow is preventable + // rather than repairable. Restoring under the same budget the trim + // enforces means the trim never has cause to fire, so nothing is + // evicted and nothing needs re-summarising. + // + // Newest-first because recency is relevance here: the last rounds' + // results are the ones the model has not seen (the cap is checked + // before the request is built), and the earliest ones are most likely + // already reflected in the compression summary above. + self.restore_cleared_outcomes(request, &self.instruction); + request + .messages + .push(TaMessage::user(self.instruction.clone())); + if let Ok(mut fired) = self.fired.lock() { + fired.insert(ctx.instance_id()); + } + Ok(()) + } +} diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs new file mode 100644 index 000000000..f15799b50 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs @@ -0,0 +1,472 @@ +//! Tests for [`FinalCallWrapUpMiddleware`]. + +use std::sync::Arc; + +use super::*; +use crate::context::{RunConfig, RunContext}; +use crate::middleware::Middleware; +use tinyinference_llm::message::Message as TaMessage; +use tinyinference_llm::model::ModelRequest; +use tinyinference_llm::tool::ToolSchema; + +/// A fixed set of captured outcomes, in the shape the wrap-up reads them. +struct Sink(Vec<(String, String)>); + +impl CapturedOutcomes for Sink { + fn content_for( + &self, + call_id: &str, + ) -> std::result::Result, OutcomesUnavailable> { + Ok(self + .0 + .iter() + .find(|(id, _)| id == call_id) + .map(|(_, content)| content.clone())) + } +} + +fn sink_with(entries: &[(&str, &str)]) -> Arc { + Arc::new(Sink( + entries + .iter() + .map(|(id, content)| ((*id).to_string(), (*content).to_string())) + .collect(), + )) +} + +/// The tools the OpenHuman host keeps for the penultimate call. +fn deliverables() -> [&'static str; 2] { + ["file_write", "apply_patch"] +} + +/// A run with calls left is untouched: no instruction, and the tool belt intact. +#[tokio::test] +async fn wrap_up_leaves_a_call_with_budget_remaining_alone() { + let mw = FinalCallWrapUpMiddleware::new("CONCLUDE NOW", "WRITE NOW", sink_with(&[]), 0); + let mut ctx = RunContext::new(RunConfig::new("mw-test").with_max_model_calls(5), ()); + ctx.limits.record_model_call().unwrap(); + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + tools: vec![ToolSchema::new("echo", "echo", serde_json::json!({}))], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert_eq!( + request.messages.len(), + 1, + "no instruction should be appended" + ); + assert_eq!(request.tools.len(), 1, "the belt must stay intact mid-turn"); +} + +/// On the last permitted call the tools are withdrawn — structurally, not by +/// asking — and the wrap-up instruction is appended as the final turn. +#[tokio::test] +async fn wrap_up_withdraws_tools_and_appends_the_instruction_on_the_last_call() { + let mw = FinalCallWrapUpMiddleware::new("CONCLUDE NOW", "WRITE NOW", sink_with(&[]), 0); + let mut ctx = RunContext::new(RunConfig::new("mw-test").with_max_model_calls(2), ()); + ctx.limits.record_model_call().unwrap(); + ctx.limits.record_model_call().unwrap(); // now the final call + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + tools: vec![ToolSchema::new("echo", "echo", serde_json::json!({}))], + tool_choice: tinyinference_llm::model::ToolChoice::Required, + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert!( + request.tools.is_empty(), + "tools must be withdrawn, not discouraged" + ); + assert!( + matches!( + request.tool_choice, + tinyinference_llm::model::ToolChoice::None + ), + "a Required choice with no tools is a provider 400" + ); + assert_eq!( + request.messages.last().map(|m| m.text()), + Some("CONCLUDE NOW".to_string()), + "the instruction must be the final turn of the request" + ); + assert!(mw.fired(&ctx)); +} + +/// The concluding call gets back the results microcompact blanked — otherwise +/// it is asked to report findings it cannot read, which is the same +/// empty-handed answer the whole mechanism exists to prevent. +#[tokio::test] +async fn wrap_up_restores_tool_results_microcompact_cleared() { + let mw = FinalCallWrapUpMiddleware::new( + "CONCLUDE NOW", + "WRITE NOW", + sink_with(&[ + ("call-old", "issue #41: auth bypass"), + ("call-new", "issue #42: leak"), + ]), + // Unbounded: this case is about restoring what was cleared, not about + // the budget that stops it (covered by its own test below). + 0, + ); + let mut ctx = RunContext::new(RunConfig::new("mw-test").with_max_model_calls(2), ()); + ctx.limits.record_model_call().unwrap(); + ctx.limits.record_model_call().unwrap(); + let mut request = ModelRequest { + messages: vec![ + // The shape microcompact leaves behind: an older result blanked to + // the placeholder, a recent one kept verbatim. + TaMessage::tool("call-old", DEFAULT_CLEARED_PLACEHOLDER), + TaMessage::tool("call-new", "issue #42: leak"), + ], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let bodies: Vec = request.messages.iter().map(|m| m.text()).collect(); + assert!( + bodies.iter().any(|b| b.contains("issue #41: auth bypass")), + "the cleared result should be restored from the capture sink: {bodies:?}" + ); + assert!( + !bodies + .iter() + .any(|b| b.trim() == DEFAULT_CLEARED_PLACEHOLDER), + "no placeholder should survive into the concluding call: {bodies:?}" + ); +} + +/// A result the model legitimately saw in full is never rewritten, even when +/// the sink holds a different (e.g. later-truncated) copy for that id. +#[tokio::test] +async fn wrap_up_does_not_rewrite_a_result_that_was_never_cleared() { + let mw = FinalCallWrapUpMiddleware::new( + "CONCLUDE NOW", + "WRITE NOW", + sink_with(&[("call-1", "FROM SINK")]), + 0, + ); + let mut ctx = RunContext::new(RunConfig::new("mw-test").with_max_model_calls(2), ()); + ctx.limits.record_model_call().unwrap(); + ctx.limits.record_model_call().unwrap(); + let mut request = ModelRequest { + messages: vec![TaMessage::tool("call-1", "IN THE TRANSCRIPT")], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert_eq!( + request.messages[0].text(), + "IN THE TRANSCRIPT", + "only a placeholder body may be replaced" + ); +} + +fn mw(sink: Arc) -> FinalCallWrapUpMiddleware { + FinalCallWrapUpMiddleware::new("CONCLUDE NOW", "WRITE NOW", sink, 0) + .with_deliverable_tools(deliverables()) +} + +/// A context sitting on the Nth call of an N-call budget, minus `back`. +fn ctx_at(max: usize, back: usize) -> RunContext { + let mut ctx = RunContext::new(RunConfig::new("mw-test").with_max_model_calls(max), ()); + for _ in 0..(max - back) { + ctx.limits.record_model_call().unwrap(); + } + ctx +} + +/// A mixed belt: two gatherers and the two writers. +fn mixed_belt() -> Vec { + ["web_fetch", "file_write", "shell", "apply_patch"] + .iter() + .map(|n| ToolSchema::new(*n, *n, serde_json::json!({}))) + .collect() +} + +fn names(request: &ModelRequest) -> Vec { + request.tools.iter().map(|t| t.name.clone()).collect() +} + +/// The defect this exists for: on the call before the conclusion the writers +/// survive, so a turn that owes a file can still write it. +#[tokio::test] +async fn penultimate_call_keeps_the_writers_and_drops_the_gatherers() { + let mw = mw(sink_with(&[])); + let mut ctx = ctx_at(15, 1); + let mut request = ModelRequest { + messages: vec![TaMessage::user("summarise the baggage rules into a file")], + tools: mixed_belt(), + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert_eq!( + names(&request), + vec!["file_write", "apply_patch"], + "only the tools that can emit a deliverable may survive" + ); + assert_eq!( + request.messages.last().map(|m| m.text()), + Some("WRITE NOW".to_string()), + "the write instruction must be the final turn of the request" + ); + assert!( + !mw.fired(&ctx), + "this is not the conclusion: the turn has one more call and must not \ + be reported as capped yet" + ); +} + +/// `Auto`, never `Required`. A turn that has already written its file, or was +/// only ever asked for an answer, must be free to spend this call on text — +/// forcing a call would make it invent a write. +#[tokio::test] +async fn penultimate_call_leaves_the_model_free_not_to_write() { + let mw = mw(sink_with(&[])); + let mut ctx = ctx_at(15, 1); + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + tools: mixed_belt(), + tool_choice: tinyinference_llm::model::ToolChoice::Required, + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert!(matches!( + request.tool_choice, + tinyinference_llm::model::ToolChoice::Auto + )); +} + +/// Asked to write its findings into a file, this call has to be able to read +/// them — so it gets the same restoration the conclusion does. +#[tokio::test] +async fn penultimate_call_restores_what_microcompact_cleared() { + let mw = mw(sink_with(&[("call-old", "carry-on max 22 x 14 x 9 in")])); + let mut ctx = ctx_at(15, 1); + let mut request = ModelRequest { + messages: vec![ + TaMessage::tool("call-old", DEFAULT_CLEARED_PLACEHOLDER), + TaMessage::tool("call-new", "checked bag 50 lb"), + ], + tools: mixed_belt(), + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let bodies: Vec = request.messages.iter().map(|m| m.text()).collect(); + assert!( + bodies.iter().any(|b| b.contains("22 x 14 x 9")), + "a call asked to write findings down must be able to read them: {bodies:?}" + ); +} + +/// A belt with no writer on it gets no narrowing and no instruction: telling a +/// read-only agent that "the only tools left are the ones that write files" +/// would simply be false. +#[tokio::test] +async fn a_belt_without_a_writer_is_left_alone() { + let mw = mw(sink_with(&[])); + let mut ctx = ctx_at(15, 1); + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + tools: vec![ + ToolSchema::new("web_fetch", "web_fetch", serde_json::json!({})), + ToolSchema::new("shell", "shell", serde_json::json!({})), + ], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert_eq!(names(&request), vec!["web_fetch", "shell"]); + assert_eq!( + request.messages.len(), + 1, + "no instruction should be appended" + ); +} + +/// Reserving out of a two-call budget would leave no round in which anything +/// could be gathered to write, so it is skipped: one tool round, then the +/// conclusion, exactly as before. +#[tokio::test] +async fn a_two_call_budget_keeps_its_one_gathering_round() { + let mw = mw(sink_with(&[])); + let mut ctx = ctx_at(2, 1); + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + tools: mixed_belt(), + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert_eq!( + names(&request), + vec!["web_fetch", "file_write", "shell", "apply_patch"], + "the belt must stay whole when there is no room to reserve a call" + ); + assert_eq!( + request.messages.len(), + 1, + "no instruction should be appended" + ); +} + +/// Two calls earlier nothing has happened yet — the reservation is for the +/// penultimate call alone, not for the tail of the turn. +#[tokio::test] +async fn the_call_before_the_penultimate_one_is_untouched() { + let mw = mw(sink_with(&[])); + let mut ctx = ctx_at(15, 2); + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + tools: mixed_belt(), + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert_eq!(request.tools.len(), 4, "the belt must stay intact mid-turn"); + assert_eq!(request.messages.len(), 1); +} + +/// And the call after it still clears the belt outright: narrowing buys the +/// artifact, the conclusion still has to be text. +#[tokio::test] +async fn the_conclusion_still_withdraws_everything() { + let mw = mw(sink_with(&[])); + let mut ctx = ctx_at(15, 0); + let mut request = ModelRequest { + messages: vec![TaMessage::user("hi")], + tools: mixed_belt(), + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + assert!( + request.tools.is_empty(), + "the final call keeps nothing, writers included" + ); + assert_eq!( + request.messages.last().map(|m| m.text()), + Some("CONCLUDE NOW".to_string()) + ); + assert!(mw.fired(&ctx)); +} + +// ── the bounds CodeRabbit asked for on #6068 ──────────────────────────────── + +/// Restoration stops at the input allowance instead of overflowing it. +/// +/// It runs after compression and before the trim, so an unbounded restore +/// pushes the request over and the trim evicts whole messages — strictly more +/// destructive than the blanking being undone, and able to discard the very +/// results just restored. Preventing the overflow beats repairing it. +#[tokio::test] +async fn wrap_up_stops_restoring_at_the_input_budget() { + let big = "x".repeat(4_000); + let mw = FinalCallWrapUpMiddleware::new( + "CONCLUDE NOW", + "WRITE NOW", + sink_with(&[("call-1", &big), ("call-2", &big), ("call-3", &big)]), + // Room for roughly one of them, not three. + 1_200, + ); + let mut ctx = RunContext::new(RunConfig::new("mw-test").with_max_model_calls(2), ()); + ctx.limits.record_model_call().unwrap(); + ctx.limits.record_model_call().unwrap(); + let mut request = ModelRequest { + messages: vec![ + TaMessage::tool("call-1", DEFAULT_CLEARED_PLACEHOLDER), + TaMessage::tool("call-2", DEFAULT_CLEARED_PLACEHOLDER), + TaMessage::tool("call-3", DEFAULT_CLEARED_PLACEHOLDER), + ], + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let restored = request + .messages + .iter() + .filter(|m| m.text().starts_with("xxx")) + .count(); + assert!( + (1..3).contains(&restored), + "some restored, not all — budget 1200 cannot hold three 4k bodies, got {restored}" + ); + // Newest-first: the last tool result is the one the model never saw (the + // cap is checked before the request is built). Not `messages.last()` — + // that is the wrap-up instruction this middleware just appended. + let newest_tool = request + .messages + .iter() + .rev() + .find(|m| matches!(m, TaMessage::Tool(_))) + .map(|m| m.text()) + .unwrap_or_default(); + assert!( + newest_tool.starts_with("xxx"), + "the newest cleared result is restored first, got: {}", + &newest_tool[..newest_tool.len().min(60)] + ); +} + +/// The other half of the no-window split: restoration must stop too. Without a +/// window there is no trim middleware behind it, so an unbounded wrap-up would +/// hand the provider a request it rejects — losing the in-loop conclusion the +/// whole cap-checkpoint path exists to deliver (CodeRabbit on #6068). +#[tokio::test] +async fn wrap_up_restoration_stays_bounded_when_no_window_is_advertised() { + let big = "x".repeat(8_000); + let entries: Vec<(String, String)> = (0..40) + .map(|i| (format!("call-{i}"), big.clone())) + .collect(); + let borrowed: Vec<(&str, &str)> = entries + .iter() + .map(|(a, b)| (a.as_str(), b.as_str())) + .collect(); + let (_toc, restore) = split_input_allowance(0); + let mw = + FinalCallWrapUpMiddleware::new("CONCLUDE NOW", "WRITE NOW", sink_with(&borrowed), restore); + let mut ctx = RunContext::new(RunConfig::new("mw-test").with_max_model_calls(2), ()); + ctx.limits.record_model_call().unwrap(); + ctx.limits.record_model_call().unwrap(); + let mut request = ModelRequest { + messages: borrowed + .iter() + .map(|(id, _)| TaMessage::tool(*id, DEFAULT_CLEARED_PLACEHOLDER)) + .collect(), + ..Default::default() + }; + + mw.before_model(&mut ctx, &(), &mut request).await.unwrap(); + + let used: u64 = request.messages.iter().map(estimate_message_tokens).sum(); + assert!( + used <= NO_WINDOW_ALLOWANCE, + "the no-window fallback must bound the whole request, not just the list: {used}" + ); + let restored = request + .messages + .iter() + .filter(|m| m.text().starts_with("xxx")) + .count(); + assert!( + (1..40).contains(&restored), + "some restored, not all 40 — got {restored}" + ); +} diff --git a/crates/tinyagents-harness/src/tool/discover/README.md b/crates/tinyagents-harness/src/tool/discover/README.md index 5e3f3d1dc..5ca081fe7 100644 --- a/crates/tinyagents-harness/src/tool/discover/README.md +++ b/crates/tinyagents-harness/src/tool/discover/README.md @@ -9,6 +9,7 @@ this file is the map of the module. | `types.rs` | `ToolDiscoveryPolicy` (the knobs, the host `ranker` and `DiscoveryRankMode`), `DeferredCatalog` (a run's deferred schemas, BM25-indexed, name-sorted, ranked through the policy), `DeferredTool`, `RankedSearch` | | `manifest.rs` | `render_manifest` — the budgeted listing inside `tool_search`'s description: full → names → count | | `bridge.rs` | The two intrinsic tools: `bridge_schemas`, `answer_tool_search` (async; returns a `SearchAnswer`), `unwrap_tool_call` | +| `embedding_ranker.rs` | `EmbeddingToolRanker` — a semantic `tinytools::ToolRanker` over any `EmbeddingModel`, cached in memory and optionally on disk by embedding-space signature | | `test.rs` | Unit tests for all of the above | The agent loop (`agent_loop/run_loop.rs`, `agent_loop/tools.rs`) is the only diff --git a/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs b/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs new file mode 100644 index 000000000..e90f66d83 --- /dev/null +++ b/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs @@ -0,0 +1,253 @@ +//! [`EmbeddingToolRanker`]: the semantic retriever for `tool_search`. +//! +//! A lexical retriever loses every paraphrase — "ping alex" never overlaps +//! `SLACK_SEND_MESSAGE send a message` — and a decision model can only pick +//! from what the retriever hands it. Measured on the orchestrator's catalogue +//! plus nine Composio toolkits (1,215 tools), BM25 recall@20 was 70%, and +//! that ceiling capped Jev at 67% top-3. This ranker embeds each tool's +//! summary once with the host's embedding model, embeds the intent per search, and ranks by cosine +//! similarity, so recall follows meaning rather than shared words. It is the +//! `retriever` inside `tinytools_jev::JevRanker` for every catalogue larger +//! than one Jev `Choice`. +//! +//! Catalogue embeddings are cached in memory by content hash and, when a +//! cache path is given, on disk keyed by the provider's signature, so a cold +//! process pays for the catalogue once and every later search embeds only +//! the intent (one provider call). + +use std::collections::HashMap; +use std::hash::{Hash, Hasher}; +use std::path::PathBuf; +use std::sync::{Arc, RwLock}; + +use serde::{Deserialize, Serialize}; +use tinyinference_embeddings::EmbeddingModel; +use tinytools::{RankCandidate, RankContext, RankError, RankHit, ToolRanker}; + +/// Texts per embedding request. The managed provider accepts far more, but a +/// bounded batch keeps one request under any body cap and lets a partial +/// failure cost one batch, not the catalogue. +const EMBED_BATCH: usize = 64; + +/// Ranks tools by cosine similarity between the intent and each tool's +/// summary, with catalogue embeddings cached. +pub struct EmbeddingToolRanker { + provider: Arc, + /// Identifies the embedding space (provider, model, dimensions) so a cached + /// vector is never compared against one from a different model. + signature: String, + cache: RwLock>>, + cache_path: Option, +} + +#[derive(Serialize, Deserialize, Default)] +struct DiskCache { + signature: String, + entries: HashMap>, +} + +impl std::fmt::Debug for EmbeddingToolRanker { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EmbeddingToolRanker") + .field("provider", &self.provider.name()) + .field("model", &self.provider.model_id()) + .field("cache_path", &self.cache_path) + .finish_non_exhaustive() + } +} + +impl EmbeddingToolRanker { + /// The stable [`ToolRanker::kind`] of this ranker. + pub const KIND: &'static str = "embedding"; + + /// A ranker over `provider` with an in-memory cache only. `signature` + /// identifies the provider's embedding space (for example + /// `provider=voyage;model=voyage-3;dims=1024`); the disk cache is valid + /// only for the signature it was written under. + pub fn new(provider: Arc, signature: impl Into) -> Self { + Self { + provider, + signature: signature.into(), + cache: RwLock::new(HashMap::new()), + cache_path: None, + } + } + + /// A ranker whose catalogue embeddings also persist at `path`, keyed by + /// the provider's signature so a model change invalidates them. A + /// missing or unreadable file is an empty cache, never an error. + pub fn with_disk_cache(mut self, path: PathBuf) -> Self { + let loaded = std::fs::read(&path) + .ok() + .and_then(|raw| serde_json::from_slice::(&raw).ok()) + .filter(|disk| disk.signature == self.signature); + if let Some(disk) = loaded { + tracing::debug!( + entries = disk.entries.len(), + path = %path.display(), + "[tool-search] loaded embedding cache" + ); + *self.cache.write().unwrap_or_else(|p| p.into_inner()) = disk.entries; + } + self.cache_path = Some(path); + self + } + + fn key(candidate: &RankCandidate) -> u64 { + let mut hasher = std::collections::hash_map::DefaultHasher::new(); + candidate.summary.hash(&mut hasher); + candidate.family.hash(&mut hasher); + hasher.finish() + } + + fn text(candidate: &RankCandidate) -> String { + match &candidate.family { + Some(family) => format!("{} ({family})", candidate.summary), + None => candidate.summary.clone(), + } + } + + /// Embeds every candidate not already cached, in batches. + async fn ensure_cached(&self, candidates: &[RankCandidate]) -> Result<(), RankError> { + let missing: Vec<(u64, String)> = { + let cache = self.cache.read().unwrap_or_else(|p| p.into_inner()); + let mut seen = std::collections::HashSet::new(); + candidates + .iter() + .map(|c| (Self::key(c), c)) + .filter(|(k, _)| !cache.contains_key(k) && seen.insert(*k)) + .map(|(k, c)| (k, Self::text(c))) + .collect() + }; + if missing.is_empty() { + return Ok(()); + } + tracing::info!( + missing = missing.len(), + provider = self.provider.name(), + "[tool-search] embedding catalogue entries" + ); + for batch in missing.chunks(EMBED_BATCH) { + let texts: Vec = batch.iter().map(|(_, t)| t.clone()).collect(); + let vectors = + self.provider + .embed(&texts) + .await + .map_err(|error| RankError::Backend { + reason: format!("embedding failed: {error:#}"), + })?; + if vectors.len() != batch.len() { + return Err(RankError::Backend { + reason: format!( + "embedding returned {} vectors for {} texts", + vectors.len(), + batch.len() + ), + }); + } + let mut cache = self.cache.write().unwrap_or_else(|p| p.into_inner()); + for ((key, _), vector) in batch.iter().zip(vectors) { + cache.insert(*key, vector); + } + } + self.persist(); + Ok(()) + } + + fn persist(&self) { + let Some(path) = &self.cache_path else { + return; + }; + let disk = DiskCache { + signature: self.signature.clone(), + entries: self.cache.read().unwrap_or_else(|p| p.into_inner()).clone(), + }; + let write = || -> std::io::Result<()> { + if let Some(parent) = path.parent() { + std::fs::create_dir_all(parent)?; + } + let tmp = path.with_extension("json.tmp"); + std::fs::write(&tmp, serde_json::to_vec(&disk)?)?; + std::fs::rename(&tmp, path) + }; + if let Err(error) = write() { + tracing::warn!( + path = %path.display(), + error = %error, + "[tool-search] could not persist embedding cache" + ); + } + } +} + +fn cosine(a: &[f32], b: &[f32]) -> f64 { + let (mut dot, mut na, mut nb) = (0.0_f64, 0.0_f64, 0.0_f64); + for (x, y) in a.iter().zip(b) { + dot += f64::from(*x) * f64::from(*y); + na += f64::from(*x) * f64::from(*x); + nb += f64::from(*y) * f64::from(*y); + } + if na == 0.0 || nb == 0.0 { + return 0.0; + } + dot / (na.sqrt() * nb.sqrt()) +} + +#[async_trait::async_trait] +impl ToolRanker for EmbeddingToolRanker { + fn kind(&self) -> &'static str { + Self::KIND + } + + async fn rank( + &self, + intent: &str, + _context: &RankContext, + candidates: &[RankCandidate], + limit: usize, + ) -> Result, RankError> { + let intent = intent.trim(); + if intent.is_empty() { + return Err(RankError::InvalidInput { + reason: "intent is empty".to_owned(), + }); + } + if candidates.is_empty() || limit == 0 { + return Ok(Vec::new()); + } + self.ensure_cached(candidates).await?; + let query = self + .provider + .embed(&[intent.to_owned()]) + .await + .map_err(|error| RankError::Backend { + reason: format!("embedding failed: {error:#}"), + })? + .into_iter() + .next() + .ok_or_else(|| RankError::Backend { + reason: "embedding returned no vector for the intent".to_owned(), + })?; + let cache = self.cache.read().unwrap_or_else(|p| p.into_inner()); + let mut scored: Vec = candidates + .iter() + .filter_map(|c| { + cache + .get(&Self::key(c)) + .map(|v| RankHit::new(c.key.clone(), cosine(&query, v))) + }) + .collect(); + scored.sort_by(|a, b| { + b.score + .partial_cmp(&a.score) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.key.cmp(&b.key)) + }); + scored.truncate(limit); + Ok(scored) + } +} + +#[cfg(test)] +#[path = "embedding_ranker_test.rs"] +mod test; diff --git a/crates/tinyagents-harness/src/tool/discover/embedding_ranker_test.rs b/crates/tinyagents-harness/src/tool/discover/embedding_ranker_test.rs new file mode 100644 index 000000000..2bba99cc9 --- /dev/null +++ b/crates/tinyagents-harness/src/tool/discover/embedding_ranker_test.rs @@ -0,0 +1,187 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use tinytools::{RankCandidate, RankContext, ToolRanker}; + +use super::*; +use tinyinference_embeddings::{EmbeddingModel, Result as EmbedResult}; + +/// Embeds a text as a bag of three hand-picked words, so similarity is +/// deterministic and readable. +struct BagEmbedder { + calls: AtomicUsize, +} + +#[async_trait::async_trait] +impl EmbeddingModel for BagEmbedder { + fn name(&self) -> &str { + "bag" + } + fn model_id(&self) -> &str { + "bag-v1" + } + fn dimensions(&self) -> usize { + 3 + } + async fn embed(&self, texts: &[String]) -> EmbedResult>> { + self.calls.fetch_add(1, Ordering::SeqCst); + Ok(texts + .iter() + .map(|t| { + let t = t.to_ascii_lowercase(); + vec![ + f32::from(u8::from(t.contains("message") || t.contains("ping"))), + f32::from(u8::from(t.contains("email") || t.contains("mail"))), + f32::from(u8::from(t.contains("file"))), + ] + }) + .collect()) + } +} + +const SIGNATURE: &str = "provider=bag;model=bag-v1;dims=3"; + +fn candidates() -> Vec { + vec![ + RankCandidate::new("SLACK_SEND_MESSAGE", "send a message to a channel") + .with_family("slack"), + RankCandidate::new("GMAIL_SEND_EMAIL", "send an email").with_family("gmail"), + RankCandidate::new("file_read", "read a file"), + ] +} + +#[tokio::test] +async fn ranks_by_cosine_and_embeds_the_catalogue_once() { + let embedder = Arc::new(BagEmbedder { + calls: AtomicUsize::new(0), + }); + let ranker = EmbeddingToolRanker::new(embedder.clone(), SIGNATURE); + assert_eq!(ranker.kind(), "embedding"); + + let hits = ranker + .rank("ping alex", &RankContext::empty(), &candidates(), 2) + .await + .unwrap(); + assert_eq!(hits[0].key, "SLACK_SEND_MESSAGE"); + assert!(hits[0].confidence.is_none()); + assert_eq!(hits.len(), 2); + // One batch for the catalogue plus one for the intent. + assert_eq!(embedder.calls.load(Ordering::SeqCst), 2); + + let hits = ranker + .rank("mail the report", &RankContext::empty(), &candidates(), 1) + .await + .unwrap(); + assert_eq!(hits[0].key, "GMAIL_SEND_EMAIL"); + // Only the intent was embedded this time. + assert_eq!(embedder.calls.load(Ordering::SeqCst), 3); +} + +#[tokio::test] +async fn disk_cache_round_trips_and_is_keyed_by_signature() { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("cache").join("tool_search_embeddings.json"); + let embedder = Arc::new(BagEmbedder { + calls: AtomicUsize::new(0), + }); + let ranker = + EmbeddingToolRanker::new(embedder.clone(), SIGNATURE).with_disk_cache(path.clone()); + ranker + .rank("ping", &RankContext::empty(), &candidates(), 1) + .await + .unwrap(); + assert!(path.exists()); + + let embedder2 = Arc::new(BagEmbedder { + calls: AtomicUsize::new(0), + }); + let warm = EmbeddingToolRanker::new(embedder2.clone(), SIGNATURE).with_disk_cache(path.clone()); + warm.rank("ping", &RankContext::empty(), &candidates(), 1) + .await + .unwrap(); + assert_eq!( + embedder2.calls.load(Ordering::SeqCst), + 1, + "a warm cache embeds only the intent" + ); +} + +#[tokio::test] +async fn empty_intent_is_rejected() { + let ranker = EmbeddingToolRanker::new( + Arc::new(BagEmbedder { + calls: AtomicUsize::new(0), + }), + SIGNATURE, + ); + assert!( + ranker + .rank(" ", &RankContext::empty(), &candidates(), 1) + .await + .is_err() + ); +} + +/// A tool that appears later — a newly connected toolkit's actions, a +/// rewritten description — is embedded on its own; the rest is a cache hit. +#[tokio::test] +async fn a_new_or_changed_tool_is_embedded_incrementally() { + let embedder = Arc::new(BagEmbedder { + calls: AtomicUsize::new(0), + }); + let ranker = EmbeddingToolRanker::new(embedder.clone(), SIGNATURE); + ranker + .rank("ping", &RankContext::empty(), &candidates(), 1) + .await + .unwrap(); + assert_eq!( + embedder.calls.load(Ordering::SeqCst), + 2, + "catalogue + intent" + ); + + let mut grown = candidates(); + grown.push(RankCandidate::new("NOTION_CREATE_PAGE", "create a page").with_family("notion")); + grown[2] = RankCandidate::new("file_read", "read a file from disk"); + ranker + .rank("ping", &RankContext::empty(), &grown, 1) + .await + .unwrap(); + // One batch for the two unseen texts (the new tool and the changed one), + // plus the intent — never the whole catalogue again. + assert_eq!(embedder.calls.load(Ordering::SeqCst), 4); + assert_eq!( + ranker.cache.read().unwrap().len(), + 5, + "old and new descriptions both cached; a stale entry is harmless" + ); +} + +/// A cache written under one embedding space is ignored under another. +#[tokio::test] +async fn a_disk_cache_written_under_another_signature_is_ignored() { + let tmp = tempfile::tempdir().unwrap(); + let path = tmp.path().join("cache.json"); + let embedder = Arc::new(BagEmbedder { + calls: AtomicUsize::new(0), + }); + EmbeddingToolRanker::new(embedder, SIGNATURE) + .with_disk_cache(path.clone()) + .rank("ping", &RankContext::empty(), &candidates(), 1) + .await + .unwrap(); + + let other = Arc::new(BagEmbedder { + calls: AtomicUsize::new(0), + }); + EmbeddingToolRanker::new(other.clone(), "provider=bag;model=bag-v2;dims=3") + .with_disk_cache(path) + .rank("ping", &RankContext::empty(), &candidates(), 1) + .await + .unwrap(); + assert_eq!( + other.calls.load(Ordering::SeqCst), + 2, + "a different signature re-embeds the catalogue and the intent" + ); +} diff --git a/crates/tinyagents-harness/src/tool/discover/mod.rs b/crates/tinyagents-harness/src/tool/discover/mod.rs index 2ca1136b6..2a1ee9305 100644 --- a/crates/tinyagents-harness/src/tool/discover/mod.rs +++ b/crates/tinyagents-harness/src/tool/discover/mod.rs @@ -11,10 +11,14 @@ //! [`DiscoveryRankMode`] that says how it is used. //! - [`bridge_schemas`] / [`answer_tool_search`] / [`unwrap_tool_call`] — the //! two intrinsic bridge tools the agent loop advertises and answers. +//! - [`EmbeddingToolRanker`] — a semantic [`tinytools::ToolRanker`] over any +//! [`tinyinference_embeddings::EmbeddingModel`], with an in-memory and +//! optional on-disk embedding cache. //! - [`render_manifest`] — the budgeted listing inside `tool_search`'s //! description. mod bridge; +mod embedding_ranker; mod manifest; mod types; @@ -22,6 +26,7 @@ pub use bridge::{ SearchAnswer, TOOL_CALL_NAME, TOOL_SEARCH_NAME, answer_tool_search, bridge_schemas, unwrap_tool_call, }; +pub use embedding_ranker::EmbeddingToolRanker; // The BM25 arithmetic lives in `tinytools::rank` now, so a host ranks with the // same index the bridge does; the old paths keep resolving. pub use manifest::{MANIFEST_DESCRIPTION_CHARS, first_sentence, render_manifest};