klieo-core 0.41.1

use super::*;
use crate::llm::{FinishReason, LlmClient};
use crate::memory::EpisodicMemory;
use crate::test_utils::{
    fake_context, FakeLlmClient, FakeStreamStep, FakeToolInvoker, InMemoryEpisodic,
};
use std::sync::Arc;
use tokio_stream::StreamExt;

fn ctx_with(llm: Arc<dyn LlmClient>) -> (AgentContext, Arc<InMemoryEpisodic>) {
    let episodic = Arc::new(InMemoryEpisodic::default());
    let mut ctx = fake_context("stream-test");
    ctx.llm = llm;
    ctx.episodic = episodic.clone();
    (ctx, episodic)
}

fn ctx_with_tools(
    llm: Arc<dyn LlmClient>,
    tools: Arc<FakeToolInvoker>,
) -> (AgentContext, Arc<InMemoryEpisodic>) {
    let episodic = Arc::new(InMemoryEpisodic::default());
    let mut ctx = fake_context("stream-test");
    ctx.llm = llm;
    ctx.episodic = episodic.clone();
    ctx.tools = tools;
    (ctx, episodic)
}

fn delta(s: &str) -> ChatChunk {
    ChatChunk {
        delta: s.into(),
        tool_calls: vec![],
        finish_reason: None,
        usage: None,
    }
}

fn final_stop_chunk() -> ChatChunk {
    ChatChunk {
        delta: String::new(),
        tool_calls: vec![],
        finish_reason: Some(FinishReason::Stop),
        usage: None,
    }
}

#[tokio::test(start_paused = true)]
async fn streaming_single_step_forwards_three_chunks() {
    let llm = Arc::new(
        FakeLlmClient::new("fake").with_stream_steps(vec![FakeStreamStep::Chunks(vec![
            Ok(delta("hel")),
            Ok(delta("lo ")),
            Ok(delta("world")),
            Ok(final_stop_chunk()),
        ])]),
    );
    let (ctx, ep) = ctx_with(llm);

    let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), RunOptions::default())
        .await
        .expect("stream opens");
    let mut chunks = Vec::new();
    while let Some(item) = s.next().await {
        chunks.push(item.expect("ok chunk"));
    }
    assert_eq!(chunks.len(), 4, "3 deltas + terminal");
    assert_eq!(chunks[0].delta, "hel");
    assert_eq!(chunks[1].delta, "lo ");
    assert_eq!(chunks[2].delta, "world");
    assert_eq!(chunks[3].finish_reason, Some(FinishReason::Stop));

    tokio::task::yield_now().await;
    tokio::task::yield_now().await;

    let events = ep.replay(ctx.run_id).await.unwrap();
    let llm_calls = events
        .iter()
        .filter(|e| matches!(e, Episode::LlmCall { .. }))
        .count();
    let completed = events
        .iter()
        .filter(|e| matches!(e, Episode::Completed))
        .count();
    assert_eq!(llm_calls, 1);
    assert_eq!(completed, 1);
}

#[tokio::test(start_paused = true)]
async fn streaming_tool_call_dispatch_starts_second_cycle() {
    let tool_call = crate::llm::ToolCall {
        id: "tc1".into(),
        name: "echo".into(),
        args: serde_json::json!({"x": 1}),
    };
    let stream1 = vec![Ok(ChatChunk {
        delta: String::new(),
        tool_calls: vec![tool_call.clone()],
        finish_reason: Some(FinishReason::ToolCalls),
        usage: None,
    })];
    let stream2 = vec![Ok(delta("done")), Ok(final_stop_chunk())];

    let llm = Arc::new(FakeLlmClient::new("fake").with_stream_steps(vec![
        FakeStreamStep::Chunks(stream1),
        FakeStreamStep::Chunks(stream2),
    ]));
    let tools = Arc::new(FakeToolInvoker::new().with_tool("echo", "echoes", Ok));
    let (ctx, ep) = ctx_with_tools(llm, tools);

    let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), RunOptions::default())
        .await
        .expect("stream opens");
    let mut deltas = Vec::new();
    while let Some(item) = s.next().await {
        deltas.push(item.expect("ok chunk"));
    }

    assert!(deltas.len() >= 3);
    let joined: String = deltas.iter().map(|c| c.delta.clone()).collect();
    assert_eq!(joined, "done");

    let events = ep.replay(ctx.run_id).await.unwrap();
    let llm_calls = events
        .iter()
        .filter(|e| matches!(e, Episode::LlmCall { .. }))
        .count();
    let tool_calls = events
        .iter()
        .filter(|e| matches!(e, Episode::ToolCall { .. }))
        .count();
    assert_eq!(llm_calls, 2, "two LLM cycles recorded");
    assert_eq!(tool_calls, 1, "one tool dispatched");
}

#[tokio::test(start_paused = true)]
async fn streaming_mid_stream_cancellation_emits_cancelled_error() {
    let mut chunks: Vec<Result<ChatChunk, LlmError>> = Vec::new();
    for i in 0..50 {
        chunks.push(Ok(delta(&format!("c{i} "))));
    }
    chunks.push(Ok(final_stop_chunk()));
    let llm = Arc::new(
        FakeLlmClient::new("fake").with_stream_steps(vec![FakeStreamStep::Chunks(chunks)]),
    );
    let (ctx, ep) = ctx_with(llm);

    let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), RunOptions::default())
        .await
        .expect("stream opens");

    let first = s.next().await.expect("first chunk").expect("ok");
    assert!(!first.delta.is_empty());
    ctx.cancel.cancel();

    let mut got_cancel_err = false;
    let mut tail_count = 0usize;
    while let Some(item) = s.next().await {
        tail_count += 1;
        match item {
            Err(LlmError::Cancelled) => {
                got_cancel_err = true;
                break;
            }
            Ok(_) => continue,
            Err(other) => panic!("unexpected error: {other:?}"),
        }
    }
    assert!(
        got_cancel_err,
        "expected terminal LlmError::Cancelled item, drained {tail_count} items"
    );

    tokio::task::yield_now().await;
    tokio::task::yield_now().await;

    let events = ep.replay(ctx.run_id).await.unwrap();
    let llm_calls = events
        .iter()
        .filter(|e| matches!(e, Episode::LlmCall { .. }))
        .count();
    assert_eq!(
        llm_calls, 0,
        "partial streaming cycle must not be recorded as a successful LlmCall"
    );
    let completed = events
        .iter()
        .filter(|e| matches!(e, Episode::Completed))
        .count();
    assert_eq!(completed, 0);
}

#[tokio::test(start_paused = true)]
async fn streaming_init_retries_server_then_succeeds() {
    let llm = Arc::new(FakeLlmClient::new("fake").with_stream_steps(vec![
        FakeStreamStep::init_err(LlmError::Server("503".into())),
        FakeStreamStep::Chunks(vec![Ok(delta("ok")), Ok(final_stop_chunk())]),
    ]));
    let (ctx, _ep) = ctx_with(llm.clone());

    let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), RunOptions::default())
        .await
        .expect("stream opens after retry");
    while let Some(item) = s.next().await {
        item.expect("ok");
    }
    assert_eq!(llm.stream_call_count(), 2, "one failure + one success");
}

#[tokio::test(start_paused = true)]
async fn streaming_init_unauthorized_propagates_after_one_attempt() {
    let llm = Arc::new(
        FakeLlmClient::new("fake")
            .with_stream_steps(vec![FakeStreamStep::init_err(LlmError::Unauthorized)]),
    );
    let (ctx, _ep) = ctx_with(llm.clone());

    let res = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), RunOptions::default()).await;
    match res {
        Ok(_) => panic!("expected Unauthorized, got Ok"),
        Err(Error::Llm(LlmError::Unauthorized)) => {}
        Err(other) => panic!("expected Llm(Unauthorized), got {other:?}"),
    }
    assert_eq!(
        llm.stream_call_count(),
        1,
        "non-retryable: exactly one attempt"
    );
}

#[tokio::test(start_paused = true)]
async fn streaming_max_steps_exceeded_when_tools_loop() {
    let tool_call = crate::llm::ToolCall {
        id: "tc1".into(),
        name: "echo".into(),
        args: serde_json::json!({}),
    };
    let mk_tool_stream = || {
        FakeStreamStep::Chunks(vec![Ok(ChatChunk {
            delta: String::new(),
            tool_calls: vec![tool_call.clone()],
            finish_reason: Some(FinishReason::ToolCalls),
            usage: None,
        })])
    };
    let llm = Arc::new(FakeLlmClient::new("fake").with_stream_steps(vec![
        mk_tool_stream(),
        mk_tool_stream(),
        mk_tool_stream(),
        mk_tool_stream(),
        mk_tool_stream(),
    ]));
    let tools = Arc::new(
        FakeToolInvoker::new().with_tool("echo", "echoes", |_| Ok(serde_json::json!("ok"))),
    );
    let (ctx, ep) = ctx_with_tools(llm, tools);

    let opts = RunOptions {
        max_steps: 2,
        ..RunOptions::default()
    };
    let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), opts)
        .await
        .expect("stream opens");

    let mut got_max_steps_err = false;
    while let Some(item) = s.next().await {
        if let Err(LlmError::Server(ref m)) = item {
            if m.contains("max steps exceeded") {
                got_max_steps_err = true;
            }
        }
    }
    assert!(
        got_max_steps_err,
        "expected terminal Err(Server(\"max steps exceeded: ...\")) before stream-close"
    );

    tokio::task::yield_now().await;
    tokio::task::yield_now().await;

    let events = ep.replay(ctx.run_id).await.unwrap();
    let failed = events
        .iter()
        .find(|e| matches!(e, Episode::Failed { .. }))
        .expect("must record Failed");
    match failed {
        Episode::Failed { error } => {
            assert!(
                error.contains("max steps"),
                "expected max-steps message, got {error}"
            );
        }
        _ => unreachable!(),
    }
}

#[tokio::test(start_paused = true)]
async fn streaming_response_byte_cap_aborts_oversized_stream() {
    let mut chunks: Vec<Result<ChatChunk, LlmError>> = Vec::new();
    for i in 0..20 {
        chunks.push(Ok(delta(&format!("payload-block-{i:02}"))));
    }
    chunks.push(Ok(final_stop_chunk()));
    let llm = Arc::new(
        FakeLlmClient::new("fake").with_stream_steps(vec![FakeStreamStep::Chunks(chunks)]),
    );
    let (ctx, ep) = ctx_with(llm);

    let opts = RunOptions {
        max_response_bytes: 64,
        ..RunOptions::default()
    };
    let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), opts)
        .await
        .expect("stream opens");

    let mut saw_cap_err = false;
    while let Some(item) = s.next().await {
        if let Err(LlmError::Server(ref m)) = item {
            if m.contains("max_response_bytes cap") {
                saw_cap_err = true;
                break;
            }
        }
    }
    assert!(
        saw_cap_err,
        "expected terminal Err carrying the byte-cap message"
    );

    tokio::task::yield_now().await;
    tokio::task::yield_now().await;

    let events = ep.replay(ctx.run_id).await.unwrap();
    let saw_failed = events.iter().any(
        |e| matches!(e, Episode::Failed { error } if error.contains("max_response_bytes cap")),
    );
    assert!(
        saw_failed,
        "expected Failed episode mentioning max_response_bytes cap, events: {events:?}"
    );
}

// ── AgentEvent emission tests ─────────────────────────────────────────────────

/// Scenario labels for the shared `build_ctx_for_events` helper.
enum LlmScenario {
    /// LLM returns one final text response immediately.
    FinalText,
    /// LLM returns one tool-call step then a final text response.
    OneToolCall,
    /// LLM always returns a tool-call step (infinite loop until max_steps).
    InfiniteToolLoop,
}

/// Which tool outcome to simulate.
enum ToolBehaviour {
    Succeeds,
    Fails,
}

/// Build an `AgentContext` wired to `progress` tx and program the
/// `FakeLlmClient` + `FakeToolInvoker` according to the supplied
/// scenario. Returns `(ctx, broadcast_receiver)`.
fn build_ctx_for_events(
    scenario: LlmScenario,
    tool_behaviour: ToolBehaviour,
    tx: tokio::sync::broadcast::Sender<crate::agent::AgentEvent>,
) -> AgentContext {
    use crate::error::ToolError;
    use crate::llm::ToolCall;
    use crate::test_utils::FakeLlmStep;

    let echo_call = ToolCall {
        id: "tc-evt".into(),
        name: "echo".into(),
        args: serde_json::json!({}),
    };

    let llm = match scenario {
        LlmScenario::FinalText => {
            Arc::new(FakeLlmClient::new("fake").with_steps(vec![FakeLlmStep::Text("done".into())]))
        }
        LlmScenario::OneToolCall => Arc::new(FakeLlmClient::new("fake").with_steps(vec![
            FakeLlmStep::ToolCalls(vec![echo_call]),
            FakeLlmStep::Text("done".into()),
        ])),
        LlmScenario::InfiniteToolLoop => {
            let calls = vec![
                FakeLlmStep::ToolCalls(vec![ToolCall {
                    id: "tc1".into(),
                    name: "echo".into(),
                    args: serde_json::json!({}),
                }]),
                FakeLlmStep::ToolCalls(vec![ToolCall {
                    id: "tc2".into(),
                    name: "echo".into(),
                    args: serde_json::json!({}),
                }]),
                FakeLlmStep::ToolCalls(vec![ToolCall {
                    id: "tc3".into(),
                    name: "echo".into(),
                    args: serde_json::json!({}),
                }]),
            ];
            Arc::new(FakeLlmClient::new("fake").with_steps(calls))
        }
    };

    let tools = Arc::new(match tool_behaviour {
        ToolBehaviour::Succeeds => {
            FakeToolInvoker::new().with_tool("echo", "echoes", |_| Ok(serde_json::json!("ok")))
        }
        ToolBehaviour::Fails => FakeToolInvoker::new().with_tool("echo", "echoes", |_| {
            Err(ToolError::Retryable {
                message: "transient".into(),
                retry_after_secs: 0,
            })
        }),
    });

    let mut ctx = fake_context("evt-test");
    ctx.llm = llm;
    ctx.tools = tools;
    ctx.progress = Some(tx);
    ctx
}

/// Drive `fut` to completion and drain every [`crate::agent::AgentEvent`]
/// that was sent to `rx` before the sender was dropped.
///
/// A brief yield after the future returns lets the runtime's final
/// `emit` calls land before we drain the channel.
async fn collect_agent_events<F>(
    tx: tokio::sync::broadcast::Sender<crate::agent::AgentEvent>,
    mut rx: tokio::sync::broadcast::Receiver<crate::agent::AgentEvent>,
    fut: F,
) -> Vec<crate::agent::AgentEvent>
where
    F: std::future::Future,
{
    let _ = fut.await;
    // One yield so any spawned-task emit calls land before we drain.
    tokio::task::yield_now().await;
    drop(tx);
    let mut events = Vec::new();
    while let Ok(ev) = rx.try_recv() {
        events.push(ev);
    }
    events
}

fn event_kinds(events: &[crate::agent::AgentEvent]) -> Vec<&'static str> {
    use crate::agent::AgentEvent;
    events
        .iter()
        .map(|e| match e {
            AgentEvent::LlmCallStarted => "LlmCallStarted",
            AgentEvent::LlmCallCompleted { .. } => "LlmCallCompleted",
            AgentEvent::ToolCallStarted { .. } => "ToolCallStarted",
            AgentEvent::ToolCallCompleted { .. } => "ToolCallCompleted",
            AgentEvent::Completed => "Completed",
            AgentEvent::Failed { .. } => "Failed",
        })
        .collect()
}

#[tokio::test(start_paused = true)]
async fn run_steps_emits_event_sequence_on_clean_run() {
    use crate::runtime::run_steps;

    let (tx, rx) = tokio::sync::broadcast::channel(64);
    let ctx = build_ctx_for_events(LlmScenario::FinalText, ToolBehaviour::Succeeds, tx.clone());

    let events = collect_agent_events(tx, rx, async {
        let _ = run_steps(&ctx, "sys", ThreadId::new("t"), RunOptions::default()).await;
    })
    .await;

    let kinds = event_kinds(&events);
    assert!(
        kinds.first() == Some(&"LlmCallStarted"),
        "first event must be LlmCallStarted; got: {kinds:?}"
    );
    assert!(
        kinds.contains(&"LlmCallCompleted"),
        "expected LlmCallCompleted; got: {kinds:?}"
    );
    assert!(
        kinds.last() == Some(&"Completed"),
        "last event must be Completed; got: {kinds:?}"
    );
}

#[tokio::test(start_paused = true)]
async fn run_steps_emits_tool_call_events() {
    use crate::agent::AgentEvent;
    use crate::runtime::run_steps;

    let (tx, rx) = tokio::sync::broadcast::channel(64);
    let ctx = build_ctx_for_events(
        LlmScenario::OneToolCall,
        ToolBehaviour::Succeeds,
        tx.clone(),
    );

    let events = collect_agent_events(tx, rx, async {
        let _ = run_steps(&ctx, "sys", ThreadId::new("t"), RunOptions::default()).await;
    })
    .await;

    assert!(
        events
            .iter()
            .any(|e| matches!(e, AgentEvent::ToolCallStarted { name } if name == "echo")),
        "expected ToolCallStarted{{name:echo}}; got: {:?}",
        event_kinds(&events)
    );
    assert!(
        events
            .iter()
            .any(|e| matches!(e, AgentEvent::ToolCallCompleted { ok: true, .. })),
        "expected ToolCallCompleted{{ok:true}}; got: {:?}",
        event_kinds(&events)
    );
}

#[tokio::test(start_paused = true)]
async fn run_steps_emits_tool_failure() {
    use crate::agent::AgentEvent;
    use crate::runtime::run_steps;

    let (tx, rx) = tokio::sync::broadcast::channel(64);
    let ctx = build_ctx_for_events(LlmScenario::OneToolCall, ToolBehaviour::Fails, tx.clone());

    let events = collect_agent_events(tx, rx, async {
        let _ = run_steps(&ctx, "sys", ThreadId::new("t"), RunOptions::default()).await;
    })
    .await;

    assert!(
        events
            .iter()
            .any(|e| matches!(e, AgentEvent::ToolCallCompleted { ok: false, .. })),
        "expected ToolCallCompleted{{ok:false}}; got: {:?}",
        event_kinds(&events)
    );
}

#[tokio::test(start_paused = true)]
async fn run_steps_emits_failed_on_max_steps() {
    use crate::agent::AgentEvent;
    use crate::runtime::run_steps;

    let (tx, rx) = tokio::sync::broadcast::channel(64);
    let ctx = build_ctx_for_events(
        LlmScenario::InfiniteToolLoop,
        ToolBehaviour::Succeeds,
        tx.clone(),
    );
    let opts = RunOptions {
        max_steps: 1,
        ..RunOptions::default()
    };

    let events = collect_agent_events(tx, rx, async {
        let _ = run_steps(&ctx, "sys", ThreadId::new("t"), opts).await;
    })
    .await;

    match events.last() {
        Some(AgentEvent::Failed { reason }) => {
            assert!(
                reason.to_lowercase().contains("max steps"),
                "expected 'max steps' in reason; got: {reason}"
            );
        }
        other => panic!("expected last event Failed; got: {other:?}"),
    }
}

#[tokio::test(start_paused = true)]
async fn run_steps_emits_failed_on_cancel() {
    use crate::agent::AgentEvent;
    use crate::runtime::run_steps;

    let (tx, rx) = tokio::sync::broadcast::channel(64);
    let ctx = build_ctx_for_events(LlmScenario::FinalText, ToolBehaviour::Succeeds, tx.clone());
    ctx.cancel.cancel();

    let events = collect_agent_events(tx, rx, async {
        let _ = run_steps(&ctx, "sys", ThreadId::new("t"), RunOptions::default()).await;
    })
    .await;

    match events.last() {
        Some(AgentEvent::Failed { reason }) => {
            assert!(
                reason.to_lowercase().contains("cancel"),
                "expected 'cancel' in reason; got: {reason}"
            );
        }
        other => panic!("expected Failed last; got: {other:?}"),
    }
}

#[tokio::test(start_paused = true)]
async fn run_steps_streaming_emits_same_event_kind_sequence_as_run_steps() {
    use crate::runtime::run_steps;

    // Blocking path — one tool call then final text.
    let tool_call_chunk = crate::llm::ToolCall {
        id: "tc-par".into(),
        name: "echo".into(),
        args: serde_json::json!({}),
    };
    let stream1 = vec![Ok(ChatChunk {
        delta: String::new(),
        tool_calls: vec![tool_call_chunk],
        finish_reason: Some(FinishReason::ToolCalls),
        usage: None,
    })];
    let stream2 = vec![
        Ok(ChatChunk {
            delta: "done".into(),
            tool_calls: vec![],
            finish_reason: None,
            usage: None,
        }),
        Ok(ChatChunk {
            delta: String::new(),
            tool_calls: vec![],
            finish_reason: Some(FinishReason::Stop),
            usage: None,
        }),
    ];

    let (tx1, rx1) = tokio::sync::broadcast::channel(64);
    let mut ctx1 = fake_context("parity-blk");
    ctx1.llm = Arc::new(FakeLlmClient::new("fake").with_steps(vec![
        crate::test_utils::FakeLlmStep::ToolCalls(vec![crate::llm::ToolCall {
            id: "tc-blk".into(),
            name: "echo".into(),
            args: serde_json::json!({}),
        }]),
        crate::test_utils::FakeLlmStep::Text("done".into()),
    ]));
    ctx1.tools = Arc::new(
        FakeToolInvoker::new().with_tool("echo", "echoes", |_| Ok(serde_json::json!("ok"))),
    );
    ctx1.progress = Some(tx1.clone());
    let events_blocking = collect_agent_events(tx1, rx1, async {
        let _ = run_steps(&ctx1, "sys", ThreadId::new("t"), RunOptions::default()).await;
    })
    .await;

    // Streaming path — same trajectory.
    let (tx2, rx2) = tokio::sync::broadcast::channel(64);
    let mut ctx2 = fake_context("parity-str");
    ctx2.llm = Arc::new(FakeLlmClient::new("fake").with_stream_steps(vec![
        FakeStreamStep::Chunks(stream1),
        FakeStreamStep::Chunks(stream2),
    ]));
    ctx2.tools = Arc::new(
        FakeToolInvoker::new().with_tool("echo", "echoes", |_| Ok(serde_json::json!("ok"))),
    );
    ctx2.progress = Some(tx2.clone());
    let events_streaming = collect_agent_events(tx2, rx2, async {
        let mut s = run_steps_streaming(&ctx2, "sys", ThreadId::new("t"), RunOptions::default())
            .await
            .expect("stream opens");
        while s.next().await.is_some() {}
    })
    .await;

    let kinds_blocking = event_kinds(&events_blocking);
    let kinds_streaming = event_kinds(&events_streaming);
    assert_eq!(
        kinds_blocking, kinds_streaming,
        "blocking vs streaming event-kind sequence mismatch:\n  blocking:  {kinds_blocking:?}\n  streaming: {kinds_streaming:?}"
    );
}

#[tokio::test(start_paused = true)]
async fn emit_is_noop_when_progress_unset() {
    use crate::runtime::run_steps;

    // ctx.progress stays None — run must complete without panic.
    let mut ctx = fake_context("noop-progress");
    ctx.llm = Arc::new(
        FakeLlmClient::new("fake")
            .with_steps(vec![crate::test_utils::FakeLlmStep::Text("ok".into())]),
    );
    let result = run_steps(&ctx, "sys", ThreadId::new("t"), RunOptions::default()).await;
    assert!(
        result.is_ok(),
        "progress=None must not panic; got: {result:?}"
    );
}

mod terminal_chunk_for_tests {
    //! Per-arm coverage of [`terminal_chunk_for`].
    //!
    //! The streaming spawn wrapper relies on `terminal_chunk_for` to
    //! translate every loop-driver `Error` into the pair forwarded to
    //! the consumer + persisted to `Episode::Failed`. Streaming
    //! integration tests indirectly cover the `Cancelled` arms only.
    //! These unit tests pin the remaining arms against their
    //! documented prefix shape so accidental rewording is caught.
    use super::super::terminal_chunk_for;
    use crate::error::{Error, LlmError, MemoryError, ToolError};

    #[test]
    fn cancelled_with_cancel_observed_yields_cancelled_episode() {
        let (chunk, episode) = terminal_chunk_for(Error::Cancelled, true);
        assert!(matches!(chunk, LlmError::Cancelled));
        assert_eq!(episode, "cancelled");
    }

    #[test]
    fn cancelled_without_cancel_observed_yields_consumer_dropped_episode() {
        let (chunk, episode) = terminal_chunk_for(Error::Cancelled, false);
        assert!(matches!(chunk, LlmError::Cancelled));
        assert_eq!(episode, "consumer-dropped");
    }

    #[test]
    fn max_steps_exceeded_renders_step_count() {
        let (chunk, episode) = terminal_chunk_for(Error::MaxStepsExceeded { steps: 7 }, false);
        let message = match chunk {
            LlmError::Server(m) => m,
            other => panic!("expected Server, got {other:?}"),
        };
        assert_eq!(message, "max steps exceeded: 7");
        assert_eq!(episode, "max steps exceeded: 7");
    }

    #[test]
    fn tool_error_prefixed_with_tool() {
        let (chunk, episode) =
            terminal_chunk_for(Error::Tool(ToolError::Permanent("boom".into())), false);
        let message = match chunk {
            LlmError::Server(m) => m,
            other => panic!("expected Server, got {other:?}"),
        };
        assert!(message.starts_with("tool: "), "got {message}");
        assert!(message.contains("boom"));
        assert_eq!(message, episode);
    }

    #[test]
    fn llm_variant_preserved_end_to_end_not_rewrapped_as_server() {
        let (chunk, episode) = terminal_chunk_for(Error::Llm(LlmError::Unauthorized), false);
        assert!(
            matches!(chunk, LlmError::Unauthorized),
            "Llm-wrapped variants must round-trip unchanged"
        );
        assert!(episode.to_lowercase().contains("unauthorized"));
    }

    #[test]
    fn llm_rate_limit_preserves_retry_after_seconds() {
        let (chunk, _) = terminal_chunk_for(
            Error::Llm(LlmError::RateLimit {
                retry_after_secs: 11,
            }),
            false,
        );
        match chunk {
            LlmError::RateLimit { retry_after_secs } => assert_eq!(retry_after_secs, 11),
            other => panic!("expected RateLimit, got {other:?}"),
        }
    }

    #[test]
    fn refused_prefixed_with_refused() {
        let (chunk, episode) = terminal_chunk_for(
            Error::Refused {
                reason: "policy".into(),
            },
            false,
        );
        let message = match chunk {
            LlmError::Server(m) => m,
            other => panic!("expected Server, got {other:?}"),
        };
        assert_eq!(message, "refused: policy");
        assert_eq!(message, episode);
    }

    #[test]
    fn handoff_includes_target_agent_and_reason() {
        let (chunk, episode) = terminal_chunk_for(
            Error::Handoff {
                agent: "safety".into(),
                reason: "needs human review".into(),
            },
            false,
        );
        let message = match chunk {
            LlmError::Server(m) => m,
            other => panic!("expected Server, got {other:?}"),
        };
        assert_eq!(message, "handoff to safety: needs human review");
        assert_eq!(message, episode);
    }

    #[test]
    fn unmatched_variant_falls_through_to_server_with_display() {
        // Memory is one of the Error variants not enumerated by the
        // explicit arms. The catch-all should wrap its Display.
        let (chunk, episode) =
            terminal_chunk_for(Error::Memory(MemoryError::Store("disk full".into())), false);
        let message = match chunk {
            LlmError::Server(m) => m,
            other => panic!("expected Server, got {other:?}"),
        };
        assert!(message.contains("disk full"), "got {message}");
        assert_eq!(message, episode);
    }
}

#[tokio::test(start_paused = true)]
async fn streaming_llm_call_completed_emits_zero_tokens_when_provider_omits_usage() {
    use crate::agent::AgentEvent;

    // All chunks carry usage: None — no streaming_usage configured on the
    // fake provider. The contract (ADR-015): LlmCallCompleted must carry
    // tokens=0 AND the next event must be Completed (not Failed), proving
    // the dropped-sentinel contract.
    let (tx, rx) = tokio::sync::broadcast::channel(64);
    let mut ctx = fake_context("usage-none-test");
    ctx.llm = Arc::new(
        FakeLlmClient::new("fake").with_stream_steps(vec![FakeStreamStep::Chunks(vec![
            Ok(delta("hi")),
            Ok(final_stop_chunk()),
        ])]),
    );
    ctx.progress = Some(tx.clone());

    let events = collect_agent_events(tx, rx, async {
        let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), RunOptions::default())
            .await
            .expect("stream opens");
        while s.next().await.is_some() {}
    })
    .await;

    let completed_idx = events
        .iter()
        .position(|e| matches!(e, AgentEvent::LlmCallCompleted { .. }))
        .expect("LlmCallCompleted must be emitted");
    match &events[completed_idx] {
        AgentEvent::LlmCallCompleted { tokens, .. } => {
            assert_eq!(
                *tokens, 0,
                "expected tokens=0 (provider omitted usage); got {tokens}"
            );
        }
        _ => unreachable!(),
    }
    // Next event must be Completed, not Failed — proves tokens=0 doesn't
    // trigger a retry or error path.
    let next = events.get(completed_idx + 1);
    assert!(
        matches!(next, Some(AgentEvent::Completed)),
        "expected next event Completed (not Failed); got: {next:?}"
    );
}

#[tokio::test(start_paused = true)]
async fn streaming_llm_call_completed_emits_real_token_count() {
    use crate::agent::AgentEvent;
    use crate::llm::Usage;

    // Final chunk carries usage so forward_chunks can surface it.
    let final_chunk_with_usage = ChatChunk {
        delta: String::new(),
        tool_calls: vec![],
        finish_reason: Some(FinishReason::Stop),
        usage: Some(Usage {
            prompt_tokens: 12,
            completion_tokens: 7,
        }),
    };
    let (tx, rx) = tokio::sync::broadcast::channel(64);
    let mut ctx = fake_context("usage-test");
    ctx.llm = Arc::new(
        FakeLlmClient::new("fake").with_stream_steps(vec![FakeStreamStep::Chunks(vec![
            Ok(delta("hi")),
            Ok(final_chunk_with_usage),
        ])]),
    );
    ctx.progress = Some(tx.clone());

    let events = collect_agent_events(tx, rx, async {
        let mut s = run_steps_streaming(&ctx, "sys", ThreadId::new("t"), RunOptions::default())
            .await
            .expect("stream opens");
        while s.next().await.is_some() {}
    })
    .await;

    let completed_event = events
        .iter()
        .find(|e| matches!(e, AgentEvent::LlmCallCompleted { .. }))
        .expect("LlmCallCompleted must be emitted");
    match completed_event {
        AgentEvent::LlmCallCompleted { tokens, .. } => {
            assert_eq!(*tokens, 19, "tokens must be prompt(12)+completion(7)=19");
        }
        _ => unreachable!(),
    }
}