From dbb5c99498647369e58c1610c291fdf90cbb7eea Mon Sep 17 00:00:00 2001 From: Yuanhao Li Date: Sun, 5 Jul 2026 19:25:11 +0200 Subject: [PATCH 1/3] feat: adoption funnel and trust fixes (roadmap phase A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Funnel: - crate-level lib.rs docs with a quick-start example (the docs.rs landing page was empty); Cargo.toml documentation now points at docs.rs - README: stale install version fixed, honest per-provider capability caveats (thinking 3/7, client-side caching Anthropic-only), registry claim corrected - API keys resolve from provider-conventional env vars when with_api_key is not called (ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY, ..., AWS composite for bedrock); explicit keys always win Trust (advertised features now actually wired): - CostConfig::cost_usd + Agent::session_cost_usd — pricing data is consumed instead of decorative - ContextTracker drives compaction sizing in the loop: real provider usage calibrates the estimation-based budget (scaled when the char/4 estimate runs low, e.g. images), tracker re-baselines after compaction - ModelConfig.reasoning now warns when thinking is requested on a model not marked reasoning-capable - CLAUDE.md claims corrected (registry dispatch, openapi factories) Papercuts: - with_temperature on Agent and SubAgentTool (SubAgentTool previously hardcoded None) - Retry-After header parsed into RateLimited.retry_after_ms so server-directed backoff actually engages - thinking_level now logs a warning on the four providers that ignore it (Azure, Gemini, Vertex, Bedrock) instead of silently no-oping - panic contracts documented on prompt*/continue_loop* (Result variants planned for 0.10) CI/deps: - macOS added to the test matrix; Windows compile-check job; MSRV (1.85) check job - serde_yaml (archived upstream) replaced with serde_yaml_ng - library tokio features trimmed from "full" to the seven actually used Co-Authored-By: Claude Fable 5 --- .github/workflows/ci.yml | 29 ++++++++++- CLAUDE.md | 6 +-- Cargo.toml | 8 ++-- README.md | 7 +-- src/agent.rs | 90 ++++++++++++++++++++++++++++++++++- src/agent_loop.rs | 35 ++++++++++++-- src/lib.rs | 51 ++++++++++++++++++++ src/openapi/adapter.rs | 2 +- src/openapi/types.rs | 2 +- src/provider/azure_openai.rs | 5 ++ src/provider/bedrock.rs | 5 ++ src/provider/google.rs | 5 ++ src/provider/google_vertex.rs | 5 ++ src/provider/mod.rs | 2 +- src/provider/model.rs | 39 ++++++++++++++- src/provider/registry.rs | 77 ++++++++++++++++++++++++++++++ src/provider/traits.rs | 31 ++++++++++-- src/sub_agent.rs | 22 ++++++++- 18 files changed, 395 insertions(+), 26 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c61d89b..c9c0155 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -12,8 +12,12 @@ env: jobs: check: - name: Check & Test - runs-on: ubuntu-latest + name: Check & Test (${{ matrix.os }}) + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, macos-latest] steps: - uses: actions/checkout@v4 @@ -34,3 +38,24 @@ jobs: - name: Build run: cargo build --all-features + + # Built-in tools shell out to unix commands, so Windows only checks compilation. + windows-check: + name: Check (windows) + runs-on: windows-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + - uses: Swatinem/rust-cache@v2 + - name: Check + run: cargo check --all-targets --all-features + + msrv: + name: MSRV (1.85) + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@1.85 + - uses: Swatinem/rust-cache@v2 + - name: Check + run: cargo check --all-features diff --git a/CLAUDE.md b/CLAUDE.md index 73da5b5..9f168bc 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,7 +37,7 @@ The loop: stream assistant response → extract tool calls → execute tools (pa ### Provider System -7 provider implementations behind `StreamProvider`, dispatched by `ApiProtocol` enum via `ProviderRegistry`: +7 provider implementations behind `StreamProvider`. `Agent` holds a concrete provider directly; `ProviderRegistry` maps `ApiProtocol` → provider for registry-based dispatch (e.g. custom multi-protocol routers): | Protocol | File | Covers | |----------|------|--------| @@ -61,7 +61,7 @@ The loop: stream assistant response → extract tool calls → execute tools (pa ### Context Management (`context.rs`) -- **`ContextTracker`** — hybrid real-usage + estimation for token tracking +- **`ContextTracker`** — hybrid real-usage + estimation; the loop uses it to calibrate the compaction budget against real provider usage - **`compact_messages()`** — tiered compaction: Level 1 (truncate tool outputs) → Level 2 (summarize old turns) → Level 3 (drop middle turns) - **`ExecutionLimits`/`ExecutionTracker`** — max turns (50), max tokens (1M), max duration (10 min) @@ -74,7 +74,7 @@ The loop: stream assistant response → extract tool calls → execute tools (pa ### OpenAPI Integration (`openapi/`, feature-gated) -Behind the `openapi` Cargo feature. `OpenApiToolAdapter` parses an OpenAPI 3.0 spec and creates one `AgentTool` per operation. Factory methods: `from_str`, `from_file`, `from_url`, `from_spec`. `OperationFilter` controls which operations become tools. Added to `Agent` via `with_openapi_file()` / `with_openapi_url()` / `with_openapi_spec()`. +Behind the `openapi` Cargo feature. `OpenApiToolAdapter` parses an OpenAPI 3.0 spec and creates one `AgentTool` per operation. Factory methods: `from_str`, `from_file`, `from_url`, `from_spec`. `OperationFilter` controls which operations become tools. Added to `Agent` via `with_openapi_spec()`. ### MCP Integration (`mcp/`) diff --git a/Cargo.toml b/Cargo.toml index 97590c3..86e4829 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,7 +6,7 @@ description = "Simple, effective agent loop with tool execution and event stream license = "MIT" repository = "https://github.com/yologdev/yoagent" homepage = "https://yologdev.github.io/yoagent/" -documentation = "https://yologdev.github.io/yoagent/" +documentation = "https://docs.rs/yoagent" readme = "README.md" keywords = ["agent", "llm", "ai", "tools", "streaming"] categories = ["api-bindings", "asynchronous"] @@ -15,7 +15,7 @@ rust-version = "1.85" exclude = ["docs/images/*", "docs/theme/*", ".github/*", "scripts/*"] [dependencies] -tokio = { version = "1", features = ["full"] } +tokio = { version = "1", features = ["rt", "sync", "fs", "macros", "process", "time", "io-util"] } tokio-util = "0.7" tokio-stream = "0.1" serde = { version = "1", features = ["derive"] } @@ -30,10 +30,10 @@ futures = "0.3" rand = "0.10.0" base64 = "0.22" openapiv3 = { version = "2", optional = true } -serde_yaml = { version = "0.9", optional = true } +serde_yaml_ng = { version = "0.10", optional = true } [features] -openapi = ["dep:openapiv3", "dep:serde_yaml", "reqwest/query"] +openapi = ["dep:openapiv3", "dep:serde_yaml_ng", "reqwest/query"] [dev-dependencies] tokio = { version = "1", features = ["full", "test-util"] } diff --git a/README.md b/README.md index 395646c..b8ec5f6 100644 --- a/README.md +++ b/README.md @@ -46,9 +46,10 @@ Everything is observable via events. Supports 7 API protocols covering 20+ LLM p - [AgentSkills](https://agentskills.io)-compatible skills — load skill directories, inject into system prompt, agent activates on demand **Multi-Provider** -- 7 API protocols, 20+ providers through a modular registry +- 7 API protocols, 20+ providers behind one `StreamProvider` trait - One OpenAI-compatible implementation covers OpenAI, xAI, Groq, Cerebras, OpenRouter, Mistral, and more -- Per-provider quirk flags (`OpenAiCompat`) handle auth, reasoning format, and tool handling differences +- Per-provider quirk flags (`OpenAiCompat`, `AnthropicCompat`) handle auth, reasoning format, and tool handling differences +- Capability notes: thinking/reasoning controls are wired for Anthropic and the OpenAI family today (Gemini, Vertex, Bedrock, and Azure planned — setting `thinking_level` there logs a warning); client-side prompt-cache breakpoints are Anthropic-specific (other providers cache server-side automatically) **Built-in Tools** - `bash` — Shell execution with timeout, output truncation, command deny patterns @@ -82,7 +83,7 @@ Or add to `Cargo.toml`: ```toml [dependencies] -yoagent = "0.6" +yoagent = "0.9" tokio = { version = "1", features = ["full"] } ``` diff --git a/src/agent.rs b/src/agent.rs index b707cc8..01a8a66 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -142,6 +142,14 @@ impl Agent { self } + /// Set the sampling temperature. Note: the newest reasoning models + /// (e.g. Claude Fable 5 / Opus 4.7+) reject sampling parameters — leave + /// unset for those. + pub fn with_temperature(mut self, temperature: f32) -> Self { + self.temperature = Some(temperature); + self + } + pub fn with_context_config(mut self, config: ContextConfig) -> Self { self.context_config = Some(config); self @@ -476,6 +484,12 @@ impl Agent { /// Call [`finish()`](Self::finish) after draining the receiver to restore /// agent state (messages, tools). `finish()` is also called automatically /// at the start of the next `prompt` / `continue_loop` call. + /// + /// # Panics + /// + /// Panics if the agent is already streaming (call [`Agent::finish`] or + /// await the previous run first). A misuse-`Result` variant is planned + /// for 0.10. pub async fn prompt(&mut self, text: impl Into) -> mpsc::UnboundedReceiver { let msg = AgentMessage::Llm(Message::user(text)); self.prompt_messages(vec![msg]).await @@ -485,6 +499,12 @@ impl Agent { /// agent loop running concurrently for true streaming. /// /// Call [`finish()`](Self::finish) after draining events to restore state. + /// + /// # Panics + /// + /// Panics if the agent is already streaming (call [`Agent::finish`] or + /// await the previous run first). A misuse-`Result` variant is planned + /// for 0.10. pub async fn prompt_messages( &mut self, messages: Vec, @@ -539,6 +559,12 @@ impl Agent { /// agent.prompt_with_sender("hello", tx).await; /// # } /// ``` + /// + /// # Panics + /// + /// Panics if the agent is already streaming (call [`Agent::finish`] or + /// await the previous run first). A misuse-`Result` variant is planned + /// for 0.10. pub async fn prompt_with_sender( &mut self, text: impl Into, @@ -550,6 +576,12 @@ impl Agent { /// Send messages as a prompt, streaming events to a caller-provided sender. /// Blocks until the loop finishes and state is restored. + /// + /// # Panics + /// + /// Panics if the agent is already streaming (call [`Agent::finish`] or + /// await the previous run first). A misuse-`Result` variant is planned + /// for 0.10. pub async fn prompt_messages_with_sender( &mut self, messages: Vec, @@ -587,6 +619,11 @@ impl Agent { /// receiver immediately with the loop running concurrently. /// /// Call [`finish()`](Self::finish) after draining events to restore state. + /// + /// # Panics + /// + /// Panics if the agent is already streaming or if there are no messages + /// to continue from. A misuse-`Result` variant is planned for 0.10. pub async fn continue_loop(&mut self) -> mpsc::UnboundedReceiver { self.finish().await; // restore from previous if needed @@ -618,6 +655,11 @@ impl Agent { /// Continue from current context, streaming events to a caller-provided sender. /// Blocks until the loop finishes and state is restored. + /// + /// # Panics + /// + /// Panics if the agent is already streaming or if there are no messages + /// to continue from. A misuse-`Result` variant is planned for 0.10. pub async fn continue_loop_with_sender(&mut self, tx: mpsc::UnboundedSender) { self.finish().await; // restore from previous if needed @@ -674,7 +716,53 @@ impl Agent { // -- Internal -- + /// The explicit key if set, else the provider-conventional env var + /// (see [`crate::provider::resolve_api_key`]). + fn resolved_api_key(&self) -> String { + if !self.api_key.is_empty() { + return self.api_key.clone(); + } + let provider = self + .model_config + .as_ref() + .map(|m| m.provider.as_str()) + .unwrap_or("anthropic"); + crate::provider::resolve_api_key(provider).unwrap_or_default() + } + + /// Total dollar cost of the assistant turns currently in history, using + /// the model's [`CostConfig`](crate::provider::CostConfig) rates. + /// + /// Returns `None` when no `ModelConfig` is set. Rates come from the + /// *current* model config; sessions that switched models mid-way are + /// priced entirely at the current rates. + pub fn session_cost_usd(&self) -> Option { + let cost = &self.model_config.as_ref()?.cost; + Some( + self.messages + .iter() + .filter_map(|m| match m { + AgentMessage::Llm(Message::Assistant { usage, .. }) => { + Some(cost.cost_usd(usage)) + } + _ => None, + }) + .sum(), + ) + } + fn build_config(&self) -> AgentLoopConfig { + if self.thinking_level != ThinkingLevel::Off { + if let Some(mc) = &self.model_config { + if !mc.reasoning { + tracing::warn!( + "thinking_level is set but model '{}' is not marked \ + reasoning-capable (ModelConfig.reasoning = false)", + mc.id + ); + } + } + } let steering_queue = self.steering_queue.clone(); let steering_mode = self.steering_mode; @@ -684,7 +772,7 @@ impl Agent { AgentLoopConfig { provider: self.provider.clone(), model: self.model.clone(), - api_key: self.api_key.clone(), + api_key: self.resolved_api_key(), thinking_level: self.thinking_level, max_tokens: self.max_tokens, temperature: self.temperature, diff --git a/src/agent_loop.rs b/src/agent_loop.rs index 0a2e3f0..e3bd1ad 100644 --- a/src/agent_loop.rs +++ b/src/agent_loop.rs @@ -8,7 +8,8 @@ //! Both return a stream of `AgentEvent`s. use crate::context::{ - self, CompactionStrategy, ContextConfig, DefaultCompaction, ExecutionLimits, ExecutionTracker, + self, CompactionStrategy, ContextConfig, ContextTracker, DefaultCompaction, ExecutionLimits, + ExecutionTracker, }; use crate::provider::{ModelConfig, StreamConfig, StreamEvent, StreamProvider, ToolDefinition}; use crate::types::*; @@ -253,6 +254,8 @@ async fn run_loop( ) { let mut first_turn = true; let mut turn_number: usize = 0; + // Blends real provider usage with estimation for compaction sizing. + let mut context_tracker = ContextTracker::new(); let mut tracker = config .execution_limits .as_ref() @@ -342,14 +345,37 @@ async fn run_loop( turn_number += 1; - // Compact context if configured (tiered: tool outputs → summarize → drop) + // Compact context if configured (tiered: tool outputs → summarize → drop). + // The context tracker blends real provider usage with estimation; + // when real usage shows the char-based estimate runs low (images, + // non-Latin text), scale the budget proportionally so compaction + // triggers at the true context size. if let Some(ref ctx_config) = config.context_config { + let estimated = context::total_tokens(&context.messages); + let hybrid = context_tracker.estimate_context_tokens(&context.messages); + let calibrated; + let effective_config = if hybrid > estimated && estimated > 0 { + calibrated = ContextConfig { + max_context_tokens: ((ctx_config.max_context_tokens as u64 + * estimated as u64) + / hybrid as u64) as usize, + ..ctx_config.clone() + }; + &calibrated + } else { + ctx_config + }; let strategy: &dyn CompactionStrategy = config .compaction_strategy .as_deref() .unwrap_or(&DefaultCompaction); + let before_len = context.messages.len(); context.messages = - strategy.compact(std::mem::take(&mut context.messages), ctx_config); + strategy.compact(std::mem::take(&mut context.messages), effective_config); + if context.messages.len() != before_len { + // Messages shifted; re-baseline from the next real usage. + context_tracker.reset(); + } } // Stream assistant response @@ -358,6 +384,9 @@ async fn run_loop( let agent_msg: AgentMessage = message.clone().into(); context.messages.push(agent_msg.clone()); new_messages.push(agent_msg.clone()); + if let Message::Assistant { usage, .. } = &message { + context_tracker.record_usage(usage, context.messages.len() - 1); + } // Check for error/abort if let Message::Assistant { diff --git a/src/lib.rs b/src/lib.rs index 1a6cb6e..50a0174 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,54 @@ +//! **yoagent** — the agent runtime for Rust. +//! +//! A simple, effective agent loop with tool execution and event streaming: +//! `Prompt → LLM stream → tool execution → loop`. The loop is the product; +//! everything else is optional layers on top of it. +//! +//! # Quick start +//! +//! ```no_run +//! use yoagent::{Agent, AgentMessage, provider::AnthropicProvider, tools}; +//! +//! # #[tokio::main] +//! # async fn main() { +//! let mut agent = Agent::new(AnthropicProvider) +//! .with_model("claude-sonnet-5") +//! .with_api_key(std::env::var("ANTHROPIC_API_KEY").unwrap()) +//! .with_system_prompt("You are a helpful coding assistant.") +//! .with_tools(tools::default_tools()); +//! +//! let mut events = agent.prompt("List the files in the current directory").await; +//! while let Some(event) = events.recv().await { +//! // stream text deltas, tool calls, usage — render however you like +//! } +//! agent.finish().await; +//! # } +//! ``` +//! +//! # What's in the box +//! +//! - **The loop** ([`agent_loop()`](agent_loop())) — a stateless free function; [`Agent`] is an +//! optional stateful wrapper (history, tool registry, steering queues). +//! - **7 provider protocols, 20+ providers** ([`provider`]) — Anthropic, +//! OpenAI (Completions + Responses), Azure, Gemini, Vertex, Bedrock, plus +//! OpenAI-compatible gateways (Groq, DeepSeek, xAI, OpenCode, Ollama, ...) +//! with per-provider quirk flags. +//! - **Tools** ([`tools`]) — bash, read/write/edit file, search; add your own +//! via the [`AgentTool`] trait. [MCP](mcp) servers and +//! [OpenAPI specs](https://docs.rs/yoagent/latest/yoagent/openapi/index.html) +//! (feature `openapi`) become tools transparently. +//! - **Steering** — interrupt a running agent mid-tool-execution +//! ([`Agent::steer`]), queue follow-ups, inspect/edit the queues. +//! - **Sub-agents** ([`SubAgentTool`]) — delegation with per-sub-agent models +//! and [`SharedState`] for passing artifacts by reference. +//! - **Context management** ([`context`]) — token tracking and tiered +//! compaction so long sessions keep running. +//! - **Skills** ([`skills`]) — load `SKILL.md` files per the +//! [AgentSkills](https://agentskills.io) standard. +//! +//! The [book](https://yologdev.github.io/yoagent/) covers concepts and +//! provider-specific guides. + pub mod agent; pub mod agent_loop; pub mod context; diff --git a/src/openapi/adapter.rs b/src/openapi/adapter.rs index e04de57..eeceeb1 100644 --- a/src/openapi/adapter.rs +++ b/src/openapi/adapter.rs @@ -349,7 +349,7 @@ fn parse_spec(input: &str) -> Result { serde_json::from_str::(input) .map_err(|e| OpenApiError::ParseError(format!("JSON: {}", e))) } else { - serde_yaml::from_str::(input) + serde_yaml_ng::from_str::(input) .map_err(|e| OpenApiError::ParseError(format!("YAML: {}", e))) } } diff --git a/src/openapi/types.rs b/src/openapi/types.rs index f51ab88..8767b75 100644 --- a/src/openapi/types.rs +++ b/src/openapi/types.rs @@ -7,7 +7,7 @@ pub enum OpenApiAuth { /// No authentication. #[default] None, - /// Bearer token (Authorization: Bearer ). + /// Bearer token (`Authorization: Bearer `). Bearer(String), /// API key in a custom header. ApiKey { header: String, value: String }, diff --git a/src/provider/azure_openai.rs b/src/provider/azure_openai.rs index d201d50..a0a8a87 100644 --- a/src/provider/azure_openai.rs +++ b/src/provider/azure_openai.rs @@ -25,6 +25,11 @@ impl StreamProvider for AzureOpenAiProvider { tx: mpsc::UnboundedSender, cancel: tokio_util::sync::CancellationToken, ) -> Result { + if config.thinking_level != ThinkingLevel::Off { + warn!( + "thinking_level is not yet wired for the Azure OpenAI provider and will be ignored" + ); + } let model_config = config .model_config .as_ref() diff --git a/src/provider/bedrock.rs b/src/provider/bedrock.rs index f7e2855..6e288ff 100644 --- a/src/provider/bedrock.rs +++ b/src/provider/bedrock.rs @@ -28,6 +28,11 @@ impl StreamProvider for BedrockProvider { tx: mpsc::UnboundedSender, cancel: tokio_util::sync::CancellationToken, ) -> Result { + if config.thinking_level != ThinkingLevel::Off { + warn!( + "thinking_level is not yet wired for the Amazon Bedrock provider and will be ignored" + ); + } let model_config = config .model_config .as_ref() diff --git a/src/provider/google.rs b/src/provider/google.rs index 33e57dd..5a67bf3 100644 --- a/src/provider/google.rs +++ b/src/provider/google.rs @@ -21,6 +21,11 @@ impl StreamProvider for GoogleProvider { tx: mpsc::UnboundedSender, cancel: tokio_util::sync::CancellationToken, ) -> Result { + if config.thinking_level != ThinkingLevel::Off { + warn!( + "thinking_level is not yet wired for the Google Gemini provider and will be ignored" + ); + } let model_config = config .model_config .as_ref() diff --git a/src/provider/google_vertex.rs b/src/provider/google_vertex.rs index b28a8bd..45f25eb 100644 --- a/src/provider/google_vertex.rs +++ b/src/provider/google_vertex.rs @@ -33,6 +33,11 @@ impl StreamProvider for GoogleVertexProvider { tx: mpsc::UnboundedSender, cancel: tokio_util::sync::CancellationToken, ) -> Result { + if config.thinking_level != ThinkingLevel::Off { + tracing::warn!( + "thinking_level is not yet wired for the Google Vertex provider and will be ignored" + ); + } let model_config = config .model_config .as_ref() diff --git a/src/provider/mod.rs b/src/provider/mod.rs index 5d8387a..002ca60 100644 --- a/src/provider/mod.rs +++ b/src/provider/mod.rs @@ -20,5 +20,5 @@ pub use mock::MockProvider; pub use model::{AnthropicCompat, ApiProtocol, CostConfig, ModelConfig, OpenAiCompat}; pub use openai_compat::OpenAiCompatProvider; pub use openai_responses::OpenAiResponsesProvider; -pub use registry::ProviderRegistry; +pub use registry::{resolve_api_key, ProviderRegistry}; pub use traits::*; diff --git a/src/provider/model.rs b/src/provider/model.rs index 523f9bf..ff17107 100644 --- a/src/provider/model.rs +++ b/src/provider/model.rs @@ -41,6 +41,20 @@ pub struct CostConfig { pub cache_write_per_million: f64, } +impl CostConfig { + /// Dollar cost of a usage record at these per-million-token rates. + /// + /// Consumed by [`crate::Agent::session_cost_usd`]; also usable directly + /// in `after_turn` callbacks for per-turn cost tracking. + pub fn cost_usd(&self, usage: &crate::types::Usage) -> f64 { + (usage.input as f64 * self.input_per_million + + usage.output as f64 * self.output_per_million + + usage.cache_read as f64 * self.cache_read_per_million + + usage.cache_write as f64 * self.cache_write_per_million) + / 1_000_000.0 + } +} + impl Default for CostConfig { fn default() -> Self { Self { @@ -297,7 +311,9 @@ pub struct ModelConfig { pub provider: String, /// Base URL for API requests (without trailing slash). pub base_url: String, - /// Whether this model supports reasoning/thinking. + /// Whether this model supports reasoning/thinking. When `false` and a + /// `thinking_level` is requested, the agent logs a warning (the request + /// is still sent — gate behavior stays with the caller). pub reasoning: bool, /// Context window size in tokens. pub context_window: u32, @@ -788,6 +804,27 @@ mod tests { assert!(config.anthropic.is_none()); } + #[test] + fn test_cost_usd() { + let cost = CostConfig { + input_per_million: 3.0, + output_per_million: 15.0, + cache_read_per_million: 0.3, + cache_write_per_million: 3.75, + }; + let usage = crate::types::Usage { + input: 1_000_000, + output: 100_000, + cache_read: 2_000_000, + cache_write: 400_000, + total_tokens: 0, + }; + // 3.0 + 1.5 + 0.6 + 1.5 = 6.6 + assert!((cost.cost_usd(&usage) - 6.6).abs() < 1e-9); + // zero rates (default) => zero cost + assert_eq!(CostConfig::default().cost_usd(&usage), 0.0); + } + #[test] fn test_new_generation_presets() { let fable = ModelConfig::claude_fable_5(); diff --git a/src/provider/registry.rs b/src/provider/registry.rs index ccd34f3..0765c9e 100644 --- a/src/provider/registry.rs +++ b/src/provider/registry.rs @@ -115,3 +115,80 @@ mod tests { assert!(registry.has(&ApiProtocol::AnthropicMessages)); } } + +/// Resolve an API key from the conventional environment variable(s) for a +/// provider name (the `ModelConfig.provider` string). +/// +/// Used as a fallback when no explicit key is set on the agent — an explicit +/// `with_api_key` always wins. Conventions: +/// +/// | provider | env var(s), first match wins | +/// |---|---| +/// | `anthropic` | `ANTHROPIC_API_KEY` | +/// | `openai` | `OPENAI_API_KEY` | +/// | `google` | `GEMINI_API_KEY`, `GOOGLE_API_KEY` | +/// | `xai` / `groq` / `deepseek` / `mistral` / `zai` / `minimax` / `openrouter` / `cerebras` | `_API_KEY` | +/// | `qwen` | `DASHSCOPE_API_KEY` | +/// | `opencode-zen` / `opencode-go` | `OPENCODE_API_KEY` | +/// | `azure` | `AZURE_OPENAI_API_KEY` | +/// | `bedrock` | `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` (+ `AWS_SESSION_TOKEN`), composed as `access:secret[:token]` | +/// | `vertex` | none — pass a short-lived OAuth token via `with_api_key` | +/// | `local` / `ollama` | no key needed (empty) | +/// | anything else | `YOAGENT_API_KEY`, then `API_KEY` | +pub fn resolve_api_key(provider: &str) -> Option { + use std::env::var; + let first = |names: &[&str]| names.iter().find_map(|n| var(n).ok()); + match provider { + "anthropic" => first(&["ANTHROPIC_API_KEY"]), + "openai" => first(&["OPENAI_API_KEY"]), + "google" => first(&["GEMINI_API_KEY", "GOOGLE_API_KEY"]), + "xai" => first(&["XAI_API_KEY"]), + "groq" => first(&["GROQ_API_KEY"]), + "deepseek" => first(&["DEEPSEEK_API_KEY"]), + "mistral" => first(&["MISTRAL_API_KEY"]), + "zai" => first(&["ZAI_API_KEY"]), + "minimax" => first(&["MINIMAX_API_KEY"]), + "openrouter" => first(&["OPENROUTER_API_KEY"]), + "cerebras" => first(&["CEREBRAS_API_KEY"]), + "qwen" => first(&["DASHSCOPE_API_KEY"]), + "opencode-zen" | "opencode-go" => first(&["OPENCODE_API_KEY"]), + "azure" => first(&["AZURE_OPENAI_API_KEY"]), + "bedrock" => { + let access = var("AWS_ACCESS_KEY_ID").ok()?; + let secret = var("AWS_SECRET_ACCESS_KEY").ok()?; + Some(match var("AWS_SESSION_TOKEN") { + Ok(token) => format!("{}:{}:{}", access, secret, token), + Err(_) => format!("{}:{}", access, secret), + }) + } + "vertex" => None, + "local" | "ollama" => Some(String::new()), + _ => first(&["YOAGENT_API_KEY", "API_KEY"]), + } +} + +#[cfg(test)] +mod resolve_key_tests { + use super::resolve_api_key; + + #[test] + fn test_deterministic_branches() { + // local/ollama need no key + assert_eq!(resolve_api_key("local").as_deref(), Some("")); + assert_eq!(resolve_api_key("ollama").as_deref(), Some("")); + // vertex expects an explicit OAuth token — never env-resolved + assert_eq!(resolve_api_key("vertex"), None); + } + + #[test] + fn test_env_resolution() { + // Use a provider name unique to this test to avoid env races with + // parallel tests. + std::env::set_var("YOAGENT_API_KEY", "from-generic-fallback"); + assert_eq!( + resolve_api_key("some-unknown-gateway").as_deref(), + Some("from-generic-fallback") + ); + std::env::remove_var("YOAGENT_API_KEY"); + } +} diff --git a/src/provider/traits.rs b/src/provider/traits.rs index 24e8a2c..5c526a1 100644 --- a/src/provider/traits.rs +++ b/src/provider/traits.rs @@ -97,14 +97,22 @@ impl ProviderError { /// Detects context overflow, rate limits, auth errors, and general API errors /// from the HTTP status code and response body. pub fn classify(status: u16, message: &str) -> Self { + Self::classify_with_retry_after(status, message, None) + } + + /// Like [`classify`](Self::classify), carrying a parsed `Retry-After` + /// value (milliseconds) into the `RateLimited` variant when present. + pub fn classify_with_retry_after( + status: u16, + message: &str, + retry_after_ms: Option, + ) -> Self { if is_context_overflow(status, message) { Self::ContextOverflow { message: message.to_string(), } } else if status == 429 { - Self::RateLimited { - retry_after_ms: None, - } + Self::RateLimited { retry_after_ms } } else if status == 401 || status == 403 { Self::Auth(message.to_string()) } else { @@ -130,8 +138,9 @@ pub async fn classify_eventsource_error(error: reqwest_eventsource::Error) -> Pr match error { reqwest_eventsource::Error::InvalidStatusCode(status, response) => { let status_code = status.as_u16(); + let retry_after_ms = parse_retry_after(response.headers()); let body = response.text().await.unwrap_or_default(); - ProviderError::classify( + ProviderError::classify_with_retry_after( status_code, &format!( "HTTP {} {}: {}", @@ -139,6 +148,7 @@ pub async fn classify_eventsource_error(error: reqwest_eventsource::Error) -> Pr status.canonical_reason().unwrap_or(""), body ), + retry_after_ms, ) } reqwest_eventsource::Error::Transport(e) => ProviderError::Network(format!("{:?}", e)), @@ -146,6 +156,19 @@ pub async fn classify_eventsource_error(error: reqwest_eventsource::Error) -> Pr } } +/// Parse a `Retry-After` header (seconds form) into milliseconds. +pub(crate) fn parse_retry_after(headers: &reqwest::header::HeaderMap) -> Option { + headers + .get(reqwest::header::RETRY_AFTER)? + .to_str() + .ok()? + .trim() + .parse::() + .ok() + .filter(|s| *s >= 0.0) + .map(|secs| (secs * 1000.0) as u64) +} + /// Classify an SSE-embedded error event message into a [`ProviderError`]. /// /// Checks the error text for known patterns (context overflow, etc.). diff --git a/src/sub_agent.rs b/src/sub_agent.rs index ade8e25..ce6bdd6 100644 --- a/src/sub_agent.rs +++ b/src/sub_agent.rs @@ -54,6 +54,7 @@ pub struct SubAgentTool { tools: Vec>, thinking_level: ThinkingLevel, max_tokens: Option, + temperature: Option, cache_config: CacheConfig, tool_execution: ToolExecutionStrategy, retry_config: crate::retry::RetryConfig, @@ -78,6 +79,7 @@ impl SubAgentTool { tools: Vec::new(), thinking_level: ThinkingLevel::Off, max_tokens: None, + temperature: None, cache_config: CacheConfig::default(), tool_execution: ToolExecutionStrategy::default(), retry_config: crate::retry::RetryConfig::default(), @@ -135,6 +137,12 @@ impl SubAgentTool { self } + /// Set the sampling temperature for the sub-agent. + pub fn with_temperature(mut self, temperature: f32) -> Self { + self.temperature = Some(temperature); + self + } + pub fn with_cache_config(mut self, config: CacheConfig) -> Self { self.cache_config = config; self @@ -287,10 +295,20 @@ impl AgentTool for SubAgentTool { let config = AgentLoopConfig { provider: self.provider.clone(), model: self.model.clone(), - api_key: self.api_key.clone(), + api_key: if self.api_key.is_empty() { + crate::provider::resolve_api_key( + self.model_config + .as_ref() + .map(|m| m.provider.as_str()) + .unwrap_or("anthropic"), + ) + .unwrap_or_default() + } else { + self.api_key.clone() + }, thinking_level: self.thinking_level, max_tokens: self.max_tokens, - temperature: None, + temperature: self.temperature, model_config: self.model_config.clone(), convert_to_llm: None, transform_context: None, From e19e8e128e5e77f2854cb0eb15be2f1f97268e21 Mon Sep 17 00:00:00 2001 From: Yuanhao Li Date: Sun, 5 Jul 2026 19:30:09 +0200 Subject: [PATCH 2/3] =?UTF-8?q?fix:=20raise=20MSRV=20to=201.86=20=E2=80=94?= =?UTF-8?q?=20the=20declared=201.85=20floor=20was=20already=20broken?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The new MSRV CI job caught it on its first run: the icu/idna stack (pulled transitively via url) requires rustc 1.86, so 1.85 users could not build yoagent from a fresh resolution anyway. Declare the real floor. Co-Authored-By: Claude Fable 5 --- .github/workflows/ci.yml | 4 ++-- Cargo.toml | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c9c0155..6eecff5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -51,11 +51,11 @@ jobs: run: cargo check --all-targets --all-features msrv: - name: MSRV (1.85) + name: MSRV (1.86) runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - - uses: dtolnay/rust-toolchain@1.85 + - uses: dtolnay/rust-toolchain@1.86 - uses: Swatinem/rust-cache@v2 - name: Check run: cargo check --all-features diff --git a/Cargo.toml b/Cargo.toml index 86e4829..cee4303 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ readme = "README.md" keywords = ["agent", "llm", "ai", "tools", "streaming"] categories = ["api-bindings", "asynchronous"] authors = ["Yuanhao "] -rust-version = "1.85" +rust-version = "1.86" exclude = ["docs/images/*", "docs/theme/*", ".github/*", "scripts/*"] [dependencies] From 8b5bd7375939982c5f6bbff0b8fc86e71951659d Mon Sep 17 00:00:00 2001 From: Yuanhao Li Date: Sun, 5 Jul 2026 22:39:06 +0200 Subject: [PATCH 3/3] fix: address PR #58 review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Multi-agent review of the Phase A funnel/trust work surfaced defects in the newly wired features. Fixes, most severe first: - Compaction budget calibration (agent_loop.rs): the ratio scaling collapsed the budget toward zero because provider usage counts the system prompt and tool schemas while the message estimate does not — the ratio fired every session. Replace with additive overhead subtraction (measured overhead = hybrid − estimate), zero the now-double-counted system_prompt reserve, and floor the result at 10% of the configured budget so a mis-measured overhead can never wipe history. Adds a tracing::debug! line. - Missing-key resolution (registry.rs, agent.rs, sub_agent.rs): empty-string fallback erased "no credentials found". Add resolve_api_key_or_warn that tracing::warn!s the exact env var to set before falling back, and a debug-log of which var resolved. - Google/Vertex usage double-counted cached tokens: promptTokenCount already includes cachedContentTokenCount, so input + cache_read over-counted. Keep input as the uncached remainder (Vertex also now parses/reports cache_read). - Retry-After bypassed max_delay_ms: clamp server-provided delay so a bad header can't stall the loop uncancellably. - session_cost_usd returned Some(0.0) for unpriced models: add CostConfig::is_configured and return None when all rates are zero (can't price, not free). Doc corrections: restore with_openapi_file/url to CLAUDE.md (they exist, pub async); rewrite the # Panics sections (finish() cannot recover a dropped _with_sender future); scope the reasoning-warning doc to the Agent wrapper; fix lib.rs steering wording and drop the unused doctest import; tighten README provider caveats (OpenAI-family thinking is opt-in per ModelConfig; Bedrock has no automatic caching); add the sub-agent temperature caveat. Tests: parse_retry_after units + a wiremock 429 asserting retry_after_ms; calibration loop tests (overhead subtraction + floor guard) via a recording compaction strategy; session_cost_usd (None cases + assistant-only sum); explicit-vs-env key precedence; sub-agent temperature/env-key plumbing; a Google cached-token regression test. Co-Authored-By: Claude Opus 4.8 (1M context) --- CLAUDE.md | 2 +- README.md | 2 +- src/agent.rs | 57 ++++++++---- src/agent_loop.rs | 37 ++++++-- src/lib.rs | 7 +- src/provider/google.rs | 8 +- src/provider/google_vertex.rs | 11 ++- src/provider/mod.rs | 1 + src/provider/model.rs | 15 +++- src/provider/registry.rs | 54 +++++++++++- src/provider/traits.rs | 42 +++++++++ src/sub_agent.rs | 7 +- tests/agent_loop_test.rs | 136 +++++++++++++++++++++++++++++ tests/agent_test.rs | 154 +++++++++++++++++++++++++++++++++ tests/anthropic_stream_test.rs | 28 ++++++ tests/google_stream_test.rs | 31 +++++++ tests/sub_agent_test.rs | 92 ++++++++++++++++++++ 17 files changed, 645 insertions(+), 39 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 9f168bc..f3a8ad3 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -74,7 +74,7 @@ The loop: stream assistant response → extract tool calls → execute tools (pa ### OpenAPI Integration (`openapi/`, feature-gated) -Behind the `openapi` Cargo feature. `OpenApiToolAdapter` parses an OpenAPI 3.0 spec and creates one `AgentTool` per operation. Factory methods: `from_str`, `from_file`, `from_url`, `from_spec`. `OperationFilter` controls which operations become tools. Added to `Agent` via `with_openapi_spec()`. +Behind the `openapi` Cargo feature. `OpenApiToolAdapter` parses an OpenAPI 3.0 spec and creates one `AgentTool` per operation. Factory methods: `from_str`, `from_file`, `from_url`, `from_spec`. `OperationFilter` controls which operations become tools. Added to `Agent` via `with_openapi_file()` / `with_openapi_url()` / `with_openapi_spec()`. ### MCP Integration (`mcp/`) diff --git a/README.md b/README.md index b8ec5f6..c446160 100644 --- a/README.md +++ b/README.md @@ -49,7 +49,7 @@ Everything is observable via events. Supports 7 API protocols covering 20+ LLM p - 7 API protocols, 20+ providers behind one `StreamProvider` trait - One OpenAI-compatible implementation covers OpenAI, xAI, Groq, Cerebras, OpenRouter, Mistral, and more - Per-provider quirk flags (`OpenAiCompat`, `AnthropicCompat`) handle auth, reasoning format, and tool handling differences -- Capability notes: thinking/reasoning controls are wired for Anthropic and the OpenAI family today (Gemini, Vertex, Bedrock, and Azure planned — setting `thinking_level` there logs a warning); client-side prompt-cache breakpoints are Anthropic-specific (other providers cache server-side automatically) +- Capability notes: thinking/reasoning controls are wired for Anthropic and for OpenAI-compatible providers whose `ModelConfig` opts in (the `openai()`/`deepseek()` presets do; other compat presets silently drop `thinking_level`). Gemini, Vertex, Bedrock, and Azure are planned — setting `thinking_level` there logs a warning. Client-side prompt-cache breakpoints are Anthropic-specific; most other providers cache server-side automatically, but Bedrock has no automatic caching **Built-in Tools** - `bash` — Shell execution with timeout, output truncation, command deny patterns diff --git a/src/agent.rs b/src/agent.rs index 01a8a66..d4c8734 100644 --- a/src/agent.rs +++ b/src/agent.rs @@ -487,9 +487,12 @@ impl Agent { /// /// # Panics /// - /// Panics if the agent is already streaming (call [`Agent::finish`] or - /// await the previous run first). A misuse-`Result` variant is planned - /// for 0.10. + /// Panics if the agent still counts as streaming — in practice only when + /// a previous `*_with_sender` future was dropped mid-run, which leaves + /// the agent stuck in the streaming state ([`Agent::finish`] cannot + /// recover it; recreate the agent). Runs started by the + /// receiver-returning methods are joined automatically. A + /// misuse-`Result` variant is planned for 0.10. pub async fn prompt(&mut self, text: impl Into) -> mpsc::UnboundedReceiver { let msg = AgentMessage::Llm(Message::user(text)); self.prompt_messages(vec![msg]).await @@ -502,9 +505,12 @@ impl Agent { /// /// # Panics /// - /// Panics if the agent is already streaming (call [`Agent::finish`] or - /// await the previous run first). A misuse-`Result` variant is planned - /// for 0.10. + /// Panics if the agent still counts as streaming — in practice only when + /// a previous `*_with_sender` future was dropped mid-run, which leaves + /// the agent stuck in the streaming state ([`Agent::finish`] cannot + /// recover it; recreate the agent). Runs started by the + /// receiver-returning methods are joined automatically. A + /// misuse-`Result` variant is planned for 0.10. pub async fn prompt_messages( &mut self, messages: Vec, @@ -562,9 +568,12 @@ impl Agent { /// /// # Panics /// - /// Panics if the agent is already streaming (call [`Agent::finish`] or - /// await the previous run first). A misuse-`Result` variant is planned - /// for 0.10. + /// Panics if the agent still counts as streaming — in practice only when + /// a previous `*_with_sender` future was dropped mid-run, which leaves + /// the agent stuck in the streaming state ([`Agent::finish`] cannot + /// recover it; recreate the agent). Runs started by the + /// receiver-returning methods are joined automatically. A + /// misuse-`Result` variant is planned for 0.10. pub async fn prompt_with_sender( &mut self, text: impl Into, @@ -579,9 +588,12 @@ impl Agent { /// /// # Panics /// - /// Panics if the agent is already streaming (call [`Agent::finish`] or - /// await the previous run first). A misuse-`Result` variant is planned - /// for 0.10. + /// Panics if the agent still counts as streaming — in practice only when + /// a previous `*_with_sender` future was dropped mid-run, which leaves + /// the agent stuck in the streaming state ([`Agent::finish`] cannot + /// recover it; recreate the agent). Runs started by the + /// receiver-returning methods are joined automatically. A + /// misuse-`Result` variant is planned for 0.10. pub async fn prompt_messages_with_sender( &mut self, messages: Vec, @@ -622,8 +634,10 @@ impl Agent { /// /// # Panics /// - /// Panics if the agent is already streaming or if there are no messages - /// to continue from. A misuse-`Result` variant is planned for 0.10. + /// Panics if there are no messages to continue from, or if a previous + /// `*_with_sender` future was dropped mid-run (the agent is stuck in the + /// streaming state; [`Agent::finish`] cannot recover it — recreate the + /// agent). A misuse-`Result` variant is planned for 0.10. pub async fn continue_loop(&mut self) -> mpsc::UnboundedReceiver { self.finish().await; // restore from previous if needed @@ -658,8 +672,10 @@ impl Agent { /// /// # Panics /// - /// Panics if the agent is already streaming or if there are no messages - /// to continue from. A misuse-`Result` variant is planned for 0.10. + /// Panics if there are no messages to continue from, or if a previous + /// `*_with_sender` future was dropped mid-run (the agent is stuck in the + /// streaming state; [`Agent::finish`] cannot recover it — recreate the + /// agent). A misuse-`Result` variant is planned for 0.10. pub async fn continue_loop_with_sender(&mut self, tx: mpsc::UnboundedSender) { self.finish().await; // restore from previous if needed @@ -727,17 +743,22 @@ impl Agent { .as_ref() .map(|m| m.provider.as_str()) .unwrap_or("anthropic"); - crate::provider::resolve_api_key(provider).unwrap_or_default() + crate::provider::resolve_api_key_or_warn(provider) } /// Total dollar cost of the assistant turns currently in history, using /// the model's [`CostConfig`](crate::provider::CostConfig) rates. /// - /// Returns `None` when no `ModelConfig` is set. Rates come from the + /// Returns `None` when no `ModelConfig` is set or when the config's + /// rates are all zero (pricing unknown — e.g. custom or local models), + /// so `None` means "can't price this", never "free". Rates come from the /// *current* model config; sessions that switched models mid-way are /// priced entirely at the current rates. pub fn session_cost_usd(&self) -> Option { let cost = &self.model_config.as_ref()?.cost; + if !cost.is_configured() { + return None; + } Some( self.messages .iter() diff --git a/src/agent_loop.rs b/src/agent_loop.rs index e3bd1ad..66b54a0 100644 --- a/src/agent_loop.rs +++ b/src/agent_loop.rs @@ -346,21 +346,39 @@ async fn run_loop( turn_number += 1; // Compact context if configured (tiered: tool outputs → summarize → drop). - // The context tracker blends real provider usage with estimation; - // when real usage shows the char-based estimate runs low (images, - // non-Latin text), scale the budget proportionally so compaction - // triggers at the true context size. + // + // Calibration: the tracker's hybrid figure is anchored on provider + // usage, which counts the FULL request (system prompt + tool + // schemas + any estimation shortfall on messages), while + // `total_tokens` counts messages only. The difference is a + // measured overhead; subtracting it from the budget makes the + // strategy's message-token checks equivalent to "true window + // occupancy <= max_context_tokens". The static + // `system_prompt_tokens` reserve is zeroed in the calibrated + // config because the measured overhead already includes the real + // system prompt. A floor of 10% of the configured budget + // guarantees a mis-measured overhead can never wipe the history. if let Some(ref ctx_config) = config.context_config { let estimated = context::total_tokens(&context.messages); let hybrid = context_tracker.estimate_context_tokens(&context.messages); + let overhead = hybrid.saturating_sub(estimated); let calibrated; - let effective_config = if hybrid > estimated && estimated > 0 { + let effective_config = if overhead > 0 { + let floor = ctx_config.max_context_tokens / 10; calibrated = ContextConfig { - max_context_tokens: ((ctx_config.max_context_tokens as u64 - * estimated as u64) - / hybrid as u64) as usize, + max_context_tokens: ctx_config + .max_context_tokens + .saturating_sub(overhead) + .max(floor), + system_prompt_tokens: 0, ..ctx_config.clone() }; + tracing::debug!( + "compaction budget calibrated: {} -> {} (measured overhead: {} tokens)", + ctx_config.max_context_tokens, + calibrated.max_context_tokens, + overhead + ); &calibrated } else { ctx_config @@ -665,8 +683,11 @@ async fn stream_assistant_response( // Abort forwarder to prevent forwarding events from failed attempt forward_handle.abort(); attempt += 1; + // Server-provided Retry-After wins over backoff, but is + // clamped to max_delay_ms so a bad header can't stall the loop. let delay = e .retry_after() + .map(|d| d.min(std::time::Duration::from_millis(retry.max_delay_ms))) .unwrap_or_else(|| retry.delay_for_attempt(attempt)); crate::retry::log_retry(attempt, retry.max_retries, &delay, e); tokio::time::sleep(delay).await; diff --git a/src/lib.rs b/src/lib.rs index 50a0174..b0b5c2e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,7 +7,7 @@ //! # Quick start //! //! ```no_run -//! use yoagent::{Agent, AgentMessage, provider::AnthropicProvider, tools}; +//! use yoagent::{Agent, provider::AnthropicProvider, tools}; //! //! # #[tokio::main] //! # async fn main() { @@ -37,8 +37,9 @@ //! via the [`AgentTool`] trait. [MCP](mcp) servers and //! [OpenAPI specs](https://docs.rs/yoagent/latest/yoagent/openapi/index.html) //! (feature `openapi`) become tools transparently. -//! - **Steering** — interrupt a running agent mid-tool-execution -//! ([`Agent::steer`]), queue follow-ups, inspect/edit the queues. +//! - **Steering** — inject guidance into a running agent ([`Agent::steer`]); +//! picked up between tool executions (per batch under the default parallel +//! strategy). Queue follow-ups, inspect/edit the queues. //! - **Sub-agents** ([`SubAgentTool`]) — delegation with per-sub-agent models //! and [`SharedState`] for passing artifacts by reference. //! - **Context management** ([`context`]) — token tracking and tiered diff --git a/src/provider/google.rs b/src/provider/google.rs index 5a67bf3..ab8cd90 100644 --- a/src/provider/google.rs +++ b/src/provider/google.rs @@ -187,7 +187,13 @@ impl StreamProvider for GoogleProvider { // Process usage if let Some(u) = &chunk.usage_metadata { - usage.input = u.prompt_token_count.unwrap_or(0); + // promptTokenCount includes cached tokens; + // keep `input` as the uncached remainder so + // downstream sums don't double-count. + usage.input = u + .prompt_token_count + .unwrap_or(0) + .saturating_sub(u.cached_content_token_count.unwrap_or(0)); usage.output = u.candidates_token_count.unwrap_or(0); usage.total_tokens = u.total_token_count.unwrap_or(0); usage.cache_read = u.cached_content_token_count.unwrap_or(0); diff --git a/src/provider/google_vertex.rs b/src/provider/google_vertex.rs index 45f25eb..f8fad3f 100644 --- a/src/provider/google_vertex.rs +++ b/src/provider/google_vertex.rs @@ -191,6 +191,8 @@ async fn parse_google_sse_response( candidates_token_count: Option, #[serde(default, rename = "totalTokenCount")] total_token_count: Option, + #[serde(default, rename = "cachedContentTokenCount")] + cached_content_token_count: Option, } let parsed: Chunk = match serde_json::from_str(data) { @@ -250,8 +252,15 @@ async fn parse_google_sse_response( } if let Some(u) = parsed.usage_metadata { - usage.input = u.prompt_token_count.unwrap_or(0); + // promptTokenCount includes cached tokens; + // keep `input` as the uncached remainder so + // downstream sums don't double-count. + usage.input = u + .prompt_token_count + .unwrap_or(0) + .saturating_sub(u.cached_content_token_count.unwrap_or(0)); usage.output = u.candidates_token_count.unwrap_or(0); + usage.cache_read = u.cached_content_token_count.unwrap_or(0); usage.total_tokens = u.total_token_count.unwrap_or(0); } } diff --git a/src/provider/mod.rs b/src/provider/mod.rs index 002ca60..4d5178d 100644 --- a/src/provider/mod.rs +++ b/src/provider/mod.rs @@ -20,5 +20,6 @@ pub use mock::MockProvider; pub use model::{AnthropicCompat, ApiProtocol, CostConfig, ModelConfig, OpenAiCompat}; pub use openai_compat::OpenAiCompatProvider; pub use openai_responses::OpenAiResponsesProvider; +pub(crate) use registry::resolve_api_key_or_warn; pub use registry::{resolve_api_key, ProviderRegistry}; pub use traits::*; diff --git a/src/provider/model.rs b/src/provider/model.rs index ff17107..a867fc9 100644 --- a/src/provider/model.rs +++ b/src/provider/model.rs @@ -42,6 +42,15 @@ pub struct CostConfig { } impl CostConfig { + /// Whether any rate is set. All-zero rates mean pricing is unknown + /// (custom/local models), not that the model is free. + pub fn is_configured(&self) -> bool { + self.input_per_million != 0.0 + || self.output_per_million != 0.0 + || self.cache_read_per_million != 0.0 + || self.cache_write_per_million != 0.0 + } + /// Dollar cost of a usage record at these per-million-token rates. /// /// Consumed by [`crate::Agent::session_cost_usd`]; also usable directly @@ -312,8 +321,10 @@ pub struct ModelConfig { /// Base URL for API requests (without trailing slash). pub base_url: String, /// Whether this model supports reasoning/thinking. When `false` and a - /// `thinking_level` is requested, the agent logs a warning (the request - /// is still sent — gate behavior stays with the caller). + /// `thinking_level` is requested, the [`Agent`](crate::Agent) wrapper + /// logs a warning; sub-agents and direct `agent_loop` calls do not. The + /// request is still sent either way — gate behavior stays with the + /// caller. pub reasoning: bool, /// Context window size in tokens. pub context_window: u32, diff --git a/src/provider/registry.rs b/src/provider/registry.rs index 0765c9e..7364b39 100644 --- a/src/provider/registry.rs +++ b/src/provider/registry.rs @@ -137,7 +137,13 @@ mod tests { /// | anything else | `YOAGENT_API_KEY`, then `API_KEY` | pub fn resolve_api_key(provider: &str) -> Option { use std::env::var; - let first = |names: &[&str]| names.iter().find_map(|n| var(n).ok()); + let first = |names: &[&str]| { + names.iter().find_map(|n| { + var(n).ok().inspect(|_| { + tracing::debug!("resolved API key for provider '{}' from ${}", provider, n) + }) + }) + }; match provider { "anthropic" => first(&["ANTHROPIC_API_KEY"]), "openai" => first(&["OPENAI_API_KEY"]), @@ -167,6 +173,52 @@ pub fn resolve_api_key(provider: &str) -> Option { } } +/// Like [`resolve_api_key`], but warns (via `tracing`) when resolution fails +/// for a provider that needs a key, then falls back to an empty string so +/// the provider returns a clear authentication error instead of the failure +/// being invisible until the first request. +pub(crate) fn resolve_api_key_or_warn(provider: &str) -> String { + match resolve_api_key(provider) { + Some(key) => key, + None => { + tracing::warn!( + "no API key found for provider '{}': {}; requests will fail \ + with an authentication error", + provider, + api_key_env_hint(provider) + ); + String::new() + } + } +} + +/// Remedy fragment for the missing-key warning, matching the +/// [`resolve_api_key`] table. +fn api_key_env_hint(provider: &str) -> &'static str { + match provider { + "anthropic" => "set ANTHROPIC_API_KEY or call .with_api_key(...)", + "openai" => "set OPENAI_API_KEY or call .with_api_key(...)", + "google" => "set GEMINI_API_KEY (or GOOGLE_API_KEY) or call .with_api_key(...)", + "xai" => "set XAI_API_KEY or call .with_api_key(...)", + "groq" => "set GROQ_API_KEY or call .with_api_key(...)", + "deepseek" => "set DEEPSEEK_API_KEY or call .with_api_key(...)", + "mistral" => "set MISTRAL_API_KEY or call .with_api_key(...)", + "zai" => "set ZAI_API_KEY or call .with_api_key(...)", + "minimax" => "set MINIMAX_API_KEY or call .with_api_key(...)", + "openrouter" => "set OPENROUTER_API_KEY or call .with_api_key(...)", + "cerebras" => "set CEREBRAS_API_KEY or call .with_api_key(...)", + "qwen" => "set DASHSCOPE_API_KEY or call .with_api_key(...)", + "opencode-zen" | "opencode-go" => "set OPENCODE_API_KEY or call .with_api_key(...)", + "azure" => "set AZURE_OPENAI_API_KEY or call .with_api_key(...)", + "bedrock" => { + "set AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY (+ AWS_SESSION_TOKEN) \ + or call .with_api_key(\"access:secret[:token]\")" + } + "vertex" => "pass a short-lived OAuth token via .with_api_key(...)", + _ => "set YOAGENT_API_KEY (or API_KEY) or call .with_api_key(...)", + } +} + #[cfg(test)] mod resolve_key_tests { use super::resolve_api_key; diff --git a/src/provider/traits.rs b/src/provider/traits.rs index 5c526a1..bf48c54 100644 --- a/src/provider/traits.rs +++ b/src/provider/traits.rs @@ -227,6 +227,48 @@ fn is_context_overflow(status: u16, message: &str) -> bool { mod tests { use super::*; + fn headers_with_retry_after(value: &str) -> reqwest::header::HeaderMap { + let mut h = reqwest::header::HeaderMap::new(); + h.insert(reqwest::header::RETRY_AFTER, value.parse().unwrap()); + h + } + + #[test] + fn parse_retry_after_whole_seconds() { + assert_eq!( + parse_retry_after(&headers_with_retry_after("5")), + Some(5000) + ); + } + + #[test] + fn parse_retry_after_fractional_seconds() { + assert_eq!( + parse_retry_after(&headers_with_retry_after("1.5")), + Some(1500) + ); + } + + #[test] + fn parse_retry_after_rejects_negative() { + assert_eq!(parse_retry_after(&headers_with_retry_after("-1")), None); + } + + #[test] + fn parse_retry_after_rejects_http_date() { + // The HTTP-date form of Retry-After is not supported; must not + // misparse as a huge delay. + assert_eq!( + parse_retry_after(&headers_with_retry_after("Wed, 21 Oct 2015 07:28:00 GMT")), + None + ); + } + + #[test] + fn parse_retry_after_missing_header() { + assert_eq!(parse_retry_after(&reqwest::header::HeaderMap::new()), None); + } + #[test] fn classify_anthropic_overflow() { let err = diff --git a/src/sub_agent.rs b/src/sub_agent.rs index ce6bdd6..8536078 100644 --- a/src/sub_agent.rs +++ b/src/sub_agent.rs @@ -137,7 +137,9 @@ impl SubAgentTool { self } - /// Set the sampling temperature for the sub-agent. + /// Set the sampling temperature for the sub-agent. Note: the newest + /// reasoning models (e.g. Claude Fable 5 / Opus 4.7+) reject sampling + /// parameters — leave unset for those. pub fn with_temperature(mut self, temperature: f32) -> Self { self.temperature = Some(temperature); self @@ -296,13 +298,12 @@ impl AgentTool for SubAgentTool { provider: self.provider.clone(), model: self.model.clone(), api_key: if self.api_key.is_empty() { - crate::provider::resolve_api_key( + crate::provider::resolve_api_key_or_warn( self.model_config .as_ref() .map(|m| m.provider.as_str()) .unwrap_or("anthropic"), ) - .unwrap_or_default() } else { self.api_key.clone() }, diff --git a/tests/agent_loop_test.rs b/tests/agent_loop_test.rs index 518fa5a..c77065b 100644 --- a/tests/agent_loop_test.rs +++ b/tests/agent_loop_test.rs @@ -2202,3 +2202,139 @@ async fn test_tool_call_with_provider_metadata_executes() { }); assert!(has_tool_result, "Tool result should contain 'hello'"); } + +// --------------------------------------------------------------------------- +// Budget calibration: measured overhead (system prompt, tool schemas, estimate +// shortfall) is subtracted from the compaction budget once real usage arrives +// --------------------------------------------------------------------------- + +/// Records the ContextConfig each compact call receives; never modifies +/// messages, so the loop's behavior is otherwise unaffected. +struct RecordingCompaction { + calls: std::sync::Mutex>, // (max_context_tokens, system_prompt_tokens) +} + +impl yoagent::CompactionStrategy for RecordingCompaction { + fn compact( + &self, + messages: Vec, + config: &yoagent::context::ContextConfig, + ) -> Vec { + self.calls + .lock() + .unwrap() + .push((config.max_context_tokens, config.system_prompt_tokens)); + messages + } +} + +fn calibration_config( + provider: std::sync::Arc, + strategy: std::sync::Arc, + max_context_tokens: usize, +) -> AgentLoopConfig { + AgentLoopConfig { + provider, + model: "usage-test".into(), + api_key: "test".into(), + thinking_level: ThinkingLevel::Off, + max_tokens: None, + temperature: None, + model_config: None, + convert_to_llm: None, + transform_context: None, + get_steering_messages: None, + get_follow_up_messages: Some(Box::new(|| { + vec![AgentMessage::Llm(Message::user("follow up"))] + })), + context_config: Some(yoagent::context::ContextConfig { + max_context_tokens, + system_prompt_tokens: 500, + keep_recent: 1, + keep_first: 1, + tool_output_max_lines: 10, + }), + compaction_strategy: Some(strategy), + execution_limits: Some(ExecutionLimits { + max_turns: 2, + max_total_tokens: 1_000_000, + max_duration: std::time::Duration::from_secs(60), + }), + cache_config: CacheConfig::default(), + tool_execution: ToolExecutionStrategy::default(), + retry_config: yoagent::RetryConfig::none(), + before_turn: None, + after_turn: None, + on_error: None, + input_filters: vec![], + turn_delay: None, + } +} + +async fn run_calibration_loop(max_context_tokens: usize) -> Vec<(usize, usize)> { + // Real usage (5010 tokens) dwarfs the char-based estimate of the tiny + // messages, so the measured overhead is ~5000 tokens. + let provider = std::sync::Arc::new(UsageProvider { + usage: Usage { + input: 5000, + output: 10, + cache_read: 0, + cache_write: 0, + total_tokens: 5010, + }, + calls: std::sync::atomic::AtomicUsize::new(0), + }); + let strategy = std::sync::Arc::new(RecordingCompaction { + calls: std::sync::Mutex::new(Vec::new()), + }); + let config = calibration_config(provider, strategy.clone(), max_context_tokens); + + let mut context = AgentContext { + system_prompt: "test".into(), + messages: Vec::new(), + tools: Vec::new(), + }; + let (tx, _rx) = mpsc::unbounded_channel(); + agent_loop( + vec![AgentMessage::Llm(Message::user("start"))], + &mut context, + &config, + tx, + CancellationToken::new(), + ) + .await; + + let calls = strategy.calls.lock().unwrap().clone(); + calls +} + +#[tokio::test] +async fn test_calibration_subtracts_measured_overhead() { + let calls = run_calibration_loop(8000).await; + assert!(calls.len() >= 2, "expected 2 turns, got {:?}", calls); + + // Turn 1: no real usage yet — config passes through unchanged. + assert_eq!(calls[0], (8000, 500)); + + // Turn 2: usage anchored at ~5010 vs a tiny estimate, so ~5000 tokens of + // overhead are subtracted and the static reserve is zeroed (the measured + // overhead already includes the real system prompt). + let (max, reserve) = calls[1]; + assert_eq!(reserve, 0); + assert!( + (2500..=3500).contains(&max), + "expected ~3000 calibrated budget, got {}", + max + ); +} + +#[tokio::test] +async fn test_calibration_floor_prevents_budget_collapse() { + // Overhead (~5000) exceeds the whole budget (4000). Without the floor the + // budget would hit 0 and compaction would wipe the conversation; the + // calibrated budget must never drop below 10% of the configured one. + let calls = run_calibration_loop(4000).await; + assert!(calls.len() >= 2, "expected 2 turns, got {:?}", calls); + assert_eq!(calls[0], (4000, 500)); + assert_eq!(calls[1], (400, 0)); +} diff --git a/tests/agent_test.rs b/tests/agent_test.rs index 8fbd6d3..9b574fa 100644 --- a/tests/agent_test.rs +++ b/tests/agent_test.rs @@ -439,3 +439,157 @@ async fn test_queue_inspection_and_take() { assert_eq!(agent.follow_up_queue_len(), 0); assert!(agent.follow_up_queue_snapshot().is_empty()); } + +// --------------------------------------------------------------------------- +// session_cost_usd +// --------------------------------------------------------------------------- + +fn assistant_with_usage(usage: Usage) -> AgentMessage { + AgentMessage::Llm(Message::assistant( + vec![Content::Text { text: "ok".into() }], + StopReason::Stop, + "id", + "model", + usage, + )) +} + +fn some_usage() -> Usage { + Usage { + input: 1_000_000, + output: 500_000, + cache_read: 0, + cache_write: 0, + total_tokens: 1_500_000, + } +} + +#[test] +fn session_cost_usd_none_without_model_config() { + let agent = + Agent::new(MockProvider::text("x")).with_messages(vec![assistant_with_usage(some_usage())]); + assert_eq!(agent.session_cost_usd(), None); +} + +#[test] +fn session_cost_usd_none_when_rates_unconfigured() { + // ModelConfig::custom has all-zero cost rates: pricing is unknown, so + // the answer must be None ("can't price"), not Some(0.0) ("free"). + let mc = yoagent::provider::ModelConfig::custom( + yoagent::provider::ApiProtocol::OpenAiCompletions, + "local", + "http://localhost:8080/v1", + "m", + "M", + ); + let agent = Agent::new(MockProvider::text("x")) + .with_model_config(mc) + .with_messages(vec![assistant_with_usage(some_usage())]); + assert_eq!(agent.session_cost_usd(), None); +} + +#[test] +fn session_cost_usd_sums_assistant_turns_only() { + let mut mc = yoagent::provider::ModelConfig::custom( + yoagent::provider::ApiProtocol::OpenAiCompletions, + "local", + "http://localhost:8080/v1", + "m", + "M", + ); + mc.cost.input_per_million = 3.0; + mc.cost.output_per_million = 15.0; + let expected_per_turn = mc.cost.cost_usd(&some_usage()); + + let agent = Agent::new(MockProvider::text("x")) + .with_model_config(mc) + .with_messages(vec![ + AgentMessage::Llm(Message::user("hi")), + assistant_with_usage(some_usage()), + assistant_with_usage(some_usage()), + ]); + let total = agent.session_cost_usd().expect("rates are configured"); + assert!(total > 0.0); + assert!((total - 2.0 * expected_per_turn).abs() < 1e-9); +} + +// --------------------------------------------------------------------------- +// API-key resolution: explicit key wins; env var is the fallback +// --------------------------------------------------------------------------- + +/// Records the api_key each stream call receives. +struct KeyCapturingProvider { + captured: Arc>, +} + +#[async_trait::async_trait] +impl yoagent::provider::StreamProvider for KeyCapturingProvider { + async fn stream( + &self, + config: yoagent::provider::StreamConfig, + tx: mpsc::UnboundedSender, + _cancel: tokio_util::sync::CancellationToken, + ) -> Result { + *self.captured.lock().unwrap() = config.api_key.clone(); + let msg = Message::assistant( + vec![Content::Text { text: "ok".into() }], + 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) + } +} + +/// Run one prompt to completion so the provider records the resolved key. +async fn run_one_prompt(agent: &mut Agent) { + let mut rx = agent.prompt("hi").await; + while rx.recv().await.is_some() {} + agent.finish().await; +} + +#[tokio::test] +async fn test_explicit_api_key_wins_over_env() { + // Each key-resolution test uses its own env var to stay race-free under + // parallel test execution. + std::env::set_var("ZAI_API_KEY", "env-key-should-lose"); + let captured = Arc::new(std::sync::Mutex::new(String::new())); + let mut agent = Agent::new(KeyCapturingProvider { + captured: captured.clone(), + }) + .with_model("m") + .with_api_key("explicit-key") + .with_model_config(yoagent::provider::ModelConfig::custom( + yoagent::provider::ApiProtocol::OpenAiCompletions, + "zai", + "http://localhost:8080/v1", + "m", + "M", + )); + run_one_prompt(&mut agent).await; + assert_eq!(*captured.lock().unwrap(), "explicit-key"); +} + +#[tokio::test] +async fn test_env_var_fallback_resolves_api_key() { + std::env::set_var("CEREBRAS_API_KEY", "cerebras-env-key"); + let captured = Arc::new(std::sync::Mutex::new(String::new())); + let mut agent = Agent::new(KeyCapturingProvider { + captured: captured.clone(), + }) + .with_model("m") + .with_model_config(yoagent::provider::ModelConfig::custom( + yoagent::provider::ApiProtocol::OpenAiCompletions, + "cerebras", + "http://localhost:8080/v1", + "m", + "M", + )); + run_one_prompt(&mut agent).await; + assert_eq!(*captured.lock().unwrap(), "cerebras-env-key"); +} diff --git a/tests/anthropic_stream_test.rs b/tests/anthropic_stream_test.rs index 2f3eeba..ac468b5 100644 --- a/tests/anthropic_stream_test.rs +++ b/tests/anthropic_stream_test.rs @@ -180,3 +180,31 @@ async fn user_authorization_header_suppresses_x_api_key() { } run_stream(config).await.expect("stream should succeed"); } + +#[tokio::test] +async fn rate_limit_carries_retry_after_from_header() { + let server = MockServer::start().await; + + Mock::given(method("POST")) + .and(path("/messages")) + .respond_with( + ResponseTemplate::new(429) + .insert_header("retry-after", "7") + .set_body_string( + r#"{"type":"error","error":{"type":"rate_limit_error","message":"rate limited"}}"#, + ), + ) + .mount(&server) + .await; + + let err = run_stream(stream_config(&server.uri(), None)) + .await + .expect_err("429 must surface as an error"); + + match err { + yoagent::provider::ProviderError::RateLimited { retry_after_ms } => { + assert_eq!(retry_after_ms, Some(7000)); + } + other => panic!("expected RateLimited, got: {:?}", other), + } +} diff --git a/tests/google_stream_test.rs b/tests/google_stream_test.rs index 39affc9..17615b8 100644 --- a/tests/google_stream_test.rs +++ b/tests/google_stream_test.rs @@ -314,3 +314,34 @@ async fn safety_finish_reason_maps_to_refusal() { "error_message should explain the block, got {error_message:?}" ); } + +/// Gemini's promptTokenCount INCLUDES cachedContentTokenCount; the mapping +/// must keep `input` as the uncached remainder so `input + cache_read` +/// doesn't double-count cached tokens. +#[tokio::test] +async fn cached_tokens_are_not_double_counted_in_usage() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path(format!( + "/v1beta/models/{}:streamGenerateContent", + MODEL + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + sse(&[ + r#"{"candidates":[{"content":{"parts":[{"text":"hi"}],"role":"model"},"finishReason":"STOP","index":0}],"usageMetadata":{"promptTokenCount":100,"cachedContentTokenCount":80,"candidatesTokenCount":5,"totalTokenCount":105}}"#, + ]), + "text/event-stream", + )) + .mount(&server) + .await; + + let message = run_stream(stream_config(&server.uri(), vec![Message::user("hi")])).await; + + let Message::Assistant { usage, .. } = &message else { + panic!("expected assistant message"); + }; + assert_eq!(usage.input, 20, "input must exclude cached tokens"); + assert_eq!(usage.cache_read, 80); + assert_eq!(usage.output, 5); + assert_eq!(usage.input + usage.cache_read + usage.output, 105); +} diff --git a/tests/sub_agent_test.rs b/tests/sub_agent_test.rs index 9785b50..3c69d0b 100644 --- a/tests/sub_agent_test.rs +++ b/tests/sub_agent_test.rs @@ -742,3 +742,95 @@ async fn test_sub_agent_in_parent_loop() { assert!(has_tool_start); assert!(has_tool_end); } + +// --------------------------------------------------------------------------- +// Config plumbing: temperature and env-var key fallback reach the provider +// --------------------------------------------------------------------------- + +/// Records the api_key and temperature each stream call receives. +struct StreamConfigCapture { + captured: Arc)>>, +} + +#[async_trait::async_trait] +impl yoagent::provider::StreamProvider for StreamConfigCapture { + async fn stream( + &self, + config: yoagent::provider::StreamConfig, + tx: mpsc::UnboundedSender, + _cancel: CancellationToken, + ) -> Result { + *self.captured.lock().unwrap() = (config.api_key.clone(), config.temperature); + let msg = Message::assistant( + vec![Content::Text { + text: "done".into(), + }], + 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) + } +} + +async fn run_sub_agent(tool: &SubAgentTool) { + tool.execute( + serde_json::json!({"task": "go"}), + ToolContext { + tool_call_id: "tc-cfg".into(), + tool_name: "cfg".into(), + cancel: CancellationToken::new(), + on_update: None, + on_progress: None, + }, + ) + .await + .expect("sub-agent should succeed"); +} + +#[tokio::test] +async fn test_sub_agent_temperature_reaches_provider() { + let captured = Arc::new(std::sync::Mutex::new((String::new(), None))); + let tool = SubAgentTool::new( + "cfg", + Arc::new(StreamConfigCapture { + captured: captured.clone(), + }), + ) + .with_model("mock") + .with_api_key("k") + .with_temperature(0.3); + + run_sub_agent(&tool).await; + assert_eq!(captured.lock().unwrap().1, Some(0.3)); +} + +#[tokio::test] +async fn test_sub_agent_env_key_fallback() { + // Own env var (not shared with other tests) to stay race-free under + // parallel execution. + std::env::set_var("MINIMAX_API_KEY", "minimax-env-key"); + let captured = Arc::new(std::sync::Mutex::new((String::new(), None))); + let tool = SubAgentTool::new( + "cfg", + Arc::new(StreamConfigCapture { + captured: captured.clone(), + }), + ) + .with_model("m") + .with_model_config(yoagent::provider::ModelConfig::custom( + yoagent::provider::ApiProtocol::OpenAiCompletions, + "minimax", + "http://localhost:8080/v1", + "m", + "M", + )); + + run_sub_agent(&tool).await; + assert_eq!(captured.lock().unwrap().0, "minimax-env-key"); +}