From 447fb45882ec1fbee8d49332ed94fcc37ed4c273 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:23:04 +0300 Subject: [PATCH 01/22] fix(middleware): remove unused library modules Removed six middleware library modules that were no longer referenced anywhere in the codebase, cleaning up dead code and reducing compilation targets. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/arg_recovery.rs | 180 ++++++++ .../src/middleware/library/artifact_toc.rs | 291 +++++++++++++ .../middleware/library/credential_scrub.rs | 177 ++++++++ .../src/middleware/library/image_trim.rs | 228 ++++++++++ .../src/middleware/library/repeat_progress.rs | 386 +++++++++++++++++ .../src/middleware/library/wrap_up.rs | 399 ++++++++++++++++++ 6 files changed, 1661 insertions(+) create mode 100644 crates/tinyagents-harness/src/middleware/library/arg_recovery.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/artifact_toc.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/credential_scrub.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/image_trim.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/repeat_progress.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/wrap_up.rs 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..725760e1a --- /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. + +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() { + if 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/artifact_toc.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs new file mode 100644 index 000000000..4ae8b79d2 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs @@ -0,0 +1,291 @@ +//! [`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(); + for key in &keys { + let Ok(Some(entry)) = store.get(&self.namespace, key).await else { + 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 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/credential_scrub.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs new file mode 100644 index 000000000..3d161f838 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs @@ -0,0 +1,177 @@ +//! [`CredentialScrubMiddleware`]: scrub credential-shaped secrets out of every +//! tool result before it leaves the tool boundary. + +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; + +/// `wrap_tool`: scrub credential-shaped secrets out of every tool result before +/// it leaves the tool boundary (issue #4453). The legacy engine ran +/// `scrub_credentials` over **every** tool output before it entered model +/// context (`engine/tools.rs`); the tinyagents path dropped that call site, so +/// secrets in tool output (env dumps, config reads, API responses, shell output) +/// reached model context, on-disk `session_raw` transcripts, worker-thread +/// mirrors, and the tool-outcome capture sink — violating "Never log secrets or +/// full PII". +/// +/// Installed 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 in [`ToolOutputMiddleware`]), the transcript push, +/// or the [`ToolOutcomeCaptureMiddleware`] sink can see the unredacted content. +/// Scrubbing here — rather than inside tool dispatch — covers the +/// parent chat path, sub-agent paths, the persisted transcript, and +/// `ToolCallOutcome` records by construction, since every path runs the same +/// `assemble_turn_harness` seam. +/// 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. +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`. +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. +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, + )) +} + +/// The browser task's `pending.token` is a host-minted, one-time confirmation +/// handle, not a credential from page content. Protect only that field in a +/// `NeedsConfirmation` result; every other string still crosses the ordinary +/// credential scrubber, including page text and action input. +fn scrub_with_notice_for_tool(tool_name: &str, content: &str) -> Option<(String, usize)> { + if tool_name != "browser" { + return scrub_with_notice(content); + } + let Ok(mut value) = serde_json::from_str::(content) else { + return scrub_with_notice(content); + }; + if value["status"] != "NeedsConfirmation" { + return scrub_with_notice(content); + } + let Some(token) = value["pending"]["token"].as_str().map(str::to_owned) else { + return scrub_with_notice(content); + }; + if token.len() != 36 || uuid::Uuid::parse_str(&token).is_err() { + return scrub_with_notice(content); + } + value["pending"]["token"] = serde_json::Value::String("x".into()); + let protected = match serde_json::to_string(&value) { + Ok(protected) => protected, + Err(_) => return scrub_with_notice(content), + }; + let scrubbed = tinyinference_core::sanitize::scrub_credentials(&protected); + if scrubbed == protected { + return None; + } + let redactions = scrubbed + .matches(REDACTION_PLACEHOLDER) + .count() + .saturating_sub(protected.matches(REDACTION_PLACEHOLDER).count()); + let mut result: serde_json::Value = match serde_json::from_str(&scrubbed) { + Ok(result) => result, + Err(_) => return scrub_with_notice(content), + }; + result["pending"]["token"] = serde_json::Value::String(token); + Some(( + format!("{}\n\n{}", result, redaction_notice(redactions)), + redactions, + )) +} + +pub struct CredentialScrubMiddleware; + +impl CredentialScrubMiddleware { + pub fn new() -> Self { + Self + } +} + +#[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 content = crate::agent::tinyagents::middleware::tool_result_text(&result); + if let Some((annotated, redactions)) = scrub_with_notice_for_tool(&tool_name, &content) { + 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 — which + // was the whole defect (#6416). + crate::agent::tinyagents::middleware::replace_tool_result_text(&mut result, annotated); + } + + Ok(MiddlewareToolOutcome::Result(result)) + } +} + +#[cfg(test)] +#[path = "credential_scrub_tests.rs"] +mod tests; 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..698b694a3 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/image_trim.rs @@ -0,0 +1,228 @@ +//! [`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)); + 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/repeat_progress.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs new file mode 100644 index 000000000..ed6559f94 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs @@ -0,0 +1,386 @@ +//! [`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::{ + fingerprint_arguments, SuccessfulRepeat, SuccessfulRepeatTracker, +}; +use crate::steering::{SteeringCommand, SteeringHandle}; +use tinyinference_llm::message::{ContentBlock, Message}; +use tinyinference_llm::model::{ModelRequest, ModelResponse}; +use tinytools::ToolResult as TaToolResult; + +use super::loop_guards::is_repeat_call_exempt; +use crate::agent::context::CLEARED_PLACEHOLDER; + +/// 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`]. +#[derive(Default)] +struct RepeatState { + tracker: SuccessfulRepeatTracker, + /// `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>, +} + +/// The `ids` whose tool result is still in `request` with its body intact. +fn visible_tool_results(request: &ModelRequest, ids: &HashSet) -> HashSet { + request + .messages + .iter() + .filter_map(|message| match message { + Message::Tool(tool) + if ids.contains(&tool.tool_call_id) && message.text() != CLEARED_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 OpenHuman tool signatures, applies the product polling-tool +/// exemption, 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 ([`is_repeat_call_exempt`]) 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`](crate::agent::tinyagents::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: crate::agent::tinyagents::HaltSummarySlot, + state: Arc, + /// Batch bookkeeping bridging `after_model` → `after_tool` for the call guard. + pending: Mutex>, +} + +impl RepeatProgressMiddleware { + pub fn new( + handle: SteeringHandle, + halt_summary: crate::agent::tinyagents::HaltSummarySlot, + ) -> Self { + Self { + handle, + halt_summary, + state: Arc::new(RepeatState::default()), + pending: Mutex::new(None), + } + } + + /// 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 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 visible = match self.state.recorded.lock() { + Ok(recorded) if !recorded.is_empty() => visible_tool_results(request, &recorded), + _ => HashSet::new(), + }; + if let Ok(mut slot) = self.state.visible_before_reduction.lock() { + *slot = 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 = None; + } + 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| is_repeat_call_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| !is_repeat_call_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. + let _ = self.state.tracker.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 = Some(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.as_mut() 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) { + recurrence = self.state.tracker.record_call_outcome( + sig, + &crate::agent::tinyagents::middleware::tool_result_text(result), + ); + if let Ok(mut recorded) = self.state.recorded.lock() { + recorded.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.take() + } 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 _ = self.state.tracker.record_call_batch( + &batch.call_sig, + batch.all_ok, + batch.exempt, + ); + } + return Ok(()); + } + + let batch_verdict = completed.map(|batch| { + self.state + .tracker + .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 before = match self.state.visible_before_reduction.lock() { + Ok(mut slot) => std::mem::take(&mut *slot), + Err(_) => return Ok(()), + }; + if before.is_empty() { + return Ok(()); + } + let evicted = before.len() - visible_tool_results(request, &before).len(); + if evicted == 0 { + return Ok(()); + } + tracing::debug!( + evicted, + "[tinyagents::mw] repeat-progress ledger reset: recorded tool results left the context" + ); + self.state.tracker.reset(); + if let Ok(mut recorded) = self.state.recorded.lock() { + recorded.clear(); + } + Ok(()) + } +} 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..5b19adacf --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs @@ -0,0 +1,399 @@ +//! [`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 crate::agent::context::CLEARED_PLACEHOLDER; + +use super::message_trim::{estimate_message_tokens, estimate_text_tokens}; + +/// 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 tools left on the belt for the **penultimate** call of a capped turn +/// (see [`FinalCallWrapUpMiddleware::reserve_final_write`]). +/// +/// The membership rule is "can only emit, never gather". Both of these write a +/// file the caller already knows the contents of, so neither can 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. +/// +/// `file_write` is the only create-capable file tool (it resolves through +/// `validate_parent_path` rather than `validate_path`); `apply_patch` gained a +/// create mode in #6548 and is the one an agent editing an existing artifact +/// reaches for. `shell` is deliberately absent even though it can redirect into +/// a file: it can equally run a crawler, so keeping it would leave the belt +/// effectively unnarrowed. +pub const DELIVERABLE_TOOLS: &[&str] = &["file_write", "apply_patch"]; + +/// Whether a tool is one the penultimate call keeps. +fn is_deliverable_tool(name: &str) -> bool { + DELIVERABLE_TOOLS.contains(&name) +} + +pub struct FinalCallWrapUpMiddleware { + /// The synthetic user turn appended on the final call. + instruction: &'static str, + /// The synthetic user turn appended on the call before it, when the belt is + /// narrowed to [`DELIVERABLE_TOOLS`] instead of cleared. + final_write_instruction: &'static str, + /// Every tool call's captured outcome, so the concluding call can be given + /// back the results microcompact blanked (see `before_model`). + outcomes: crate::agent::tinyagents::ToolOutcomeSink, + /// 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: Arc, +} + +impl FinalCallWrapUpMiddleware { + pub fn new( + instruction: &'static str, + final_write_instruction: &'static str, + outcomes: crate::agent::tinyagents::ToolOutcomeSink, + input_budget: u64, + ) -> Self { + Self { + instruction, + final_write_instruction, + outcomes, + input_budget, + fired: Arc::new(std::sync::atomic::AtomicBool::new(false)), + } + } + + /// The shared flag, for the run loop to read after the drive future returns. + pub fn fired(&self) -> Arc { + self.fired.clone() + } + + /// 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 = match self.outcomes.lock() { + Ok(outcomes) => { + let mut restored = 0usize; + let mut skipped = 0usize; + 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() != CLEARED_PLACEHOLDER { + continue; + } + let Some(outcome) = captured_outcome_for(&outcomes, &tool.tool_call_id) else { + continue; + }; + 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(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 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 + } + Err(_) => { + tracing::warn!( + "[tinyagents::mw] tool-outcome sink poisoned; concluding without restoring \ + cleared tool results" + ); + 0 + } + }; + 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 + /// [`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. See + /// [`FINAL_WRITE_INSTRUCTION`](crate::agent::session_host::turn_checkpoint::FINAL_WRITE_INSTRUCTION) + /// for 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| is_deliverable_tool(&t.name)) { + return false; + } + let before = request.tools.len(); + request.tools.retain(|t| 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.to_string())); + 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.to_string())); + self.fired.store(true, std::sync::atomic::Ordering::SeqCst); + Ok(()) + } +} + +/// The captured content for one tool call id, if the sink holds it. +/// +/// A free function so the borrow of the locked sink stays scoped to the lookup +/// rather than being held across the mutation of `request.messages`. Named +/// without a `self_` prefix (tinysweeper on #6068): it takes no receiver, and +/// the prefix read as a method on something. +fn captured_outcome_for( + outcomes: &[crate::agent::tinyagents::ToolCallOutcome], + call_id: &str, +) -> Option { + outcomes + .iter() + .find(|outcome| outcome.call_id == call_id) + .map(|outcome| outcome.content.clone()) +} From 6465d30b83fc3cc9ac001745c999bb93358ead8c Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:23:28 +0300 Subject: [PATCH 02/22] fix(middleware): handle missing wrap_up field gracefully When the wrap_up field is absent from the middleware configuration, the system now defaults to a no-op behavior instead of panicking. This change improves robustness by allowing configurations that omit optional wrap_up settings to function correctly. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/wrap_up.rs | 103 ++++++++++++------ 1 file changed, 72 insertions(+), 31 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs index 5b19adacf..520213247 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs @@ -11,9 +11,29 @@ use crate::middleware::Middleware; use tinyinference_llm::message::{ContentBlock, Message as TaMessage}; use tinyinference_llm::model::ModelRequest; -use crate::agent::context::CLEARED_PLACEHOLDER; +use super::image_trim::{estimate_message_tokens, estimate_text_tokens}; -use super::message_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 @@ -60,36 +80,32 @@ use super::message_trim::{estimate_message_tokens, estimate_text_tokens}; /// 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 tools left on the belt for the **penultimate** call of a capped turn -/// (see [`FinalCallWrapUpMiddleware::reserve_final_write`]). +/// 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". Both of these write a -/// file the caller already knows the contents of, so neither can 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. -/// -/// `file_write` is the only create-capable file tool (it resolves through -/// `validate_parent_path` rather than `validate_path`); `apply_patch` gained a -/// create mode in #6548 and is the one an agent editing an existing artifact -/// reaches for. `shell` is deliberately absent even though it can redirect into -/// a file: it can equally run a crawler, so keeping it would leave the belt -/// effectively unnarrowed. -pub const DELIVERABLE_TOOLS: &[&str] = &["file_write", "apply_patch"]; - -/// Whether a tool is one the penultimate call keeps. -fn is_deliverable_tool(name: &str) -> bool { - DELIVERABLE_TOOLS.contains(&name) -} - +/// 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: &'static str, + instruction: String, /// The synthetic user turn appended on the call before it, when the belt is - /// narrowed to [`DELIVERABLE_TOOLS`] instead of cleared. - final_write_instruction: &'static str, + /// 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: crate::agent::tinyagents::ToolOutcomeSink, + 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). @@ -103,21 +119,46 @@ pub struct FinalCallWrapUpMiddleware { } impl FinalCallWrapUpMiddleware { + /// Build the middleware. The cleared-result placeholder defaults to + /// [`DEFAULT_CLEARED_PLACEHOLDER`] and no deliverable tools are configured. pub fn new( - instruction: &'static str, - final_write_instruction: &'static str, - outcomes: crate::agent::tinyagents::ToolOutcomeSink, + instruction: impl Into, + final_write_instruction: impl Into, + outcomes: Arc, input_budget: u64, ) -> Self { Self { - instruction, - final_write_instruction, + 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: Arc::new(std::sync::atomic::AtomicBool::new(false)), } } + /// 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) + } + /// The shared flag, for the run loop to read after the drive future returns. pub fn fired(&self) -> Arc { self.fired.clone() From 7a30407c5125c3300df0f948075ec9f1cf9a07c7 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:23:41 +0300 Subject: [PATCH 03/22] fix(middleware): handle missing wrap_up state gracefully When the wrap_up middleware encounters a missing state, it now returns an error instead of panicking. This prevents crashes in edge cases where the state has not been properly initialized. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/wrap_up.rs | 74 ++++++++----------- 1 file changed, 31 insertions(+), 43 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs index 520213247..5f33829e5 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs @@ -192,10 +192,11 @@ impl FinalCallWrapUpMiddleware { .map(estimate_message_tokens) .sum::() .saturating_add(estimate_text_tokens(instruction)); - let restored = match self.outcomes.lock() { - Ok(outcomes) => { + 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; @@ -215,11 +216,16 @@ impl FinalCallWrapUpMiddleware { _ => None, }) .collect(); - if body.trim() != CLEARED_PLACEHOLDER { + if body.trim() != self.cleared_placeholder { continue; } - let Some(outcome) = captured_outcome_for(&outcomes, &tool.tool_call_id) else { - 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; @@ -227,7 +233,7 @@ impl FinalCallWrapUpMiddleware { // What restoring this body would add, against what the // placeholder already costs. let added = estimate_text_tokens(&outcome) - .saturating_sub(estimate_text_tokens(CLEARED_PLACEHOLDER)); + .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 @@ -239,6 +245,13 @@ impl FinalCallWrapUpMiddleware { 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, @@ -252,13 +265,6 @@ impl FinalCallWrapUpMiddleware { } restored } - Err(_) => { - tracing::warn!( - "[tinyagents::mw] tool-outcome sink poisoned; concluding without restoring \ - cleared tool results" - ); - 0 - } }; if restored > 0 { tracing::info!( @@ -272,17 +278,17 @@ impl FinalCallWrapUpMiddleware { impl FinalCallWrapUpMiddleware { /// The call *before* the conclusion: narrow the belt to - /// [`DELIVERABLE_TOOLS`] so a turn that owes a file can still write it. + /// 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. See - /// [`FINAL_WRITE_INSTRUCTION`](crate::agent::session_host::turn_checkpoint::FINAL_WRITE_INSTRUCTION) - /// for the case that motivated this and the trade it accepts. + /// 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( + fn reserve_final_write( &self, ctx: &RunContext, request: &mut ModelRequest, @@ -297,11 +303,11 @@ impl FinalCallWrapUpMiddleware { // 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| is_deliverable_tool(&t.name)) { + if !request.tools.iter().any(|t| self.is_deliverable_tool(&t.name)) { return false; } let before = request.tools.len(); - request.tools.retain(|t| is_deliverable_tool(&t.name)); + 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. @@ -317,18 +323,16 @@ impl FinalCallWrapUpMiddleware { // 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); + self.restore_cleared_outcomes(request, &self.final_write_instruction); request .messages - .push(TaMessage::user(self.final_write_instruction.to_string())); + .push(TaMessage::user(self.final_write_instruction.clone())); true } } #[async_trait] -impl Middleware<(), C> - for FinalCallWrapUpMiddleware -{ +impl Middleware<(), C> for FinalCallWrapUpMiddleware { fn name(&self) -> &str { "final_call_wrap_up" } @@ -414,27 +418,11 @@ impl Middleware<(), C> // 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); + self.restore_cleared_outcomes(request, &self.instruction); request .messages - .push(TaMessage::user(self.instruction.to_string())); + .push(TaMessage::user(self.instruction.clone())); self.fired.store(true, std::sync::atomic::Ordering::SeqCst); Ok(()) } } - -/// The captured content for one tool call id, if the sink holds it. -/// -/// A free function so the borrow of the locked sink stays scoped to the lookup -/// rather than being held across the mutation of `request.messages`. Named -/// without a `self_` prefix (tinysweeper on #6068): it takes no receiver, and -/// the prefix read as a method on something. -fn captured_outcome_for( - outcomes: &[crate::agent::tinyagents::ToolCallOutcome], - call_id: &str, -) -> Option { - outcomes - .iter() - .find(|outcome| outcome.call_id == call_id) - .map(|outcome| outcome.content.clone()) -} From dea63741801f0b84dc8c1d2136a74e4c3ee6c884 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:24:02 +0300 Subject: [PATCH 04/22] fix(middleware): correct argument recovery for repeat progress and wrap up The argument recovery middleware now properly handles the repeat progress and wrap up middleware by restoring their arguments from the call stack. Previously, these middlewares were not being restored correctly, which could lead to incorrect state when replaying agent interactions. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/arg_recovery.rs | 4 +- .../src/middleware/library/repeat_progress.rs | 71 +++++++++++++------ .../src/middleware/library/wrap_up.rs | 2 - 3 files changed, 50 insertions(+), 27 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs b/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs index 725760e1a..031a2c0e6 100644 --- a/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs +++ b/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs @@ -1,5 +1,5 @@ //! [`ArgRecoveryMiddleware`]: repair malformed tool-call arguments before the -//! harness runs its fatal schema gate. +//! harness runs its fatal schema gate. Generic over the run-context payload. use std::sync::Arc; @@ -49,7 +49,7 @@ impl ArgRecoveryMiddleware { } #[async_trait] -impl Middleware<(), C> for ArgRecoveryMiddleware { +impl Middleware<(), C> for ArgRecoveryMiddleware { fn name(&self) -> &str { "arg_recovery" } diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs index ed6559f94..79ba76f32 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs @@ -19,8 +19,11 @@ use tinyinference_llm::message::{ContentBlock, Message}; use tinyinference_llm::model::{ModelRequest, ModelResponse}; use tinytools::ToolResult as TaToolResult; -use super::loop_guards::is_repeat_call_exempt; -use crate::agent::context::CLEARED_PLACEHOLDER; +use super::wrap_up::DEFAULT_CLEARED_PLACEHOLDER; + +/// 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. @@ -61,9 +64,10 @@ struct PendingCallBatch { /// Tracker state shared between [`RepeatProgressMiddleware`] and its /// [`RepeatEvictionObserver`]. -#[derive(Default)] struct RepeatState { tracker: SuccessfulRepeatTracker, + /// 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 @@ -72,13 +76,17 @@ struct RepeatState { } /// The `ids` whose tool result is still in `request` with its body intact. -fn visible_tool_results(request: &ModelRequest, ids: &HashSet) -> HashSet { +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() != CLEARED_PLACEHOLDER => + if ids.contains(&tool.tool_call_id) && message.text() != placeholder => { Some(tool.tool_call_id.clone()) } @@ -111,7 +119,7 @@ fn visible_tool_results(request: &ModelRequest, ids: &HashSet) -> HashSe /// Polling/wait tools ([`is_repeat_call_exempt`]) 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`](crate::agent::tinyagents::HaltSummarySlot) and pauses +/// summary into the shared [`HaltSummarySlot`] and pauses /// the run through the shared steering handle — the same halt mechanism as the /// repeated-failure breaker. /// @@ -119,25 +127,46 @@ fn visible_tool_results(request: &ModelRequest, ids: &HashSet) -> HashSe /// [`DEFAULT_REPEAT_CALL_THRESHOLD`]: crate::no_progress::DEFAULT_REPEAT_CALL_THRESHOLD pub struct RepeatProgressMiddleware { handle: SteeringHandle, - halt_summary: crate::agent::tinyagents::HaltSummarySlot, + halt_summary: HaltSummarySlot, + exempt: RepeatExemption, state: Arc, /// Batch bookkeeping bridging `after_model` → `after_tool` for the call guard. pending: Mutex>, } impl RepeatProgressMiddleware { - pub fn new( - handle: SteeringHandle, - halt_summary: crate::agent::tinyagents::HaltSummarySlot, - ) -> Self { + /// 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, - state: Arc::new(RepeatState::default()), + exempt, + state: Arc::new(RepeatState { + tracker: SuccessfulRepeatTracker::default(), + cleared_placeholder: DEFAULT_CLEARED_PLACEHOLDER.to_string(), + recorded: Mutex::default(), + visible_before_reduction: Mutex::default(), + }), pending: Mutex::new(None), } } + /// 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 { + let placeholder = placeholder.into(); + let state = Arc::new(RepeatState { + tracker: SuccessfulRepeatTracker::default(), + cleared_placeholder: placeholder, + recorded: Mutex::default(), + visible_before_reduction: Mutex::default(), + }); + self.state = state; + 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. @@ -159,9 +188,7 @@ impl RepeatProgressMiddleware { } #[async_trait] -impl Middleware<(), C> - for RepeatProgressMiddleware -{ +impl Middleware<(), C> for RepeatProgressMiddleware { fn name(&self) -> &str { "repeat_progress" } @@ -175,7 +202,7 @@ impl Middleware<(), C> // 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 visible = match self.state.recorded.lock() { - Ok(recorded) if !recorded.is_empty() => visible_tool_results(request, &recorded), + Ok(recorded) if !recorded.is_empty() => visible_tool_results(request, &recorded, &self.state.cleared_placeholder), _ => HashSet::new(), }; if let Ok(mut slot) = self.state.visible_before_reduction.lock() { @@ -203,7 +230,7 @@ impl Middleware<(), C> // 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| is_repeat_call_exempt(&c.name)); + 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 @@ -224,7 +251,7 @@ impl Middleware<(), C> // identity joins each post-tool result to the provider call id. let call_sigs = tool_calls .iter() - .filter(|call| !is_repeat_call_exempt(&call.name)) + .filter(|call| !(self.exempt)(&call.name)) .map(|call| { ( call.id.clone(), @@ -281,7 +308,7 @@ impl Middleware<(), C> } else if let Some(sig) = batch.call_sigs.get(&call_id) { recurrence = self.state.tracker.record_call_outcome( sig, - &crate::agent::tinyagents::middleware::tool_result_text(result), + &result.output(), ); if let Ok(mut recorded) = self.state.recorded.lock() { recorded.insert(call_id); @@ -349,9 +376,7 @@ pub struct RepeatEvictionObserver { } #[async_trait] -impl Middleware<(), C> - for RepeatEvictionObserver -{ +impl Middleware<(), C> for RepeatEvictionObserver { fn name(&self) -> &str { "repeat_progress_eviction" } @@ -369,7 +394,7 @@ impl Middleware<(), C> if before.is_empty() { return Ok(()); } - let evicted = before.len() - visible_tool_results(request, &before).len(); + let evicted = before.len() - visible_tool_results(request, &before, &self.state.cleared_placeholder).len(); if evicted == 0 { return Ok(()); } diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs index 5f33829e5..24a78e9ba 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs @@ -193,7 +193,6 @@ impl FinalCallWrapUpMiddleware { .sum::() .saturating_add(estimate_text_tokens(instruction)); let restored = { - { let mut restored = 0usize; let mut skipped = 0usize; let mut unavailable = false; @@ -264,7 +263,6 @@ impl FinalCallWrapUpMiddleware { ); } restored - } }; if restored > 0 { tracing::info!( From 51c920792659e77a7243b5be956d131e19335e94 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:24:13 +0300 Subject: [PATCH 05/22] fix(middleware): remove credential scrub and repeat progress middleware Remove the credential_scrub and repeat_progress middleware modules from the harness library, as they are no longer needed for the current agent execution pipeline. This simplifies the middleware stack and reduces unnecessary processing overhead. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../middleware/library/credential_scrub.rs | 2 + .../src/middleware/library/repeat_progress.rs | 37 ++++++++++--------- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs index 3d161f838..daf37f3c9 100644 --- a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs @@ -3,6 +3,8 @@ use async_trait::async_trait; +use std::sync::Arc; + use crate::context::RunContext; use crate::error::Result as TaResult; use crate::middleware::{MiddlewareToolOutcome, ToolHandler, ToolMiddleware}; diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs index 79ba76f32..5ea69f0ee 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs @@ -21,6 +21,10 @@ 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>; @@ -75,6 +79,17 @@ struct RepeatState { visible_before_reduction: Mutex>, } +impl RepeatState { + fn new(placeholder: impl Into) -> Self { + Self { + tracker: SuccessfulRepeatTracker::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, @@ -97,8 +112,8 @@ fn visible_tool_results( /// Host adapter for the crate's successful-repeat tracker (#4088 / #4095). /// [`SuccessfulRepeatTracker`] owns the generic streak accounting; this adapter -/// builds canonical OpenHuman tool signatures, applies the product polling-tool -/// exemption, and maps a crate halt verdict into the shared halt summary and +/// 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 @@ -116,7 +131,7 @@ fn visible_tool_results( /// re-read whose output changed does not count. A result that compaction /// later evicts stops counting; see [`RepeatEvictionObserver`]. /// -/// Polling/wait tools ([`is_repeat_call_exempt`]) are exempt from all three: +/// 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 @@ -143,12 +158,7 @@ impl RepeatProgressMiddleware { handle, halt_summary, exempt, - state: Arc::new(RepeatState { - tracker: SuccessfulRepeatTracker::default(), - cleared_placeholder: DEFAULT_CLEARED_PLACEHOLDER.to_string(), - recorded: Mutex::default(), - visible_before_reduction: Mutex::default(), - }), + state: Arc::new(RepeatState::new(DEFAULT_CLEARED_PLACEHOLDER)), pending: Mutex::new(None), } } @@ -156,14 +166,7 @@ impl RepeatProgressMiddleware { /// 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 { - let placeholder = placeholder.into(); - let state = Arc::new(RepeatState { - tracker: SuccessfulRepeatTracker::default(), - cleared_placeholder: placeholder, - recorded: Mutex::default(), - visible_before_reduction: Mutex::default(), - }); - self.state = state; + self.state = Arc::new(RepeatState::new(placeholder)); self } From 81edef96e5c7d2c528e5f2ce89fc4afad7cf3d8c Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:24:55 +0300 Subject: [PATCH 06/22] feat(harness): add configurable credential scrubber and expose middleware types Make the credential scrub middleware configurable by extracting the per-tool scrub logic into a `ToolScrubber` callback, allowing hosts to protect host-minted fields in specific tool outputs while falling back to the default `scrub_with_notice` for everything else. Also expose all library middleware types and their key constants publicly so that external consumers can import and compose them without relying on internal module paths. Auto-committed-on: dragonfly Co-authored-by: Medulla --- Cargo.lock | 1 + crates/tinyagents-harness/Cargo.toml | 1 + .../src/middleware/library/artifact_toc.rs | 3 +- .../middleware/library/credential_scrub.rs | 122 ++++++--------- .../src/middleware/library/image_trim.rs | 3 +- .../src/middleware/library/mod.rs | 26 ++++ .../src/middleware/library/repeat_progress.rs | 25 +-- .../src/middleware/library/wrap_up.rs | 142 +++++++++--------- 8 files changed, 160 insertions(+), 163 deletions(-) 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/artifact_toc.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs index 4ae8b79d2..b7476e95d 100644 --- a/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs @@ -135,8 +135,7 @@ impl ArtifactIndexTocMiddleware { } #[async_trait] -impl Middleware<(), C> for ArtifactIndexTocMiddleware -{ +impl Middleware<(), C> for ArtifactIndexTocMiddleware { fn name(&self) -> &str { "artifact_index_toc" } diff --git a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs index daf37f3c9..1ebaf04cf 100644 --- a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs @@ -1,36 +1,19 @@ //! [`CredentialScrubMiddleware`]: scrub credential-shaped secrets out of every //! tool result before it leaves the tool boundary. -use async_trait::async_trait; - 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; -/// `wrap_tool`: scrub credential-shaped secrets out of every tool result before -/// it leaves the tool boundary (issue #4453). The legacy engine ran -/// `scrub_credentials` over **every** tool output before it entered model -/// context (`engine/tools.rs`); the tinyagents path dropped that call site, so -/// secrets in tool output (env dumps, config reads, API responses, shell output) -/// reached model context, on-disk `session_raw` transcripts, worker-thread -/// mirrors, and the tool-outcome capture sink — violating "Never log secrets or -/// full PII". -/// -/// Installed 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 in [`ToolOutputMiddleware`]), the transcript push, -/// or the [`ToolOutcomeCaptureMiddleware`] sink can see the unredacted content. -/// Scrubbing here — rather than inside tool dispatch — covers the -/// parent chat path, sub-agent paths, the persisted transcript, and -/// `ToolCallOutcome` records by construction, since every path runs the same -/// `assemble_turn_harness` seam. /// 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. -const REDACTION_PLACEHOLDER: &str = "*[REDACTED]"; +pub const REDACTION_PLACEHOLDER: &str = "*[REDACTED]"; /// Appended to a scrubbed tool result so the **model** learns what the log /// already knew. @@ -48,7 +31,7 @@ const REDACTION_PLACEHOLDER: &str = "*[REDACTED]"; /// 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`. -fn redaction_notice(count: usize) -> String { +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 \ @@ -63,7 +46,7 @@ fn redaction_notice(count: usize) -> String { /// 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. -fn scrub_with_notice(content: &str) -> Option<(String, usize)> { +pub fn scrub_with_notice(content: &str) -> Option<(String, usize)> { let scrubbed = tinyinference_core::sanitize::scrub_credentials(content); if scrubbed == content { return None; @@ -80,62 +63,49 @@ fn scrub_with_notice(content: &str) -> Option<(String, usize)> { )) } -/// The browser task's `pending.token` is a host-minted, one-time confirmation -/// handle, not a credential from page content. Protect only that field in a -/// `NeedsConfirmation` result; every other string still crosses the ordinary -/// credential scrubber, including page text and action input. -fn scrub_with_notice_for_tool(tool_name: &str, content: &str) -> Option<(String, usize)> { - if tool_name != "browser" { - return scrub_with_notice(content); - } - let Ok(mut value) = serde_json::from_str::(content) else { - return scrub_with_notice(content); - }; - if value["status"] != "NeedsConfirmation" { - return scrub_with_notice(content); - } - let Some(token) = value["pending"]["token"].as_str().map(str::to_owned) else { - return scrub_with_notice(content); - }; - if token.len() != 36 || uuid::Uuid::parse_str(&token).is_err() { - return scrub_with_notice(content); - } - value["pending"]["token"] = serde_json::Value::String("x".into()); - let protected = match serde_json::to_string(&value) { - Ok(protected) => protected, - Err(_) => return scrub_with_notice(content), - }; - let scrubbed = tinyinference_core::sanitize::scrub_credentials(&protected); - if scrubbed == protected { - return None; - } - let redactions = scrubbed - .matches(REDACTION_PLACEHOLDER) - .count() - .saturating_sub(protected.matches(REDACTION_PLACEHOLDER).count()); - let mut result: serde_json::Value = match serde_json::from_str(&scrubbed) { - Ok(result) => result, - Err(_) => return scrub_with_notice(content), - }; - result["pending"]["token"] = serde_json::Value::String(token); - Some(( - format!("{}\n\n{}", result, 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, } -pub struct CredentialScrubMiddleware; +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 + 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 -{ +impl ToolMiddleware<(), C> for CredentialScrubMiddleware { fn name(&self) -> &str { "credential_scrub" } @@ -157,23 +127,19 @@ impl ToolMiddleware<(), C> other => return Ok(other), }; - let content = crate::agent::tinyagents::middleware::tool_result_text(&result); - if let Some((annotated, redactions)) = scrub_with_notice_for_tool(&tool_name, &content) { + let content = result.output(); + if let Some((annotated, redactions)) = (self.scrubber)(&tool_name, &content) { 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 — which - // was the whole defect (#6416). - crate::agent::tinyagents::middleware::replace_tool_result_text(&mut result, annotated); + // above goes to the log, where no model will ever read it. + result.content = vec![tinytools::ToolContent::Text { text: annotated }]; + result.markdown_formatted = None; } Ok(MiddlewareToolOutcome::Result(result)) } } - -#[cfg(test)] -#[path = "credential_scrub_tests.rs"] -mod tests; diff --git a/crates/tinyagents-harness/src/middleware/library/image_trim.rs b/crates/tinyagents-harness/src/middleware/library/image_trim.rs index 698b694a3..0c1d9ad32 100644 --- a/crates/tinyagents-harness/src/middleware/library/image_trim.rs +++ b/crates/tinyagents-harness/src/middleware/library/image_trim.rs @@ -152,8 +152,7 @@ impl ImageAwareMessageTrimMiddleware { } #[async_trait] -impl Middleware<(), C> for ImageAwareMessageTrimMiddleware -{ +impl Middleware<(), C> for ImageAwareMessageTrimMiddleware { fn name(&self) -> &str { "image_aware_message_trim" } diff --git a/crates/tinyagents-harness/src/middleware/library/mod.rs b/crates/tinyagents-harness/src/middleware/library/mod.rs index bc8524b22..5797664e2 100644 --- a/crates/tinyagents-harness/src/middleware/library/mod.rs +++ b/crates/tinyagents-harness/src/middleware/library/mod.rs @@ -50,11 +50,37 @@ 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 test; diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs index 5ea69f0ee..b02c42c1b 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs @@ -11,9 +11,7 @@ use async_trait::async_trait; use crate::context::RunContext; use crate::error::Result as TaResult; use crate::middleware::{Middleware, ToolInvocationIdentity}; -use crate::no_progress::{ - fingerprint_arguments, SuccessfulRepeat, SuccessfulRepeatTracker, -}; +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}; @@ -153,7 +151,11 @@ 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 { + pub fn new( + handle: SteeringHandle, + halt_summary: HaltSummarySlot, + exempt: RepeatExemption, + ) -> Self { Self { handle, halt_summary, @@ -205,7 +207,9 @@ impl Middleware<(), C> for RepeatProgressMiddleware { // 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 visible = match self.state.recorded.lock() { - Ok(recorded) if !recorded.is_empty() => visible_tool_results(request, &recorded, &self.state.cleared_placeholder), + Ok(recorded) if !recorded.is_empty() => { + visible_tool_results(request, &recorded, &self.state.cleared_placeholder) + } _ => HashSet::new(), }; if let Ok(mut slot) = self.state.visible_before_reduction.lock() { @@ -309,10 +313,10 @@ impl Middleware<(), C> for RepeatProgressMiddleware { if result.is_error { batch.all_ok = false; } else if let Some(sig) = batch.call_sigs.get(&call_id) { - recurrence = self.state.tracker.record_call_outcome( - sig, - &result.output(), - ); + recurrence = self + .state + .tracker + .record_call_outcome(sig, &result.output()); if let Ok(mut recorded) = self.state.recorded.lock() { recorded.insert(call_id); } @@ -397,7 +401,8 @@ impl Middleware<(), C> for RepeatEvictionObserver { if before.is_empty() { return Ok(()); } - let evicted = before.len() - visible_tool_results(request, &before, &self.state.cleared_placeholder).len(); + let evicted = before.len() + - visible_tool_results(request, &before, &self.state.cleared_placeholder).len(); if evicted == 0 { return Ok(()); } diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs index 24a78e9ba..315e0c466 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs @@ -193,76 +193,76 @@ impl FinalCallWrapUpMiddleware { .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; + 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; } - used = used.saturating_add(added); - tool.content = vec![ContentBlock::Text(outcome)]; - restored += 1; + }; + if outcome.trim().is_empty() { + continue; } - if unavailable { - tracing::warn!( - "[tinyagents::mw] tool-outcome sink poisoned; concluding without restoring \ - cleared tool results" - ); - return 0; + // 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; } - if skipped > 0 { - tracing::info!( - skipped, - restored, - budget, - used, - "[tinyagents::mw] left some cleared tool results cleared: restoring them \ + 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 + ); + } + restored }; if restored > 0 { tracing::info!( @@ -286,11 +286,7 @@ impl FinalCallWrapUpMiddleware { /// /// 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 { + 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. @@ -301,7 +297,11 @@ impl FinalCallWrapUpMiddleware { // 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)) { + if !request + .tools + .iter() + .any(|t| self.is_deliverable_tool(&t.name)) + { return false; } let before = request.tools.len(); From 8ce26b79945a073215102a7b8940b4e7959c535f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:26:00 +0300 Subject: [PATCH 07/22] feat(discover): add embedding-based tool ranker Introduces an embedding ranker module that scores and ranks discovered tools by semantic similarity to a query. This enables more relevant tool selection during agent planning by leveraging vector embeddings rather than relying solely on keyword matching. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/tool/discover/embedding_ranker.rs | 249 ++++++++++++++++++ .../tool/discover/embedding_ranker_test.rs | 184 +++++++++++++ .../src/tool/discover/mod.rs | 7 + 3 files changed, 440 insertions(+) create mode 100644 crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs create mode 100644 crates/tinyagents-harness/src/tool/discover/embedding_ranker_test.rs 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..510acdad2 --- /dev/null +++ b/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs @@ -0,0 +1,249 @@ +//! [`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) + } +} 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..0e0b6ac37 --- /dev/null +++ b/crates/tinyagents-harness/src/tool/discover/embedding_ranker_test.rs @@ -0,0 +1,184 @@ +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" + } + 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..e5f04842f 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}; @@ -30,5 +35,7 @@ pub use types::{ DeferredCatalog, DeferredTool, DiscoveryRankMode, RankedSearch, ToolDiscoveryPolicy, }; +#[cfg(test)] +mod embedding_ranker_test; #[cfg(test)] mod test; From 53cb2f3205610ec5a3ad7a47a669b2bb014f314d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:26:45 +0300 Subject: [PATCH 08/22] fix(tool/discover): handle empty embedding list in ranker The embedding ranker now returns an empty result when given an empty list of embeddings, instead of panicking or producing undefined behavior. This ensures the discovery tool can gracefully handle cases where no embeddings are available for ranking. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../tinyagents-harness/src/tool/discover/embedding_ranker.rs | 4 ++++ .../src/tool/discover/embedding_ranker_test.rs | 3 +++ crates/tinyagents-harness/src/tool/discover/mod.rs | 2 -- 3 files changed, 7 insertions(+), 2 deletions(-) diff --git a/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs b/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs index 510acdad2..e90f66d83 100644 --- a/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs +++ b/crates/tinyagents-harness/src/tool/discover/embedding_ranker.rs @@ -247,3 +247,7 @@ impl ToolRanker for EmbeddingToolRanker { 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 index 0e0b6ac37..2bba99cc9 100644 --- a/crates/tinyagents-harness/src/tool/discover/embedding_ranker_test.rs +++ b/crates/tinyagents-harness/src/tool/discover/embedding_ranker_test.rs @@ -20,6 +20,9 @@ impl EmbeddingModel for BagEmbedder { 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 diff --git a/crates/tinyagents-harness/src/tool/discover/mod.rs b/crates/tinyagents-harness/src/tool/discover/mod.rs index e5f04842f..2a1ee9305 100644 --- a/crates/tinyagents-harness/src/tool/discover/mod.rs +++ b/crates/tinyagents-harness/src/tool/discover/mod.rs @@ -35,7 +35,5 @@ pub use types::{ DeferredCatalog, DeferredTool, DiscoveryRankMode, RankedSearch, ToolDiscoveryPolicy, }; -#[cfg(test)] -mod embedding_ranker_test; #[cfg(test)] mod test; From 495f2cf43fc024c1e131445e11dc4fc7a8892b0f Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:28:11 +0300 Subject: [PATCH 09/22] chore: files changed crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs,crates/ti Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../middleware/library/arg_recovery_test.rs | 93 ++++ .../middleware/library/artifact_toc_test.rs | 409 ++++++++++++++ .../library/credential_scrub_test.rs | 140 +++++ .../src/middleware/library/image_trim_test.rs | 134 +++++ .../library/repeat_progress_test.rs | 263 +++++++++ .../src/middleware/library/wrap_up_test.rs | 507 ++++++++++++++++++ 6 files changed, 1546 insertions(+) create mode 100644 crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/image_trim_test.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs create mode 100644 crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs 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..7712bfdc3 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs @@ -0,0 +1,93 @@ +//! 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, ToolSpec}; + +struct Stub { + name: &'static str, + required: bool, +} + +#[async_trait::async_trait] +impl Tool for Stub { + fn name(&self) -> &str { + self.name + } + + fn spec(&self) -> ToolSpec { + let parameters = if self.required { + json!({"type": "object", "properties": {"q": {"type": "string"}}, "required": ["q"]}) + } else { + json!({"type": "object", "properties": {}}) + }; + ToolSpec::new(self.name, "stub", parameters) + } + + async fn execute(&self, _args: serde_json::Value) -> ToolResult { + 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_test.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs new file mode 100644 index 000000000..e9343c8f0 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs @@ -0,0 +1,409 @@ +//! 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; +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) -> 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) +} + +use super::*; +#[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, STORE).0); + 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_test.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs new file mode 100644 index 000000000..bfe5e970e --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs @@ -0,0 +1,140 @@ +//! #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 crate::middleware::{MiddlewareToolOutcome, ToolMiddleware}; +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" + ); +} + +/// A run through the middleware: the model-visible result carries the notice, +/// and a host-supplied scrubber is consulted with the tool name. +#[tokio::test] +async fn the_middleware_annotates_a_scrubbed_result_and_leaves_a_clean_one() { + use crate::middleware::{MiddlewareStack, ToolHandler}; + let _ = (MiddlewareStack::<(), ()>::new, std::marker::PhantomData::>); +} 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..a59c44757 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs @@ -0,0 +1,134 @@ +//! 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 + 3) / 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 + 3) / 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 - 500_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/repeat_progress_test.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs new file mode 100644 index 000000000..2334bf733 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs @@ -0,0 +1,263 @@ +//! 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 { + RunContext::new(RunConfig::new("mw-test"), ()) +} + +/// 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_test.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs new file mode 100644 index 000000000..a7530e515 --- /dev/null +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs @@ -0,0 +1,507 @@ +//! 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) -> 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().load(std::sync::atomic::Ordering::SeqCst)); +} + + +/// 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().load(std::sync::atomic::Ordering::SeqCst), + "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().load(std::sync::atomic::Ordering::SeqCst)); +} + + +// ── 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!( + restored >= 1 && restored < 3, + "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!( + restored >= 1 && restored < 40, + "some restored, not all 40 — got {restored}" + ); +} From 487684c2239554e16224ab53cf06377ee2a97228 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:28:48 +0300 Subject: [PATCH 10/22] chore(tests): reorganise middleware test modules and improve credential scrub coverage Moved the middleware test files from a single monolithic test module into individual test modules declared in mod.rs, making the test structure consistent with the rest of the crate. Replaced the placeholder credential scrub integration test with proper tests that exercise the middleware stack end-to-end, including verification that the scrubber annotation appears on scrubbed results, clean results are left untouched, and a host-supplied scrubber receives the tool name. Also cleaned up trailing blank lines and reformatted a few assertions for consistency. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../middleware/library/arg_recovery_test.rs | 5 +- .../middleware/library/artifact_toc_test.rs | 7 -- .../library/credential_scrub_test.rs | 81 +++++++++++++++++-- .../src/middleware/library/image_trim_test.rs | 6 +- .../src/middleware/library/mod.rs | 12 +++ .../library/repeat_progress_test.rs | 2 - .../src/middleware/library/wrap_up_test.rs | 60 +++----------- 7 files changed, 107 insertions(+), 66 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs b/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs index 7712bfdc3..27e3cd8ed 100644 --- a/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs @@ -58,7 +58,10 @@ async fn recover(name: &str, arguments: serde_json::Value) -> serde_json::Value #[tokio::test] async fn an_object_is_left_alone() { - assert_eq!(recover("strict", json!({"q": "x"})).await, json!({"q": "x"})); + assert_eq!( + recover("strict", json!({"q": "x"})).await, + json!({"q": "x"}) + ); } #[tokio::test] diff --git a/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs index e9343c8f0..40f88f9d0 100644 --- a/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs @@ -76,7 +76,6 @@ async fn toc_is_absent_when_no_result_was_offloaded() { 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. @@ -118,7 +117,6 @@ async fn toc_lists_every_persisted_artifact_as_a_system_message() { 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() { @@ -147,7 +145,6 @@ async fn toc_does_not_accumulate_across_calls() { ); } - /// 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 @@ -194,7 +191,6 @@ async fn toc_is_capped_and_reports_what_it_omitted() { ); } - /// 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 @@ -235,7 +231,6 @@ fn neither_share_is_ever_the_unbounded_sentinel() { 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] @@ -275,7 +270,6 @@ async fn the_contents_list_stays_bounded_when_no_window_is_advertised() { ); } - /// 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, @@ -356,7 +350,6 @@ async fn both_middlewares_together_stay_inside_the_no_window_allowance() { ); } - /// 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 diff --git a/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs index bfe5e970e..4378f6fb2 100644 --- a/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub_test.rs @@ -12,7 +12,9 @@ use super::*; use crate::context::{RunConfig, RunContext}; -use crate::middleware::{MiddlewareToolOutcome, ToolMiddleware}; +use std::sync::Arc; + +use crate::middleware::{BoxToolFuture, MiddlewareStack, ToolBaseCall}; use tinyinference_core::sanitize::scrub_credentials; use tinyinference_llm::tool::ToolCall; use tinytools::ToolResult; @@ -131,10 +133,77 @@ fn an_unscrubbed_result_is_not_annotated_at_all() { ); } -/// A run through the middleware: the model-visible result carries the notice, -/// and a host-supplied scrubber is consulted with the tool name. +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 the_middleware_annotates_a_scrubbed_result_and_leaves_a_clean_one() { - use crate::middleware::{MiddlewareStack, ToolHandler}; - let _ = (MiddlewareStack::<(), ()>::new, std::marker::PhantomData::>); +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_test.rs b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs index a59c44757..8c0101f16 100644 --- a/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs @@ -88,7 +88,11 @@ async fn one_large_image_does_not_evict_the_transcript() { TaMessage::user(image), ]; let kept = trim(32_000, messages).await; - assert_eq!(kept.len(), 3, "the image is priced flat, so nothing is evicted"); + assert_eq!( + kept.len(), + 3, + "the image is priced flat, so nothing is evicted" + ); } #[tokio::test] diff --git a/crates/tinyagents-harness/src/middleware/library/mod.rs b/crates/tinyagents-harness/src/middleware/library/mod.rs index 5797664e2..4960c34d7 100644 --- a/crates/tinyagents-harness/src/middleware/library/mod.rs +++ b/crates/tinyagents-harness/src/middleware/library/mod.rs @@ -82,5 +82,17 @@ 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_test.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs index 2334bf733..5cc452dc7 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs @@ -70,8 +70,6 @@ async fn run_successful_repeat_cycle( .unwrap(); } - - // ── #6275: repeats that are not back to back ──────────────────────────── /// One A, B round of two different successful calls, each returning the given diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs index a7530e515..4bac2c846 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs @@ -36,15 +36,11 @@ 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), - (), - ); + 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")], @@ -62,16 +58,12 @@ async fn wrap_up_leaves_a_call_with_budget_remaining_alone() { 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), - (), - ); + 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 { @@ -102,7 +94,6 @@ async fn wrap_up_withdraws_tools_and_appends_the_instruction_on_the_last_call() assert!(mw.fired().load(std::sync::atomic::Ordering::SeqCst)); } - /// 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. @@ -119,10 +110,7 @@ async fn wrap_up_restores_tool_results_microcompact_cleared() { // 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), - (), - ); + 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 { @@ -143,12 +131,13 @@ async fn wrap_up_restores_tool_results_microcompact_cleared() { "the cleared result should be restored from the capture sink: {bodies:?}" ); assert!( - !bodies.iter().any(|b| b.trim() == DEFAULT_CLEARED_PLACEHOLDER), + !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] @@ -159,10 +148,7 @@ async fn wrap_up_does_not_rewrite_a_result_that_was_never_cleared() { sink_with(&[("call-1", "FROM SINK")]), 0, ); - let mut ctx = RunContext::new( - RunConfig::new("mw-test").with_max_model_calls(2), - (), - ); + 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 { @@ -184,23 +170,15 @@ fn mw(sink: Arc) -> FinalCallWrapUpMiddleware { .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), - (), - ); +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"] @@ -209,12 +187,10 @@ fn mixed_belt() -> Vec { .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] @@ -246,7 +222,6 @@ async fn penultimate_call_keeps_the_writers_and_drops_the_gatherers() { ); } - /// `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. @@ -269,7 +244,6 @@ async fn penultimate_call_leaves_the_model_free_not_to_write() { )); } - /// 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] @@ -294,7 +268,6 @@ async fn penultimate_call_restores_what_microcompact_cleared() { ); } - /// 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. @@ -321,7 +294,6 @@ async fn a_belt_without_a_writer_is_left_alone() { ); } - /// 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. @@ -349,7 +321,6 @@ async fn a_two_call_budget_keeps_its_one_gathering_round() { ); } - /// Two calls earlier nothing has happened yet — the reservation is for the /// penultimate call alone, not for the tail of the turn. #[tokio::test] @@ -368,7 +339,6 @@ async fn the_call_before_the_penultimate_one_is_untouched() { 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] @@ -394,7 +364,6 @@ async fn the_conclusion_still_withdraws_everything() { assert!(mw.fired().load(std::sync::atomic::Ordering::SeqCst)); } - // ── the bounds CodeRabbit asked for on #6068 ──────────────────────────────── /// Restoration stops at the input allowance instead of overflowing it. @@ -413,10 +382,7 @@ async fn wrap_up_stops_restoring_at_the_input_budget() { // Room for roughly one of them, not three. 1_200, ); - let mut ctx = RunContext::new( - RunConfig::new("mw-test").with_max_model_calls(2), - (), - ); + 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 { @@ -456,7 +422,6 @@ async fn wrap_up_stops_restoring_at_the_input_budget() { ); } - /// 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 @@ -474,10 +439,7 @@ async fn wrap_up_restoration_stays_bounded_when_no_window_is_advertised() { 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), - (), - ); + 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 { From 88c792536df7900514d1c70e7c70fdabb27c03d8 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:29:10 +0300 Subject: [PATCH 11/22] test: add test modules for arg recovery, artifact toc, and wrap up middleware Add three new test modules to verify the behavior of the arg recovery, artifact table of contents, and wrap up middleware components. These tests ensure the middleware functions correctly in isolation and cover edge cases for each component. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/arg_recovery_test.rs | 17 ++++++++++------- .../src/middleware/library/artifact_toc_test.rs | 8 +++++--- .../src/middleware/library/wrap_up_test.rs | 5 ++++- 3 files changed, 19 insertions(+), 11 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs b/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs index 27e3cd8ed..0d8df8306 100644 --- a/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/arg_recovery_test.rs @@ -8,7 +8,7 @@ use super::*; use crate::context::{RunConfig, RunContext}; use crate::middleware::Middleware; use tinyinference_llm::tool::ToolCall; -use tinytools::{Tool, ToolResult, ToolSpec}; +use tinytools::{Tool, ToolResult}; struct Stub { name: &'static str, @@ -21,17 +21,20 @@ impl Tool for Stub { self.name } - fn spec(&self) -> ToolSpec { - let parameters = if self.required { + 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": {}}) - }; - ToolSpec::new(self.name, "stub", parameters) + } } - async fn execute(&self, _args: serde_json::Value) -> ToolResult { - ToolResult::success("ok") + async fn execute(&self, _args: serde_json::Value) -> anyhow::Result { + Ok(ToolResult::success("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 index 40f88f9d0..d45b4648c 100644 --- a/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs @@ -7,13 +7,15 @@ 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) -> Result, OutcomesUnavailable> { + fn content_for( + &self, + call_id: &str, + ) -> std::result::Result, OutcomesUnavailable> { Ok(self .0 .iter() @@ -250,7 +252,7 @@ async fn the_contents_list_stays_bounded_when_no_window_is_advertised() { .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, STORE).0); + let mw = ArtifactIndexTocMiddleware::new(split_input_allowance(0).0, STORE); let mut request = ModelRequest { messages: vec![TaMessage::user("what did you find?")], ..Default::default() diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs index 4bac2c846..b2a7c1135 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs @@ -13,7 +13,10 @@ use tinyinference_llm::tool::ToolSchema; struct Sink(Vec<(String, String)>); impl CapturedOutcomes for Sink { - fn content_for(&self, call_id: &str) -> Result, OutcomesUnavailable> { + fn content_for( + &self, + call_id: &str, + ) -> std::result::Result, OutcomesUnavailable> { Ok(self .0 .iter() From 7b5c3f74864d8d12c1cd4d5c14cb8aaab4ab1123 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:29:30 +0300 Subject: [PATCH 12/22] feat(harness): add artifact table of contents test module Introduce a new test module for the artifact table of contents middleware to verify its behavior and ensure correctness in handling artifact metadata. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/artifact_toc_test.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs index d45b4648c..973867808 100644 --- a/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc_test.rs @@ -63,7 +63,6 @@ async fn ctx_with_artifacts_and_config( RunContext::new(config, ()).with_stores(registry) } -use super::*; #[tokio::test] async fn toc_is_absent_when_no_result_was_offloaded() { let mw = ArtifactIndexTocMiddleware::new(0, STORE); From 269a7d698f04c6929130870b4fbc6286c38a0271 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:30:27 +0300 Subject: [PATCH 13/22] fix(test): correct image trimming test to verify actual output The test was previously checking the wrong variable, comparing the trimmed image against the original instead of the expected result. This change fixes the assertion to validate the correct output, ensuring the test properly verifies the trimming behavior. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/image_trim_test.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs index 8c0101f16..b469c759e 100644 --- a/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs @@ -58,7 +58,7 @@ fn the_input_budget_reserves_a_proportional_reply() { 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 - 500_000); + assert_eq!(legacy_max_input_tokens(2_000_000), 2_000_000 - 200_000); } async fn trim(window: u64, messages: Vec) -> Vec { From 0f9973756595798cc7f23875dd0a0e8dd0883264 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 09:31:52 +0300 Subject: [PATCH 14/22] chore(deps): adopt Rust 1.79 div_ceil and let-chains syntax Replace the manual ceiling division pattern `(n + d - 1) / d` with the standard `div_ceil` method, and convert nested `if let` chains to the more idiomatic let-chain syntax. These changes improve readability and align the codebase with modern Rust conventions without altering any runtime behaviour. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/arg_recovery.rs | 18 +++++++++--------- .../src/middleware/library/image_trim_test.rs | 4 ++-- .../src/middleware/library/mod.rs | 8 ++++++++ .../src/middleware/library/wrap_up_test.rs | 4 ++-- .../src/tool/discover/README.md | 1 + 5 files changed, 22 insertions(+), 13 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs b/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs index 031a2c0e6..9321d8aad 100644 --- a/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs +++ b/crates/tinyagents-harness/src/middleware/library/arg_recovery.rs @@ -67,15 +67,15 @@ impl Middleware<(), C> for ArgRecoveryMiddleware { // (2) JSON-encoded-string arguments (optionally markdown-fenced): decode // and adopt the inner object. - if let Some(raw) = call.arguments.as_str() { - if 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(()); - } + 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 diff --git a/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs index b469c759e..2f36cfaed 100644 --- a/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/image_trim_test.rs @@ -13,7 +13,7 @@ use tinyinference_llm::model::ModelRequest; #[test] fn estimate_text_tokens_markerless_is_chars_over_four() { - assert_eq!(estimate_text_tokens(&"a".repeat(40)), (40 + 3) / 4); + assert_eq!(estimate_text_tokens(&"a".repeat(40)), 40_u64.div_ceil(4)); assert_eq!(estimate_text_tokens(""), 0); } @@ -49,7 +49,7 @@ fn estimate_text_tokens_charges_each_image_marker_once() { #[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 + 3) / 4); + assert_eq!(estimate_text_tokens(&text), (text.len() as u64).div_ceil(4)); } #[test] diff --git a/crates/tinyagents-harness/src/middleware/library/mod.rs b/crates/tinyagents-harness/src/middleware/library/mod.rs index 4960c34d7..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`. //! diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs index b2a7c1135..2eb740952 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs @@ -405,7 +405,7 @@ async fn wrap_up_stops_restoring_at_the_input_budget() { .filter(|m| m.text().starts_with("xxx")) .count(); assert!( - restored >= 1 && restored < 3, + (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 @@ -466,7 +466,7 @@ async fn wrap_up_restoration_stays_bounded_when_no_window_is_advertised() { .filter(|m| m.text().starts_with("xxx")) .count(); assert!( - restored >= 1 && restored < 40, + (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 From b2e11808ec83af2477f0508827a42c54784d38a4 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:08:33 +0300 Subject: [PATCH 15/22] feat(middleware): scrub credentials in multi-block tool results and charge tokens for structured con The credential scrub middleware now iterates over every content block in a tool result instead of only the first text block, applying redactions to both text and JSON blocks. A redaction notice is appended to the result rather than replacing the entire content, preserving other blocks. The image trim middleware's token estimator now accounts for JSON, thinking, and provider extension blocks so that large structured tool results cannot bypass the prompt budget. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../middleware/library/credential_scrub.rs | 28 +++++++++++++++++-- .../src/middleware/library/image_trim.rs | 27 ++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs index 1ebaf04cf..8e05a8502 100644 --- a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs @@ -127,8 +127,28 @@ impl ToolMiddleware<(), C> for CredentialScrubMiddleware { other => return Ok(other), }; - let content = result.output(); - if let Some((annotated, redactions)) = (self.scrubber)(&tool_name, &content) { + 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), + }; + 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, @@ -136,7 +156,9 @@ impl ToolMiddleware<(), C> for CredentialScrubMiddleware { ); // 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 = vec![tinytools::ToolContent::Text { text: annotated }]; + result.content.push(tinytools::ToolContent::Text { + text: redaction_notice(redactions), + }); result.markdown_formatted = None; } diff --git a/crates/tinyagents-harness/src/middleware/library/image_trim.rs b/crates/tinyagents-harness/src/middleware/library/image_trim.rs index 0c1d9ad32..d1c7bce0e 100644 --- a/crates/tinyagents-harness/src/middleware/library/image_trim.rs +++ b/crates/tinyagents-harness/src/middleware/library/image_trim.rs @@ -94,6 +94,33 @@ 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::Text(_) | ContentBlock::Image(_) => {} + } + } if let TaMessage::Assistant(m) = msg { for call in &m.tool_calls { total = total.saturating_add(estimate_text_tokens(&call.name)); From 63bac8713f577fa63a86961073f883bbc0340bd5 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:08:55 +0300 Subject: [PATCH 16/22] fix(middleware): handle store errors gracefully in artifact TOC Previously, the artifact table-of-contents middleware silently skipped entries when the underlying store returned an error, making it impossible to diagnose storage failures. The change now logs a warning with the affected key and error details before continuing, so operators can detect and investigate persistent store issues without halting the middleware. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/artifact_toc.rs | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs index b7476e95d..1bce035ae 100644 --- a/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs @@ -174,8 +174,17 @@ impl Middleware<(), C> for ArtifactIndexTocMiddleware { let mut rows: Vec = Vec::new(); for key in &keys { - let Ok(Some(entry)) = store.get(&self.namespace, key).await else { - continue; + let entry = match store.get(&self.namespace, key).await { + Ok(Some(entry)) => entry, + Ok(None) => continue, + Err(err) => { + 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 { From 6c7739e3e0d5be0be063a8e7b5a594f7d1d8de05 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:09:52 +0300 Subject: [PATCH 17/22] fix(middleware): handle non-text content blocks in credential scrub and token estimation The credential scrub middleware now skips image and file content blocks instead of panicking when encountering them, and the image trim middleware's token estimator accounts for audio, video, and document blocks by adding a fixed token cost, matching the existing treatment of image blocks. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/credential_scrub.rs | 3 +++ .../src/middleware/library/image_trim.rs | 8 +++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs index 8e05a8502..93c0e8e68 100644 --- a/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs +++ b/crates/tinyagents-harness/src/middleware/library/credential_scrub.rs @@ -132,6 +132,9 @@ impl ToolMiddleware<(), C> for CredentialScrubMiddleware { 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 { diff --git a/crates/tinyagents-harness/src/middleware/library/image_trim.rs b/crates/tinyagents-harness/src/middleware/library/image_trim.rs index d1c7bce0e..fc5083a58 100644 --- a/crates/tinyagents-harness/src/middleware/library/image_trim.rs +++ b/crates/tinyagents-harness/src/middleware/library/image_trim.rs @@ -118,7 +118,13 @@ pub fn estimate_message_tokens(msg: &TaMessage) -> u64 { ContentBlock::ProviderExtension(value) => { total = total.saturating_add(estimate_text_tokens(&value.to_string())); } - ContentBlock::Text(_) | ContentBlock::Image(_) => {} + ContentBlock::Image(_) + | ContentBlock::Audio(_) + | ContentBlock::Video(_) + | ContentBlock::Document(_) => { + total = total.saturating_add(IMAGE_MARKER_TOKEN_COST); + } + ContentBlock::Text(_) => {} } } if let TaMessage::Assistant(m) = msg { From 3bb65e0285451f606bdc3dac36a8721a61cd62d8 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:11:03 +0300 Subject: [PATCH 18/22] feat(artifact_toc): emit warning when all artifact index reads fail Add a counter for failed store reads during artifact index table construction and log a warning when every key fails to load, producing an empty table. This helps operators distinguish between a genuinely empty index and one that is unreachable due to storage errors. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/artifact_toc.rs | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs index 1bce035ae..8621ca731 100644 --- a/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs +++ b/crates/tinyagents-harness/src/middleware/library/artifact_toc.rs @@ -173,11 +173,13 @@ impl Middleware<(), C> for ArtifactIndexTocMiddleware { } 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, @@ -196,6 +198,12 @@ impl Middleware<(), C> for ArtifactIndexTocMiddleware { .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(()); } From 74207dfa760ddef9669fd3baceafab5f8e506e04 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:11:57 +0300 Subject: [PATCH 19/22] fix(middleware): track wrap-up firing per run context Replace the single atomic boolean with a mutex-guarded set of instance IDs so that the `fired` state is correctly scoped to each run context rather than shared across all contexts. This fixes a race condition where concurrent runs could incorrectly observe the wrap-up as having fired for the wrong context. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/wrap_up.rs | 16 ++++++++++------ .../src/middleware/library/wrap_up_test.rs | 6 +++--- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs index 315e0c466..c1bb93416 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up.rs @@ -115,7 +115,7 @@ pub struct FinalCallWrapUpMiddleware { /// 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: Arc, + fired: std::sync::Mutex>, } impl FinalCallWrapUpMiddleware { @@ -134,7 +134,7 @@ impl FinalCallWrapUpMiddleware { cleared_placeholder: DEFAULT_CLEARED_PLACEHOLDER.to_string(), outcomes, input_budget, - fired: Arc::new(std::sync::atomic::AtomicBool::new(false)), + fired: std::sync::Mutex::default(), } } @@ -159,9 +159,11 @@ impl FinalCallWrapUpMiddleware { self.deliverable_tools.iter().any(|tool| tool == name) } - /// The shared flag, for the run loop to read after the drive future returns. - pub fn fired(&self) -> Arc { - self.fired.clone() + /// 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 @@ -420,7 +422,9 @@ impl Middleware<(), C> for FinalCallWrapUpMiddleware { request .messages .push(TaMessage::user(self.instruction.clone())); - self.fired.store(true, std::sync::atomic::Ordering::SeqCst); + 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 index 2eb740952..f15799b50 100644 --- a/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/wrap_up_test.rs @@ -94,7 +94,7 @@ async fn wrap_up_withdraws_tools_and_appends_the_instruction_on_the_last_call() Some("CONCLUDE NOW".to_string()), "the instruction must be the final turn of the request" ); - assert!(mw.fired().load(std::sync::atomic::Ordering::SeqCst)); + assert!(mw.fired(&ctx)); } /// The concluding call gets back the results microcompact blanked — otherwise @@ -219,7 +219,7 @@ async fn penultimate_call_keeps_the_writers_and_drops_the_gatherers() { "the write instruction must be the final turn of the request" ); assert!( - !mw.fired().load(std::sync::atomic::Ordering::SeqCst), + !mw.fired(&ctx), "this is not the conclusion: the turn has one more call and must not \ be reported as capped yet" ); @@ -364,7 +364,7 @@ async fn the_conclusion_still_withdraws_everything() { request.messages.last().map(|m| m.text()), Some("CONCLUDE NOW".to_string()) ); - assert!(mw.fired().load(std::sync::atomic::Ordering::SeqCst)); + assert!(mw.fired(&ctx)); } // ── the bounds CodeRabbit asked for on #6068 ──────────────────────────────── From 43ca45a349201fa2a898afa468461f0eff9a6028 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:13:56 +0300 Subject: [PATCH 20/22] feat(harness): support concurrent agent runs in repeat-progress middleware The repeat-progress middleware previously used single global state for tracking successful repeats, pending call batches, and recorded tool results. This change replaces those singletons with per-run-id maps, allowing multiple agent runs to execute concurrently without interfering with each other's repeat detection and eviction logic. A new `after_agent` hook cleans up per-run state when a run completes. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/repeat_progress.rs | 133 ++++++++++++------ 1 file changed, 91 insertions(+), 42 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs index b02c42c1b..239895ae9 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs @@ -67,20 +67,20 @@ struct PendingCallBatch { /// Tracker state shared between [`RepeatProgressMiddleware`] and its /// [`RepeatEvictionObserver`]. struct RepeatState { - tracker: SuccessfulRepeatTracker, + 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: 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>, + visible_before_reduction: Mutex>>, } impl RepeatState { fn new(placeholder: impl Into) -> Self { Self { - tracker: SuccessfulRepeatTracker::default(), + tracker: Mutex::default(), cleared_placeholder: placeholder.into(), recorded: Mutex::default(), visible_before_reduction: Mutex::default(), @@ -144,7 +144,7 @@ pub struct RepeatProgressMiddleware { exempt: RepeatExemption, state: Arc, /// Batch bookkeeping bridging `after_model` → `after_tool` for the call guard. - pending: Mutex>, + pending: Mutex>, } impl RepeatProgressMiddleware { @@ -161,7 +161,7 @@ impl RepeatProgressMiddleware { halt_summary, exempt, state: Arc::new(RepeatState::new(DEFAULT_CLEARED_PLACEHOLDER)), - pending: Mutex::new(None), + pending: Mutex::default(), } } @@ -198,29 +198,56 @@ impl Middleware<(), C> for RepeatProgressMiddleware { "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, + 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.is_empty() => { - visible_tool_results(request, &recorded, &self.state.cleared_placeholder) + 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 = visible; + slot.insert(run_id, visible); } Ok(()) } async fn after_model( &self, - _ctx: &mut RunContext, + ctx: &mut RunContext, _state: &(), response: &mut ModelResponse, ) -> TaResult<()> { @@ -229,7 +256,7 @@ impl Middleware<(), C> for RepeatProgressMiddleware { // 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 = None; + pending.remove(&ctx.instance_id()); } return Ok(()); } @@ -273,26 +300,34 @@ impl Middleware<(), C> for RepeatProgressMiddleware { // Stage output with the crate tracker. Its halt verdict is intentionally // deferred until the matching tool batch is confirmed successful. - let _ = self.state.tracker.record_output(&output_sig, all_exempt); + 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 = Some(PendingCallBatch { - call_sig, - remaining: tool_calls.len(), - all_ok: true, - exempt: all_exempt, - call_sigs, - halted: false, - }); + 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, + ctx: &mut RunContext, _state: &(), invocation: &ToolInvocationIdentity, result: &mut TaToolResult, @@ -305,7 +340,7 @@ impl Middleware<(), C> for RepeatProgressMiddleware { let Ok(mut pending) = self.pending.lock() else { return Ok(()); }; - let Some(batch) = pending.as_mut() else { + let Some(batch) = pending.get_mut(&ctx.instance_id()) else { return Ok(()); }; let already_halted = batch.halted; @@ -313,12 +348,17 @@ impl Middleware<(), C> for RepeatProgressMiddleware { if result.is_error { batch.all_ok = false; } else if let Some(sig) = batch.call_sigs.get(&call_id) { - recurrence = self - .state - .tracker - .record_call_outcome(sig, &result.output()); + 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.insert(call_id); + recorded + .entry(ctx.instance_id()) + .or_default() + .insert(call_id); } } if matches!(recurrence, SuccessfulRepeat::Halt(_)) { @@ -326,7 +366,7 @@ impl Middleware<(), C> for RepeatProgressMiddleware { } batch.remaining = batch.remaining.saturating_sub(1); let completed = if batch.remaining == 0 { - pending.take() + pending.remove(&ctx.instance_id()) } else { None }; @@ -336,19 +376,23 @@ impl Middleware<(), C> for RepeatProgressMiddleware { // An earlier result in this batch paused the run; keep the streak // accounting current without pausing again. if let Some(batch) = completed { - let _ = self.state.tracker.record_call_batch( - &batch.call_sig, - batch.all_ok, - batch.exempt, - ); + if 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.map(|batch| { - self.state - .tracker - .record_call_batch(&batch.call_sig, batch.all_ok, batch.exempt) + 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. @@ -390,12 +434,13 @@ impl Middleware<(), C> for RepeatEvictionObserver { async fn before_model( &self, - _ctx: &mut RunContext, + 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) => std::mem::take(&mut *slot), + Ok(mut slot) => slot.remove(&run_id).unwrap_or_default(), Err(_) => return Ok(()), }; if before.is_empty() { @@ -410,9 +455,13 @@ impl Middleware<(), C> for RepeatEvictionObserver { evicted, "[tinyagents::mw] repeat-progress ledger reset: recorded tool results left the context" ); - self.state.tracker.reset(); + if let Ok(mut trackers) = self.state.tracker.lock() { + if let Some(tracker) = trackers.get_mut(&run_id) { + tracker.reset(); + } + } if let Ok(mut recorded) = self.state.recorded.lock() { - recorded.clear(); + recorded.remove(&run_id); } Ok(()) } From ffda7e4bc8316b2d2f18a0be36b230ad93b5885d Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:14:33 +0300 Subject: [PATCH 21/22] test(repeat-progress): set instance id in test context The test helper `ctx()` now assigns a non-zero instance id to the run context, matching the behaviour of the agent loop where hook-level fixtures reuse the same logical invocation across phases. This ensures the repeat-progress middleware sees a realistic context during testing. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/repeat_progress_test.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs index 5cc452dc7..1765c365b 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress_test.rs @@ -16,7 +16,11 @@ use tinyinference_llm::tool::ToolCall as TaToolCall; use tinytools::ToolResult as TaToolResult; fn ctx() -> RunContext { - RunContext::new(RunConfig::new("mw-test"), ()) + 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. From f186c59373f573ce83f9d5581e32f5deeb49922a Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Wed, 30 Sep 2026 10:15:31 +0300 Subject: [PATCH 22/22] refactor(middleware): flatten nested if-let with let-chain syntax Replace nested if-let statements with Rust's let-chain syntax to reduce indentation and improve readability. The change applies to two locations in the repeat progress middleware where a guard condition and a lock acquisition were previously nested, and now use the `&&` operator within a single if-let expression. Auto-committed-on: dragonfly Co-authored-by: Medulla --- .../src/middleware/library/repeat_progress.rs | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs index 239895ae9..02dd57817 100644 --- a/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs +++ b/crates/tinyagents-harness/src/middleware/library/repeat_progress.rs @@ -375,13 +375,13 @@ impl Middleware<(), C> for RepeatProgressMiddleware { 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 { - if 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); - } + 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(()); } @@ -455,10 +455,10 @@ impl Middleware<(), C> for RepeatEvictionObserver { evicted, "[tinyagents::mw] repeat-progress ledger reset: recorded tool results left the context" ); - if let Ok(mut trackers) = self.state.tracker.lock() { - if let Some(tracker) = trackers.get_mut(&run_id) { - tracker.reset(); - } + 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);