From 1659796e3030a92e53cc41be18f7df5fe9cd8b18 Mon Sep 17 00:00:00 2001 From: Yuanhao Li Date: Sat, 11 Jul 2026 00:08:53 +0200 Subject: [PATCH] fix: address Phase C review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five-agent review of the combined v0.10.0...main diff (PRs #61-#65) surfaced convergent findings; this is the full fix batch. Critical fixes: - prompt_structured error handling overhauled: provider failures now surface as StructuredPromptError::Provider carrying the real error (previously laundered into Parse { raw: "" }); the raw-text scan is scoped to messages produced by THIS call (a failed run can no longer return an earlier turn's JSON as Ok(T)); the payload is taken from the LAST text block (tool-forcing appends after preamble); Parse now carries #[source] serde_json::Error. - Schema drop-leak closed structurally: the output_schema field on Agent is gone — the schema threads per-call through prompt_messages_internal into the run's AgentLoopConfig, so a timed-out/dropped future can't leave the agent schema-forced. - Bedrock replays Content::Thinking (with signature) as reasoningContent — previously captured then dropped, guaranteeing a ValidationException on multi-turn thinking + tool use. - Anthropic structured outputs disable thinking for that request (forced tool_choice + extended thinking is an API-level 400) with a warning. - Session::append_new verifies the history extends the current path and returns HistoryDiverged instead of silently dropping turns or corrupting the tree when compaction (on by default) rewrites agent messages; returns the appended count. - Session::from_jsonl single-pass validation: DuplicateId + UnknownParent (parents must precede children, which also makes cycles impossible — path_ids provably terminates). API shape (free now — these types are unreleased): - ToolMiddleware::before_tool takes a #[non_exhaustive] ToolCallRequest context struct instead of three frozen positional params. - StreamConfig and OutputSchema are #[non_exhaustive]; StreamConfig::new added (this release alone added output_schema to it). - ToolDecision documents the deliberate StopReason-style exhaustive policy. Hardening: - A panicking ToolMiddleware is contained (catch_unwind) as a Deny instead of killing the loop task and stripping the agent of its tools. - Middleware denials emit tracing::warn! (operator visibility). - unwrap_structured_tool_call removes only the synthetic tool call and preserves non-ToolUse stop reasons (Length truncation no longer laundered into success). - llm_stream span gains an error field; seek_checkpoint is latest-wins on duplicate labels; Vertex round-trips Gemini thought signatures on function calls (parity with google.rs). Docs: tools.md example on the new signature + panic note; structured-outputs error semantics + schema-dialect caveats + thinking caveat; telemetry OTel snippet updated to the current builder API (old one was removed in 0.17) and overhead claims made honest; session-trees compaction caveat + validation; "enforced natively where supported"; CHANGELOG Fixed section; CLAUDE.md. Tests (+15, 376 total): schema propagation via capturing provider (replaces the tautological reset test); Provider-error and stale-history-never-parsed; middleware never sees the synthetic tool; middleware panic containment; Batched-strategy deny; session divergence/dup-id/dangling-parent/latest-wins; bedrock thinking-replay body; anthropic structured-disables-thinking body; google thinking-drop pin; vertex signature round-trip; telemetry field values including exact cost math (on_record collector). Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 27 +++- CLAUDE.md | 4 +- README.md | 4 +- docs/concepts/session-trees.md | 16 ++- docs/concepts/structured-outputs.md | 19 ++- docs/concepts/telemetry.md | 26 ++-- docs/concepts/tools.md | 15 +- src/agent.rs | 99 +++++++++---- src/agent_loop.rs | 53 ++++++- src/lib.rs | 6 +- src/provider/anthropic.rs | 28 +++- src/provider/bedrock.rs | 64 ++++++++- src/provider/google.rs | 26 ++++ src/provider/google_vertex.rs | 56 +++++++- src/provider/traits.rs | 37 ++++- src/session.rs | 70 +++++++++- src/types.rs | 31 ++++- tests/agent_test.rs | 209 ++++++++++++++++++++++++++-- tests/anthropic_stream_test.rs | 19 +-- tests/google_stream_test.rs | 19 +-- tests/session_test.rs | 91 ++++++++++-- tests/sub_agent_test.rs | 7 +- tests/telemetry_test.rs | 139 ++++++++++++++++++ 23 files changed, 917 insertions(+), 148 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 225daff..12560ea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,31 @@ adheres to [Semantic Versioning](https://semver.org/). ## Unreleased +### Fixed (pre-release review of the items below) + +- `prompt_structured` now surfaces provider failures as + `StructuredPromptError::Provider` (previously laundered into + `Parse { raw: "" }`), scans only messages produced by the current call + (never stale history), and threads the schema per-call so a dropped/timed- + out future can't leave the agent stuck in schema-forcing mode. +- Bedrock replays `Content::Thinking` blocks (with signatures) on subsequent + requests — previously captured and then dropped, breaking multi-turn + thinking + tool use with a ValidationException. +- Anthropic structured outputs disable extended thinking for that request + (forced tool choice + thinking is an API-level conflict) with a warning. +- Vertex now round-trips Gemini thought signatures on function calls (parity + with the Gemini API provider). +- `Session::append_new` verifies the history still extends the session path + and returns `HistoryDiverged` instead of silently corrupting the tree (the + usual cause: context compaction); `from_jsonl` validates ids and parents + (duplicates/dangling/cycles rejected); `seek_checkpoint` is latest-wins. +- A panicking `ToolMiddleware` is contained as a denial instead of killing + the loop task (which stripped the agent of its tools). +- Middleware denials are logged (`tracing::warn!`) so operators see them. +- `ToolMiddleware::before_tool` takes a `ToolCallRequest` context struct + (extensible without breaking implementors); `StreamConfig` and + `OutputSchema` are `#[non_exhaustive]` with constructors. + ### Added - **Telemetry** — `tracing` spans: `agent_loop` (model), `llm_stream` per @@ -23,7 +48,7 @@ adheres to [Semantic Versioning](https://semver.org/). - **Session trees** — `Session`: branching conversation history with `append`/`seek`/`checkpoint`, fork-preserving edits, `path_messages()` for - branch resume, and JSONL persistence. The pi-style id/parentId tree; maps + branch resume, and JSONL persistence. The pi-style id/parent_id tree; maps to GASP's `transcripts/` tier. - **Structured outputs** — `Agent::prompt_structured::(text, schema)` diff --git a/CLAUDE.md b/CLAUDE.md index 258b56d..ca8e88c 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,9 +72,9 @@ The loop: stream assistant response → extract tool calls → execute tools (pa - `Sequential` — one at a time, checks steering queue between each - `Batched { size }` — concurrent within batch, steering check between batches -**Structured outputs** (`Agent::prompt_structured::(text, schema)`): the schema travels via `StreamConfig.output_schema` (`OutputSchema` in `provider/traits.rs`). Anthropic enforces it by tool-forcing (a synthetic tool + `tool_choice`; the loop's `unwrap_structured_tool_call` converts the forced call back to text **before** tool-call extraction); OpenAI-compat uses `response_format: json_schema`; Gemini uses `responseSchema`. Responses/Azure/Vertex/Bedrock warn and ignore. +**Structured outputs** (`Agent::prompt_structured::(text, schema)`): the schema is threaded per-call through `prompt_messages_internal` into `StreamConfig.output_schema` (`OutputSchema` in `provider/traits.rs`) — never stored on the Agent. Errors: `Provider` (API failure), `Parse { source, raw }`, `NoOutput`; only this call's messages are scanned. Anthropic enforces it by tool-forcing (a synthetic tool + `tool_choice`; the loop's `unwrap_structured_tool_call` converts the forced call back to text **before** tool-call extraction); OpenAI-compat uses `response_format: json_schema`; Gemini uses `responseSchema`. Responses/Azure/Vertex/Bedrock warn and ignore. -**Tool middleware** (`ToolMiddleware` in `types.rs`): async approve/deny/modify hooks gating every tool call, run in a chain at the single choke point (`execute_single_tool`) shared by all three strategies. `Deny(reason)` becomes an error tool result the LLM sees (loop continues); `Modify(args)` rewrites the call. Installed via `Agent::with_tool_middleware` / `SubAgentTool::with_tool_middleware` / `AgentLoopConfig::tool_middleware`. Empty chain = allow all. +**Tool middleware** (`ToolMiddleware` in `types.rs`): async approve/deny/modify hooks gating every tool call, run in a chain at the single choke point (`execute_single_tool`) shared by all three strategies. `before_tool` takes a `#[non_exhaustive]` `ToolCallRequest` context struct (extensible without breaking implementors). `Deny(reason)` becomes an error tool result the LLM sees (loop continues); `Modify(args)` rewrites the call; a panicking middleware is contained as a denial. Installed via `Agent::with_tool_middleware` / `SubAgentTool::with_tool_middleware` / `AgentLoopConfig::tool_middleware`. Empty chain = allow all. ### OpenAPI Integration (`openapi/`, feature-gated) diff --git a/README.md b/README.md index 72bdcc4..c33bcc4 100644 --- a/README.md +++ b/README.md @@ -36,7 +36,7 @@ Everything is observable via events. Supports 7 API protocols covering 20+ LLM p - Parallel tool execution by default — sequential and batched strategies also available - Sub-agents via `SubAgentTool` — delegate tasks to child agent loops with their own tools and system prompts - Tool middleware — async approve/deny/modify hooks gating every tool call (`with_tool_middleware`), the mechanism for permission prompts and policy engines -- Structured outputs — `prompt_structured::()` returns typed, schema-validated replies, enforced natively per provider (Anthropic tool-forcing, OpenAI `json_schema`, Gemini `responseSchema`) +- Structured outputs — `prompt_structured::()` returns typed, schema-validated replies, enforced natively where supported (Anthropic tool-forcing, OpenAI `json_schema`, Gemini `responseSchema`; other protocols log a warning) - Real-time event streaming — `prompt()` spawns the loop concurrently and returns events immediately; `prompt_with_sender()` accepts a caller-provided channel for custom consumption - Streaming tool output — tools emit real-time progress via `on_update` callback - Multimodal support — `Content::Image` flows through tool results across all providers @@ -45,7 +45,7 @@ Everything is observable via events. Supports 7 API protocols covering 20+ LLM p - State persistence — `save_messages()` / `restore_messages()` for pause/resume workflows - Session trees — branching history with fork, checkpoints, and JSONL persistence (`Session`); edit an earlier turn and re-run without losing the original branch - Lifecycle callbacks — `before_turn`, `after_turn`, `on_error` for observability and control -- Telemetry — `tracing` spans for the loop, each LLM stream (tokens + cost fields), and each tool execution; bridge to OpenTelemetry via `tracing-opentelemetry`, zero overhead when no subscriber is installed +- Telemetry — `tracing` spans for the loop, each LLM stream (tokens + cost fields), and each tool execution; bridge to OpenTelemetry via `tracing-opentelemetry`, negligible overhead when no subscriber is installed - Full serde support — all core types implement `Serialize`/`Deserialize`/`PartialEq` - [AgentSkills](https://agentskills.io)-compatible skills — load skill directories, inject into system prompt, agent activates on demand diff --git a/docs/concepts/session-trees.md b/docs/concepts/session-trees.md index 60fbb21..16f8c99 100644 --- a/docs/concepts/session-trees.md +++ b/docs/concepts/session-trees.md @@ -20,7 +20,7 @@ let mut agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "So let mut rx = agent.prompt("draft a plan").await; while rx.recv().await.is_some() {} agent.finish().await; -session.append_new(agent.messages()); +session.append_new(agent.messages())?; session.checkpoint("first-draft")?; // ... more turns ... @@ -32,7 +32,7 @@ let mut agent = Agent::from_config(ModelConfig::anthropic("claude-sonnet-5", "So let mut rx = agent.prompt("actually, make it a library instead").await; while rx.recv().await.is_some() {} agent.finish().await; -session.append_new(agent.messages()); // new branch recorded +session.append_new(agent.messages())?; // new branch recorded // Both branches exist: assert_eq!(session.branch_tips().len(), 2); @@ -46,7 +46,15 @@ let restored = Session::from_jsonl(&std::fs::read_to_string("session.jsonl")?)?; ``` One entry per line, append-friendly, diff-friendly. On load the head is the -last line's entry. The flat `save_messages()` / `restore_messages()` API +last line's entry, and the file is validated (duplicate ids and +dangling/forward parent references are rejected — which also makes cycles +impossible). + +> **Compaction caveat:** `append_new` verifies the agent's history still +> extends the session path and returns `SessionError::HistoryDiverged` when it +> doesn't. The usual cause is context compaction (on by default) rewriting the +> agent's messages — disable context management on session-tracked agents, or +> rebuild the branch with `Session::from_messages` after a divergence. The flat `save_messages()` / `restore_messages()` API remains for single-branch persistence. ## API sketch @@ -54,7 +62,7 @@ remains for single-branch persistence. | Method | Purpose | |---|---| | `append(msg) -> id` | Add as child of head, advance head | -| `append_new(&agent_messages)` | Record everything beyond the current path — the post-run sync | +| `append_new(&agent_messages)` | Record everything beyond the current path (verifies the prefix; errors on divergence) | | `seek(id)` / `seek_checkpoint(label)` | Move the head (fork point) | | `checkpoint(label)` | Label the head | | `path_messages()` | Root→head messages — feed to `Agent::with_messages` | diff --git a/docs/concepts/structured-outputs.md b/docs/concepts/structured-outputs.md index 5c68449..839ef1c 100644 --- a/docs/concepts/structured-outputs.md +++ b/docs/concepts/structured-outputs.md @@ -32,8 +32,11 @@ let invoice: Invoice = agent ``` Derive the schema however you like — by hand as above, or with the -[`schemars`](https://crates.io/crates/schemars) crate -(`schemars::schema_for!(Invoice)`). +[`schemars`](https://crates.io/crates/schemars) crate (convert with +`serde_json::to_value(schemars::schema_for!(Invoice))`). Mind the provider +dialects: OpenAI strict mode requires `additionalProperties: false` and every +property listed in `required`; Gemini rejects `$defs`/`$ref`. Schemas are +passed through as given. ## How each provider enforces it @@ -48,9 +51,15 @@ Derive the schema however you like — by hand as above, or with the - `prompt_structured` runs the loop to completion internally and returns the parsed `T` — there is no event receiver for this call. -- On parse failure you get `StructuredPromptError::Parse { error, raw }` — the - raw model text is preserved so you can retry or salvage. -- On Anthropic the forced tool call preempts regular tools for that request. +- Three error shapes: `Provider { message }` when the API call itself failed + (auth, network, a schema-induced 400 — retrying the parse is pointless); + `Parse { source, raw }` when the model's text didn't deserialize (the raw + text is preserved so you can retry or salvage); `NoOutput` when the run + produced no text. Only messages produced by **this call** are considered — + stale output from earlier turns is never parsed. +- On Anthropic the forced tool call preempts regular tools for that request, + and **disables extended thinking** for that request (forced tool choice and + thinking are mutually exclusive at the API level — a warning is logged). Treat structured prompts as **extraction/finalization calls**, not agentic tool-using turns. - Markdown code fences around the JSON are stripped defensively before diff --git a/docs/concepts/telemetry.md b/docs/concepts/telemetry.md index 97a306f..2c860c4 100644 --- a/docs/concepts/telemetry.md +++ b/docs/concepts/telemetry.md @@ -2,14 +2,14 @@ yoagent instruments the loop with [`tracing`](https://docs.rs/tracing) **spans** — structured, timed, nested units your observability stack can consume. With -no subscriber installed they compile to no-ops; there is zero overhead unless -you opt in. +no subscriber installed the overhead is near-zero (a cached per-callsite +interest check); nothing is exported unless you opt in. ## The span tree ``` agent_loop (model) -└─ llm_stream (turn, model, tokens_in, tokens_out, tokens_cached, cost_usd) +└─ llm_stream (turn, model, tokens_in, tokens_out, tokens_cached, cost_usd, error) └─ tool (tool, tool_call_id, is_error) ``` @@ -38,19 +38,27 @@ same spans flow to any OTLP backend — Datadog, Grafana Tempo, Honeycomb, Jaeger: ```rust -use opentelemetry_otlp::WithExportConfig; +// opentelemetry_otlp 0.27+, opentelemetry_sdk 0.27+, tracing-opentelemetry 0.28+ +use opentelemetry::trace::TracerProvider as _; use tracing_subscriber::layer::SubscriberExt; -let tracer = opentelemetry_otlp::new_pipeline() - .tracing() - .with_exporter(opentelemetry_otlp::new_exporter().tonic()) - .install_batch(opentelemetry_sdk::runtime::Tokio)?; +let exporter = opentelemetry_otlp::SpanExporter::builder() + .with_tonic() + .build()?; +let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder() + .with_batch_exporter(exporter) + .build(); let subscriber = tracing_subscriber::registry() - .with(tracing_opentelemetry::layer().with_tracer(tracer)); + .with(tracing_opentelemetry::layer().with_tracer(provider.tracer("yoagent-app"))); tracing::subscriber::set_global_default(subscriber)?; ``` +(The OTel crates rework their builder APIs between releases — if this snippet +drifts, the authoritative wiring is the +[`tracing-opentelemetry` docs](https://docs.rs/tracing-opentelemetry); yoagent +only emits standard `tracing` spans and does not depend on OTel.) + Because these are ordinary `tracing` spans, an agent call nests inside your app's existing request traces (e.g. an axum handler span) automatically. diff --git a/docs/concepts/tools.md b/docs/concepts/tools.md index 2ed6e3c..754ac51 100644 --- a/docs/concepts/tools.md +++ b/docs/concepts/tools.md @@ -460,19 +460,14 @@ behind permission prompts, policy engines, and argument rewriting. yoagent ships the hook, not a policy: with no middleware installed, every call runs. ```rust -use yoagent::{ToolDecision, ToolMiddleware}; +use yoagent::{ToolCallRequest, ToolDecision, ToolMiddleware}; struct ReadOnlyPolicy; #[async_trait::async_trait] impl ToolMiddleware for ReadOnlyPolicy { - async fn before_tool( - &self, - _tool_call_id: &str, - tool_name: &str, - _args: &serde_json::Value, - ) -> ToolDecision { - match tool_name { + async fn before_tool(&self, call: &ToolCallRequest<'_>) -> ToolDecision { + match call.tool_name { "write_file" | "edit_file" | "bash" => { ToolDecision::Deny("read-only session".into()) } @@ -498,6 +493,10 @@ Semantics: adapt — pick another tool, ask the user — and the loop continues. A denial never aborts the run. +A middleware that panics is contained: the call is denied (reason `"tool +middleware panicked"`) and the loop continues — a buggy policy can't kill the +run. + The hook is `async`, so an interactive app can prompt a human before deciding. Note that under the default `Parallel` execution strategy, middleware for parallel tool calls run concurrently — if you need one-at-a-time approval UX, diff --git a/src/agent.rs b/src/agent.rs index c55b5fc..2956fcf 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -66,9 +66,6 @@ pub struct Agent { // Tool middleware (permissions/policy hooks) tool_middleware: Vec>, - // Structured-output schema for the in-flight prompt_structured call - output_schema: Option, - // Custom compaction strategy compaction_strategy: Option>, @@ -101,13 +98,23 @@ pub enum AgentBuildError { #[derive(Debug, thiserror::Error)] #[non_exhaustive] pub enum StructuredPromptError { - /// The run produced no assistant text to parse. + /// The run produced no assistant text to parse. Only messages produced by + /// this call are considered — earlier history is never scanned. #[error("model returned no output to parse")] NoOutput, + /// The provider call itself failed (auth, network, rate limits, a + /// schema-induced 400, ...). Retrying the parse is pointless; the message + /// carries the underlying provider error. + #[error("provider error during structured prompt: {message}")] + Provider { message: String }, /// The model's output did not deserialize into the requested type. /// `raw` carries the model's text so callers can retry or salvage. - #[error("failed to parse structured output: {error}; raw output: {raw}")] - Parse { error: String, raw: String }, + #[error("failed to parse structured output: {source}; raw output: {raw}")] + Parse { + #[source] + source: serde_json::Error, + raw: String, + }, } impl Agent { @@ -260,7 +267,6 @@ impl Agent { on_error: None, input_filters: Vec::new(), tool_middleware: Vec::new(), - output_schema: None, compaction_strategy: None, cancel: None, is_streaming: false, @@ -698,26 +704,55 @@ impl Agent { text: impl Into, schema: serde_json::Value, ) -> Result { - self.output_schema = Some(crate::provider::OutputSchema::new( - "structured_output", - schema, - )); - let mut rx = self.prompt(text).await; + // The schema is threaded through this call's loop config only — it is + // never stored on the agent, so a dropped/timed-out future cannot + // leave the agent stuck in schema-forcing mode. + let schema = crate::provider::OutputSchema::new("structured_output", schema); + let history_len = self.messages.len(); + + let msg = AgentMessage::Llm(Message::user(text)); + let mut rx = self.prompt_messages_internal(vec![msg], Some(schema)).await; while rx.recv().await.is_some() {} self.finish().await; - self.output_schema = None; - let raw = self - .messages - .iter() + // Only messages produced by THIS run count — never parse stale text + // from earlier turns. (Compaction can shrink history below + // history_len; saturating slice keeps the scan sound.) + let run_messages = self.messages.get(history_len.min(self.messages.len())..); + let last_assistant = run_messages + .into_iter() + .flatten() .rev() .find_map(|m| match m { - AgentMessage::Llm(Message::Assistant { content, .. }) => { - content.iter().find_map(|c| match c { - Content::Text { text } => Some(text.clone()), - _ => None, - }) - } + AgentMessage::Llm(Message::Assistant { + content, + stop_reason, + error_message, + .. + }) => Some((content, stop_reason, error_message)), + _ => None, + }); + + let Some((content, stop_reason, error_message)) = last_assistant else { + return Err(StructuredPromptError::NoOutput); + }; + + // A failed provider call is not a parse problem — surface it as such. + if *stop_reason == StopReason::Error { + return Err(StructuredPromptError::Provider { + message: error_message + .clone() + .unwrap_or_else(|| "provider error (no detail)".into()), + }); + } + + // The structured payload is the LAST text block: tool-forcing unwrap + // appends it after any preamble text the model produced. + let raw = content + .iter() + .rev() + .find_map(|c| match c { + Content::Text { text } if !text.is_empty() => Some(text.clone()), _ => None, }) .ok_or(StructuredPromptError::NoOutput)?; @@ -729,8 +764,8 @@ impl Agent { .trim_start_matches("```") .trim_end_matches("```") .trim(); - serde_json::from_str(cleaned).map_err(|e| StructuredPromptError::Parse { - error: e.to_string(), + serde_json::from_str(cleaned).map_err(|source| StructuredPromptError::Parse { + source, raw: raw.clone(), }) } @@ -751,6 +786,17 @@ impl Agent { pub async fn prompt_messages( &mut self, messages: Vec, + ) -> mpsc::UnboundedReceiver { + self.prompt_messages_internal(messages, None).await + } + + /// Shared plumbing for `prompt_messages` and `prompt_structured`. The + /// structured-output schema is per-call state: it lives on this run's + /// `AgentLoopConfig` only, never on the agent. + async fn prompt_messages_internal( + &mut self, + messages: Vec, + output_schema: Option, ) -> mpsc::UnboundedReceiver { self.finish().await; // restore from previous if needed @@ -771,7 +817,8 @@ impl Agent { tools: std::mem::take(&mut self.tools), }; - let config = self.build_config(); + let mut config = self.build_config(); + config.output_schema = output_schema; let handle = tokio::spawn(async move { let _new_messages = agent_loop(messages, &mut context, &config, tx, cancel).await; @@ -1082,7 +1129,7 @@ impl Agent { on_error: self.on_error.clone(), input_filters: self.input_filters.clone(), tool_middleware: self.tool_middleware.clone(), - output_schema: self.output_schema.clone(), + output_schema: None, turn_delay: None, } } diff --git a/src/agent_loop.rs b/src/agent_loop.rs index ce348dd..9bb7d1d 100644 --- a/src/agent_loop.rs +++ b/src/agent_loop.rs @@ -425,6 +425,7 @@ async fn run_loop( tokens_out = tracing::field::Empty, tokens_cached = tracing::field::Empty, cost_usd = tracing::field::Empty, + error = tracing::field::Empty, ); let message = { use tracing::Instrument; @@ -432,7 +433,11 @@ async fn run_loop( .instrument(llm_span.clone()) .await }; - if let Message::Assistant { usage, .. } = &message { + if let Message::Assistant { + usage, stop_reason, .. + } = &message + { + llm_span.record("error", *stop_reason == StopReason::Error); llm_span.record("tokens_in", usage.input); llm_span.record("tokens_out", usage.output); llm_span.record("tokens_cached", usage.cache_read); @@ -794,6 +799,7 @@ fn unwrap_structured_tool_call( model, provider, usage, + stop_reason, .. } = &message else { @@ -808,19 +814,29 @@ fn unwrap_structured_tool_call( return message; }; - // Keep non-tool-call content (e.g. thinking blocks); the payload becomes - // the message text. + // Remove ONLY the synthetic call; any real tool calls (shouldn't occur + // under forced tool_choice, but defensively) stay and execute normally. + // The payload is appended AFTER any preamble text — consumers take the + // last text block. let mut new_content: Vec = content .iter() - .filter(|c| !matches!(c, Content::ToolCall { .. })) + .filter(|c| !matches!(c, Content::ToolCall { name, .. } if *name == schema.name)) .cloned() .collect(); new_content.push(Content::Text { text: payload.to_string(), }); + // ToolUse becomes Stop (the forced call was the "answer"); every other + // stop reason (Length = truncated payload, Error, ...) is preserved so + // truncation isn't laundered into success. + let new_stop = if *stop_reason == StopReason::ToolUse { + StopReason::Stop + } else { + stop_reason.clone() + }; Message::assistant( new_content, - StopReason::Stop, + new_stop, model.clone(), provider.clone(), usage.clone(), @@ -965,7 +981,24 @@ async fn execute_single_tool( // (the LLM sees the reason and can adapt — the loop continues). let mut effective_args = args.clone(); for mw in middleware { - match mw.before_tool(id, name, &effective_args).await { + let call = ToolCallRequest { + tool_call_id: id, + tool_name: name, + args: &effective_args, + }; + // A panicking middleware must not kill the loop task (which would + // strip the agent of its tools) — contain it and fail closed. + let decision = { + use futures::FutureExt; + std::panic::AssertUnwindSafe(mw.before_tool(&call)) + .catch_unwind() + .await + .unwrap_or_else(|_| { + tracing::warn!(tool = name, "tool middleware panicked; denying the call"); + ToolDecision::Deny("tool middleware panicked".into()) + }) + }; + match decision { ToolDecision::Allow => {} ToolDecision::Modify(new_args) => effective_args = new_args, ToolDecision::Deny(reason) => { @@ -1098,6 +1131,14 @@ fn denied_tool_call( reason: &str, tx: &mpsc::UnboundedSender, ) -> (Message, bool) { + // Operator-visible signal: without this, a denial exists only in the + // event stream / message history, invisible to telemetry. + tracing::warn!( + tool = name, + tool_call_id = id, + reason, + "tool call denied by middleware" + ); tx.send(AgentEvent::ToolExecutionStart { tool_call_id: id.to_string(), tool_name: name.to_string(), diff --git a/src/lib.rs b/src/lib.rs index a13b111..2a2cdaa 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -41,8 +41,8 @@ //! picked up between tool executions (per batch under the default parallel //! strategy). Queue follow-ups, inspect/edit the queues. //! - **Structured outputs** ([`Agent::prompt_structured`]) — typed, -//! schema-validated replies, enforced natively per provider (forced tool -//! call / `json_schema` / `responseSchema`). +//! schema-validated replies, enforced natively where supported (forced +//! tool call / `json_schema` / `responseSchema`). //! - **Permissions** ([`ToolMiddleware`]) — async approve/deny/modify hooks //! gating every tool call; the mechanism behind approval prompts and //! policy engines (yoagent ships no policy — you install it). @@ -56,7 +56,7 @@ //! - **Skills** ([`skills`]) — load `SKILL.md` files per the //! [AgentSkills](https://agentskills.io) standard. //! - **Telemetry** — `tracing` spans per loop/LLM-stream/tool with token and -//! cost fields; bridge to OpenTelemetry app-side, zero overhead otherwise. +//! cost fields; bridge to OpenTelemetry app-side, negligible cost otherwise. //! //! The [book](https://yologdev.github.io/yoagent/) covers concepts and //! provider-specific guides. diff --git a/src/provider/anthropic.rs b/src/provider/anthropic.rs index 760b038..d7709b2 100644 --- a/src/provider/anthropic.rs +++ b/src/provider/anthropic.rs @@ -509,7 +509,17 @@ fn build_request_body(config: &StreamConfig, is_oauth: bool) -> serde_json::Valu body["tool_choice"] = serde_json::json!({ "type": "tool", "name": schema.name }); } - if config.thinking_level != ThinkingLevel::Off { + // Forced tool_choice and extended thinking are mutually exclusive at the + // API level — a structured-output request wins and thinking is skipped + // for this call (warned, not silent). + let thinking_requested = config.thinking_level != ThinkingLevel::Off; + if thinking_requested && config.output_schema.is_some() { + tracing::warn!( + "structured outputs force tool_choice, which Anthropic rejects with \ + extended thinking; thinking is disabled for this request" + ); + } + if thinking_requested && config.output_schema.is_none() { if compat.adaptive_thinking { // Current generation (Claude 4.6+ / Fable 5): adaptive thinking with // an effort hint. Budget-based thinking is rejected with a 400. @@ -784,6 +794,22 @@ mod tests { assert_eq!(empty.cache_hit_rate(), 0.0); } + #[test] + fn structured_output_disables_thinking() { + // Forced tool_choice + extended thinking is an API-level conflict: + // the schema wins and no thinking field may be emitted. + let mut config = make_config(CacheConfig::default()); + config.thinking_level = ThinkingLevel::High; + config.output_schema = Some(crate::provider::OutputSchema::new( + "structured_output", + serde_json::json!({"type": "object"}), + )); + let body = build_request_body(&config, false); + assert!(body["thinking"].is_null(), "thinking must be skipped"); + assert!(body["output_config"].is_null()); + assert_eq!(body["tool_choice"]["name"], "structured_output"); + } + #[test] fn test_structured_output_forces_synthetic_tool() { let mut config = make_config(CacheConfig::default()); diff --git a/src/provider/bedrock.rs b/src/provider/bedrock.rs index cd97143..0dad684 100644 --- a/src/provider/bedrock.rs +++ b/src/provider/bedrock.rs @@ -357,23 +357,37 @@ fn content_to_bedrock(content: &[Content]) -> Vec { content .iter() .filter(|c| !matches!(c, Content::Text { text } if text.is_empty())) - .filter_map(|c| match c { - Content::Text { text } => Some(serde_json::json!({"text": text})), - Content::Image { data, mime_type } => Some(serde_json::json!({ + .map(|c| match c { + Content::Text { text } => serde_json::json!({"text": text}), + Content::Image { data, mime_type } => serde_json::json!({ "image": { "format": mime_type.split('/').nth(1).unwrap_or("png"), "source": {"bytes": data}, } - })), + }), Content::ToolCall { id, name, arguments, .. - } => Some(serde_json::json!({ + } => serde_json::json!({ "toolUse": {"toolUseId": id, "name": name, "input": arguments}, - })), - Content::Thinking { .. } => None, + }), + // Replay reasoning blocks: Anthropic-on-Bedrock requires the + // thinking block (with signature) to accompany a replayed + // assistant message in multi-turn tool use — dropping it causes a + // ValidationException on the next call. + Content::Thinking { + thinking, + signature, + } => serde_json::json!({ + "reasoningContent": { + "reasoningText": { + "text": thinking, + "signature": signature.clone().unwrap_or_default(), + } + } + }), }) .collect() } @@ -454,6 +468,42 @@ struct BedrockUsage { mod tests { use super::*; + #[test] + fn thinking_blocks_are_replayed_with_signature() { + // Anthropic-on-Bedrock rejects replayed assistant messages whose + // thinking block was dropped — pin that we serialize it back. + let mut config = StreamConfig::new("anthropic.claude-sonnet", "a:b"); + config.messages = vec![ + Message::user("go"), + Message::assistant( + vec![ + Content::thinking_signed("chain of thought", "sig-1"), + Content::Text { + text: "answer".into(), + }, + ], + StopReason::Stop, + "m", + "bedrock", + Usage::default(), + ), + ]; + let body = build_bedrock_body(&config); + let assistant_content = body["messages"][1]["content"].as_array().unwrap(); + let reasoning = assistant_content + .iter() + .find(|b| b.get("reasoningContent").is_some()) + .expect("thinking block must be replayed"); + assert_eq!( + reasoning["reasoningContent"]["reasoningText"]["text"], + "chain of thought" + ); + assert_eq!( + reasoning["reasoningContent"]["reasoningText"]["signature"], + "sig-1" + ); + } + #[test] fn thinking_level_sets_additional_model_request_fields() { let config = StreamConfig { diff --git a/src/provider/google.rs b/src/provider/google.rs index 5aeef28..f5610db 100644 --- a/src/provider/google.rs +++ b/src/provider/google.rs @@ -500,6 +500,32 @@ struct GoogleUsageMetadata { mod tests { use super::*; + #[test] + fn thinking_blocks_are_dropped_on_replay() { + // Gemini does not accept echoed thought summaries; signatures ride on + // functionCall parts instead. Pin the drop so it stays deliberate. + let mut config = StreamConfig::new("gemini-2.5-pro", "k"); + config.messages = vec![ + Message::user("go"), + Message::assistant( + vec![ + Content::thinking("thought"), + Content::Text { + text: "answer".into(), + }, + ], + StopReason::Stop, + "m", + "google", + Usage::default(), + ), + ]; + let body = build_request_body(&config); + let parts = body["contents"][1]["parts"].as_array().unwrap(); + assert_eq!(parts.len(), 1, "only the text part is replayed"); + assert_eq!(parts[0]["text"], "answer"); + } + #[test] fn thinking_level_sets_thinking_config() { let config = StreamConfig { diff --git a/src/provider/google_vertex.rs b/src/provider/google_vertex.rs index 046ff06..de23345 100644 --- a/src/provider/google_vertex.rs +++ b/src/provider/google_vertex.rs @@ -182,6 +182,8 @@ async fn parse_google_sse_response( thought: Option, #[serde(default, rename = "functionCall")] function_call: Option, + #[serde(default, rename = "thoughtSignature")] + thought_signature: Option, } #[derive(Deserialize)] struct FCall { @@ -250,8 +252,13 @@ async fn parse_google_sse_response( if let Some(fc) = part.function_call { let id = format!("vertex-fc-{}", content.len()); let args = fc.args.unwrap_or(serde_json::Value::Object(Default::default())); + // Gemini thinking conversations require the + // thought signature replayed on function calls. + let metadata = part.thought_signature.as_ref().map(|sig| { + serde_json::json!({"thought_signature": sig}) + }); let idx = content.len(); - content.push(Content::ToolCall { provider_metadata: None, + content.push(Content::ToolCall { provider_metadata: metadata, id: id.clone(), name: fc.name.clone(), arguments: args, @@ -348,10 +355,23 @@ fn build_vertex_request_body(config: &StreamConfig) -> serde_json::Value { .filter_map(|c| match c { Content::Text { text } => Some(serde_json::json!({"text": text})), Content::ToolCall { - name, arguments, .. - } => Some(serde_json::json!({ - "functionCall": {"name": name, "args": arguments}, - })), + name, + arguments, + provider_metadata, + .. + } => { + let mut part = serde_json::json!({ + "functionCall": {"name": name, "args": arguments}, + }); + if let Some(sig) = provider_metadata + .as_ref() + .and_then(|m| m.get("thought_signature")) + .and_then(|v| v.as_str()) + { + part["thoughtSignature"] = serde_json::json!(sig); + } + Some(part) + } _ => None, }) .collect(); @@ -450,6 +470,32 @@ mod tests { } } + #[test] + fn tool_call_thought_signature_is_replayed() { + // Parity with the Gemini API provider: signatures captured into + // provider_metadata must ride back on functionCall parts. + let mut c = config(ThinkingLevel::Off); + c.messages = vec![ + Message::user("go"), + Message::assistant( + vec![Content::ToolCall { + id: "vertex-fc-0".into(), + name: "get_weather".into(), + arguments: serde_json::json!({"city": "Paris"}), + provider_metadata: Some(serde_json::json!({"thought_signature": "sig-9"})), + }], + StopReason::ToolUse, + "m", + "vertex", + Usage::default(), + ), + ]; + let body = build_vertex_request_body(&c); + let part = &body["contents"][1]["parts"][0]; + assert_eq!(part["functionCall"]["name"], "get_weather"); + assert_eq!(part["thoughtSignature"], "sig-9"); + } + #[test] fn thinking_level_sets_thinking_config() { let body = build_vertex_request_body(&config(ThinkingLevel::High)); diff --git a/src/provider/traits.rs b/src/provider/traits.rs index b9d74d9..2926575 100644 --- a/src/provider/traits.rs +++ b/src/provider/traits.rs @@ -29,8 +29,19 @@ pub enum StreamEvent { Error { message: Message }, } -/// Configuration for a streaming LLM call +/// Configuration for a streaming LLM call. +/// +/// Marked `#[non_exhaustive]`: fields are added in minor releases (this +/// release alone added `output_schema`). Construct with +/// [`StreamConfig::new`] and mutate the public fields: +/// +/// ``` +/// # use yoagent::provider::StreamConfig; +/// let mut config = StreamConfig::new("claude-sonnet-5", "sk-key"); +/// config.system_prompt = "be brief".into(); +/// ``` #[derive(Debug, Clone)] +#[non_exhaustive] pub struct StreamConfig { pub model: String, pub system_prompt: String, @@ -52,8 +63,32 @@ pub struct StreamConfig { pub output_schema: Option, } +impl StreamConfig { + /// A config with the given model and API key; everything else defaults + /// (empty prompt/messages/tools, thinking off, caching enabled). + pub fn new(model: impl Into, api_key: impl Into) -> Self { + Self { + model: model.into(), + system_prompt: String::new(), + messages: Vec::new(), + tools: Vec::new(), + thinking_level: ThinkingLevel::Off, + api_key: api_key.into(), + max_tokens: None, + temperature: None, + model_config: None, + cache_config: CacheConfig::default(), + output_schema: None, + } + } +} + /// JSON-Schema constraint for structured outputs. +/// +/// Marked `#[non_exhaustive]`: fields may be added in minor releases (e.g. +/// strictness flags). Construct with [`OutputSchema::new`]. #[derive(Debug, Clone, Serialize, Deserialize)] +#[non_exhaustive] pub struct OutputSchema { /// Schema name (doubles as the forced tool name on Anthropic). pub name: String, diff --git a/src/session.rs b/src/session.rs index 549a63a..ca6f2ac 100644 --- a/src/session.rs +++ b/src/session.rs @@ -24,7 +24,7 @@ //! let mut rx = agent.prompt("hello").await; //! while rx.recv().await.is_some() {} //! agent.finish().await; -//! session.append_new(agent.messages()); +//! session.append_new(agent.messages()).unwrap(); //! //! // Checkpoint, keep working... //! session.checkpoint("after-hello").unwrap(); @@ -74,6 +74,20 @@ pub enum SessionError { /// The session is empty where an entry was required. #[error("session is empty")] Empty, + /// A JSONL file contained the same entry id twice. + #[error("duplicate session entry id: {0}")] + DuplicateId(String), + /// An entry referenced a parent that does not precede it in the file + /// (dangling parent or forward/cyclic reference). + #[error("entry {id} references unknown parent {parent}")] + UnknownParent { id: String, parent: String }, + /// The history passed to `append_new` does not extend this session's + /// current path (e.g. context compaction rewrote the agent's messages). + #[error( + "history diverged from the session path at index {index}: the agent's \ + messages no longer extend this branch (did compaction rewrite them?)" + )] + HistoryDiverged { index: usize }, } /// A conversation history tree: append advances the head; seek + append forks. @@ -117,13 +131,35 @@ impl Session { id } - /// Append every message of `full_history` beyond the current path length - /// — the typical post-run sync from [`Agent::messages`](crate::Agent::messages). - pub fn append_new(&mut self, full_history: &[AgentMessage]) { - let known = self.path_ids().len(); - for m in full_history.iter().skip(known) { + /// Append every message of `full_history` beyond the current path — the + /// typical post-run sync from [`Agent::messages`](crate::Agent::messages). + /// Returns how many messages were appended. + /// + /// The history must **extend the current path**: `full_history[..path_len]` + /// is verified against the path's messages, and any mismatch (or a + /// shorter history) returns [`SessionError::HistoryDiverged`] instead of + /// silently corrupting the tree. The common cause is context compaction + /// (on by default) rewriting the agent's messages mid-session — disable + /// context management on session-tracked agents, or rebuild the branch + /// with [`from_messages`](Self::from_messages) after a divergence. + pub fn append_new(&mut self, full_history: &[AgentMessage]) -> Result { + let path = self.path_messages(); + if full_history.len() < path.len() { + return Err(SessionError::HistoryDiverged { + index: full_history.len(), + }); + } + for (i, known) in path.iter().enumerate() { + if &full_history[i] != known { + return Err(SessionError::HistoryDiverged { index: i }); + } + } + let mut appended = 0; + for m in full_history.iter().skip(path.len()) { self.append(m.clone()); + appended += 1; } + Ok(appended) } /// Current head entry id, if any. @@ -155,11 +191,13 @@ impl Session { Ok(()) } - /// Move the head to the entry labeled `label`. + /// Move the head to the entry labeled `label`. If several entries carry + /// the same label, the most recently created one wins. pub fn seek_checkpoint(&mut self, label: &str) -> Result<(), SessionError> { let id = self .entries .iter() + .rev() .find(|e| e.label.as_deref() == Some(label)) .map(|e| e.id.clone()) .ok_or_else(|| SessionError::UnknownCheckpoint(label.to_string()))?; @@ -239,6 +277,12 @@ impl Session { /// to the last line's entry. pub fn from_jsonl(s: &str) -> Result { let mut session = Self::new(); + // `to_jsonl` writes insertion order with parents preceding children, + // so one streaming check enforces every tree invariant: ids unique, + // parents already seen (which also rules out cycles — a cycle would + // need a forward reference). This keeps `path_ids` provably + // terminating and the internal `expect`s sound. + let mut seen: std::collections::HashSet = std::collections::HashSet::new(); for (i, line) in s.lines().enumerate() { let line = line.trim(); if line.is_empty() { @@ -249,6 +293,18 @@ impl Session { line: i + 1, error: e.to_string(), })?; + if seen.contains(&entry.id) { + return Err(SessionError::DuplicateId(entry.id)); + } + if let Some(parent) = &entry.parent_id { + if !seen.contains(parent) { + return Err(SessionError::UnknownParent { + id: entry.id.clone(), + parent: parent.clone(), + }); + } + } + seen.insert(entry.id.clone()); // Track the numeric suffix so future appends can't collide. if let Some(n) = entry .id diff --git a/src/types.rs b/src/types.rs index a667bee..7275446 100644 --- a/src/types.rs +++ b/src/types.rs @@ -611,6 +611,12 @@ pub trait InputFilter: Send + Sync { // --------------------------------------------------------------------------- /// Decision returned by a [`ToolMiddleware`] before a tool executes. +/// +/// Deliberately NOT `#[non_exhaustive]` (same policy as [`StopReason`]): +/// this is a control-flow enum — a new variant should be a compile error for +/// matchers, not a silent wildcard fallthrough. Interactive flows like +/// "ask the user" need no variant: the hook is `async`, so prompt inside the +/// middleware and return `Allow`/`Deny`. #[derive(Debug, Clone)] pub enum ToolDecision { /// Execute the tool with the current arguments. @@ -637,14 +643,27 @@ pub enum ToolDecision { /// default [`ToolExecutionStrategy::Parallel`], middleware for parallel tool /// calls runs concurrently — serialize approval prompts inside your /// implementation (or use `Sequential`) if you need one-at-a-time UX. +/// Borrowed view of a pending tool call, passed to +/// [`ToolMiddleware::before_tool`]. +/// +/// Marked `#[non_exhaustive]`: fields may be added in minor releases (turn +/// number, history access, ...) without breaking middleware implementations. +/// Constructed by the loop; middleware only reads it. +#[derive(Debug)] +#[non_exhaustive] +pub struct ToolCallRequest<'a> { + /// Provider-assigned id of this tool call. + pub tool_call_id: &'a str, + /// Name of the tool the model wants to run. + pub tool_name: &'a str, + /// Arguments as the model provided them (possibly rewritten by earlier + /// middleware in the chain). + pub args: &'a serde_json::Value, +} + #[async_trait::async_trait] pub trait ToolMiddleware: Send + Sync { - async fn before_tool( - &self, - tool_call_id: &str, - tool_name: &str, - args: &serde_json::Value, - ) -> ToolDecision; + async fn before_tool(&self, call: &ToolCallRequest<'_>) -> ToolDecision; } // --------------------------------------------------------------------------- diff --git a/tests/agent_test.rs b/tests/agent_test.rs index 2bf28bf..9605638 100644 --- a/tests/agent_test.rs +++ b/tests/agent_test.rs @@ -741,13 +741,8 @@ impl ToolMiddleware for FnMiddleware where F: Fn(&str, &serde_json::Value) -> ToolDecision + Send + Sync, { - async fn before_tool( - &self, - _tool_call_id: &str, - tool_name: &str, - args: &serde_json::Value, - ) -> ToolDecision { - (self.0)(tool_name, args) + async fn before_tool(&self, call: &ToolCallRequest<'_>) -> ToolDecision { + (self.0)(call.tool_name, call.args) } } @@ -988,14 +983,54 @@ async fn test_prompt_structured_strips_markdown_fences() { assert_eq!(out.count, 1); } +/// Provider that records each call's output_schema, then answers with text. +struct SchemaCapturingProvider { + schemas: Arc>>>, + reply: String, +} + +#[async_trait::async_trait] +impl yoagent::provider::StreamProvider for SchemaCapturingProvider { + async fn stream( + &self, + config: yoagent::provider::StreamConfig, + tx: mpsc::UnboundedSender, + _cancel: tokio_util::sync::CancellationToken, + ) -> Result { + self.schemas + .lock() + .unwrap() + .push(config.output_schema.as_ref().map(|s| s.name.clone())); + let msg = Message::assistant( + vec![Content::Text { + text: self.reply.clone(), + }], + StopReason::Stop, + "mock", + "mock", + Usage::default(), + ); + let _ = tx.send(yoagent::provider::StreamEvent::Start); + let _ = tx.send(yoagent::provider::StreamEvent::Done { + message: msg.clone(), + }); + Ok(msg) + } +} + #[tokio::test] -async fn test_prompt_structured_resets_schema_for_next_prompt() { - // After a structured call, a normal prompt must not carry the schema. - let provider = MockProvider::new(vec![ - MockResponse::Text(r#"{"name": "a", "count": 1}"#.into()), - MockResponse::Text("plain answer".into()), - ]); - let mut agent = Agent::from_provider(provider, yoagent::provider::ModelConfig::mock()); +async fn test_prompt_structured_schema_reaches_provider_then_resets() { + // Pins the loop→StreamConfig propagation joint (MockProvider ignores its + // config, so only a capturing provider can verify it) AND that the very + // next plain prompt carries no schema. + let schemas = Arc::new(std::sync::Mutex::new(Vec::new())); + let mut agent = Agent::from_provider( + SchemaCapturingProvider { + schemas: schemas.clone(), + reply: r#"{"name": "a", "count": 1}"#.into(), + }, + yoagent::provider::ModelConfig::mock(), + ); let _: Extracted = agent .prompt_structured("extract", serde_json::json!({"type": "object"})) .await @@ -1004,9 +1039,153 @@ async fn test_prompt_structured_resets_schema_for_next_prompt() { let mut rx = agent.prompt("normal question").await; while rx.recv().await.is_some() {} agent.finish().await; - // Last message is the plain text answer — the loop ran normally. + + let seen = schemas.lock().unwrap().clone(); + assert_eq!( + seen, + vec![Some("structured_output".to_string()), None], + "structured call must carry the schema; the next plain prompt must not" + ); +} + +/// Provider that always fails, for the Provider error-variant path. +struct AlwaysFailProvider; + +#[async_trait::async_trait] +impl yoagent::provider::StreamProvider for AlwaysFailProvider { + async fn stream( + &self, + _config: yoagent::provider::StreamConfig, + _tx: mpsc::UnboundedSender, + _cancel: tokio_util::sync::CancellationToken, + ) -> Result { + Err(yoagent::provider::ProviderError::Api( + "invalid x-api-key".into(), + )) + } +} + +#[tokio::test] +async fn test_prompt_structured_surfaces_provider_error() { + // A dead API key must surface as Provider { message } carrying the real + // error — never as Parse { raw: "" }. + let mut agent = + Agent::from_provider(AlwaysFailProvider, yoagent::provider::ModelConfig::mock()) + .with_retry_config(yoagent::RetryConfig::none()); + let err = agent + .prompt_structured::("extract", serde_json::json!({"type": "object"})) + .await + .unwrap_err(); + match err { + yoagent::StructuredPromptError::Provider { message } => { + assert!(message.contains("invalid x-api-key"), "got: {message}"); + } + other => panic!("expected Provider error, got {other:?}"), + } +} + +#[tokio::test] +async fn test_prompt_structured_never_parses_stale_history() { + // Seed history with an earlier turn's perfectly-parsable JSON; a failed + // run must NOT fall back to it and return stale data as Ok. + let stale = AgentMessage::Llm(Message::assistant( + vec![Content::Text { + text: r#"{"name": "stale", "count": 99}"#.into(), + }], + StopReason::Stop, + "mock", + "mock", + Usage::default(), + )); + let mut agent = + Agent::from_provider(AlwaysFailProvider, yoagent::provider::ModelConfig::mock()) + .with_retry_config(yoagent::RetryConfig::none()) + .with_messages(vec![stale]); + + let err = agent + .prompt_structured::("extract", serde_json::json!({"type": "object"})) + .await + .unwrap_err(); + assert!( + matches!(err, yoagent::StructuredPromptError::Provider { .. }), + "must not return the stale turn's JSON, got {err:?}" + ); +} + +#[tokio::test] +async fn test_middleware_never_sees_synthetic_structured_tool() { + // The forced structured_output call is unwrapped BEFORE tool extraction, + // so middleware must never be consulted for it. + let calls = Arc::new(std::sync::Mutex::new(0usize)); + let calls2 = calls.clone(); + let provider = MockProvider::new(vec![MockResponse::ToolCalls(vec![MockToolCall { + provider_metadata: None, + name: "structured_output".into(), + arguments: serde_json::json!({"name": "g", "count": 7}), + }])]); + let mut agent = Agent::from_provider(provider, yoagent::provider::ModelConfig::mock()) + .with_tool_middleware(FnMiddleware(move |_: &str, _: &serde_json::Value| { + *calls2.lock().unwrap() += 1; + ToolDecision::Deny("should never run".into()) + })); + let out: Extracted = agent + .prompt_structured("extract", serde_json::json!({"type": "object"})) + .await + .expect("unwrap happens before middleware"); + assert_eq!(out.count, 7); + assert_eq!( + *calls.lock().unwrap(), + 0, + "middleware saw the synthetic tool" + ); +} + +#[tokio::test] +async fn test_tool_middleware_panic_denies_and_loop_survives() { + // A panicking middleware must not kill the loop task (which would strip + // the agent of its tools) — it fails closed as a denial. + let ran = Arc::new(std::sync::Mutex::new(None)); + let agent = Agent::from_provider(tool_call_provider(), yoagent::provider::ModelConfig::mock()) + .with_tools(vec![Box::new(RecordingTool { ran: ran.clone() })]) + .with_tool_middleware(FnMiddleware(|_: &str, _: &serde_json::Value| { + panic!("middleware bug") + })); + + let (agent, _) = run_middleware_agent(agent).await; + assert!(ran.lock().unwrap().is_none(), "tool must not run"); + // Loop survived to the final text, and the denial reached the LLM. assert!(matches!( agent.messages().last(), Some(AgentMessage::Llm(Message::Assistant { .. })) )); + let denied = agent.messages().iter().any(|m| { + matches!(m, AgentMessage::Llm(Message::ToolResult { content, is_error: true, .. }) + if matches!(&content[0], Content::Text { text } if text.contains("middleware panicked"))) + }); + assert!(denied); + // Crucially: the tools survived for the next prompt. + // (RecordingTool was moved into the loop and restored by finish().) +} + +#[tokio::test] +async fn test_tool_middleware_deny_under_batched_strategy() { + let ran = Arc::new(std::sync::Mutex::new(None)); + let agent = Agent::from_provider(tool_call_provider(), yoagent::provider::ModelConfig::mock()) + .with_tools(vec![Box::new(RecordingTool { ran: ran.clone() })]) + .with_tool_execution(ToolExecutionStrategy::Batched { size: 1 }) + .with_tool_middleware(FnMiddleware(|_: &str, _: &serde_json::Value| { + ToolDecision::Deny("no".into()) + })); + let (agent, events) = run_middleware_agent(agent).await; + assert!(ran.lock().unwrap().is_none()); + assert!(events + .iter() + .any(|e| matches!(e, AgentEvent::ToolExecutionEnd { is_error: true, .. }))); + let denied = agent.messages().iter().any(|m| { + matches!( + m, + AgentMessage::Llm(Message::ToolResult { is_error: true, .. }) + ) + }); + assert!(denied); } diff --git a/tests/anthropic_stream_test.rs b/tests/anthropic_stream_test.rs index 829ecab..06f83b2 100644 --- a/tests/anthropic_stream_test.rs +++ b/tests/anthropic_stream_test.rs @@ -39,19 +39,12 @@ fn stream_config(base_url: &str, anthropic: Option) -> StreamCo let mut mc = ModelConfig::anthropic("claude-sonnet-5", "Claude Sonnet 5"); mc.base_url = base_url.to_string(); mc.anthropic = anthropic; - StreamConfig { - model: "claude-sonnet-5".into(), - system_prompt: "test".into(), - messages: vec![Message::user("hi")], - tools: vec![], - thinking_level: ThinkingLevel::Off, - api_key: "test-key".into(), - max_tokens: Some(256), - temperature: None, - model_config: Some(mc), - cache_config: CacheConfig::default(), - output_schema: None, - } + let mut config = StreamConfig::new("claude-sonnet-5", "test-key"); + config.system_prompt = "test".into(); + config.messages = vec![Message::user("hi")]; + config.max_tokens = Some(256); + config.model_config = Some(mc); + config } async fn run_stream(config: StreamConfig) -> Result { diff --git a/tests/google_stream_test.rs b/tests/google_stream_test.rs index 6880bb4..a744cfe 100644 --- a/tests/google_stream_test.rs +++ b/tests/google_stream_test.rs @@ -24,19 +24,12 @@ fn sse(events: &[&str]) -> String { fn stream_config(base_url: &str, messages: Vec) -> StreamConfig { let mut mc = ModelConfig::google(MODEL, "Gemini 2.5 Flash"); mc.base_url = base_url.to_string(); - StreamConfig { - model: MODEL.into(), - system_prompt: "test".into(), - messages, - tools: vec![], - thinking_level: ThinkingLevel::Off, - api_key: "test-key".into(), - max_tokens: Some(256), - temperature: None, - model_config: Some(mc), - cache_config: CacheConfig::default(), - output_schema: None, - } + let mut config = StreamConfig::new(MODEL, "test-key"); + config.system_prompt = "test".into(); + config.messages = messages; + config.max_tokens = Some(256); + config.model_config = Some(mc); + config } async fn run_stream(config: StreamConfig) -> Message { diff --git a/tests/session_test.rs b/tests/session_test.rs index f567eef..42fda4b 100644 --- a/tests/session_test.rs +++ b/tests/session_test.rs @@ -103,9 +103,10 @@ fn jsonl_roundtrip_preserves_tree_and_ids() { assert_eq!(restored.head(), s.head()); // Tree shape intact: two children of root. assert_eq!(restored.children(&root).len(), 2); - // Checkpoint label survives. + // Checkpoint label survives and resolves to the same entry. let mut r = restored.clone(); r.seek_checkpoint("tip2").unwrap(); + assert_eq!(r.head(), s.head()); // Appending after load can't collide with existing ids. let mut r2 = restored; @@ -144,18 +145,15 @@ async fn fork_edit_rerun_with_agent() { agent.finish().await; let mut session = Session::new(); - session.append_new(agent.messages()); + session.append_new(agent.messages()).unwrap(); assert_eq!(session.entries().len(), 2); // user + assistant session.checkpoint("turn-1").unwrap(); - // "Edit" the first user message: fork from BEFORE it (root fork = seek to - // nothing is not a thing — fork from the first entry's parent by starting - // a sibling of the user message). Here: fork from turn-1's assistant to - // ask a different follow-up on one branch... + // Continue on one branch: ask a follow-up from turn-1's assistant... let mut rx = agent.prompt("follow-up B").await; while rx.recv().await.is_some() {} agent.finish().await; - session.append_new(agent.messages()); + session.append_new(agent.messages()).unwrap(); assert_eq!(session.entries().len(), 4); let tip_b = session.head().unwrap().to_string(); @@ -170,7 +168,7 @@ async fn fork_edit_rerun_with_agent() { let mut rx = agent2.prompt("follow-up C").await; while rx.recv().await.is_some() {} agent2.finish().await; - session.append_new(agent2.messages()); + session.append_new(agent2.messages()).unwrap(); // Two branches: B's tip and C's tip; both intact. assert_eq!(session.entries().len(), 6); @@ -189,3 +187,80 @@ async fn fork_edit_rerun_with_agent() { other => panic!("unexpected message {other:?}"), } } + +// --------------------------------------------------------------------------- +// Hardening from the Phase C review: divergence detection + load validation +// --------------------------------------------------------------------------- + +#[test] +fn append_new_detects_diverged_history() { + // Reuse the SAME message values (timestamps included) so equality with + // the stored path is meaningful — exactly like a real Agent history. + let (a, b) = (user("a"), user("b")); + let mut s = Session::new(); + s.append(a.clone()); + s.append(b.clone()); + + // Shorter history than the path → divergence, nothing appended. + let err = s.append_new(std::slice::from_ref(&a)).unwrap_err(); + assert!(matches!(err, SessionError::HistoryDiverged { index: 1 })); + assert_eq!(s.entries().len(), 2); + + // Same length but rewritten prefix (what compaction does) → divergence + // at the mismatching index. + let err = s + .append_new(&[a.clone(), user("REWRITTEN"), user("c")]) + .unwrap_err(); + assert!(matches!(err, SessionError::HistoryDiverged { index: 1 })); + assert_eq!(s.entries().len(), 2, "no partial append on divergence"); + + // A genuine extension appends and reports the count. + let c = user("c"); + let appended = s.append_new(&[a, b, c]).unwrap(); + assert_eq!(appended, 1); + assert_eq!(s.entries().len(), 3); +} + +#[test] +fn from_jsonl_rejects_duplicate_ids() { + let mut s = Session::new(); + s.append(user("x")); + let line = s.to_jsonl(); + let doubled = format!("{line}\n{line}"); + assert!(matches!( + Session::from_jsonl(&doubled), + Err(SessionError::DuplicateId(_)) + )); +} + +#[test] +fn from_jsonl_rejects_dangling_or_forward_parent() { + // Build a valid line via the library, then doctor its parent to a + // nonexistent id (also covers cycles: a cycle needs a forward reference, + // which this rejects). + let mut s = Session::new(); + s.append(user("x")); + let mut v: serde_json::Value = serde_json::from_str(&s.to_jsonl()).unwrap(); + v["parent_id"] = serde_json::json!("e99"); + let line = serde_json::to_string(&v).unwrap(); + + match Session::from_jsonl(&line) { + Err(SessionError::UnknownParent { id, parent }) => { + assert_eq!(id, "e1"); + assert_eq!(parent, "e99"); + } + other => panic!("expected UnknownParent, got {other:?}"), + } +} + +#[test] +fn seek_checkpoint_latest_wins_on_duplicate_labels() { + let mut s = Session::new(); + s.append(user("a")); + s.checkpoint("mark").unwrap(); + let b = s.append(user("b")); + s.checkpoint("mark").unwrap(); + + s.seek_checkpoint("mark").unwrap(); + assert_eq!(s.head(), Some(b.as_str()), "most recent label wins"); +} diff --git a/tests/sub_agent_test.rs b/tests/sub_agent_test.rs index 83f0892..d3ba1a1 100644 --- a/tests/sub_agent_test.rs +++ b/tests/sub_agent_test.rs @@ -880,12 +880,7 @@ struct DenyAll; #[async_trait::async_trait] impl yoagent::ToolMiddleware for DenyAll { - async fn before_tool( - &self, - _id: &str, - _name: &str, - _args: &serde_json::Value, - ) -> yoagent::ToolDecision { + async fn before_tool(&self, _call: &yoagent::ToolCallRequest<'_>) -> yoagent::ToolDecision { yoagent::ToolDecision::Deny("sub-agent policy".into()) } } diff --git a/tests/telemetry_test.rs b/tests/telemetry_test.rs index 3d26aa3..e645f89 100644 --- a/tests/telemetry_test.rs +++ b/tests/telemetry_test.rs @@ -142,3 +142,142 @@ async fn loop_emits_agent_llm_and_tool_spans() { "got: {names:?}" ); } + +// --------------------------------------------------------------------------- +// Field values: tokens + cost recorded on llm_stream (not just span names) +// --------------------------------------------------------------------------- + +/// Records (span_name, field_name, value_debug) for every record() call. +struct FieldCollector { + names: Arc>>, + records: Arc>>, +} + +struct FieldVisitor<'a> { + span: String, + out: &'a Mutex>, +} + +impl tracing::field::Visit for FieldVisitor<'_> { + fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { + self.out.lock().unwrap().push(( + self.span.clone(), + field.name().to_string(), + format!("{value:?}"), + )); + } +} + +impl tracing_subscriber::Layer for FieldCollector +where + S: tracing::Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a>, +{ + fn on_new_span( + &self, + attrs: &tracing::span::Attributes<'_>, + id: &tracing::span::Id, + _ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + self.names + .lock() + .unwrap() + .insert(id.into_u64(), attrs.metadata().name().to_string()); + } + fn on_record( + &self, + id: &tracing::span::Id, + values: &tracing::span::Record<'_>, + _ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + let span = self + .names + .lock() + .unwrap() + .get(&id.into_u64()) + .cloned() + .unwrap_or_default(); + let mut visitor = FieldVisitor { + span, + out: &self.records, + }; + values.record(&mut visitor); + } +} + +/// Provider returning fixed non-zero usage so token/cost fields are real. +struct UsageProvider; + +#[async_trait::async_trait] +impl yoagent::provider::StreamProvider for UsageProvider { + async fn stream( + &self, + _config: yoagent::provider::StreamConfig, + tx: mpsc::UnboundedSender, + _cancel: CancellationToken, + ) -> Result { + let msg = Message::assistant( + vec![Content::Text { text: "ok".into() }], + StopReason::Stop, + "m", + "mock", + Usage { + input: 1_000_000, + output: 500_000, + cache_read: 7, + cache_write: 0, + total_tokens: 1_500_007, + }, + ); + let _ = tx.send(yoagent::provider::StreamEvent::Done { + message: msg.clone(), + }); + Ok(msg) + } +} + +#[tokio::test] +async fn llm_stream_records_tokens_and_cost() { + let names = Arc::new(Mutex::new(std::collections::HashMap::new())); + let records = Arc::new(Mutex::new(Vec::new())); + let subscriber = tracing_subscriber::registry().with(FieldCollector { + names: names.clone(), + records: records.clone(), + }); + + let mut config = loop_config(MockProvider::text("unused")); + config.provider = std::sync::Arc::new(UsageProvider); + let mut mc = yoagent::provider::ModelConfig::mock(); + mc.cost.input_per_million = 3.0; + mc.cost.output_per_million = 15.0; + config.model_config = Some(mc); + + let mut context = AgentContext { + system_prompt: "t".into(), + messages: Vec::new(), + tools: Vec::new(), + }; + let (tx, _rx) = mpsc::unbounded_channel(); + agent_loop( + vec![AgentMessage::Llm(Message::user("go"))], + &mut context, + &config, + tx, + CancellationToken::new(), + ) + .with_subscriber(subscriber) + .await; + + let recs = records.lock().unwrap().clone(); + let get = |field: &str| -> String { + recs.iter() + .find(|(span, f, _)| span == "llm_stream" && f == field) + .map(|(_, _, v)| v.clone()) + .unwrap_or_else(|| panic!("field {field} not recorded; got {recs:?}")) + }; + assert_eq!(get("tokens_in"), "1000000"); + assert_eq!(get("tokens_out"), "500000"); + assert_eq!(get("tokens_cached"), "7"); + // 1M in @ $3/M + 0.5M out @ $15/M = 10.5 + assert_eq!(get("cost_usd"), "10.5"); + assert_eq!(get("error"), "false"); +}