yoagent 0.25.5

//! Tests for SubAgentTool using MockProvider.

use std::sync::Arc;
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;
use yoagent::agent_loop::{agent_loop, AgentLoopConfig};
use yoagent::provider::mock::*;
use yoagent::provider::MockProvider;
use yoagent::provider::ModelConfig;
use yoagent::sub_agent::SubAgentTool;
use yoagent::*;

fn make_config(provider: MockProvider) -> AgentLoopConfig {
    let mut config = AgentLoopConfig::new(std::sync::Arc::new(provider), "mock");
    config.api_key = "test".into();
    config
}

fn collect_events(mut rx: mpsc::UnboundedReceiver<AgentEvent>) -> Vec<AgentEvent> {
    let mut events = Vec::new();
    while let Ok(e) = rx.try_recv() {
        events.push(e);
    }
    events
}

// ---------------------------------------------------------------------------
// Basic sub-agent execution
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_sub_agent_basic() {
    // The sub-agent's mock provider returns a simple text response
    let sub_provider = Arc::new(MockProvider::text("Research result: Rust is great"));

    let sub_agent = SubAgentTool::from_provider("researcher", sub_provider, ModelConfig::mock())
        .with_description("Researches topics")
        .with_system_prompt("You are a research assistant.");

    // Execute the sub-agent tool directly
    let params = serde_json::json!({"task": "Tell me about Rust"});

    let result = sub_agent
        .execute(params, ToolContext::new("tc-1", "researcher"))
        .await
        .expect("sub-agent should succeed");

    // Should contain the sub-agent's response text
    let text = match &result.content[0] {
        Content::Text { text } => text.as_str(),
        _ => panic!("Expected text content"),
    };
    assert_eq!(text, "Research result: Rust is great");

    // Details should include sub-agent metadata
    assert_eq!(result.details["sub_agent"], "researcher");
}

// ---------------------------------------------------------------------------
// Sub-agent with its own tools
// ---------------------------------------------------------------------------

struct EchoTool;

#[async_trait::async_trait]
impl AgentTool for EchoTool {
    fn name(&self) -> &str {
        "echo"
    }
    fn label(&self) -> &str {
        "Echo"
    }
    fn description(&self) -> &str {
        "Echoes input"
    }
    fn parameters_schema(&self) -> serde_json::Value {
        serde_json::json!({
            "type": "object",
            "properties": {
                "text": {"type": "string"}
            }
        })
    }
    async fn execute(
        &self,
        params: serde_json::Value,
        _ctx: ToolContext,
    ) -> Result<ToolResult, ToolError> {
        let text = params["text"].as_str().unwrap_or("(empty)");
        Ok(ToolResult {
            content: vec![Content::Text {
                text: format!("echoed: {}", text),
            }],
            details: serde_json::Value::Null,
        })
    }
}

#[tokio::test]
async fn test_sub_agent_with_tools() {
    // Sub-agent first calls the echo tool, then responds with text
    let sub_provider = Arc::new(MockProvider::new(vec![
        MockResponse::ToolCalls(vec![MockToolCall {
            provider_metadata: None,
            name: "echo".into(),
            arguments: serde_json::json!({"text": "hello"}),
        }]),
        MockResponse::Text("The echo returned: echoed: hello".into()),
    ]));

    let echo_tool: Arc<dyn AgentTool> = Arc::new(EchoTool);

    let sub_agent = SubAgentTool::from_provider("echo_agent", sub_provider, ModelConfig::mock())
        .with_description("Agent that echoes")
        .with_system_prompt("Use the echo tool.")
        .with_tools(vec![echo_tool]);

    let params = serde_json::json!({"task": "Echo hello"});

    let result = sub_agent
        .execute(params, ToolContext::new("tc-1", "echo_agent"))
        .await
        .expect("sub-agent should succeed");

    let text = match &result.content[0] {
        Content::Text { text } => text.as_str(),
        _ => panic!("Expected text content"),
    };
    assert_eq!(text, "The echo returned: echoed: hello");
}

// ---------------------------------------------------------------------------
// Cancellation propagation
// ---------------------------------------------------------------------------

/// Cancelled before it ran: a cancelled delegation, not an empty success
/// (as a cancel mid-call is, below).
#[tokio::test]
async fn test_sub_agent_cancellation() {
    let sub_provider = Arc::new(MockProvider::text("Should not appear"));

    let sub_agent =
        SubAgentTool::from_provider("cancelled_agent", sub_provider, ModelConfig::mock());

    let cancel = CancellationToken::new();
    cancel.cancel(); // Cancel immediately

    let params = serde_json::json!({"task": "Do something"});

    let result = sub_agent
        .execute(
            params,
            ToolContext::new("tc-1", "cancelled_agent").with_cancel(cancel),
        )
        .await;
    assert!(
        matches!(result, Err(ToolError::Cancelled)),
        "a cancelled delegation: {result:?}"
    );
}

#[tokio::test]
async fn test_sub_agent_cancelled_mid_call_fails_the_delegation() {
    // The child's provider call is cancelled: its run ends as `Aborted` with
    // empty text. That is a failed delegation, not an empty success.
    struct CancelledProvider;

    #[async_trait::async_trait]
    impl yoagent::provider::StreamProvider for CancelledProvider {
        async fn stream(
            &self,
            _config: yoagent::provider::StreamConfig,
            _tx: tokio::sync::mpsc::UnboundedSender<yoagent::provider::StreamEvent>,
            _cancel: CancellationToken,
        ) -> Result<Message, yoagent::provider::ProviderError> {
            Err(yoagent::provider::ProviderError::Cancelled)
        }
    }

    let sub_agent = SubAgentTool::from_provider(
        "cancelled_agent",
        Arc::new(CancelledProvider),
        ModelConfig::mock(),
    );
    let err = sub_agent
        .execute(
            serde_json::json!({"task": "Do something"}),
            ToolContext::new("tc-1", "cancelled_agent"),
        )
        .await
        .expect_err("a cancelled delegation must fail");
    let text = err.to_string();
    assert!(
        text.contains("cancelled_agent") && text.contains("Cancelled"),
        "got: {text}"
    );
}

// ---------------------------------------------------------------------------
// Max turns limit
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_sub_agent_max_turns() {
    // Sub-agent keeps calling tools indefinitely — max_turns should stop it.
    // With max_turns=1, the sub-agent gets 1 LLM call.
    // Response 1: tool call → executes tool → hits turn limit → returns limit message
    // The sub-agent won't get a second LLM call to produce text.
    let sub_provider = Arc::new(MockProvider::new(vec![
        MockResponse::ToolCalls(vec![MockToolCall {
            provider_metadata: None,
            name: "echo".into(),
            arguments: serde_json::json!({"text": "loop"}),
        }]),
        // This response won't be reached due to turn limit
        MockResponse::Text("Should not reach".into()),
    ]));

    let echo_tool: Arc<dyn AgentTool> = Arc::new(EchoTool);

    let sub_agent = SubAgentTool::from_provider("limited_agent", sub_provider, ModelConfig::mock())
        .with_tools(vec![echo_tool])
        .with_max_turns(1); // Only 1 turn allowed

    let params = serde_json::json!({"task": "Keep going"});

    let result = sub_agent
        .execute(params, ToolContext::new("tc-1", "limited_agent"))
        .await
        .expect("sub-agent should succeed");

    // The sub-agent was stopped by turn limit — it won't have the second text response
    let text = match &result.content[0] {
        Content::Text { text } => text.as_str(),
        _ => panic!("Expected text content"),
    };
    // Should NOT contain the text from the second response
    assert_ne!(text, "Should not reach");
}

// ---------------------------------------------------------------------------
// Parallel sub-agent execution (via parent agent loop)
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_sub_agent_parallel() {
    // Two sub-agents that each take ~50ms, run in parallel via the parent loop.
    // The parent's mock emits both sub-agent tool calls, then a final text.

    struct SlowProvider {
        delay_ms: u64,
        text: String,
    }

    #[async_trait::async_trait]
    impl yoagent::provider::StreamProvider for SlowProvider {
        async fn stream(
            &self,
            _config: yoagent::provider::StreamConfig,
            tx: tokio::sync::mpsc::UnboundedSender<yoagent::provider::StreamEvent>,
            cancel: tokio_util::sync::CancellationToken,
        ) -> Result<Message, yoagent::provider::ProviderError> {
            if cancel.is_cancelled() {
                return Err(yoagent::provider::ProviderError::Cancelled);
            }
            tokio::time::sleep(std::time::Duration::from_millis(self.delay_ms)).await;

            let _ = tx.send(yoagent::provider::StreamEvent::Start);
            let _ = tx.send(yoagent::provider::StreamEvent::TextDelta {
                content_index: 0,
                delta: self.text.clone(),
            });
            let msg = Message::assistant(
                vec![Content::Text {
                    text: self.text.clone(),
                }],
                StopReason::Stop,
                "slow",
                "slow",
                Usage::default(),
            );
            let _ = tx.send(yoagent::provider::StreamEvent::Done {
                message: msg.clone(),
            });
            Ok(msg)
        }
    }

    let sub_a = SubAgentTool::from_provider(
        "agent_a",
        Arc::new(SlowProvider {
            delay_ms: 50,
            text: "Result A".into(),
        }),
        ModelConfig::mock(),
    );

    let sub_b = SubAgentTool::from_provider(
        "agent_b",
        Arc::new(SlowProvider {
            delay_ms: 50,
            text: "Result B".into(),
        }),
        ModelConfig::mock(),
    );

    // Parent provider: first call triggers both sub-agents, second returns final text
    let parent_provider = MockProvider::new(vec![
        MockResponse::ToolCalls(vec![
            MockToolCall {
                provider_metadata: None,
                name: "agent_a".into(),
                arguments: serde_json::json!({"task": "Do A"}),
            },
            MockToolCall {
                provider_metadata: None,
                name: "agent_b".into(),
                arguments: serde_json::json!({"task": "Do B"}),
            },
        ]),
        MockResponse::Text("Both sub-agents completed.".into()),
    ]);

    let config = make_config(parent_provider);

    let mut context = AgentContext {
        system_prompt: "You are a coordinator.".into(),
        messages: Vec::new(),
        tools: vec![Box::new(sub_a), Box::new(sub_b)],
    };

    let prompt = AgentMessage::Llm(Message::user("Run both agents"));
    let (tx, rx) = mpsc::unbounded_channel();
    let cancel = CancellationToken::new();

    let start = std::time::Instant::now();
    let new_messages = agent_loop(vec![prompt], &mut context, &config, tx, cancel).await;
    let elapsed = start.elapsed();

    let _events = collect_events(rx);

    // Both tool results should be present
    let tool_results: Vec<_> = new_messages
        .iter()
        .filter(|m| m.role() == "toolResult")
        .collect();
    assert_eq!(tool_results.len(), 2);

    // Should complete in roughly 50-100ms (parallel), not 100ms+ (sequential)
    assert!(
        elapsed.as_millis() < 130,
        "Parallel sub-agents took {}ms, expected <130ms",
        elapsed.as_millis()
    );
}

// ---------------------------------------------------------------------------
// Event forwarding via on_update
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_sub_agent_event_forwarding() {
    let sub_provider = Arc::new(MockProvider::text("Sub-agent done"));

    let sub_agent =
        SubAgentTool::from_provider("streaming_agent", sub_provider, ModelConfig::mock());

    let params = serde_json::json!({"task": "Do work"});

    // Collect on_update calls
    let updates: Arc<std::sync::Mutex<Vec<String>>> = Arc::new(std::sync::Mutex::new(Vec::new()));
    let updates_clone = updates.clone();
    let on_update: ToolUpdateFn = Arc::new(move |result: ToolResult| {
        if let Some(Content::Text { text }) = result.content.first() {
            updates_clone.lock().unwrap().push(text.clone());
        }
    });

    let result = sub_agent
        .execute(
            params,
            ToolContext::new("tc-1", "streaming_agent").with_on_update(on_update),
        )
        .await
        .expect("sub-agent should succeed");

    // Final result should contain the sub-agent's text
    let text = match &result.content[0] {
        Content::Text { text } => text.as_str(),
        _ => panic!("Expected text content"),
    };
    assert_eq!(text, "Sub-agent done");

    // on_update should have received streaming deltas
    let collected = updates.lock().unwrap();
    assert!(
        !collected.is_empty(),
        "Expected on_update to receive streaming events"
    );
    // Should contain the text delta from the mock provider
    assert!(
        collected.iter().any(|t| t.contains("Sub-agent done")),
        "Expected text delta in updates, got: {:?}",
        *collected
    );
}

#[tokio::test]
async fn test_sub_agent_marks_a_retried_attempt_in_updates() {
    // The child's first attempt streams partial text and fails with a
    // retryable error. The parent's updates must say where the retry starts,
    // so the partial text is not read as the start of the answer.
    struct FlakyOnce {
        attempts: std::sync::atomic::AtomicUsize,
    }

    #[async_trait::async_trait]
    impl yoagent::provider::StreamProvider for FlakyOnce {
        async fn stream(
            &self,
            config: yoagent::provider::StreamConfig,
            tx: tokio::sync::mpsc::UnboundedSender<yoagent::provider::StreamEvent>,
            cancel: CancellationToken,
        ) -> Result<Message, yoagent::provider::ProviderError> {
            use std::sync::atomic::Ordering;
            if self.attempts.fetch_add(1, Ordering::SeqCst) == 0 {
                let _ = tx.send(yoagent::provider::StreamEvent::Start);
                let _ = tx.send(yoagent::provider::StreamEvent::TextDelta {
                    content_index: 0,
                    delta: "partial".into(),
                });
                return Err(yoagent::provider::ProviderError::Network(
                    "stream cut".into(),
                ));
            }
            MockProvider::text("recovered")
                .stream(config, tx, cancel)
                .await
        }
    }

    let sub_agent = SubAgentTool::from_provider(
        "flaky_agent",
        Arc::new(FlakyOnce {
            attempts: Default::default(),
        }),
        ModelConfig::mock(),
    )
    .with_retry_config(yoagent::RetryConfig {
        max_retries: 1,
        initial_delay_ms: 1,
        backoff_multiplier: 1.0,
        max_delay_ms: 1,
    });

    let updates: Arc<std::sync::Mutex<Vec<String>>> = Arc::default();
    let sink = updates.clone();
    let on_update: ToolUpdateFn = Arc::new(move |result: ToolResult| {
        if let Some(Content::Text { text }) = result.content.first() {
            sink.lock().unwrap().push(text.clone());
        }
    });

    let result = sub_agent
        .execute(
            serde_json::json!({"task": "Do work"}),
            ToolContext::new("tc-1", "flaky_agent").with_on_update(on_update),
        )
        .await
        .expect("the retry succeeds");
    assert!(matches!(&result.content[0], Content::Text { text } if text == "recovered"));

    let collected = updates.lock().unwrap().clone();
    assert_eq!(
        collected,
        [
            "partial",
            "[sub-agent retrying after: Network error: stream cut]",
            "recovered"
        ],
    );
}

// ---------------------------------------------------------------------------
// Invalid parameters
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_sub_agent_missing_task_parameter() {
    let sub_provider = Arc::new(MockProvider::text("Should not run"));

    let sub_agent = SubAgentTool::from_provider("test_agent", sub_provider, ModelConfig::mock());

    let params = serde_json::json!({}); // Missing "task"

    let result = sub_agent
        .execute(params, ToolContext::new("tc-1", "test_agent"))
        .await;
    assert!(result.is_err());

    match result.unwrap_err() {
        ToolError::InvalidArgs(msg) => assert!(msg.contains("task")),
        other => panic!("Expected InvalidArgs, got: {:?}", other),
    }
}

// ---------------------------------------------------------------------------
// Skills: with_skills injects the skills index into the sub-agent system prompt
// ---------------------------------------------------------------------------

/// Provider that records the system prompt it is dispatched with, so tests can
/// assert on the exact prompt the sub-agent assembles.
struct CapturingProvider {
    captured: Arc<std::sync::Mutex<String>>,
}

#[async_trait::async_trait]
impl yoagent::provider::StreamProvider for CapturingProvider {
    async fn stream(
        &self,
        config: yoagent::provider::StreamConfig,
        tx: mpsc::UnboundedSender<yoagent::provider::StreamEvent>,
        _cancel: CancellationToken,
    ) -> Result<Message, yoagent::provider::ProviderError> {
        *self.captured.lock().unwrap() = config.system_prompt.clone();
        let _ = tx.send(yoagent::provider::StreamEvent::Start);
        let msg = Message::assistant(
            vec![Content::Text {
                text: "done".into(),
            }],
            StopReason::Stop,
            "mock",
            "mock",
            Usage::default(),
        );
        let _ = tx.send(yoagent::provider::StreamEvent::Done {
            message: msg.clone(),
        });
        Ok(msg)
    }
}

/// RAII guard for a per-test temp skills directory. Holds a unique path
/// (avoids collisions under parallel `cargo test`) and removes it on drop,
/// so cleanup runs even if the test panics.
struct SkillsDir(std::path::PathBuf);

impl SkillsDir {
    /// Create a temp dir containing a single `<name>/SKILL.md`. `unique` must
    /// differ per test to avoid concurrent collisions on the shared temp dir.
    fn with_one_skill(unique: &str, name: &str, description: &str) -> Self {
        let dir = std::env::temp_dir().join(format!("yoagent-test-skills-{unique}"));
        let _ = std::fs::remove_dir_all(&dir);
        let skill_dir = dir.join(name);
        std::fs::create_dir_all(&skill_dir).unwrap();
        std::fs::write(
            skill_dir.join("SKILL.md"),
            format!("---\nname: {name}\ndescription: {description}\n---\n\nBody.\n"),
        )
        .unwrap();
        Self(dir)
    }

    fn load(&self) -> yoagent::skills::SkillSet {
        yoagent::skills::SkillSet::load(&[self.0.to_string_lossy().to_string()]).unwrap()
    }
}

impl Drop for SkillsDir {
    fn drop(&mut self) {
        let _ = std::fs::remove_dir_all(&self.0);
    }
}

/// Dispatch a sub-agent through a `CapturingProvider` and return the system
/// prompt the provider was called with.
async fn capture_system_prompt(
    build: impl FnOnce(Arc<CapturingProvider>) -> SubAgentTool,
) -> String {
    let captured = Arc::new(std::sync::Mutex::new(String::new()));
    let provider = Arc::new(CapturingProvider {
        captured: captured.clone(),
    });
    let sub_agent = build(provider);

    sub_agent
        .execute(
            serde_json::json!({"task": "do work"}),
            ToolContext::new("tc-1", "sub"),
        )
        .await
        .expect("sub-agent should succeed");

    let prompt = captured.lock().unwrap().clone();
    prompt
}

#[tokio::test]
async fn test_sub_agent_with_skills() {
    let skills_dir = SkillsDir::with_one_skill(
        "with-skills",
        "research",
        "How to call the search and read APIs",
    );
    let skills = skills_dir.load();
    assert_eq!(skills.len(), 1, "expected the research skill to load");

    let prompt = capture_system_prompt(|provider| {
        SubAgentTool::from_provider("researcher", provider, ModelConfig::mock())
            .with_system_prompt("You are a research assistant.")
            .with_skills(skills)
    })
    .await;

    // Base system prompt is preserved...
    assert!(
        prompt.contains("You are a research assistant."),
        "base system prompt missing, got: {prompt}"
    );
    // ...and the skills index is appended.
    assert!(
        prompt.contains("<available_skills>") && prompt.contains("<name>research</name>"),
        "skills index not injected into sub-agent system prompt, got: {prompt}"
    );
}

#[tokio::test]
async fn test_sub_agent_with_skills_empty_base_prompt() {
    // Exercises the `system_prompt.is_empty()` branch: skills become the entire
    // prompt with no leading blank line. assert_eq pins the exact output.
    let skills_dir = SkillsDir::with_one_skill("empty-base", "research", "desc");
    let skills = skills_dir.load();
    let expected = skills.format_for_prompt();
    assert!(!expected.is_empty());

    let prompt = capture_system_prompt(|provider| {
        // No with_system_prompt() call — base prompt is empty.
        SubAgentTool::from_provider("researcher", provider, ModelConfig::mock()).with_skills(skills)
    })
    .await;

    assert_eq!(
        prompt, expected,
        "with empty base prompt, the skills index should be the whole prompt verbatim"
    );
}

#[tokio::test]
async fn test_sub_agent_with_empty_skillset_is_noop() {
    // An empty SkillSet must not alter the system prompt (no trailing "\n\n").
    let prompt = capture_system_prompt(|provider| {
        SubAgentTool::from_provider("researcher", provider, ModelConfig::mock())
            .with_system_prompt("Base prompt.")
            .with_skills(yoagent::skills::SkillSet::empty())
    })
    .await;

    assert_eq!(prompt, "Base prompt.", "empty SkillSet should be a no-op");
}

#[tokio::test]
async fn test_sub_agent_skills_before_shared_state() {
    // Skills and shared-state both append to the prompt; lock in the order
    // base -> skills -> shared-state.
    let skills_dir = SkillsDir::with_one_skill("ordering", "research", "desc");
    let skills = skills_dir.load();
    let state = SharedState::new();

    let prompt = capture_system_prompt(|provider| {
        SubAgentTool::from_provider("researcher", provider, ModelConfig::mock())
            .with_system_prompt("Base prompt.")
            .with_skills(skills)
            .with_shared_state(state)
    })
    .await;

    let skills_at = prompt
        .find("<available_skills>")
        .expect("skills index present");
    let shared_at = prompt
        .find("## Shared State")
        .expect("shared-state block present");
    assert!(
        skills_at < shared_at,
        "skills index should precede the shared-state block, got: {prompt}"
    );
}

// ---------------------------------------------------------------------------
// Integration: sub-agent tool in a parent agent loop
// ---------------------------------------------------------------------------

#[tokio::test]
async fn test_sub_agent_in_parent_loop() {
    // Parent calls sub-agent, sub-agent returns text, parent summarizes
    let sub_provider = Arc::new(MockProvider::text("42 is the answer"));

    let sub_agent = SubAgentTool::from_provider("calculator", sub_provider, ModelConfig::mock())
        .with_description("Calculates things");

    let parent_provider = MockProvider::new(vec![
        MockResponse::ToolCalls(vec![MockToolCall {
            provider_metadata: None,
            name: "calculator".into(),
            arguments: serde_json::json!({"task": "What is 6*7?"}),
        }]),
        MockResponse::Text("The calculator says: 42 is the answer".into()),
    ]);

    let config = make_config(parent_provider);

    let mut context = AgentContext {
        system_prompt: "You are a coordinator.".into(),
        messages: Vec::new(),
        tools: vec![Box::new(sub_agent)],
    };

    let prompt = AgentMessage::Llm(Message::user("What is 6*7?"));
    let (tx, rx) = mpsc::unbounded_channel();
    let cancel = CancellationToken::new();

    let new_messages = agent_loop(vec![prompt], &mut context, &config, tx, cancel).await;

    let events = collect_events(rx);

    // Should have: user, assistant(tool_call), toolResult, assistant(text)
    assert_eq!(new_messages.len(), 4);
    assert_eq!(new_messages[0].role(), "user");
    assert_eq!(new_messages[1].role(), "assistant");
    assert_eq!(new_messages[2].role(), "toolResult");
    assert_eq!(new_messages[3].role(), "assistant");

    // Tool result should contain sub-agent's output
    if let AgentMessage::Llm(Message::ToolResult { content, .. }) = &new_messages[2] {
        let text = match &content[0] {
            Content::Text { text } => text.as_str(),
            _ => panic!("Expected text content"),
        };
        assert_eq!(text, "42 is the answer");
    } else {
        panic!("Expected tool result message");
    }

    // Should have tool execution events
    let has_tool_start = events
        .iter()
        .any(|e| matches!(e, AgentEvent::ToolExecutionStart { tool_name, .. } if tool_name == "calculator"));
    let has_tool_end = events
        .iter()
        .any(|e| matches!(e, AgentEvent::ToolExecutionEnd { tool_name, .. } if tool_name == "calculator"));
    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<std::sync::Mutex<(String, Option<f32>)>>,
}

#[async_trait::async_trait]
impl yoagent::provider::StreamProvider for StreamConfigCapture {
    async fn stream(
        &self,
        config: yoagent::provider::StreamConfig,
        tx: mpsc::UnboundedSender<yoagent::provider::StreamEvent>,
        _cancel: CancellationToken,
    ) -> Result<Message, yoagent::provider::ProviderError> {
        *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::new("tc-cfg", "cfg"),
    )
    .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::from_provider(
        "cfg",
        Arc::new(StreamConfigCapture {
            captured: captured.clone(),
        }),
        ModelConfig::mock(),
    )
    .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::from_provider(
        "cfg",
        Arc::new(StreamConfigCapture {
            captured: captured.clone(),
        }),
        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");
}

#[tokio::test]
async fn test_sub_agent_from_provider_construction() {
    // from_provider + ModelConfig::mock() mirrors Agent's construction path.
    let tool = SubAgentTool::from_provider(
        "researcher",
        Arc::new(MockProvider::text("Research result")),
        yoagent::provider::ModelConfig::mock(),
    )
    .with_description("Researches topics");

    let result = tool
        .execute(
            serde_json::json!({"task": "go"}),
            ToolContext::new("tc-fp", "researcher"),
        )
        .await
        .expect("sub-agent should succeed");
    let text = match &result.content[0] {
        Content::Text { text } => text,
        other => panic!("expected text, got {other:?}"),
    };
    assert!(text.contains("Research result"));
}

#[test]
fn test_sub_agent_from_config_wires_model() {
    // from_config selects a built-in provider from config.api and sets the id.
    let tool = SubAgentTool::from_config(
        "analyst",
        yoagent::provider::ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5"),
    );
    assert_eq!(tool.name(), "analyst");
}

#[test]
fn test_sub_agent_from_config_with_errors_on_empty_registry() {
    let registry = yoagent::provider::ProviderRegistry::new();
    let err = match SubAgentTool::from_config_with(
        &registry,
        "analyst",
        yoagent::provider::ModelConfig::anthropic("claude-sonnet-5", "Sonnet 5"),
    ) {
        Ok(_) => panic!("empty registry must fail"),
        Err(e) => e,
    };
    assert_eq!(
        err,
        yoagent::AgentBuildError::NoProviderForProtocol(
            yoagent::provider::ApiProtocol::AnthropicMessages
        )
    );
}

// ---------------------------------------------------------------------------
// Tool middleware wiring: sub-agent's own tool calls are gated
// ---------------------------------------------------------------------------

struct DenyAll;

#[async_trait::async_trait]
impl yoagent::ToolMiddleware for DenyAll {
    async fn before_tool(&self, _call: &yoagent::ToolCallRequest<'_>) -> yoagent::ToolDecision {
        yoagent::ToolDecision::Deny("sub-agent policy".into())
    }
}

/// A tool that must never run under DenyAll.
struct MustNotRun {
    ran: Arc<std::sync::Mutex<bool>>,
}

#[async_trait::async_trait]
impl AgentTool for MustNotRun {
    fn name(&self) -> &str {
        "must_not_run"
    }
    fn label(&self) -> &str {
        "Must Not Run"
    }
    fn description(&self) -> &str {
        "test"
    }
    fn parameters_schema(&self) -> serde_json::Value {
        serde_json::json!({"type": "object"})
    }
    async fn execute(
        &self,
        _params: serde_json::Value,
        _ctx: ToolContext,
    ) -> Result<ToolResult, ToolError> {
        *self.ran.lock().unwrap() = true;
        Ok(ToolResult {
            content: vec![Content::Text { text: "ran".into() }],
            details: serde_json::Value::Null,
        })
    }
}

#[tokio::test]
async fn test_sub_agent_tool_middleware_denies() {
    let ran = Arc::new(std::sync::Mutex::new(false));
    let provider = Arc::new(MockProvider::new(vec![
        MockResponse::ToolCalls(vec![MockToolCall {
            provider_metadata: None,
            name: "must_not_run".into(),
            arguments: serde_json::json!({}),
        }]),
        MockResponse::Text("finished".into()),
    ]));

    let tool =
        SubAgentTool::from_provider("gated", provider, yoagent::provider::ModelConfig::mock())
            .with_tools(vec![Arc::new(MustNotRun { ran: ran.clone() })])
            .with_tool_middleware(DenyAll);

    let result = tool
        .execute(
            serde_json::json!({"task": "go"}),
            ToolContext::new("tc-mw", "gated"),
        )
        .await
        .expect("sub-agent completes despite denial");

    assert!(
        !*ran.lock().unwrap(),
        "denied tool must not run in sub-agent"
    );
    // Sub-agent still produced its final text.
    let text = match &result.content[0] {
        Content::Text { text } => text,
        other => panic!("expected text, got {other:?}"),
    };
    assert!(text.contains("finished"));
}

// ---------------------------------------------------------------------------
// Scoped stash: the sink and the tool must agree on scope (issue #134)
// ---------------------------------------------------------------------------

struct BigOutputTool;

#[async_trait::async_trait]
impl AgentTool for BigOutputTool {
    fn name(&self) -> &str {
        "big_output"
    }
    fn label(&self) -> &str {
        "Big"
    }
    fn description(&self) -> &str {
        "emits many lines"
    }
    fn parameters_schema(&self) -> serde_json::Value {
        serde_json::json!({"type": "object"})
    }
    async fn execute(
        &self,
        _p: serde_json::Value,
        _c: ToolContext,
    ) -> Result<ToolResult, ToolError> {
        Ok(ToolResult {
            content: vec![Content::Text {
                text: (0..400)
                    .map(|i| format!("line {i}"))
                    .collect::<Vec<_>>()
                    .join("\n"),
            }],
            details: serde_json::Value::Null,
        })
    }
}

/// A scoped sub-agent stashes under `scope␟key` while its marker names the bare
/// key. That resolves only because the `shared_state` tool it is given carries
/// the *same* scoped handle — an agreement that currently holds by construction
/// and was asserted nowhere. A change to either side would break retrieval
/// silently.
#[tokio::test]
async fn a_scoped_sub_agent_stash_resolves_through_its_own_scoped_tool() {
    let parent_view = SharedState::new();

    let sub_provider = MockProvider::new(vec![
        MockResponse::ToolCalls(vec![MockToolCall {
            provider_metadata: None,
            name: "big_output".into(),
            arguments: serde_json::json!({}),
        }]),
        MockResponse::Text("done".into()),
    ]);

    let sub_agent = SubAgentTool::from_provider(
        "worker",
        std::sync::Arc::new(sub_provider),
        ModelConfig::mock(),
    )
    .with_tools(vec![std::sync::Arc::new(BigOutputTool)])
    .with_scoped_shared_state(parent_view.clone(), "worker-1")
    .with_context_config(yoagent::context::ContextConfig {
        tool_output_max_lines: 20,
        ..Default::default()
    });

    sub_agent
        .execute(
            serde_json::json!({"task": "produce output"}),
            ToolContext::new("tc-1", "worker"),
        )
        .await
        .expect("sub-agent should succeed");

    // The parent's unscoped view sees the entry under the scope prefix, so the
    // stash really did go through the scoped handle.
    let parent_keys = parent_view.keys().await;
    assert_eq!(
        parent_keys.len(),
        1,
        "the sub-agent's truncated output must be stashed, got {parent_keys:?}"
    );
    assert!(
        parent_keys[0].contains("worker-1"),
        "the entry must carry the sub-agent's scope, got {parent_keys:?}"
    );

    // And the sub-agent's own scoped view resolves it by the bare key the
    // marker names — which is the invariant that was never pinned.
    let scoped = parent_view.scoped("worker-1");
    let scoped_keys = scoped.keys().await;
    assert_eq!(
        scoped_keys.len(),
        1,
        "the scoped view must see exactly its own entry, got {scoped_keys:?}"
    );
    let full = scoped
        .get(&scoped_keys[0])
        .await
        .expect("the scoped key the marker names must resolve");
    assert!(
        full.contains("line 200"),
        "retrieval must return the elided middle"
    );
}

/// A provider that records the system prompt it was handed.
struct PromptRecorder {
    prompts: std::sync::Arc<std::sync::Mutex<Vec<String>>>,
    inner: MockProvider,
}

#[async_trait::async_trait]
impl yoagent::provider::StreamProvider for PromptRecorder {
    async fn stream(
        &self,
        config: yoagent::provider::StreamConfig,
        tx: tokio::sync::mpsc::UnboundedSender<yoagent::provider::StreamEvent>,
        cancel: tokio_util::sync::CancellationToken,
    ) -> Result<Message, yoagent::provider::ProviderError> {
        self.prompts
            .lock()
            .unwrap()
            .push(config.system_prompt.clone());
        self.inner.stream(config, tx, cancel).await
    }
}

/// Stashed tool output must not leak into the sub-agent's system prompt.
///
/// This asserts on the prompt the provider actually receives. An earlier
/// version called `prompt_summary()` directly and never inspected a prompt —
/// reverting `sub_agent.rs` to the leaking `summary()` left it green, so the
/// test named for the regression could not see it.
///
/// The prompt embeds a `SharedState` summary computed once per invocation, so
/// a second run would otherwise carry the first run's machine-generated
/// `tool-out-*` keys — a different system prompt per delegation, meaning no
/// sub-agent call can reuse the previous one's cached prefix.
#[tokio::test]
async fn stashed_output_does_not_leak_into_the_sub_agent_system_prompt() {
    let store = SharedState::new();
    let prompts = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));

    let run = |label: &'static str| {
        let store = store.clone();
        let prompts = prompts.clone();
        async move {
            SubAgentTool::from_provider(
                "worker",
                std::sync::Arc::new(PromptRecorder {
                    prompts,
                    inner: MockProvider::new(vec![
                        MockResponse::ToolCalls(vec![MockToolCall {
                            provider_metadata: None,
                            name: "big_output".into(),
                            arguments: serde_json::json!({}),
                        }]),
                        MockResponse::Text("done".into()),
                    ]),
                }),
                ModelConfig::mock(),
            )
            .with_tools(vec![std::sync::Arc::new(BigOutputTool)])
            .with_system_prompt("You are a worker.")
            .with_shared_state(store)
            .with_context_config(yoagent::context::ContextConfig {
                tool_output_max_lines: 20,
                ..Default::default()
            })
            .execute(
                serde_json::json!({ "task": label }),
                ToolContext::new("tc-1", "worker"),
            )
            .await
            .expect("run")
        }
    };

    run("first").await;
    assert!(
        !store.keys().await.is_empty(),
        "precondition: the first run must have stashed something"
    );
    run("second").await;

    let seen = prompts.lock().unwrap().clone();
    assert!(
        seen.len() >= 2,
        "both invocations must have reached the provider"
    );
    for (i, prompt) in seen.iter().enumerate() {
        assert!(
            !prompt.contains("tool-out-"),
            "prompt {i} carries a stash key: {prompt}"
        );
    }

    // The load-bearing property: the prompt is byte-identical across
    // invocations, so a delegation can reuse the previous one's cached prefix.
    let first_of_run_one = &seen[0];
    let first_of_run_two = seen.last().unwrap();
    assert_eq!(
        first_of_run_one, first_of_run_two,
        "the system prompt must not change between invocations"
    );

    // A user-set variable still belongs in the prompt summary, and stashes are
    // still discoverable at runtime through the tool's own listing.
    store.set("findings", "user data".into()).await.unwrap();
    assert!(store.prompt_summary().await.contains("findings"));
    assert!(store.summary().await.contains("tool-out-"));
}

/// A sub-agent stopped by loop detection must not report success.
///
/// Regression: `extract_error` matched only `StopReason::Error`, but a loop
/// abort leaves `ToolUse` on the last assistant message, and
/// `extract_final_text` scans assistant messages only — so it never saw the
/// trailing `[Agent stopped: …]` and fell through to "(sub-agent produced no
/// text output)". The parent got `is_error: false`. A sub-agent that burned its
/// whole budget looping was indistinguishable from one that had nothing to say.
///
/// Loop detection is inherited, not configured here: `SubAgentTool` takes
/// `ExecutionLimits::default()`, so every delegation runs with `Some(3)`.
#[tokio::test]
async fn a_looping_sub_agent_reports_failure_not_empty_success() {
    let responses: Vec<MockResponse> = (0..12)
        .map(|_| {
            MockResponse::ToolCalls(vec![MockToolCall {
                provider_metadata: None,
                name: "echo".into(),
                arguments: serde_json::json!({"text": "same"}),
            }])
        })
        .collect();

    let echo_tool: Arc<dyn AgentTool> = Arc::new(EchoTool);
    let sub_agent = SubAgentTool::from_provider(
        "looper",
        Arc::new(MockProvider::new(responses)),
        ModelConfig::mock(),
    )
    .with_system_prompt("Use the echo tool.")
    .with_tools(vec![echo_tool]);

    let result = sub_agent
        .execute(
            serde_json::json!({"task": "go"}),
            ToolContext::new("tc-1", "looper"),
        )
        .await;

    match result {
        Err(e) => {
            let msg = e.to_string();
            assert!(
                msg.contains("echo") || msg.contains("repeatedly"),
                "the failure must say why the sub-agent stopped, got: {msg}"
            );
        }
        Ok(r) => {
            let text = match &r.content[0] {
                Content::Text { text } => text.clone(),
                _ => String::new(),
            };
            panic!("a loop-aborted sub-agent reported success to its parent: {text:?}");
        }
    }
}

/// A turn-limited sub-agent returns its work, and says the work is partial.
///
/// The two self-stop kinds mean opposite things to a parent: `max_turns` is a
/// bound (keep the output), a loop abort is a failure (there is nothing to
/// keep). Without the notice, half-finished work is indistinguishable from a
/// complete answer and the parent's model treats it as final.
///
/// The limit fires at the top of a turn, so the fixture must call a tool —
/// a text-only reply ends the run normally and never reaches the check.
#[tokio::test]
async fn a_turn_limited_sub_agent_returns_partial_work_and_says_so() {
    let sub_provider = Arc::new(MockProvider::new(vec![
        MockResponse::ToolCalls(vec![MockToolCall {
            provider_metadata: None,
            name: "echo".into(),
            arguments: serde_json::json!({"text": "more"}),
        }]),
        MockResponse::Text("Should not reach".into()),
    ]));

    let echo_tool: Arc<dyn AgentTool> = Arc::new(EchoTool);
    let sub_agent = SubAgentTool::from_provider("limited2", sub_provider, ModelConfig::mock())
        .with_tools(vec![echo_tool])
        .with_max_turns(1);

    let result = sub_agent
        .execute(
            serde_json::json!({"task": "keep going"}),
            ToolContext::new("tc-1", "limited2"),
        )
        .await
        .expect("a turn limit is a bound, not a failure");

    let text = match &result.content[0] {
        Content::Text { text } => text.clone(),
        _ => panic!("expected text"),
    };
    assert!(
        text.contains(yoagent::agent_loop::AGENT_STOPPED_PREFIX),
        "and the parent must be told the answer is partial: {text:?}"
    );
    assert!(
        !text.contains("Should not reach"),
        "the turn limit must actually have stopped it: {text:?}"
    );
}

// ---------------------------------------------------------------------------
// Sub-agent spend reaches the parent (#173)
// ---------------------------------------------------------------------------

fn usage(input: u64, output: u64, cache_read: u64, cache_write: u64) -> Usage {
    Usage {
        input,
        output,
        cache_read,
        cache_write,
        total_tokens: 0,
    }
}

fn delegate(to: &str) -> MockResponse {
    delegate_with_usage(to, Usage::default())
}

fn delegate_with_usage(to: &str, u: Usage) -> MockResponse {
    MockResponse::ToolCallsWithUsage(
        vec![MockToolCall {
            name: to.into(),
            arguments: serde_json::json!({"task": "do it"}),
            provider_metadata: None,
        }],
        u,
    )
}

fn echo_call(u: Usage) -> MockResponse {
    MockResponse::ToolCallsWithUsage(
        vec![MockToolCall {
            name: "echo".into(),
            arguments: serde_json::json!({"text": "x"}),
            provider_metadata: None,
        }],
        u,
    )
}

/// Run a parent loop over `tools` and return its events.
async fn run_parent(provider: MockProvider, tools: Vec<Box<dyn AgentTool>>) -> Vec<AgentEvent> {
    run_parent_with(make_config(provider), tools).await
}

async fn run_parent_with(
    config: AgentLoopConfig,
    tools: Vec<Box<dyn AgentTool>>,
) -> Vec<AgentEvent> {
    let (tx, rx) = mpsc::unbounded_channel();
    let mut context = AgentContext {
        system_prompt: String::new(),
        messages: vec![],
        tools,
    };
    agent_loop(
        vec![AgentMessage::Llm(Message::user("go"))],
        &mut context,
        &config,
        tx,
        CancellationToken::new(),
    )
    .await;
    collect_events(rx)
}

fn end_stats(events: &[AgentEvent]) -> SessionStats {
    events
        .iter()
        .find_map(|e| match e {
            AgentEvent::AgentEnd { stats, .. } => Some(stats.clone()),
            _ => None,
        })
        .expect("AgentEnd must be emitted")
}

fn tool_end(events: &[AgentEvent], name: &str) -> (ToolResult, bool) {
    events
        .iter()
        .find_map(|e| match e {
            AgentEvent::ToolExecutionEnd {
                tool_name,
                result,
                is_error,
                ..
            } if tool_name == name => Some((result.clone(), *is_error)),
            _ => None,
        })
        .unwrap_or_else(|| panic!("no ToolExecutionEnd for {name}"))
}

/// The main #173 test. Each figure is distinct, so a rollup that drops a turn,
/// copies the last one, or merges the child into the parent's own bucket
/// cannot pass.
#[tokio::test]
async fn sub_agent_usage_reaches_the_parent_in_a_separate_bucket() {
    // The child spends over two turns: a tool call, then its answer.
    let child = SubAgentTool::from_provider(
        "researcher",
        Arc::new(MockProvider::new(vec![
            echo_call(usage(100, 10, 1000, 1)),
            MockResponse::TextWithUsage("found it".into(), usage(200, 20, 2000, 2)),
        ])),
        ModelConfig::mock(),
    )
    .with_tools(vec![Arc::new(EchoTool)]);

    let parent = MockProvider::new(vec![
        delegate_with_usage("researcher", usage(1, 2, 3, 4)),
        MockResponse::TextWithUsage("done".into(), usage(5, 6, 7, 8)),
    ]);
    let events = run_parent(parent, vec![Box::new(child)]).await;
    let stats = end_stats(&events);

    // Separate bucket: the parent's own figures are its own two turns only.
    assert_eq!(stats.turns, 2);
    assert_eq!(stats.usage, usage(6, 8, 10, 12));

    // The child's two turns, summed.
    assert_eq!(stats.sub_agents.usage, usage(300, 30, 3000, 3));
    assert_eq!(stats.sub_agents.runs, 1);
    assert_eq!(stats.total_usage(), usage(306, 38, 3010, 15));

    // Per delegation, a streaming consumer reads the same off ToolExecutionEnd.
    let (result, is_error) = tool_end(&events, "researcher");
    assert!(!is_error);
    let child_stats =
        SessionStats::from_sub_agent_result(&result).expect("sub-agent stats in details");
    assert_eq!(child_stats.usage, usage(300, 30, 3000, 3));
    assert_eq!(child_stats.turns, 2);
    assert!(child_stats.sub_agents.is_empty());
}

/// A run with no delegation serializes exactly as before — the bucket is
/// omitted from the wire rather than written as zeros.
#[tokio::test]
async fn no_delegation_leaves_the_bucket_empty_and_off_the_wire() {
    let events = run_parent(
        MockProvider::new(vec![MockResponse::TextWithUsage(
            "hi".into(),
            usage(1, 1, 0, 0),
        )]),
        vec![],
    )
    .await;
    let stats = end_stats(&events);
    assert!(stats.sub_agents.is_empty());
    let json = serde_json::to_value(&stats).unwrap();
    assert!(json.get("subAgents").is_none(), "{json}");
}

/// Recursion: parent → planner → worker. One top-level number covers the
/// tree, and the planner's own spend stays distinguishable from its nested
/// worker's.
#[tokio::test]
async fn nested_sub_agent_usage_sums_through_the_tree() {
    let worker = SubAgentTool::from_provider(
        "worker",
        Arc::new(MockProvider::new(vec![MockResponse::TextWithUsage(
            "worked".into(),
            usage(1000, 100, 0, 0),
        )])),
        ModelConfig::mock(),
    );
    let planner = SubAgentTool::from_provider(
        "planner",
        Arc::new(MockProvider::new(vec![
            delegate_with_usage("worker", usage(10, 1, 0, 0)),
            MockResponse::TextWithUsage("planned".into(), usage(20, 2, 0, 0)),
        ])),
        ModelConfig::mock(),
    )
    .with_tools(vec![Arc::new(worker)]);

    let parent = MockProvider::new(vec![delegate("planner"), MockResponse::Text("done".into())]);
    let events = run_parent(parent, vec![Box::new(planner)]).await;
    let stats = end_stats(&events);

    assert_eq!(stats.usage, Usage::default(), "parent spent nothing itself");
    assert_eq!(stats.sub_agents.usage, usage(1030, 103, 0, 0));
    assert_eq!(stats.sub_agents.runs, 2, "planner and its worker");

    let (result, _) = tool_end(&events, "planner");
    let planner_stats = SessionStats::from_sub_agent_result(&result).unwrap();
    assert_eq!(planner_stats.usage, usage(30, 3, 0, 0), "planner's own");
    assert_eq!(
        planner_stats.sub_agents.usage,
        usage(1000, 100, 0, 0),
        "planner's nested worker"
    );
    assert_eq!(planner_stats.sub_agents.runs, 1);
}

/// A sub-agent that fails partway still spent tokens. `SubAgentTool` returns
/// `Err`, which has nowhere to carry them, so this pins the side channel.
#[tokio::test]
async fn failed_sub_agent_still_reports_what_it_spent() {
    let child = SubAgentTool::from_provider(
        "flaky",
        Arc::new(MockProvider::new(vec![
            echo_call(usage(50, 5, 0, 0)),
            MockResponse::ErrorWithUsage("upstream exploded".into(), usage(7, 0, 0, 0)),
        ])),
        ModelConfig::mock(),
    )
    .with_tools(vec![Arc::new(EchoTool)]);

    let parent = MockProvider::new(vec![delegate("flaky"), MockResponse::Text("ok".into())]);
    let events = run_parent(parent, vec![Box::new(child)]).await;

    let (result, is_error) = tool_end(&events, "flaky");
    assert!(is_error, "the delegation must still read as failed");
    let child_stats = SessionStats::from_sub_agent_result(&result)
        .expect("a failed delegation must still carry its stats");
    assert_eq!(child_stats.usage, usage(57, 5, 0, 0));

    let stats = end_stats(&events);
    assert_eq!(stats.sub_agents.usage, usage(57, 5, 0, 0));
    assert_eq!(stats.sub_agents.runs, 1);
}

/// A sub-agent on a different model is priced at its own rates, never
/// re-priced at the parent's.
#[tokio::test]
async fn sub_agent_cost_uses_the_childs_own_pricing() {
    let mut child_config = ModelConfig::mock();
    child_config.cost = Some(yoagent::provider::CostConfig::new(1.0, 2.0));
    let mut parent_config = ModelConfig::mock();
    parent_config.cost = Some(yoagent::provider::CostConfig::new(10.0, 20.0));

    let child_usage = usage(1_000_000, 500_000, 0, 0);
    let child = SubAgentTool::from_provider(
        "cheap",
        Arc::new(MockProvider::new(vec![MockResponse::TextWithUsage(
            "done".into(),
            child_usage.clone(),
        )])),
        child_config.clone(),
    );

    let mut config = make_config(MockProvider::new(vec![
        delegate("cheap"),
        MockResponse::TextWithUsage("ok".into(), usage(1_000_000, 0, 0, 0)),
    ]));
    config.model_config = Some(parent_config);
    let stats = end_stats(&run_parent_with(config, vec![Box::new(child)]).await);

    let child_cost = child_config.cost.as_ref().unwrap().cost_usd(&child_usage);
    assert!((child_cost - 2.0).abs() < 1e-9, "sanity: {child_cost}");
    let reported = stats.sub_agents.cost_usd.expect("child is priced");
    assert!(
        (reported - child_cost).abs() < 1e-9,
        "{reported} != {child_cost}"
    );

    let own = stats.cost_usd.expect("parent is priced");
    assert!(
        (own - 10.0).abs() < 1e-9,
        "own cost unchanged by delegation: {own}"
    );
    let total = stats.total_cost_usd().unwrap();
    assert!((total - 12.0).abs() < 1e-9, "{total}");
}

/// An unpriced sub-agent makes the delegated cost unknown rather than a
/// silently low sum — and it stays unknown when a priced one follows.
#[test]
fn unpriced_spend_poisons_the_cost_instead_of_under_reporting() {
    let mut priced = SubAgentSpend::default();
    priced.usage = usage(10, 0, 0, 0);
    priced.cost_usd = Some(1.0);
    priced.runs = 1;
    let mut unpriced = SubAgentSpend::default();
    unpriced.usage = usage(10, 0, 0, 0);
    unpriced.runs = 1;

    let mut acc = SubAgentSpend::default();
    acc.merge(&priced);
    assert_eq!(acc.cost_usd, Some(1.0));
    acc.merge(&unpriced);
    assert_eq!(acc.cost_usd, None);
    acc.merge(&priced);
    assert_eq!(acc.cost_usd, None, "a later priced run must not revive it");
    assert_eq!(acc.usage, usage(30, 0, 0, 0));
    assert_eq!(acc.runs, 3);

    // Zero tokens needs no price.
    let mut free = SubAgentSpend::default();
    free.runs = 1;
    let mut acc = SubAgentSpend::default();
    acc.merge(&priced);
    acc.merge(&free);
    assert_eq!(acc.cost_usd, Some(1.0));
}

/// `Agent` keeps the bucket across runs, apart from its own spend.
#[tokio::test]
async fn agent_exposes_sub_agent_spend_across_runs() {
    let child = SubAgentTool::from_provider(
        "helper",
        Arc::new(MockProvider::new(vec![
            MockResponse::TextWithUsage("a".into(), usage(100, 1, 0, 0)),
            MockResponse::TextWithUsage("b".into(), usage(200, 2, 0, 0)),
        ])),
        ModelConfig::mock(),
    );
    let parent = MockProvider::new(vec![
        delegate("helper"),
        MockResponse::TextWithUsage("one".into(), usage(1, 0, 0, 0)),
        delegate("helper"),
        MockResponse::TextWithUsage("two".into(), usage(2, 0, 0, 0)),
    ]);
    let mut agent = Agent::from_provider(parent, ModelConfig::mock()).with_sub_agent(child);

    let mut rx = agent.prompt("first").await;
    while rx.recv().await.is_some() {}
    agent.finish().await;
    assert_eq!(agent.sub_agent_spend().usage, usage(100, 1, 0, 0));

    let (tx, mut rx) = mpsc::unbounded_channel();
    let drain = tokio::spawn(async move { while rx.recv().await.is_some() {} });
    agent.prompt_with_sender("second", tx).await;
    drain.await.unwrap();
    assert_eq!(agent.sub_agent_spend().usage, usage(300, 3, 0, 0));
    assert_eq!(agent.sub_agent_spend().runs, 2);

    // The parent's own history holds only its own turns.
    let own_input: u64 = agent
        .messages()
        .iter()
        .filter_map(|m| match m {
            AgentMessage::Llm(Message::Assistant { usage, .. }) => Some(usage.input),
            _ => None,
        })
        .sum();
    assert_eq!(own_input, 3);

    agent.reset().await;
    assert!(agent.sub_agent_spend().is_empty());
}

// ---------------------------------------------------------------------------
// Whole-bill totals, the unpriced tri-state, and custom delegation tools
// ---------------------------------------------------------------------------

/// `ModelConfig::mock()` priced at `dollars_per_k` per thousand input
/// tokens, so `usage(1_000, ..)` costs exactly that many dollars. (Kept to
/// thousands so runs stay under the default execution token limit.)
fn priced_mock(dollars_per_k: f64) -> ModelConfig {
    let mut config = ModelConfig::mock();
    config.cost = Some(yoagent::provider::CostConfig::new(
        dollars_per_k * 1000.0,
        0.0,
    ));
    config
}

fn text_turn(u: Usage) -> MockResponse {
    MockResponse::TextWithUsage("ok".into(), u)
}

async fn run_agent(agent: &mut Agent, text: &str) {
    let mut rx = agent.prompt(text).await;
    while rx.recv().await.is_some() {}
    agent.finish().await;
}

/// With no delegation the whole bill is the agent's own cost. The delegated
/// bucket's `None` here means "nothing spent", so a naive `zip` of the two
/// options — which yields `None` — gets this case wrong.
#[tokio::test]
async fn total_cost_without_delegation_is_the_agents_own() {
    let mut agent = Agent::from_provider(
        MockProvider::new(vec![text_turn(usage(1_000, 0, 0, 0))]),
        priced_mock(1.0),
    );
    assert_eq!(agent.total_cost_usd(), None, "nothing spent yet");

    run_agent(&mut agent, "go").await;
    assert!(agent.sub_agent_spend().is_empty());
    assert!(!agent.sub_agent_spend().is_unpriced());
    assert_eq!(agent.session_cost_usd(), Some(1.0));
    assert_eq!(agent.total_cost_usd(), Some(1.0));
    assert_eq!(agent.total_usage(), usage(1_000, 0, 0, 0));
}

/// The total covers runs, not history: clearing or replacing history lowers
/// the history-derived `session_cost_usd` but not the bill, which moves only
/// on `reset` — the same window as `sub_agent_spend`.
#[tokio::test]
async fn total_cost_combines_own_and_delegated_spend_over_one_window() {
    let child = SubAgentTool::from_provider(
        "helper",
        Arc::new(MockProvider::new(vec![text_turn(usage(1_000, 0, 0, 0))])),
        priced_mock(2.0),
    );
    let parent = MockProvider::new(vec![
        delegate_with_usage("helper", usage(1_000, 0, 0, 0)),
        text_turn(usage(1_000, 0, 0, 0)),
    ]);
    let mut agent = Agent::from_provider(parent, priced_mock(1.0)).with_sub_agent(child);
    run_agent(&mut agent, "go").await;

    // Own: 2k input at $1/k. Delegated: 1k at the child's $2/k.
    assert_eq!(agent.session_cost_usd(), Some(2.0));
    assert_eq!(agent.sub_agent_spend().cost_usd, Some(2.0));
    assert_eq!(agent.total_cost_usd(), Some(4.0));
    assert_eq!(agent.total_usage(), usage(3_000, 0, 0, 0));

    agent.clear_messages();
    assert_eq!(agent.session_cost_usd(), Some(0.0), "history-derived");
    assert_eq!(agent.total_cost_usd(), Some(4.0), "spent is spent");
    assert_eq!(agent.total_usage(), usage(3_000, 0, 0, 0));

    agent.reset().await;
    assert_eq!(agent.total_cost_usd(), None);
    assert_eq!(agent.total_usage(), Usage::default());
}

/// An unpriced sub-agent that spent tokens makes the bill unknown, not a sum
/// that treats its spend as free (`unwrap_or(0.0)`).
#[tokio::test]
async fn unpriced_delegation_makes_the_total_unknown() {
    // `ModelConfig::mock()` has no rates.
    let child = SubAgentTool::from_provider(
        "helper",
        Arc::new(MockProvider::new(vec![text_turn(usage(500, 0, 0, 0))])),
        ModelConfig::mock(),
    );
    let parent = MockProvider::new(vec![
        delegate_with_usage("helper", usage(1_000, 0, 0, 0)),
        text_turn(Usage::default()),
    ]);
    let mut agent = Agent::from_provider(parent, priced_mock(1.0)).with_sub_agent(child);
    run_agent(&mut agent, "go").await;

    assert_eq!(agent.session_cost_usd(), Some(1.0), "own spend is priced");
    assert!(agent.sub_agent_spend().is_unpriced());
    assert_eq!(agent.sub_agent_spend().cost_usd, None);
    assert_eq!(agent.total_cost_usd(), None);
}

#[test]
fn is_unpriced_separates_unknown_from_nothing_spent() {
    let empty = SubAgentSpend::default();
    assert_eq!(empty.cost_usd, None);
    assert!(!empty.is_unpriced(), "nothing to price");

    let mut free_run = SubAgentSpend::default();
    free_run.runs = 1;
    assert!(!free_run.is_unpriced(), "a run that spent no tokens");

    let mut unpriced = SubAgentSpend::default();
    unpriced.usage = usage(1, 0, 0, 0);
    unpriced.runs = 1;
    assert!(unpriced.is_unpriced());

    let mut priced = unpriced.clone();
    priced.cost_usd = Some(0.5);
    assert!(!priced.is_unpriced());

    // SessionStats: a zero-turn run reports `None` too, and is not unpriced.
    let zero_turns = SessionStats::default();
    assert_eq!(zero_turns.cost_usd, None);
    assert!(!zero_turns.is_unpriced());
    assert!(SessionStats::new(usage(1, 0, 0, 0), 1, None, 0).is_unpriced());
}

/// A custom delegation tool: it reports each run it "started" through
/// `ToolContext::report_delegated_run`, then succeeds or fails.
struct FanOutTool {
    runs: Vec<SessionStats>,
    fail: bool,
}

#[async_trait::async_trait]
impl AgentTool for FanOutTool {
    fn name(&self) -> &str {
        "fanout"
    }
    fn label(&self) -> &str {
        "Fan out"
    }
    fn description(&self) -> &str {
        "Delegates to several agents"
    }
    fn parameters_schema(&self) -> serde_json::Value {
        serde_json::json!({"type": "object", "properties": {}})
    }
    async fn execute(
        &self,
        _params: serde_json::Value,
        ctx: ToolContext,
    ) -> Result<ToolResult, ToolError> {
        for run in &self.runs {
            ctx.report_delegated_run(run.clone());
        }
        if self.fail {
            return Err(ToolError::Failed("a delegate failed".into()));
        }
        Ok(ToolResult {
            content: vec![Content::Text {
                text: "fanned out".into(),
            }],
            details: serde_json::Value::Null,
        })
    }
}

fn call_fanout() -> MockResponse {
    MockResponse::ToolCalls(vec![MockToolCall {
        name: "fanout".into(),
        arguments: serde_json::json!({}),
        provider_metadata: None,
    }])
}

/// A custom tool's report is counted like `SubAgentTool`'s — including when
/// the tool then fails.
#[tokio::test]
async fn custom_tool_reporting_a_delegated_run_is_counted() {
    let tool = FanOutTool {
        runs: vec![SessionStats::new(usage(40, 4, 0, 0), 2, Some(0.25), 0)],
        fail: true,
    };
    let parent = MockProvider::new(vec![call_fanout(), MockResponse::Text("ok".into())]);
    let events = run_parent(parent, vec![Box::new(tool)]).await;

    let stats = end_stats(&events);
    assert_eq!(stats.sub_agents.usage, usage(40, 4, 0, 0));
    assert_eq!(stats.sub_agents.cost_usd, Some(0.25));
    assert_eq!(stats.sub_agents.runs, 1);

    let (result, is_error) = tool_end(&events, "fanout");
    assert!(is_error);
    let reported = SessionStats::from_sub_agent_result(&result).unwrap();
    assert_eq!(reported.usage, usage(40, 4, 0, 0));
}

/// Two runs reported from one tool call: the rollup counts both, and the
/// details carry their combination rather than nothing.
#[tokio::test]
async fn two_delegations_in_one_call_attach_combined_stats() {
    let first = SessionStats::new(usage(100, 10, 0, 0), 2, Some(1.0), 0);
    let mut second = SessionStats::new(usage(200, 20, 0, 0), 3, Some(2.0), 1);
    // The second run delegated further itself.
    second.sub_agents.usage = usage(7, 0, 0, 0);
    second.sub_agents.cost_usd = Some(0.5);
    second.sub_agents.runs = 1;

    let tool = FanOutTool {
        runs: vec![first, second],
        fail: false,
    };
    let parent = MockProvider::new(vec![call_fanout(), MockResponse::Text("ok".into())]);
    let events = run_parent(parent, vec![Box::new(tool)]).await;

    let stats = end_stats(&events);
    assert_eq!(stats.sub_agents.usage, usage(307, 30, 0, 0));
    assert_eq!(stats.sub_agents.runs, 3, "two reported runs and one nested");
    assert_eq!(stats.sub_agents.cost_usd, Some(3.5));

    let (result, is_error) = tool_end(&events, "fanout");
    assert!(!is_error);
    let combined =
        SessionStats::from_sub_agent_result(&result).expect("several runs must still attach stats");
    assert_eq!(combined.usage, usage(300, 30, 0, 0), "the runs' own spend");
    assert_eq!(combined.turns, 5);
    assert_eq!(combined.cost_usd, Some(3.0));
    assert_eq!(combined.compactions, 1);
    assert_eq!(combined.sub_agents.usage, usage(7, 0, 0, 0));
    assert_eq!(combined.total_usage(), stats.sub_agents.usage);
    assert_eq!(combined.total_cost_usd(), stats.sub_agents.cost_usd);
}

/// Outside the loop there is no parent to report to; reporting is a no-op
/// rather than a panic.
#[test]
fn report_delegated_run_without_a_loop_is_a_no_op() {
    ToolContext::new("id", "tool").report_delegated_run(SessionStats::default());
}

/// A tool that says when it started and returns when released.
struct Gated {
    started: Arc<tokio::sync::Notify>,
    release: Arc<tokio::sync::Notify>,
}

#[async_trait::async_trait]
impl AgentTool for Gated {
    fn name(&self) -> &str {
        "wait"
    }
    fn label(&self) -> &str {
        "Wait"
    }
    fn description(&self) -> &str {
        "Waits until released"
    }
    fn parameters_schema(&self) -> serde_json::Value {
        serde_json::json!({"type": "object", "properties": {}})
    }
    async fn execute(
        &self,
        _params: serde_json::Value,
        _ctx: ToolContext,
    ) -> Result<ToolResult, ToolError> {
        self.started.notify_one();
        self.release.notified().await;
        Ok(ToolResult {
            content: vec![Content::Text {
                text: "done".into(),
            }],
            details: serde_json::Value::Null,
        })
    }
}

/// Cancelled while its tool ran (between provider calls), the child's run
/// ends with the cancel marker, so the delegation fails — it used to return
/// the "no text output" fallback as a success. Deterministic: the cancel
/// lands while the tool is known to be running.
#[tokio::test]
async fn test_sub_agent_cancelled_during_a_tool_fails_the_delegation() {
    let sub_provider = Arc::new(MockProvider::new(vec![
        MockResponse::ToolCalls(vec![MockToolCall {
            provider_metadata: None,
            name: "wait".into(),
            arguments: serde_json::json!({}),
        }]),
        MockResponse::Text("never reached".into()),
    ]));
    let started = Arc::new(tokio::sync::Notify::new());
    let release = Arc::new(tokio::sync::Notify::new());
    let gated: Arc<dyn AgentTool> = Arc::new(Gated {
        started: started.clone(),
        release: release.clone(),
    });
    let sub_agent = SubAgentTool::from_provider("waiter", sub_provider, ModelConfig::mock())
        .with_tools(vec![gated]);

    let cancel = CancellationToken::new();
    let trigger = cancel.clone();
    tokio::spawn(async move {
        started.notified().await;
        trigger.cancel();
        release.notify_one();
    });
    let err = sub_agent
        .execute(
            serde_json::json!({"task": "wait"}),
            ToolContext::new("tc-1", "waiter").with_cancel(cancel),
        )
        .await
        .expect_err("a cancelled delegation is a failure");
    assert!(
        err.to_string()
            .contains(yoagent::agent_loop::CANCELLED_MARKER),
        "{err}"
    );
}

/// Three tool turns that each report a 400K-token prompt (mostly cached),
/// then the answer.
fn heavy_delegation() -> Arc<MockProvider> {
    let heavy = Usage {
        input: 20_000,
        cache_read: 380_000,
        output: 500,
        ..Default::default()
    };
    let call = || MockToolCall {
        provider_metadata: None,
        name: "echo".into(),
        arguments: serde_json::json!({"text": "x"}),
    };
    Arc::new(MockProvider::new(vec![
        MockResponse::ToolCallsWithUsage(vec![call()], heavy.clone()),
        MockResponse::ToolCallsWithUsage(vec![call()], heavy.clone()),
        MockResponse::ToolCallsWithUsage(vec![call()], heavy),
        MockResponse::Text("final answer".into()),
    ]))
}

fn text_of(result: &ToolResult) -> String {
    match &result.content[0] {
        Content::Text { text } => text.clone(),
        other => panic!("expected text, got {other:?}"),
    }
}

/// The default caps (1M tokens per delegation) stop a delegation working over
/// a large context after three turns; `with_execution_limits` lifts them.
/// They used to be hard-coded, with only `with_max_turns` adjustable.
#[tokio::test]
async fn execution_limits_are_configurable_per_delegation() {
    let echo: Arc<dyn AgentTool> = Arc::new(EchoTool);

    let capped = SubAgentTool::from_provider("worker", heavy_delegation(), ModelConfig::mock())
        .with_tools(vec![echo.clone()]);
    let result = capped
        .execute(
            serde_json::json!({"task": "go"}),
            ToolContext::new("t", "worker"),
        )
        .await
        .unwrap();
    let text = text_of(&result);
    assert!(!text.contains("final answer"), "{text}");
    assert!(
        text.contains("Max tokens"),
        "stopped by the token cap: {text}"
    );

    let lifted = SubAgentTool::from_provider("worker", heavy_delegation(), ModelConfig::mock())
        .with_tools(vec![echo])
        .with_execution_limits(
            yoagent::context::ExecutionLimits::default()
                .with_max_turns(20)
                .with_max_total_tokens(usize::MAX),
        );
    let result = lifted
        .execute(
            serde_json::json!({"task": "go"}),
            ToolContext::new("t", "worker"),
        )
        .await
        .unwrap();
    assert_eq!(text_of(&result), "final answer");
}

/// `with_max_turns` after `with_execution_limits` changes only the turns: the
/// lifted token cap survives, so the three heavy turns (1.2M tokens) and the
/// answer all fit under a 5-turn cap. Had it reset the limits, the default 1M
/// cap would stop the run after turn 3.
#[tokio::test]
async fn with_max_turns_keeps_the_other_limits() {
    let echo: Arc<dyn AgentTool> = Arc::new(EchoTool);
    let tool = SubAgentTool::from_provider("worker", heavy_delegation(), ModelConfig::mock())
        .with_tools(vec![echo])
        .with_execution_limits(
            yoagent::context::ExecutionLimits::default().with_max_total_tokens(usize::MAX),
        )
        .with_max_turns(5);
    let result = tool
        .execute(
            serde_json::json!({"task": "go"}),
            ToolContext::new("t", "worker"),
        )
        .await
        .unwrap();
    assert_eq!(text_of(&result), "final answer");
}