//! Agent executor: core run loop.
//!
//! Loads session history, calls the LLM, streams results, handles tool calls,
//! appends messages to history, and marks the run complete.

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::Duration;

use serde_json::{json, Value};
use uuid::Uuid;

use crate::agent::context::{build_context, build_context_with_tagging};
use crate::agent::prompt_guard::{postflight, preflight};
use crate::agent::provider::*;
use crate::agent::tools::{self, ToolCallResult};
use crate::agent::{AgentConfig, AgentError};
use crate::plugins::hook_utils;
use crate::plugins::HookDispatchResult;
use crate::server::ws::{broadcast_agent_event, broadcast_chat_event, WsServerState};

/// Maximum wall-clock time for a single LLM turn (call + stream processing).
/// This is a safety net above the reqwest-level timeout (300s) to catch hangs
/// in stream processing or channel backpressure.
const TURN_TIMEOUT: Duration = Duration::from_secs(600);

/// Maximum time to wait for the next chunk in the SSE stream.
/// If the provider stalls without sending any data for this long,
/// we treat it as a timeout and abort the turn.
const STREAM_CHUNK_TIMEOUT: Duration = Duration::from_secs(90);
use crate::sessions::{ChatMessage, MessageRole};
use tokio::sync::mpsc;
use tokio_util::sync::CancellationToken;

/// Result of processing an LLM stream for a single turn.
struct StreamResult {
    turn_text: String,
    assistant_blocks: Vec<ContentBlock>,
    pending_tool_calls: Vec<(String, String, Value)>,
    stop_reason: StopReason,
    turn_usage: TokenUsage,
}

fn dispatch_plugin_hook(
    state: &Arc<WsServerState>,
    hook_name: &str,
    payload: &Value,
) -> Option<HookDispatchResult> {
    let plugin_registry = state.plugin_registry()?.clone();
    hook_utils::dispatch_hook(plugin_registry, hook_name, payload)
}

fn parse_hook_payload(result: &HookDispatchResult, hook_name: &str) -> Option<Value> {
    hook_utils::parse_hook_payload(result, hook_name)
}

fn apply_agent_hook_overrides(
    config: &mut AgentConfig,
    payload: &Value,
    settings: &Value,
) -> Result<(), AgentError> {
    let Some(obj) = payload.as_object() else {
        return Ok(());
    };

    if let Some(system) = obj.get("system") {
        match system {
            Value::String(value) => config.system = Some(value.clone()),
            Value::Null => config.system = None,
            _ => {}
        }
    }

    let hook_model = obj.get("model").and_then(|v| v.as_str());
    let hook_route = obj.get("route").and_then(|v| v.as_str());

    if hook_route.is_some() && hook_model.is_some() {
        return Err(AgentError::Provider(
            "cannot set both `route` and `model` in before_agent_start hook response".to_string(),
        ));
    } else if let Some(route) = hook_route {
        crate::agent::resolve_agent_model(
            config,
            settings,
            None,
            &crate::agent::ModelResolutionOverrides {
                request_route: Some(route),
                ..Default::default()
            },
        )?;
    } else if let Some(model) = hook_model {
        config.model = model.to_string();
    }

    if let Some(max_tokens) = obj.get("maxTokens").and_then(|value| value.as_u64()) {
        config.max_tokens = max_tokens.min(u32::MAX as u64) as u32;
    }

    if let Some(temperature) = obj.get("temperature").and_then(|value| value.as_f64()) {
        config.temperature = Some(temperature);
    }

    if let Some(extra) = obj.get("extra") {
        if extra.is_null() {
            config.extra = None;
        } else {
            config.extra = Some(extra.clone());
        }
    }

    Ok(())
}

fn apply_tool_input_override(tool_input: &mut Value, payload: &Value) {
    let Some(obj) = payload.as_object() else {
        return;
    };
    if let Some(input) = obj.get("input") {
        *tool_input = input.clone();
    }
}

fn apply_tool_result_override(result_content: &mut String, is_error: &mut bool, payload: &Value) {
    let Some(obj) = payload.as_object() else {
        return;
    };

    if let Some(result) = obj.get("result") {
        match result {
            Value::String(value) => *result_content = value.clone(),
            Value::Null => {}
            _ => {
                if let Ok(serialized) = serde_json::to_string(result) {
                    *result_content = serialized;
                }
            }
        }
    }

    if let Some(flag) = obj.get("isError").and_then(|value| value.as_bool()) {
        *is_error = flag;
    }
}

/// Dispatch a single stream event, updating turn accumulators.
///
/// Returns `Ok(true)` if the event was a `Stop` (caller should break),
/// `Ok(false)` to continue reading, or `Err` on a provider error event.
fn handle_stream_event(
    event: StreamEvent,
    result: &mut StreamResult,
    state: &Arc<WsServerState>,
    run_id: &str,
    session_key: &str,
    seq: &AtomicU64,
) -> Result<bool, AgentError> {
    match event {
        StreamEvent::TextDelta { text, metadata } => {
            if text.is_empty() && metadata.is_none() {
                return Ok(false);
            }

            result.turn_text.push_str(&text);
            push_text_block(&mut result.assistant_blocks, text.clone(), metadata);

            if text.is_empty() {
                return Ok(false);
            }

            broadcast_agent_event(
                state,
                run_id,
                seq.fetch_add(1, Ordering::Relaxed),
                "text",
                json!({ "delta": &text }),
            );

            broadcast_chat_event(
                state,
                run_id,
                session_key,
                seq.fetch_add(1, Ordering::Relaxed),
                "delta",
                Some(json!({ "content": &text })),
                None,
                None,
                None,
            );
            Ok(false)
        }

        StreamEvent::ToolUse {
            id,
            name,
            input,
            metadata,
        } => {
            broadcast_agent_event(
                state,
                run_id,
                seq.fetch_add(1, Ordering::Relaxed),
                "tool_use",
                json!({
                    "toolUseId": &id,
                    "name": &name,
                    "input": &input,
                }),
            );
            result.assistant_blocks.push(ContentBlock::ToolUse {
                id: id.clone(),
                name: name.clone(),
                input: input.clone(),
                metadata,
            });
            result.pending_tool_calls.push((id, name, input));
            Ok(false)
        }

        StreamEvent::Stop { reason, usage } => {
            result.stop_reason = reason;
            result.turn_usage = usage;
            if reason != StopReason::ToolUse {
                cancel_pending_tool_calls(state, run_id, seq, result, None);
            }
            Ok(true)
        }

        StreamEvent::Error { message, usage } => {
            // Preserve any tokens accumulated before the error so the caller can
            // still record billing/observability data for the partial turn.
            if let Some(usage) = usage {
                result.turn_usage = usage;
            }
            let safe_message = sanitize_provider_error(&message);
            cancel_pending_tool_calls(state, run_id, seq, result, Some(&safe_message));
            broadcast_agent_event(
                state,
                run_id,
                seq.fetch_add(1, Ordering::Relaxed),
                "error",
                json!({ "message": &safe_message }),
            );
            broadcast_chat_event(
                state,
                run_id,
                session_key,
                seq.fetch_add(1, Ordering::Relaxed),
                "error",
                None,
                Some(&safe_message),
                None,
                None,
            );
            tracing::error!(run_id = %run_id, error = %safe_message, "LLM provider error");
            Err(AgentError::Provider(safe_message))
        }
    }
}

/// Build the `tool_use_cancelled` payload for each tool call that was streamed
/// (and therefore already broadcast to clients as `tool_use`) but will not be
/// executed. Pure projection — no broadcast side effects — so it can be unit
/// tested independently of the WS state.
fn pending_tool_cancellation_payloads(
    pending_tool_calls: &[(String, String, serde_json::Value)],
    reason: Option<&str>,
) -> Vec<Value> {
    pending_tool_calls
        .iter()
        .map(|(tool_use_id, name, _input)| {
            let mut payload = json!({
                "toolUseId": tool_use_id,
                "name": name,
            });
            if let Some(reason) = reason {
                payload["reason"] = json!(reason);
            }
            payload
        })
        .collect()
}

/// Broadcast a `tool_use_cancelled` agent event for each tool call that was
/// streamed but will not be executed because the turn terminated with a
/// non-`ToolUse` stop reason, a provider Error, run cancellation, chunk
/// timeout, or a premature stream end. Without this, clients see a `tool_use`
/// event followed by nothing, leaving the UI in a "tool is running" state
/// forever.
fn broadcast_pending_tool_cancellations(
    state: &Arc<WsServerState>,
    run_id: &str,
    seq: &AtomicU64,
    pending_tool_calls: &[(String, String, serde_json::Value)],
    reason: Option<&str>,
) {
    for payload in pending_tool_cancellation_payloads(pending_tool_calls, reason) {
        broadcast_agent_event(
            state,
            run_id,
            seq.fetch_add(1, Ordering::Relaxed),
            "tool_use_cancelled",
            payload,
        );
    }
}

/// Cancel all currently pending streamed tool calls.
///
/// A pending tool was already broadcast to live clients as `tool_use`. When the
/// turn terminates before those tools are executed, live clients need exactly
/// one `tool_use_cancelled` event and persisted assistant history must not keep
/// executable `tool_use` blocks without matching `tool_result` messages.
fn cancel_pending_tool_calls(
    state: &Arc<WsServerState>,
    run_id: &str,
    seq: &AtomicU64,
    result: &mut StreamResult,
    reason: Option<&str>,
) -> usize {
    let cancelled_count = result.pending_tool_calls.len();
    if cancelled_count == 0 {
        return 0;
    }

    broadcast_pending_tool_cancellations(state, run_id, seq, &result.pending_tool_calls, reason);

    let cancelled_ids: Vec<String> = result
        .pending_tool_calls
        .iter()
        .map(|(id, _, _)| id.clone())
        .collect();
    result.assistant_blocks.retain(|block| match block {
        ContentBlock::ToolUse { id, .. } => {
            !cancelled_ids.iter().any(|cancelled_id| cancelled_id == id)
        }
        _ => true,
    });
    result.pending_tool_calls.clear();

    cancelled_count
}

/// Process the LLM event stream for a single turn.
///
/// Reads events from `rx`, accumulates text deltas and tool-use blocks,
/// broadcasts events to clients, and checks for cancellation. Always
/// returns the partial `StreamResult` alongside an `Ok(())` on a clean
/// Stop or an `Err(AgentError)` on cancellation / stream failure / provider
/// Error. Returning the partial result on Err lets the caller record any
/// tokens accumulated before the failure (`turn_usage` is updated whenever
/// a `StreamEvent::Stop` or `StreamEvent::Error { usage: Some(_), .. }`
/// arrives).
async fn process_llm_stream(
    rx: &mut mpsc::Receiver<StreamEvent>,
    state: &Arc<WsServerState>,
    run_id: &str,
    session_key: &str,
    seq: &AtomicU64,
    cancel_token: &CancellationToken,
) -> (StreamResult, Result<(), AgentError>) {
    let mut result = StreamResult {
        turn_text: String::new(),
        assistant_blocks: Vec::new(),
        pending_tool_calls: Vec::new(),
        stop_reason: StopReason::EndTurn,
        turn_usage: TokenUsage::default(),
    };

    let outcome: Result<(), AgentError> = loop {
        let event = tokio::select! {
            _ = cancel_token.cancelled() => {
                broadcast_agent_event(state, run_id, seq.fetch_add(1, Ordering::Relaxed), "cancelled", json!({}));
                break Err(AgentError::Cancelled);
            }
            recv = tokio::time::timeout(STREAM_CHUNK_TIMEOUT, rx.recv()) => {
                match recv {
                    Ok(Some(e)) => e,
                    Ok(None) => {
                        // Stream ended without a terminal Stop or Error.
                        break Err(AgentError::Stream(
                            "stream ended without stop event".to_string(),
                        ));
                    }
                    Err(_) => {
                        tracing::error!(
                            run_id = %run_id,
                            timeout_secs = STREAM_CHUNK_TIMEOUT.as_secs(),
                            "LLM stream stalled — no data received within chunk timeout"
                        );
                        break Err(AgentError::Provider(format!(
                            "LLM stream stalled — no data received for {}s",
                            STREAM_CHUNK_TIMEOUT.as_secs()
                        )));
                    }
                }
            }
        };

        match handle_stream_event(event, &mut result, state, run_id, session_key, seq) {
            Ok(true) => break Ok(()),
            Ok(false) => continue,
            Err(e) => break Err(e),
        }
    };

    // For non-`handle_stream_event` terminal paths (cancellation, chunk timeout,
    // premature stream end) we still need to clear any `tool_use` events that
    // were already broadcast for tools that will not now execute. The `Stop`/
    // `Error` arms in `handle_stream_event` do this themselves; here we cover
    // the cases where the outer loop breaks before any terminal event arrives.
    if let Err(ref err) = outcome {
        if !result.pending_tool_calls.is_empty() {
            let reason = match err {
                AgentError::Cancelled => Some("cancelled"),
                AgentError::Stream(_) => Some("stream ended without stop event"),
                AgentError::Provider(_) => Some("stream stalled"),
                _ => None,
            };
            cancel_pending_tool_calls(state, run_id, seq, &mut result, reason);
        }
    }

    (result, outcome)
}

#[derive(Default)]
struct TextSanitizationStats {
    finding_count: usize,
    blocked: bool,
    html_modified: bool,
}

fn coalesce_assistant_blocks(blocks: &[ContentBlock]) -> Vec<ContentBlock> {
    let mut coalesced = Vec::new();

    for block in blocks {
        match block {
            ContentBlock::Text { text, metadata } => {
                if text.is_empty() && metadata.is_none() {
                    continue;
                }
                push_text_block(&mut coalesced, text.clone(), metadata.clone());
            }
            _ => coalesced.push(block.clone()),
        }
    }

    coalesced
}

fn push_text_block(
    blocks: &mut Vec<ContentBlock>,
    text: String,
    metadata: Option<ContentBlockMetadata>,
) {
    let incoming_is_signature_only_boundary = text.is_empty() && metadata.is_some();
    if let Some(ContentBlock::Text {
        text: existing,
        metadata: existing_metadata,
    }) = blocks.last_mut()
    {
        let existing_is_signature_only_boundary =
            existing.is_empty() && existing_metadata.is_some();
        if *existing_metadata == metadata
            && !incoming_is_signature_only_boundary
            && !existing_is_signature_only_boundary
        {
            existing.push_str(&text);
            return;
        }
    }

    blocks.push(ContentBlock::Text { text, metadata });
}

fn metadata_boundary_marker(boundary_seed: &str, index: usize) -> String {
    // These control characters are chosen as low-likelihood separators for the
    // synthetic marked-text pass. If a sanitizer normalizes or strips them in a
    // way that prevents boundary recovery, the equality check below falls back
    // to sanitized plain text with metadata dropped for that run.
    format!("\u{001f}CARAPACE_METADATA_BOUNDARY_{boundary_seed}_{index}\u{001e}")
}

fn sanitize_text_with_stats(text: &str, config: &AgentConfig) -> (String, TextSanitizationStats) {
    let mut stats = TextSanitizationStats::default();
    let sanitized = sanitize_text_fragment(text, config, &mut stats);
    (sanitized, stats)
}

fn append_text_sanitization_stats(
    total: &mut TextSanitizationStats,
    current: &TextSanitizationStats,
) {
    total.finding_count += current.finding_count;
    total.blocked |= current.blocked;
    total.html_modified |= current.html_modified;
}

fn split_sanitized_text_run(
    text_blocks: &[(String, Option<ContentBlockMetadata>)],
    sanitized_marked_text: &str,
    boundary_seed: &str,
) -> Option<Vec<String>> {
    if text_blocks.is_empty() {
        return Some(Vec::new());
    }

    let mut segments = Vec::with_capacity(text_blocks.len());
    let mut remaining = sanitized_marked_text;

    for marker_index in 0..text_blocks.len().saturating_sub(1) {
        let marker = metadata_boundary_marker(boundary_seed, marker_index);
        let marker_pos = remaining.find(&marker)?;
        segments.push(remaining[..marker_pos].to_string());
        remaining = &remaining[marker_pos + marker.len()..];
    }

    segments.push(remaining.to_string());
    Some(segments)
}

fn sanitize_text_run(
    text_blocks: &[(String, Option<ContentBlockMetadata>)],
    config: &AgentConfig,
    stats: &mut TextSanitizationStats,
    run_id: &str,
) -> (String, Vec<ContentBlock>) {
    let original_text: String = text_blocks.iter().map(|(text, _)| text).cloned().collect();
    let (sanitized_text, run_stats) = sanitize_text_with_stats(&original_text, config);
    append_text_sanitization_stats(stats, &run_stats);

    let has_provider_metadata = text_blocks.iter().any(|(_, metadata)| metadata.is_some());
    if !has_provider_metadata {
        let mut blocks = Vec::new();
        if !sanitized_text.is_empty() {
            blocks.push(ContentBlock::Text {
                text: sanitized_text.clone(),
                metadata: None,
            });
        }
        return (sanitized_text, blocks);
    }

    let boundary_seed = Uuid::new_v4().simple().to_string();
    let mut marked_text = String::new();
    for (index, (text, _)) in text_blocks.iter().enumerate() {
        if index > 0 {
            marked_text.push_str(&metadata_boundary_marker(&boundary_seed, index - 1));
        }
        marked_text.push_str(text);
    }

    // The marker-bearing pass is only used to recover per-block boundaries
    // after sanitization. Its stats are intentionally discarded because the
    // marked string contains synthetic delimiters; the authoritative
    // audit/logging signal comes from the unmarked `original_text` pass above.
    let (sanitized_marked_text, _) = sanitize_text_with_stats(&marked_text, config);
    if let Some(sanitized_segments) =
        split_sanitized_text_run(text_blocks, &sanitized_marked_text, &boundary_seed)
    {
        let reconstructed_text: String = sanitized_segments.concat();
        if reconstructed_text == sanitized_text {
            let mut blocks = Vec::with_capacity(text_blocks.len());
            for ((_, metadata), sanitized_segment) in text_blocks.iter().zip(sanitized_segments) {
                if !sanitized_segment.is_empty() || metadata.is_some() {
                    blocks.push(ContentBlock::Text {
                        text: sanitized_segment,
                        metadata: metadata.clone(),
                    });
                }
            }
            return (sanitized_text, blocks);
        }
    }

    tracing::warn!(
        run_id = %run_id,
        "metadata-bearing text run required whole-turn sanitization fallback; dropping per-part metadata for that run"
    );

    let mut blocks = Vec::new();
    if !sanitized_text.is_empty() {
        blocks.push(ContentBlock::Text {
            text: sanitized_text.clone(),
            metadata: None,
        });
    }
    (sanitized_text, blocks)
}

fn sanitize_text_fragment(
    text: &str,
    config: &AgentConfig,
    stats: &mut TextSanitizationStats,
) -> String {
    let mut sanitized = text.to_string();

    if config.prompt_guard.enabled && config.prompt_guard.postflight.enabled {
        let postflight_result =
            postflight::filter_output(&sanitized, &config.prompt_guard.postflight);
        if !postflight_result.is_clean() {
            stats.finding_count += postflight_result.findings.len();
            stats.blocked |= postflight_result.blocked;
        }
        sanitized = postflight_result.sanitized;
    }

    if config.output_sanitizer.sanitize_html {
        let sanitized_output =
            crate::agent::output_sanitizer::sanitize_output(&sanitized, &config.output_sanitizer);
        if sanitized_output.was_modified {
            stats.html_modified = true;
        }
        sanitized = sanitized_output.content;
    }

    sanitized
}

fn log_text_sanitization_stats(run_id: &str, stats: &TextSanitizationStats) {
    if stats.finding_count > 0 {
        tracing::warn!(
            run_id = %run_id,
            findings = stats.finding_count,
            blocked = stats.blocked,
            "prompt guard post-flight detected sensitive content in output"
        );
        if stats.blocked {
            crate::logging::audit::audit(crate::logging::audit::AuditEvent::PromptGuardBlocked {
                layer: "postflight".to_string(),
                reason: format!("{} findings (output sanitized)", stats.finding_count),
                run_id: crate::logging::audit::truncate_audit_free_text_field(
                    run_id,
                    crate::logging::audit::AUDIT_FREE_TEXT_FIELD_MAX_BYTES,
                ),
            });
        }
    }

    if stats.html_modified {
        tracing::info!(
            run_id = %run_id,
            "output sanitizer modified agent response for safe rendering"
        );
    }
}

fn sanitize_assistant_turn(
    assistant_blocks: &[ContentBlock],
    config: &AgentConfig,
    run_id: &str,
) -> (String, Vec<ContentBlock>) {
    let assistant_blocks = coalesce_assistant_blocks(assistant_blocks);
    let mut stats = TextSanitizationStats::default();
    let mut sanitized_turn_text = String::new();
    let mut sanitized_blocks = Vec::with_capacity(assistant_blocks.len());
    let mut pending_text_run: Vec<(String, Option<ContentBlockMetadata>)> = Vec::new();

    let flush_text_run = |pending_text_run: &mut Vec<(String, Option<ContentBlockMetadata>)>,
                          sanitized_turn_text: &mut String,
                          sanitized_blocks: &mut Vec<ContentBlock>,
                          stats: &mut TextSanitizationStats| {
        if pending_text_run.is_empty() {
            return;
        }

        let (sanitized_text, sanitized_run_blocks) =
            sanitize_text_run(pending_text_run, config, stats, run_id);
        sanitized_turn_text.push_str(&sanitized_text);
        sanitized_blocks.extend(sanitized_run_blocks);
        pending_text_run.clear();
    };

    for block in assistant_blocks {
        match block {
            ContentBlock::Text { text, metadata } => {
                pending_text_run.push((text, metadata));
            }
            other => {
                flush_text_run(
                    &mut pending_text_run,
                    &mut sanitized_turn_text,
                    &mut sanitized_blocks,
                    &mut stats,
                );
                sanitized_blocks.push(other);
            }
        }
    }

    flush_text_run(
        &mut pending_text_run,
        &mut sanitized_turn_text,
        &mut sanitized_blocks,
        &mut stats,
    );

    log_text_sanitization_stats(run_id, &stats);
    (sanitized_turn_text, sanitized_blocks)
}

fn combined_assistant_text(assistant_blocks: &[ContentBlock]) -> String {
    assistant_blocks
        .iter()
        .filter_map(|block| match block {
            ContentBlock::Text { text, .. } => Some(text.as_str()),
            _ => None,
        })
        .collect()
}

/// Build an assistant history message from accumulated turn text and tool-use blocks.
///
/// Returns the `ChatMessage` ready for persistence, or `None` if there is
/// nothing to record (empty text and no tool calls).
fn build_assistant_message(
    session_id: &str,
    assistant_blocks: &[ContentBlock],
    output_tokens: u64,
) -> Option<ChatMessage> {
    if assistant_blocks.is_empty() {
        return None;
    }

    let has_tool_use = assistant_blocks
        .iter()
        .any(|block| matches!(block, ContentBlock::ToolUse { .. }));
    let has_provider_metadata = assistant_blocks
        .iter()
        .any(ContentBlock::has_provider_metadata);
    let turn_text = combined_assistant_text(assistant_blocks);

    let content = if !has_tool_use && !has_provider_metadata {
        turn_text.clone()
    } else if has_provider_metadata {
        crate::agent::context::serialize_assistant_blocks(assistant_blocks)
            .unwrap_or_else(|_| turn_text.clone())
    } else {
        let mut blocks: Vec<ContentBlock> = Vec::new();
        if !turn_text.is_empty() {
            blocks.push(ContentBlock::Text {
                text: turn_text.clone(),
                metadata: None,
            });
        }
        for block in assistant_blocks {
            if let ContentBlock::ToolUse { .. } = block {
                blocks.push(block.clone());
            }
        }
        crate::agent::context::serialize_assistant_blocks(&blocks)
            .unwrap_or_else(|_| turn_text.clone())
    };

    Some(ChatMessage::assistant(session_id, &content).with_tokens(output_tokens))
}

/// Execute pending tool calls with exfiltration guard and tool-policy checks,
/// broadcast results, and return the corresponding history messages.
///
/// SECURITY/FEATURE: `agent_id` is threaded from `session.metadata.agent_id`
/// (available at the `execute_single_turn` caller) down to
/// `tools::execute_tool_call_with_sandbox`. Without this, the builtin
/// tools that gate fail-closed on `ctx.agent_id.is_none()` (i.e.
/// `session_list_tool` and `session_read_tool`) refuse every invocation
/// from a production agent run — a feature regression hidden behind
/// the prior security fix. The agent_id is also routed to
/// `memory_store_path` so each agent gets its own memory namespace
/// rather than colliding on the shared `"default.json"` fallback.
#[allow(clippy::too_many_arguments)]
fn execute_tools_with_guards(
    pending_tool_calls: &[(String, String, Value)],
    config: &AgentConfig,
    state: &Arc<WsServerState>,
    session_id: &str,
    session_key: &str,
    agent_id: Option<&str>,
    message_channel: Option<&str>,
    run_id: &str,
    seq: &AtomicU64,
) -> Vec<ChatMessage> {
    let mut tool_msgs = Vec::with_capacity(pending_tool_calls.len());

    for (tool_id, tool_name, tool_input) in pending_tool_calls {
        let mut tool_input = tool_input.clone();
        let original_tool_input = tool_input.clone();

        // Check exfiltration guard before tool policy (defence-in-depth)
        let tool_result = if config.exfiltration_guard
            && crate::agent::exfiltration::is_exfiltration_sensitive(tool_name)
        {
            ToolCallResult::Error {
                message: format!(
                    "Tool \"{}\" is blocked by the exfiltration guard. \
                     This tool sends data externally and requires explicit approval. \
                     Set exfiltration_guard: false in agent config to allow.",
                    tool_name
                ),
            }
        } else if !config.tool_policy.is_allowed(tool_name) {
            ToolCallResult::Error {
                message: format!("Tool \"{}\" is not available for this agent", tool_name),
            }
        } else if let Some(result) = dispatch_plugin_hook(
            state,
            "before_tool_call",
            &json!({
                "runId": run_id,
                "sessionKey": session_key,
                "toolUseId": tool_id,
                "name": tool_name,
                "input": &tool_input,
                "messageChannel": message_channel,
            }),
        ) {
            if result.cancelled {
                ToolCallResult::Error {
                    message: format!("Tool \"{}\" cancelled by hook", tool_name),
                }
            } else {
                if let Some(payload) = parse_hook_payload(&result, "before_tool_call") {
                    apply_tool_input_override(&mut tool_input, &payload);
                    if tool_input != original_tool_input {
                        tracing::info!(
                            run_id = %run_id,
                            tool = %tool_name,
                            tool_use_id = %tool_id,
                            "tool input modified by hook"
                        );
                    }
                }

                if let Some(tools_registry) = state.tools_registry() {
                    let sandbox = if config.process_sandbox.enabled {
                        Some(&config.process_sandbox)
                    } else {
                        None
                    };
                    tools::execute_tool_call_with_sandbox(
                        tool_name,
                        tool_input.clone(),
                        tools_registry,
                        session_key,
                        agent_id,
                        message_channel,
                        sandbox,
                    )
                } else {
                    ToolCallResult::Error {
                        message: "no tools registry available".to_string(),
                    }
                }
            }
        } else if let Some(tools_registry) = state.tools_registry() {
            let sandbox = if config.process_sandbox.enabled {
                Some(&config.process_sandbox)
            } else {
                None
            };
            // Thread `agent_id` on the non-hook fallback too. This branch
            // is practically unreachable in production (`dispatch_plugin_hook`
            // returns `Some(HookDispatchResult { handler_count: 0, .. })`
            // whenever the plugin registry is wired up, which it always is
            // outside tests), but the prior `None` argument here diverged
            // from the hook-success branch above and would silently route
            // `session_list_tool` / `session_read_tool` to the shared
            // default namespace (or fail-closed) if reached. Threading
            // agent_id on both branches keeps agent-scoped builtin tools
            // consistent regardless of which path runs.
            tools::execute_tool_call_with_sandbox(
                tool_name,
                tool_input.clone(),
                tools_registry,
                session_key,
                agent_id,
                message_channel,
                sandbox,
            )
        } else {
            ToolCallResult::Error {
                message: "no tools registry available".to_string(),
            }
        };

        let (mut result_content, mut is_error) = match &tool_result {
            ToolCallResult::Ok { output } => (output.clone(), false),
            ToolCallResult::Error { message } => (message.clone(), true),
        };

        let _ = dispatch_plugin_hook(
            state,
            "after_tool_call",
            &json!({
                "runId": run_id,
                "sessionKey": session_key,
                "toolUseId": tool_id,
                "name": tool_name,
                "originalInput": &original_tool_input,
                "input": &tool_input,
                "result": &result_content,
                "isError": is_error,
                "messageChannel": message_channel,
            }),
        );

        if let Some(result) = dispatch_plugin_hook(
            state,
            "tool_result_persist",
            &json!({
                "runId": run_id,
                "sessionKey": session_key,
                "toolUseId": tool_id,
                "name": tool_name,
                "originalInput": &original_tool_input,
                "input": &tool_input,
                "result": &result_content,
                "isError": is_error,
                "messageChannel": message_channel,
            }),
        ) {
            if result.cancelled {
                result_content = format!("Tool \"{}\" result suppressed by hook", tool_name);
                is_error = true;
            } else if let Some(payload) = parse_hook_payload(&result, "tool_result_persist") {
                apply_tool_result_override(&mut result_content, &mut is_error, &payload);
            }
        }

        // Broadcast tool_result event
        broadcast_agent_event(
            state,
            run_id,
            seq.fetch_add(1, Ordering::Relaxed),
            "tool_result",
            json!({
                "toolUseId": tool_id,
                "name": tool_name,
                "result": &result_content,
                "isError": is_error,
            }),
        );

        let mut tool_msg = ChatMessage::tool(session_id, tool_name, tool_id, &result_content);
        if is_error {
            tool_msg.metadata = Some(json!({"is_error": true}));
        }
        tool_msgs.push(tool_msg);
    }

    tool_msgs
}

fn usage_provider_for_model(model: &str) -> Option<&'static str> {
    if crate::agent::claude_cli::is_claude_cli_model(model) {
        Some("claude-cli")
    } else if crate::agent::ollama::is_ollama_model(model) {
        Some("ollama")
    } else if crate::agent::nearai::is_nearai_model(model) {
        Some("nearai")
    } else if crate::agent::venice::is_venice_model(model) {
        Some("venice")
    } else if crate::agent::vertex::is_vertex_model(model) {
        Some("vertex")
    } else if crate::agent::gemini::is_gemini_model(model) {
        Some("gemini")
    } else if crate::agent::codex::is_codex_model(model) {
        Some("codex")
    } else if crate::agent::openai::is_openai_model(model) {
        Some("openai")
    } else if crate::agent::bedrock::is_bedrock_model(model) {
        Some("bedrock")
    } else if crate::agent::anthropic::is_anthropic_model(model) {
        Some("anthropic")
    } else {
        None
    }
}

/// Record token usage for a single turn via the usage tracker.
fn record_turn_usage(session_key: &str, model: &str, usage: &TokenUsage) {
    let provider_name = match usage_provider_for_model(model) {
        Some(provider) => provider,
        None => {
            tracing::warn!(model, "unrecognized model prefix for usage recording");
            "unknown"
        }
    };
    crate::server::ws::record_usage(
        session_key,
        provider_name,
        model,
        usage.input_tokens,
        usage.output_tokens,
    );
}

/// Build the `CompletionRequest` for a single turn from in-memory history.
fn build_turn_request(
    history: &[ChatMessage],
    config: &AgentConfig,
    state: &Arc<WsServerState>,
    message_channel: Option<&str>,
) -> CompletionRequest {
    let (system, messages) = if config.prompt_guard.enabled && config.prompt_guard.tagging.enabled {
        build_context_with_tagging(
            history,
            config.system.as_deref(),
            &config.prompt_guard.tagging,
        )
    } else {
        build_context(history, config.system.as_deref())
    };

    let tools = if let Some(tools_registry) = state.tools_registry() {
        let all_tools = tools::list_provider_tools(tools_registry, message_channel);
        config
            .tool_policy
            .filter_tools_with_guard(all_tools, config.exfiltration_guard)
    } else {
        vec![]
    };

    CompletionRequest {
        model: config.model.clone(),
        messages,
        system,
        tools,
        max_tokens: config.max_tokens,
        temperature: config.temperature,
        extra: config.extra.clone(),
    }
}

/// Broadcast run-completion events and mark the run as completed in the registry.
#[allow(clippy::too_many_arguments)]
fn finalize_run(
    state: &Arc<WsServerState>,
    run_id: &str,
    session_key: &str,
    seq: &AtomicU64,
    final_stop_reason: StopReason,
    total_input_tokens: u64,
    total_output_tokens: u64,
    accumulated_text: String,
    csp_policy: &str,
) {
    let stop_reason_str = match final_stop_reason {
        StopReason::EndTurn => "end_turn",
        StopReason::MaxTokens => "max_tokens",
        StopReason::ToolUse => "tool_use",
    };

    broadcast_agent_event(
        state,
        run_id,
        seq.fetch_add(1, Ordering::Relaxed),
        "complete",
        json!({
            "stopReason": stop_reason_str,
            "usage": {
                "inputTokens": total_input_tokens,
                "outputTokens": total_output_tokens,
            },
            "contentPolicy": {
                "csp": csp_policy,
            }
        }),
    );

    broadcast_chat_event(
        state,
        run_id,
        session_key,
        seq.fetch_add(1, Ordering::Relaxed),
        "final",
        Some(json!({
            "content": &accumulated_text,
            "contentPolicy": {
                "csp": csp_policy,
            }
        })),
        None,
        Some(json!({
            "inputTokens": total_input_tokens,
            "outputTokens": total_output_tokens,
        })),
        Some(stop_reason_str),
    );

    let mut registry = state.agent_run_registry.lock();
    registry.mark_completed(run_id, accumulated_text);
}

fn delivery_context_from_registry(state: &WsServerState, run_id: &str) -> Option<String> {
    let registry = state.agent_run_registry.lock();
    registry
        .get(run_id)
        .map(|run| run.delivery_recipient_id.clone())
        .unwrap_or_else(|| {
            tracing::warn!(
                run_id = %run_id,
                "agent run missing from registry before delivery finalization"
            );
            None
        })
}

/// Execute a single LLM turn: call the provider, stream the response,
/// record usage, persist the assistant message, and optionally execute tools.
///
/// Returns `Ok(true)` if the loop should continue (tool calls pending),
/// `Ok(false)` if the run is done (end-turn / max-tokens).
#[allow(clippy::too_many_arguments)]
async fn execute_single_turn(
    config: &AgentConfig,
    state: &Arc<WsServerState>,
    provider: &Arc<dyn LlmProvider>,
    cancel_token: &CancellationToken,
    run_id: &str,
    session_key: &str,
    session_id: &str,
    agent_id: Option<&str>,
    message_channel: Option<&str>,
    seq: &AtomicU64,
    history: &mut Vec<ChatMessage>,
    accumulated_text: &mut String,
    total_input_tokens: &mut u64,
    total_output_tokens: &mut u64,
    final_stop_reason: &mut StopReason,
) -> Result<bool, AgentError> {
    // Check cancellation
    if cancel_token.is_cancelled() {
        broadcast_agent_event(
            state,
            run_id,
            seq.fetch_add(1, Ordering::Relaxed),
            "cancelled",
            json!({}),
        );
        return Err(AgentError::Cancelled);
    }

    let request = build_turn_request(history, config, state, message_channel);

    // Call LLM (with per-turn timeout)
    let mut rx = match tokio::time::timeout(
        TURN_TIMEOUT,
        provider.complete(request, cancel_token.clone()),
    )
    .await
    {
        Ok(Ok(rx)) => rx,
        Ok(Err(AgentError::Cancelled)) => return Err(AgentError::Cancelled),
        Ok(Err(e)) => return Err(AgentError::Provider(e.to_string())),
        Err(_) => {
            return Err(AgentError::Provider(format!(
                "LLM turn timed out after {}s",
                TURN_TIMEOUT.as_secs()
            )));
        }
    };

    let (
        StreamResult {
            turn_text: _,
            assistant_blocks,
            pending_tool_calls,
            stop_reason,
            turn_usage,
        },
        outcome,
    ) = process_llm_stream(&mut rx, state, run_id, session_key, seq, cancel_token).await;

    // Track usage when a terminal event actually populated it, so the operator
    // can account for tokens consumed by partial turns (e.g. a Google stream
    // that terminated with MALFORMED_FUNCTION_CALL after several content
    // chunks). Skip the zero-token case — cancellation / chunk-timeout /
    // premature-stream-end terminate before any usage was reported, and
    // calling `record_turn_usage` with zero would still bump the per-day /
    // per-month / per-session `requests` counters and append an empty
    // `UsageRecord` to disk on every cancelled run.
    if turn_usage.input_tokens > 0 || turn_usage.output_tokens > 0 {
        *total_input_tokens += turn_usage.input_tokens;
        *total_output_tokens += turn_usage.output_tokens;
        record_turn_usage(session_key, &config.model, &turn_usage);
    }

    // Propagate stream error after recording usage.
    outcome?;

    let (turn_text, assistant_blocks) = sanitize_assistant_turn(&assistant_blocks, config, run_id);

    // Append assistant message to history (after post-flight filtering)
    if let Some(msg) =
        build_assistant_message(session_id, &assistant_blocks, turn_usage.output_tokens)
    {
        crate::sessions::append_message_blocking(state.session_store().clone(), msg.clone())
            .await
            .map_err(|e| AgentError::SessionStore(e.to_string()))?;
        history.push(msg);
    }

    accumulated_text.push_str(&turn_text);
    *final_stop_reason = stop_reason;

    // If there are tool calls, execute them and signal continuation
    if !pending_tool_calls.is_empty() && stop_reason == StopReason::ToolUse {
        let tool_msgs = execute_tools_with_guards(
            &pending_tool_calls,
            config,
            state,
            session_id,
            session_key,
            agent_id,
            message_channel,
            run_id,
            seq,
        );
        crate::sessions::append_messages_blocking(state.session_store().clone(), tool_msgs.clone())
            .await
            .map_err(|e| AgentError::SessionStore(e.to_string()))?;
        history.extend(tool_msgs);
        return Ok(true);
    }

    Ok(false)
}

/// Execute an agent run to completion.
///
/// This is the core loop: load history → call LLM → stream results →
/// handle tool calls → append to history → mark complete.
pub async fn execute_run(
    run_id: String,
    session_key: String,
    mut config: AgentConfig,
    state: Arc<WsServerState>,
    provider: Arc<dyn LlmProvider>,
    cancel_token: CancellationToken,
) -> Result<(), AgentError> {
    let seq = AtomicU64::new(0);

    // 1. Mark run as Running (returns false if already cancelled)
    let run_started = {
        let mut registry = state.agent_run_registry.lock();
        registry.mark_started(&run_id)
    };
    if !run_started {
        return Err(AgentError::Cancelled);
    }

    let result = async {
        broadcast_agent_event(
            &state,
            &run_id,
            seq.fetch_add(1, Ordering::Relaxed),
            "started",
            json!({ "sessionKey": &session_key, "model": &config.model }),
        );

        // 2. Load or create session
        let session = crate::sessions::get_session_by_key_blocking(
            state.session_store().clone(),
            session_key.clone(),
        )
        .await
        .map_err(|e| AgentError::SessionNotFound(format!("{session_key}: {e}")))?;
        let message_channel = session.metadata.channel.clone();

        if let Some(result) = dispatch_plugin_hook(
            &state,
            "before_agent_start",
            &json!({
                "runId": run_id,
                "sessionKey": session_key,
                "model": &config.model,
                "system": &config.system,
                "maxTokens": config.max_tokens,
                "temperature": config.temperature,
                "deliver": config.deliver,
                "messageChannel": &message_channel,
                "extra": &config.extra,
            }),
        ) {
            if result.cancelled {
                return Err(AgentError::Cancelled);
            }
            if let Some(payload) = parse_hook_payload(&result, "before_agent_start") {
                let cfg =
                    crate::config::load_config().unwrap_or(Value::Object(serde_json::Map::new()));
                apply_agent_hook_overrides(&mut config, &payload, &cfg)?;
            }
        }

        // 2b. Pre-flight system prompt check
        if config.prompt_guard.enabled && config.prompt_guard.preflight.enabled {
            if let Some(ref system) = config.system {
                let preflight_result =
                    preflight::analyze_system_prompt(system, &config.prompt_guard.preflight);
                if preflight_result.has_critical() {
                    let reasons: Vec<String> = preflight_result
                        .findings
                        .iter()
                        .map(|f| f.description.clone())
                        .collect();
                    let reason = reasons.join("; ");
                    tracing::warn!(
                        run_id = %run_id,
                        findings = %reason,
                        "prompt guard pre-flight blocked system prompt"
                    );
                    crate::logging::audit::audit(
                        crate::logging::audit::AuditEvent::PromptGuardBlocked {
                            layer: "preflight".to_string(),
                            reason: crate::logging::audit::truncate_audit_free_text_field(
                                &reason,
                                crate::logging::audit::AUDIT_FREE_TEXT_FIELD_MAX_BYTES,
                            ),
                            run_id: crate::logging::audit::truncate_audit_free_text_field(
                                &run_id,
                                crate::logging::audit::AUDIT_FREE_TEXT_FIELD_MAX_BYTES,
                            ),
                        },
                    );
                    return Err(AgentError::Provider(format!(
                        "Prompt guard pre-flight blocked: {reason}"
                    )));
                }
            }
        }

        // 2c. Classifier (optional, fail-open)
        if let Some(ref clf_config) = config.classifier {
            if clf_config.enabled
                && clf_config.mode != crate::agent::classifier::ClassifierMode::Off
            {
                let last_user_msg = crate::sessions::get_history_blocking(
                    state.session_store().clone(),
                    session.id.clone(),
                    None,
                    None,
                )
                .await
                .ok()
                .and_then(|history| {
                    history
                        .iter()
                        .rev()
                        .find(|m| m.role == MessageRole::User)
                        .map(|m| m.content.clone())
                });

                if let Some(user_message) = last_user_msg {
                    match crate::agent::classifier::classify_message(
                        &user_message,
                        clf_config,
                        &config.model,
                        provider.as_ref(),
                    )
                    .await
                    {
                        Ok(verdict) if verdict.should_block(clf_config) => {
                            crate::logging::audit::audit(
                                crate::logging::audit::AuditEvent::ClassifierBlocked {
                                    category: verdict.category.to_string(),
                                    confidence: verdict.confidence as f64,
                                    reasoning:
                                        crate::logging::audit::truncate_audit_free_text_field(
                                            &verdict.reasoning,
                                            crate::logging::audit::AUDIT_FREE_TEXT_FIELD_MAX_BYTES,
                                        ),
                                    run_id: crate::logging::audit::truncate_audit_free_text_field(
                                        &run_id,
                                        crate::logging::audit::AUDIT_FREE_TEXT_FIELD_MAX_BYTES,
                                    ),
                                },
                            );
                            return Err(AgentError::ClassifierBlocked(
                                verdict.category.to_string(),
                                verdict.reasoning,
                            ));
                        }
                        Ok(verdict) if verdict.should_warn(clf_config) => {
                            crate::logging::audit::audit(
                                crate::logging::audit::AuditEvent::ClassifierWarned {
                                    category: verdict.category.to_string(),
                                    confidence: verdict.confidence as f64,
                                    reasoning:
                                        crate::logging::audit::truncate_audit_free_text_field(
                                            &verdict.reasoning,
                                            crate::logging::audit::AUDIT_FREE_TEXT_FIELD_MAX_BYTES,
                                        ),
                                    run_id: crate::logging::audit::truncate_audit_free_text_field(
                                        &run_id,
                                        crate::logging::audit::AUDIT_FREE_TEXT_FIELD_MAX_BYTES,
                                    ),
                                },
                            );
                            tracing::warn!(
                                run_id = %run_id,
                                category = %verdict.category,
                                confidence = verdict.confidence,
                                "classifier warned: {}",
                                verdict.reasoning
                            );
                        }
                        Ok(_) => {}
                        Err(e) => {
                            tracing::warn!(
                                run_id = %run_id,
                                error = %e,
                                "classifier error (fail-open)"
                            );
                        }
                    }
                }
            }
        }

        let delivery_channel_id = if config.deliver {
            message_channel.as_deref()
        } else {
            None
        };
        let channel_activity_policy = if let Some(channel_id) = delivery_channel_id {
            Some(crate::channels::activity::load_channel_activity_policy_async(channel_id).await)
        } else {
            None
        };
        let (channel_capabilities, capability_probe_failed) =
            if let (Some(channel_id), Some(plugin_registry)) =
                (delivery_channel_id, state.plugin_registry())
            {
                if let Some(plugin) = plugin_registry.get_channel(channel_id) {
                    match tokio::task::spawn_blocking(move || plugin.get_capabilities()).await {
                        Ok(Ok(capabilities)) => (Some(capabilities), false),
                        Ok(Err(err)) => {
                            tracing::warn!(
                                run_id = %run_id,
                                channel = %channel_id,
                                error = %err,
                                "failed to load channel capabilities for activity features"
                            );
                            (None, true)
                        }
                        Err(err) => {
                            tracing::warn!(
                                run_id = %run_id,
                                channel = %channel_id,
                                error = %err,
                                "activity capability worker failed"
                            );
                            (None, true)
                        }
                    }
                } else {
                    (None, false)
                }
            } else {
                (None, false)
            };

        let mut typing_handle = if let (Some(channel_id), Some(plugin_registry), Some(policy)) = (
            delivery_channel_id,
            state.plugin_registry(),
            channel_activity_policy.as_ref(),
        ) {
            if !channel_capabilities
                .as_ref()
                .map(|capabilities| capabilities.typing_indicators)
                .unwrap_or(false)
            {
                if policy.typing.enabled && !capability_probe_failed {
                    state
                        .activity_service()
                        .warn_unsupported_feature(channel_id, "typing");
                }
                None
            } else {
                let typing_context = {
                    let registry = state.agent_run_registry.lock();
                    registry
                        .get(&run_id)
                        .and_then(|run| run.typing_context.clone())
                        .or_else(|| {
                            session.metadata.chat_id.clone().map(|to| {
                                crate::plugins::TypingContext {
                                    to,
                                    ..Default::default()
                                }
                            })
                        })
                };

                match typing_context {
                    Some(ctx) => {
                        // Typing policy is snapshotted for the lifetime of this
                        // run. Reload only affects future runs; do not re-check the
                        // policy mid-run and break the receive-time snapshot
                        // contract for activity side effects.
                        crate::channels::activity::maybe_start_typing_loop(
                            plugin_registry.clone(),
                            channel_id,
                            policy,
                            channel_capabilities.clone(),
                            state.activity_service().clone(),
                            ctx,
                        )
                        .await
                    }
                    None => None,
                }
            }
        } else {
            None
        };

        let execution_result = async {
            // 3. Main agentic loop
            let mut accumulated_text = String::new();
            let mut total_input_tokens: u64 = 0;
            let mut total_output_tokens: u64 = 0;
            let mut final_stop_reason = StopReason::EndTurn;

            let mut history = crate::sessions::get_history_blocking(
                state.session_store().clone(),
                session.id.clone(),
                None,
                None,
            )
            .await
            .map_err(|e| AgentError::SessionStore(e.to_string()))?;

            for _turn in 0..config.max_turns {
                let should_continue = execute_single_turn(
                    &config,
                    &state,
                    &provider,
                    &cancel_token,
                    &run_id,
                    &session_key,
                    &session.id,
                    session.metadata.agent_id.as_deref(),
                    message_channel.as_deref(),
                    &seq,
                    &mut history,
                    &mut accumulated_text,
                    &mut total_input_tokens,
                    &mut total_output_tokens,
                    &mut final_stop_reason,
                )
                .await?;

                if !should_continue {
                    break;
                }
            }

            let delivery_text = if config.deliver {
                Some(accumulated_text.clone())
            } else {
                None
            };
            let delivery_recipient_id = delivery_context_from_registry(&state, &run_id);

            if let Some(handle) = typing_handle.take() {
                handle.stop().await;
            }

            if let Some(text) = delivery_text {
                if !text.is_empty() {
                    let recipient_id =
                        delivery_recipient_id.or_else(|| session.metadata.chat_id.clone());
                    if let (Some(channel_id), Some(chat_id)) = (&message_channel, recipient_id) {
                        let metadata = crate::messages::outbound::MessageMetadata {
                            recipient_id: Some(chat_id),
                            ..Default::default()
                        };
                        let outbound = crate::messages::outbound::OutboundMessage::new(
                            channel_id.clone(),
                            crate::messages::outbound::MessageContent::text(text),
                        )
                        .with_metadata(metadata);
                        let mut ctx = crate::messages::outbound::OutboundContext::new()
                            .with_trace_id(&run_id)
                            .with_source("agent");
                        if channel_id == crate::channels::matrix::MATRIX_CHANNEL_ID {
                            ctx =
                                ctx.with_retries(crate::channels::matrix::MATRIX_OUTBOUND_RETRIES);
                        }
                        if let Err(err) = state.message_pipeline().queue(outbound, ctx) {
                            tracing::warn!(
                                run_id = %run_id,
                                channel = %channel_id,
                                error = %err,
                                "failed to queue agent response for delivery"
                            );
                        }
                    }
                }
            }

            finalize_run(
                &state,
                &run_id,
                &session_key,
                &seq,
                final_stop_reason,
                total_input_tokens,
                total_output_tokens,
                accumulated_text,
                &config.output_sanitizer.csp_policy,
            );

            Ok(())
        }
        .await;

        if let Some(handle) = typing_handle {
            handle.stop().await;
        }

        execution_result
    }
    .await;

    result
}

/// Sanitize a provider error message before sending to clients.
///
/// Strips potential secrets (API keys, internal URLs with auth) from error
/// messages while preserving the error type and human-readable portion.
fn sanitize_provider_error(message: &str) -> String {
    use std::sync::LazyLock;

    static API_KEY_RE: LazyLock<regex::Regex> = LazyLock::new(|| {
        regex::Regex::new(r"(sk-ant-|sk-|key-)[A-Za-z0-9_-]{10,}")
            .expect("failed to compile regex: api_key_pattern")
    });
    static AUTH_HEADER_RE: LazyLock<regex::Regex> = LazyLock::new(|| {
        regex::Regex::new(r"(?i)(authorization|x-api-key):\s*(bearer\s+)?\S+")
            .expect("failed to compile regex: auth_header_pattern")
    });

    // Strip anything that looks like an API key (sk-ant-..., sk-..., key-...)
    let sanitized = API_KEY_RE.replace_all(message, "[REDACTED]");

    // Strip Authorization and x-api-key header values
    let sanitized = AUTH_HEADER_RE
        .replace_all(&sanitized, "$1: [REDACTED]")
        .into_owned();

    // Cap length to prevent huge error payloads (char-boundary safe)
    if sanitized.len() > 500 {
        let boundary = sanitized
            .char_indices()
            .map(|(i, _)| i)
            .take_while(|&i| i <= 500)
            .last()
            .unwrap_or(0);
        format!("{}... (truncated)", &sanitized[..boundary])
    } else {
        sanitized
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::agent::gemini::GeminiProvider;
    use crate::agent::provider::{LlmProvider, StopReason, StreamEvent, TokenUsage};
    use crate::agent::vertex::build_gemini_body as build_vertex_gemini_body;
    use crate::agent::AgentConfig;
    use crate::plugins::{
        BindingError, ChannelCapabilities, ChannelPluginInstance, DeliveryResult, OutboundContext,
        PluginRegistry, ReadReceiptContext, TypingContext,
    };
    use crate::server::ws::{WsServerConfig, WsServerState};
    use crate::sessions;
    use async_trait::async_trait;
    use parking_lot::Mutex;
    use std::sync::Arc;
    use tokio::sync::mpsc;
    use tokio::sync::Notify;
    use tokio_util::sync::CancellationToken;

    /// Mock LLM provider that returns canned responses.
    struct MockProvider {
        /// Each call to `complete` pops the next response sequence.
        responses: parking_lot::Mutex<Vec<Vec<StreamEvent>>>,
    }

    impl MockProvider {
        /// Create a provider that returns the given sequences in order.
        fn new(responses: Vec<Vec<StreamEvent>>) -> Self {
            // Reverse so we can pop from the back (FIFO via pop from reversed vec)
            let mut reversed = responses;
            reversed.reverse();
            Self {
                responses: parking_lot::Mutex::new(reversed),
            }
        }

        /// Convenience: single-turn text response.
        fn text(text: &str) -> Self {
            Self::new(vec![vec![
                StreamEvent::TextDelta {
                    text: text.to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ]])
        }

        /// Always returns a tool_use, useful for testing max_turns.
        fn always_tool_use(turns: usize) -> Self {
            let mut responses = Vec::new();
            for _ in 0..turns {
                responses.push(vec![
                    StreamEvent::ToolUse {
                        id: "tool_1".to_string(),
                        name: "time".to_string(),
                        input: serde_json::json!({}),
                        metadata: None,
                    },
                    StreamEvent::Stop {
                        reason: StopReason::ToolUse,
                        usage: TokenUsage {
                            input_tokens: 10,
                            output_tokens: 5,
                        },
                    },
                ]);
            }
            Self::new(responses)
        }
    }

    #[async_trait]
    impl LlmProvider for MockProvider {
        async fn complete(
            &self,
            _request: crate::agent::provider::CompletionRequest,
            _cancel_token: CancellationToken,
        ) -> Result<mpsc::Receiver<StreamEvent>, crate::agent::AgentError> {
            let events = {
                let mut responses = self.responses.lock();
                responses.pop().unwrap_or_default()
            };
            let (tx, rx) = mpsc::channel(64);
            tokio::spawn(async move {
                for event in events {
                    let _ = tx.send(event).await;
                }
            });
            Ok(rx)
        }
    }

    struct RecordingProvider {
        responses: parking_lot::Mutex<Vec<Vec<StreamEvent>>>,
        requests: parking_lot::Mutex<Vec<CompletionRequest>>,
    }

    impl RecordingProvider {
        fn new(responses: Vec<Vec<StreamEvent>>) -> Self {
            let mut reversed = responses;
            reversed.reverse();
            Self {
                responses: parking_lot::Mutex::new(reversed),
                requests: parking_lot::Mutex::new(Vec::new()),
            }
        }

        fn requests(&self) -> Vec<CompletionRequest> {
            self.requests.lock().clone()
        }
    }

    #[async_trait]
    impl LlmProvider for RecordingProvider {
        async fn complete(
            &self,
            request: crate::agent::provider::CompletionRequest,
            _cancel_token: CancellationToken,
        ) -> Result<mpsc::Receiver<StreamEvent>, crate::agent::AgentError> {
            self.requests.lock().push(request);
            let events = {
                let mut responses = self.responses.lock();
                responses.pop().unwrap_or_default()
            };
            let (tx, rx) = mpsc::channel(64);
            tokio::spawn(async move {
                for event in events {
                    let _ = tx.send(event).await;
                }
            });
            Ok(rx)
        }
    }

    struct ActivityRecordingChannel {
        events: Mutex<Vec<&'static str>>,
        send_notify: Notify,
        mark_read_notify: Notify,
        start_notify: Notify,
        stop_notify: Notify,
    }

    impl ActivityRecordingChannel {
        fn new() -> Self {
            Self {
                events: Mutex::new(Vec::new()),
                send_notify: Notify::new(),
                mark_read_notify: Notify::new(),
                start_notify: Notify::new(),
                stop_notify: Notify::new(),
            }
        }

        fn record(&self, event: &'static str) {
            self.events.lock().push(event);
        }

        fn events(&self) -> Vec<&'static str> {
            self.events.lock().clone()
        }
    }

    impl ChannelPluginInstance for ActivityRecordingChannel {
        fn get_info(&self) -> Result<crate::plugins::ChannelInfo, BindingError> {
            Ok(crate::plugins::ChannelInfo {
                id: "signal".to_string(),
                label: "Signal".to_string(),
                selection_label: "Signal Channel".to_string(),
                docs_path: "".to_string(),
                blurb: "".to_string(),
                order: 0,
            })
        }

        fn get_capabilities(&self) -> Result<ChannelCapabilities, BindingError> {
            Ok(ChannelCapabilities {
                typing_indicators: true,
                read_receipts: true,
                ..Default::default()
            })
        }

        fn send_text(&self, _ctx: OutboundContext) -> Result<DeliveryResult, BindingError> {
            self.record("send");
            self.send_notify.notify_one();
            Ok(DeliveryResult {
                ok: true,
                message_id: Some("sent-1".to_string()),
                error: None,
                retryability: crate::plugins::Retryability::Terminal,
                conversation_id: None,
                to_jid: None,
                poll_id: None,
                error_kind: None,
            })
        }

        fn send_media(&self, _ctx: OutboundContext) -> Result<DeliveryResult, BindingError> {
            unreachable!("executor lifecycle test only sends text")
        }

        fn start_typing(&self, _ctx: TypingContext) -> Result<(), BindingError> {
            self.record("start");
            self.start_notify.notify_one();
            Ok(())
        }

        fn stop_typing(&self, _ctx: TypingContext) -> Result<(), BindingError> {
            self.record("stop");
            self.stop_notify.notify_one();
            Ok(())
        }

        fn mark_read(&self, _ctx: ReadReceiptContext) -> Result<(), BindingError> {
            self.record("read");
            self.mark_read_notify.notify_one();
            Ok(())
        }
    }

    #[test]
    fn test_build_assistant_message_preserves_text_metadata_in_history() {
        let assistant_message = build_assistant_message(
            "sess-metadata",
            &[ContentBlock::Text {
                text: "Hello".to_string(),
                metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                    "sig-text".to_string(),
                )),
            }],
            7,
        )
        .expect("assistant message should be created");

        assert!(
            assistant_message.content.starts_with('['),
            "metadata-bearing assistant turns should be stored as structured blocks"
        );

        let (_system, messages) = crate::agent::context::build_context(&[assistant_message], None);
        match &messages[0].content[0] {
            ContentBlock::Text { text, metadata } => {
                assert_eq!(text, "Hello");
                assert_eq!(
                    metadata
                        .as_ref()
                        .and_then(ContentBlockMetadata::gemini_thought_signature),
                    Some("sig-text")
                );
            }
            other => panic!("expected text block, got {other:?}"),
        }
    }

    #[test]
    fn test_build_assistant_message_preserves_tool_use_metadata_in_history() {
        let assistant_message = build_assistant_message(
            "sess-tool-metadata",
            &[ContentBlock::ToolUse {
                id: "tool_1".to_string(),
                name: "get_weather".to_string(),
                input: serde_json::json!({"city": "SF"}),
                metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                    "sig-tool".to_string(),
                )),
            }],
            5,
        )
        .expect("assistant message should be created");

        let (_system, messages) = crate::agent::context::build_context(&[assistant_message], None);
        match &messages[0].content[0] {
            ContentBlock::ToolUse {
                id,
                name,
                input,
                metadata,
            } => {
                assert_eq!(id, "tool_1");
                assert_eq!(name, "get_weather");
                assert_eq!(input["city"], "SF");
                assert_eq!(
                    metadata
                        .as_ref()
                        .and_then(ContentBlockMetadata::gemini_thought_signature),
                    Some("sig-tool")
                );
            }
            other => panic!("expected tool-use block, got {other:?}"),
        }
    }

    #[test]
    fn test_build_assistant_message_ignores_empty_metadata_objects() {
        let assistant_message = build_assistant_message(
            "sess-empty-metadata",
            &[ContentBlock::Text {
                text: "Hello".to_string(),
                metadata: Some(ContentBlockMetadata::default()),
            }],
            3,
        )
        .expect("assistant message should be created");

        assert_eq!(
            assistant_message.content, "Hello",
            "empty metadata should not force structured assistant-history persistence"
        );
    }

    #[test]
    fn test_sanitize_assistant_turn_preserves_signature_only_part() {
        let config = AgentConfig::default();
        let (turn_text, assistant_blocks) = sanitize_assistant_turn(
            &[
                ContentBlock::Text {
                    text: "Hel".to_string(),
                    metadata: None,
                },
                ContentBlock::Text {
                    text: "lo".to_string(),
                    metadata: None,
                },
                ContentBlock::Text {
                    text: "".to_string(),
                    metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                        "sig-empty".to_string(),
                    )),
                },
            ],
            &config,
            "run-sanitize-signature-only",
        );

        assert_eq!(turn_text, "Hello");
        assert_eq!(assistant_blocks.len(), 2);
        match &assistant_blocks[0] {
            ContentBlock::Text { text, metadata } => {
                assert_eq!(text, "Hello");
                assert!(metadata.is_none());
            }
            other => panic!("expected coalesced text block, got {other:?}"),
        }
        match &assistant_blocks[1] {
            ContentBlock::Text { text, metadata } => {
                assert_eq!(text, "");
                assert_eq!(
                    metadata
                        .as_ref()
                        .and_then(ContentBlockMetadata::gemini_thought_signature),
                    Some("sig-empty")
                );
            }
            other => panic!("expected signature-only text block, got {other:?}"),
        }
    }

    #[test]
    fn test_sanitize_assistant_turn_sanitizes_metadata_text_blocks() {
        let config = AgentConfig::default();
        let (turn_text, assistant_blocks) = sanitize_assistant_turn(
            &[
                ContentBlock::Text {
                    text: "<script>alert(1)</script>Hello".to_string(),
                    metadata: None,
                },
                ContentBlock::Text {
                    text: "".to_string(),
                    metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                        "sig-empty".to_string(),
                    )),
                },
            ],
            &config,
            "run-sanitize-metadata-text",
        );

        assert!(
            !turn_text.contains("<script"),
            "sanitized turn text should not preserve dangerous HTML"
        );
        match &assistant_blocks[0] {
            ContentBlock::Text { text, metadata } => {
                assert!(
                    !text.contains("<script"),
                    "persisted metadata-bearing history should use sanitized text"
                );
                assert!(metadata.is_none());
            }
            other => panic!("expected sanitized text block, got {other:?}"),
        }
    }

    #[test]
    fn test_sanitize_assistant_turn_coalesces_identical_metadata_text_chunks() {
        let config = AgentConfig::default();
        let shared_metadata =
            ContentBlockMetadata::with_gemini_thought_signature(Some("sig-shared".to_string()));
        let (turn_text, assistant_blocks) = sanitize_assistant_turn(
            &[
                ContentBlock::Text {
                    text: "Hello ".to_string(),
                    metadata: shared_metadata.clone(),
                },
                ContentBlock::Text {
                    text: "world".to_string(),
                    metadata: shared_metadata.clone(),
                },
            ],
            &config,
            "run-coalesce-identical-metadata",
        );

        assert_eq!(turn_text, "Hello world");
        assert_eq!(assistant_blocks.len(), 1);
        match &assistant_blocks[0] {
            ContentBlock::Text { text, metadata } => {
                assert_eq!(text, "Hello world");
                assert_eq!(metadata, &shared_metadata);
            }
            other => panic!("expected coalesced metadata-bearing text block, got {other:?}"),
        }
    }

    #[test]
    fn test_sanitize_assistant_turn_preserves_empty_signature_only_boundary() {
        let config = AgentConfig::default();
        let shared_metadata =
            ContentBlockMetadata::with_gemini_thought_signature(Some("sig-shared".to_string()));
        let (turn_text, assistant_blocks) = sanitize_assistant_turn(
            &[
                ContentBlock::Text {
                    text: "".to_string(),
                    metadata: shared_metadata.clone(),
                },
                ContentBlock::Text {
                    text: "Hello".to_string(),
                    metadata: shared_metadata.clone(),
                },
            ],
            &config,
            "run-preserve-empty-signature-boundary",
        );

        assert_eq!(turn_text, "Hello");
        assert_eq!(assistant_blocks.len(), 2);
        match &assistant_blocks[0] {
            ContentBlock::Text { text, metadata } => {
                assert_eq!(text, "");
                assert_eq!(metadata, &shared_metadata);
            }
            other => panic!("expected signature-only boundary block, got {other:?}"),
        }
        match &assistant_blocks[1] {
            ContentBlock::Text { text, metadata } => {
                assert_eq!(text, "Hello");
                assert_eq!(metadata, &shared_metadata);
            }
            other => panic!("expected follow-on text block, got {other:?}"),
        }
    }

    #[test]
    fn test_sanitize_assistant_turn_falls_back_when_boundary_changes_sanitization() {
        let config = AgentConfig::default();
        let (turn_text, assistant_blocks) = sanitize_assistant_turn(
            &[
                ContentBlock::Text {
                    text: "<scr".to_string(),
                    metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                        "sig-a".to_string(),
                    )),
                },
                ContentBlock::Text {
                    text: "ipt>alert(1)</script>Hello".to_string(),
                    metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                        "sig-b".to_string(),
                    )),
                },
            ],
            &config,
            "run-cross-boundary-fallback",
        );

        assert!(
            !turn_text.contains("<script"),
            "whole-turn sanitization should remove dangerous HTML even when the tag spans metadata block boundaries"
        );
        assert_eq!(turn_text, "Hello");
        assert_eq!(assistant_blocks.len(), 1);
        match &assistant_blocks[0] {
            ContentBlock::Text { text, metadata } => {
                assert!(
                    !text.contains("<script"),
                    "fallback block should persist sanitized content"
                );
                assert_eq!(text, "Hello");
                assert!(
                    metadata.is_none(),
                    "fallback path should drop per-part metadata when whole-turn sanitization crosses block boundaries"
                );
            }
            other => panic!("expected fallback plain text block, got {other:?}"),
        }
    }

    /// Helper to set up test state with a temp session store.
    fn make_test_state() -> (Arc<WsServerState>, tempfile::TempDir) {
        let tmp = tempfile::tempdir().unwrap();
        let store = Arc::new(sessions::SessionStore::with_base_path(
            tmp.path().join("sessions"),
        ));
        let state = WsServerState::new(WsServerConfig::default()).with_session_store(store);
        (Arc::new(state), tmp)
    }

    /// Helper to set up test state with a tools registry.
    fn make_test_state_with_tools() -> (Arc<WsServerState>, tempfile::TempDir) {
        let tmp = tempfile::tempdir().unwrap();
        let store = Arc::new(sessions::SessionStore::with_base_path(
            tmp.path().join("sessions"),
        ));
        let tools_registry = Arc::new(crate::plugins::tools::ToolsRegistry::new());
        let state = WsServerState::new(WsServerConfig::default())
            .with_session_store(store)
            .with_tools_registry(tools_registry);
        (Arc::new(state), tmp)
    }

    fn make_test_state_with_plugin_registry(
        plugin_registry: Arc<PluginRegistry>,
    ) -> (Arc<WsServerState>, tempfile::TempDir) {
        let tmp = tempfile::tempdir().unwrap();
        let store = Arc::new(sessions::SessionStore::with_base_path(
            tmp.path().join("sessions"),
        ));
        let state = WsServerState::new(WsServerConfig::default())
            .with_session_store(store)
            .with_plugin_registry(plugin_registry);
        (Arc::new(state), tmp)
    }

    /// Helper to set up a session and register an agent run.
    fn setup_session_and_run(
        state: &WsServerState,
        session_key: &str,
        run_id: &str,
    ) -> sessions::Session {
        let session = state
            .session_store()
            .get_or_create_session(session_key, sessions::SessionMetadata::default())
            .unwrap();
        // Append a user message so the executor has context
        state
            .session_store()
            .append_message(sessions::ChatMessage::user(&session.id, "Hello"))
            .unwrap();
        // Register the run
        {
            use crate::server::ws::{AgentRun, AgentRunStatus};
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis() as u64;
            let mut registry = state.agent_run_registry.lock();
            registry.register(AgentRun {
                run_id: run_id.to_string(),
                session_key: session_key.to_string(),
                delivery_recipient_id: None,
                typing_context: None,
                status: AgentRunStatus::Queued,
                message: "Hello".to_string(),
                response: String::new(),
                error: None,
                created_at: now,
                started_at: None,
                completed_at: None,
                cancel_token: CancellationToken::new(),
                waiters: Vec::new(),
            });
        }
        session
    }

    fn setup_session_and_run_with_channel_activity(
        state: &WsServerState,
        session_key: &str,
        run_id: &str,
        channel: &str,
        chat_id: &str,
    ) -> sessions::Session {
        let metadata = sessions::SessionMetadata {
            channel: Some(channel.to_string()),
            chat_id: Some(chat_id.to_string()),
            ..Default::default()
        };
        let session = state
            .session_store()
            .get_or_create_session(session_key, metadata)
            .unwrap();
        state
            .session_store()
            .append_message(sessions::ChatMessage::user(&session.id, "Hello"))
            .unwrap();
        {
            use crate::server::ws::{AgentRun, AgentRunStatus};
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis() as u64;
            let mut registry = state.agent_run_registry.lock();
            registry.register(AgentRun {
                run_id: run_id.to_string(),
                session_key: session_key.to_string(),
                delivery_recipient_id: Some(chat_id.to_string()),
                typing_context: Some(TypingContext {
                    to: chat_id.to_string(),
                    ..Default::default()
                }),
                status: AgentRunStatus::Queued,
                message: "Hello".to_string(),
                response: String::new(),
                error: None,
                created_at: now,
                started_at: None,
                completed_at: None,
                cancel_token: CancellationToken::new(),
                waiters: Vec::new(),
            });
        }
        session
    }

    #[test]
    fn test_delivery_context_from_registry_returns_none_when_run_missing() {
        let (state, _tmp) = make_test_state();
        let recipient = delivery_context_from_registry(&state, "missing-run");
        assert!(recipient.is_none());
    }

    #[tokio::test]
    async fn test_single_turn_text_response() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-text-1";
        let session_key = "test-session-1";
        setup_session_and_run(&state, session_key, run_id);

        let provider = Arc::new(MockProvider::text("Hello world!"));
        let config = AgentConfig {
            max_turns: 5,
            deliver: true,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        // Verify the run is marked completed
        let registry = state.agent_run_registry.lock();
        let run = registry.get(run_id).unwrap();
        assert_eq!(run.status, crate::server::ws::AgentRunStatus::Completed);
        assert_eq!(run.response, "Hello world!");
    }

    #[tokio::test]
    async fn test_max_turns_reached() {
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-max-turns";
        let session_key = "test-session-max";
        setup_session_and_run(&state, session_key, run_id);

        let max_turns = 3;
        let provider = Arc::new(MockProvider::always_tool_use(max_turns as usize + 1));
        let config = AgentConfig {
            max_turns,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        // Should complete without error — the loop simply exits after max_turns
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());
    }

    #[tokio::test]
    async fn test_tool_use_loop() {
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-tool-loop";
        let session_key = "test-session-tool";
        setup_session_and_run(&state, session_key, run_id);

        // Turn 1: tool use (time), Turn 2: text response
        let provider = Arc::new(MockProvider::new(vec![
            // Turn 1: tool use
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            // Turn 2: text response
            vec![
                StreamEvent::TextDelta {
                    text: "The time is now.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 10,
                    },
                },
            ],
        ]));
        let config = AgentConfig {
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let registry = state.agent_run_registry.lock();
        let run = registry.get(run_id).unwrap();
        assert_eq!(run.status, crate::server::ws::AgentRunStatus::Completed);
        assert_eq!(run.response, "The time is now.");
    }

    #[tokio::test(flavor = "current_thread")]
    async fn test_execute_run_channel_activity_lifecycle_orders_stop_before_delivery() {
        let _fixture = crate::test_support::config::StableConfigFixture::new(serde_json::json!({
            "channels": {
                "signal": {
                    "features": {
                        "typing": {
                            "enabled": true,
                            "intervalSeconds": 30
                        },
                        "readReceipts": {
                            "enabled": true
                        }
                    }
                }
            }
        }));

        let plugin = Arc::new(ActivityRecordingChannel::new());
        let plugin_registry = Arc::new(PluginRegistry::new());
        plugin_registry.register_channel("signal".to_string(), plugin.clone());
        let (state, _tmp) = make_test_state_with_plugin_registry(plugin_registry.clone());
        state.channel_registry().register(
            crate::channels::ChannelInfo::new("signal", "Signal")
                .with_status(crate::channels::ChannelStatus::Connected),
        );

        let run_id = "run-channel-activity";
        let session_key = "test-channel-activity";
        let chat_id = "+15551234567";
        setup_session_and_run_with_channel_activity(&state, session_key, run_id, "signal", chat_id);

        let provider = Arc::new(MockProvider::text("Hello world!"));
        let config = AgentConfig {
            max_turns: 5,
            deliver: true,
            ..Default::default()
        };
        let (delivery_shutdown_tx, delivery_shutdown_rx) = tokio::sync::watch::channel(false);
        let delivery_handle = tokio::spawn(crate::messages::delivery::delivery_loop(
            state.message_pipeline().clone(),
            plugin_registry.clone(),
            state.channel_registry().clone(),
            delivery_shutdown_rx,
        ));

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        // Structural lifecycle awaits: each `Notify` stores a permit on
        // `notify_one`, so this is order-independent for correctness but
        // ordered here for readability. By the time `execute_run` returns,
        // typing has already started and stopped; only delivery is still in
        // flight on the spawned `delivery_loop`.
        tokio::time::timeout(
            std::time::Duration::from_secs(2),
            plugin.start_notify.notified(),
        )
        .await
        .expect("typing start should be observed before assertions");
        tokio::time::timeout(
            std::time::Duration::from_secs(2),
            plugin.stop_notify.notified(),
        )
        .await
        .expect("typing stop should be observed before delivery is queued");
        tokio::time::timeout(
            std::time::Duration::from_secs(2),
            plugin.send_notify.notified(),
        )
        .await
        .expect("delivery loop should send the queued outbound message");
        delivery_shutdown_tx
            .send(true)
            .expect("delivery loop shutdown signal should send");
        state.message_pipeline().notifier().notify_one();
        tokio::time::timeout(std::time::Duration::from_secs(2), delivery_handle)
            .await
            .expect("delivery loop should exit")
            .expect("delivery loop task should succeed");
        state.shutdown_activity_service().await;

        assert_eq!(plugin.events(), vec!["start", "stop", "send"]);
    }

    #[tokio::test(flavor = "current_thread")]
    async fn test_inbound_signal_read_receipt_is_sent_before_typing_starts() {
        let _fixture = crate::test_support::config::StableConfigFixture::new(serde_json::json!({
            "channels": {
                "signal": {
                    "features": {
                        "typing": {
                            "enabled": true,
                            "intervalSeconds": 30
                        },
                        "readReceipts": {
                            "enabled": true
                        }
                    }
                }
            }
        }));

        let plugin = Arc::new(ActivityRecordingChannel::new());
        let plugin_registry = Arc::new(PluginRegistry::new());
        plugin_registry.register_channel("signal".to_string(), plugin.clone());
        let (state, _tmp) = make_test_state_with_plugin_registry(plugin_registry);
        let (read_receipt_shutdown_tx, read_receipt_shutdown_rx) =
            tokio::sync::watch::channel(false);
        state
            .activity_service()
            .spawn_read_receipt_worker(state.clone(), read_receipt_shutdown_rx);

        let chat_id = "+15551234567";
        let claimed_read_receipt = state
            .activity_service()
            .try_claim_read_receipt(
                "signal",
                ReadReceiptContext {
                    recipient: chat_id.to_string(),
                    timestamp: Some(1706745600000),
                    ..Default::default()
                },
            )
            .expect("receipt ownership should be claimable before inbound append");

        let dispatch = crate::channels::inbound::dispatch_inbound_text_with_options(
            &state,
            "signal",
            chat_id,
            chat_id,
            "Hello",
            Some(chat_id.to_string()),
            crate::channels::inbound::InboundDispatchOptions {
                claimed_read_receipt: Some(claimed_read_receipt),
                ..Default::default()
            },
        )
        .await
        .expect("inbound dispatch should succeed");

        // Wait for the read receipt to be observed before asserting ordering.
        tokio::time::timeout(
            std::time::Duration::from_secs(2),
            plugin.mark_read_notify.notified(),
        )
        .await
        .expect("read receipt should fire before run starts");
        assert_eq!(
            plugin.events(),
            vec!["read"],
            "the claimed Signal receipt should be sent before the run can start typing"
        );

        // No-provider dispatch doesn't register the run — manually register
        // it so the controlled `execute_run` below has a registry entry to
        // update. Read the session key off the persisted session rather than
        // re-deriving it; this stays aligned with dispatch's internal session
        // resolution even if the fixture grows session-scoping config later.
        let user_sessions = state
            .session_store()
            .list_sessions(crate::sessions::SessionFilter {
                user_id: Some(chat_id.to_string()),
                ..Default::default()
            })
            .expect("dispatch should have persisted exactly one session");
        let session_key = user_sessions
            .iter()
            .find(|s| s.metadata.channel.as_deref() == Some("signal"))
            .map(|s| s.session_key.clone())
            .expect("dispatch should have created a signal session for chat_id");
        let now_ms = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap_or_default()
            .as_millis() as u64;
        {
            let mut registry = state.agent_run_registry.lock();
            registry.register(crate::server::ws::AgentRun {
                run_id: dispatch.run_id.clone(),
                session_key: session_key.clone(),
                delivery_recipient_id: Some(chat_id.to_string()),
                typing_context: Some(TypingContext {
                    to: chat_id.to_string(),
                    ..Default::default()
                }),
                status: crate::server::ws::AgentRunStatus::Queued,
                message: "Hello".to_string(),
                response: String::new(),
                error: None,
                created_at: now_ms,
                started_at: None,
                completed_at: None,
                cancel_token: CancellationToken::new(),
                waiters: Vec::new(),
            });
        }

        let provider = Arc::new(MockProvider::new(vec![vec![StreamEvent::Stop {
            reason: StopReason::EndTurn,
            usage: TokenUsage {
                input_tokens: 5,
                output_tokens: 0,
            },
        }]]));
        let config = AgentConfig {
            max_turns: 5,
            deliver: true,
            ..Default::default()
        };

        let result = execute_run(
            dispatch.run_id,
            session_key,
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        // Wait for the run's typing lifecycle to be observed before asserting.
        tokio::time::timeout(
            std::time::Duration::from_secs(2),
            plugin.start_notify.notified(),
        )
        .await
        .expect("typing start should be observed");
        tokio::time::timeout(
            std::time::Duration::from_secs(2),
            plugin.stop_notify.notified(),
        )
        .await
        .expect("typing stop should be observed");
        assert!(
            state.message_pipeline().channels_with_messages().is_empty(),
            "empty provider response should not queue outbound delivery"
        );
        assert_eq!(
            plugin.events(),
            vec!["read", "start", "stop"],
            "Signal should be marked read before typing begins"
        );

        let receipt_tasks = state.activity_service().read_receipt_queue().list();
        assert_eq!(receipt_tasks.len(), 1);
        assert_eq!(receipt_tasks[0].state, crate::tasks::TaskState::Done);

        read_receipt_shutdown_tx
            .send(true)
            .expect("read receipt worker shutdown signal should send");
        state.shutdown_activity_service().await;
    }

    #[tokio::test(flavor = "current_thread")]
    async fn test_execute_run_skips_channel_activity_when_delivery_disabled() {
        let _fixture = crate::test_support::config::StableConfigFixture::new(serde_json::json!({
            "channels": {
                "signal": {
                    "features": {
                        "typing": {
                            "enabled": true,
                            "intervalSeconds": 30
                        },
                        "readReceipts": {
                            "enabled": true
                        }
                    }
                }
            }
        }));

        let plugin = Arc::new(ActivityRecordingChannel::new());
        let plugin_registry = Arc::new(PluginRegistry::new());
        plugin_registry.register_channel("signal".to_string(), plugin.clone());
        let (state, _tmp) = make_test_state_with_plugin_registry(plugin_registry);
        state.channel_registry().register(
            crate::channels::ChannelInfo::new("signal", "Signal")
                .with_status(crate::channels::ChannelStatus::Connected),
        );

        let run_id = "run-channel-activity-no-delivery";
        let session_key = "test-channel-activity-no-delivery";
        let chat_id = "+15551234567";
        setup_session_and_run_with_channel_activity(&state, session_key, run_id, "signal", chat_id);

        let provider = Arc::new(MockProvider::text("Hello world!"));
        let config = AgentConfig {
            max_turns: 5,
            deliver: false,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        assert!(
            state.message_pipeline().channels_with_messages().is_empty(),
            "deliver=false should not queue outbound channel messages"
        );
        assert!(
            plugin.events().is_empty(),
            "deliver=false should not emit typing or read-receipt channel activity"
        );
    }

    #[tokio::test]
    async fn test_execute_run_replays_gemini_tool_signature_in_next_request_body() {
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-gemini-thought-signature";
        let session_key = "test-gemini-thought-signature";
        setup_session_and_run(&state, session_key, run_id);

        let provider = Arc::new(RecordingProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                        "sig-tool".to_string(),
                    )),
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "Done.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 5,
                    },
                },
            ],
        ]));
        let config = AgentConfig {
            model: "gemini-2.5-flash".to_string(),
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider.clone(),
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let requests = provider.requests();
        assert_eq!(requests.len(), 2, "expected two LLM turns");

        let body = GeminiProvider::new("test-key".to_string())
            .expect("provider")
            .build_body(&requests[1]);
        let contents = body["contents"].as_array().expect("contents array");
        let assistant_part = contents
            .iter()
            .find(|item| item["role"] == "model")
            .expect("assistant model message");
        assert_eq!(assistant_part["parts"][0]["functionCall"]["name"], "time");
        assert_eq!(assistant_part["parts"][0]["thoughtSignature"], "sig-tool");
    }

    #[tokio::test]
    async fn test_execute_run_replays_vertex_tool_signature_in_next_request_body() {
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-vertex-thought-signature";
        let session_key = "test-vertex-thought-signature";
        setup_session_and_run(&state, session_key, run_id);

        let provider = Arc::new(RecordingProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: ContentBlockMetadata::with_gemini_thought_signature(Some(
                        "sig-tool".to_string(),
                    )),
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "Done.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 5,
                    },
                },
            ],
        ]));
        let config = AgentConfig {
            model: "vertex:gemini-2.5-flash".to_string(),
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider.clone(),
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let requests = provider.requests();
        assert_eq!(requests.len(), 2, "expected two LLM turns");

        let body = build_vertex_gemini_body(&requests[1]);
        let contents = body["contents"].as_array().expect("contents array");
        let assistant_part = contents
            .iter()
            .find(|item| item["role"] == "model")
            .expect("assistant model message");
        assert_eq!(assistant_part["parts"][0]["functionCall"]["name"], "time");
        assert_eq!(assistant_part["parts"][0]["thoughtSignature"], "sig-tool");
    }

    #[tokio::test]
    async fn test_empty_response_handling() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-empty";
        let session_key = "test-session-empty";
        setup_session_and_run(&state, session_key, run_id);

        // Provider returns stop immediately with no text
        let provider = Arc::new(MockProvider::new(vec![vec![StreamEvent::Stop {
            reason: StopReason::EndTurn,
            usage: TokenUsage {
                input_tokens: 5,
                output_tokens: 0,
            },
        }]]));
        let config = AgentConfig {
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let registry = state.agent_run_registry.lock();
        let run = registry.get(run_id).unwrap();
        assert_eq!(run.status, crate::server::ws::AgentRunStatus::Completed);
        assert!(run.response.is_empty(), "response should be empty");
    }

    #[tokio::test]
    async fn test_delivery_uses_run_recipient_override() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-signal-delivery-override";
        let session_key = "signal:global";
        let session = state
            .session_store()
            .get_or_create_session(
                session_key,
                sessions::SessionMetadata {
                    channel: Some("signal".to_string()),
                    chat_id: Some("stale-peer".to_string()),
                    ..Default::default()
                },
            )
            .unwrap();
        state
            .session_store()
            .append_message(sessions::ChatMessage::user(&session.id, "Hello"))
            .unwrap();

        {
            use crate::server::ws::{AgentRun, AgentRunStatus};
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis() as u64;
            let mut registry = state.agent_run_registry.lock();
            registry.register(AgentRun {
                run_id: run_id.to_string(),
                session_key: session_key.to_string(),
                delivery_recipient_id: Some("fresh-peer".to_string()),
                typing_context: None,
                status: AgentRunStatus::Queued,
                message: "Hello".to_string(),
                response: String::new(),
                error: None,
                created_at: now,
                started_at: None,
                completed_at: None,
                cancel_token: CancellationToken::new(),
                waiters: Vec::new(),
            });
        }

        let provider = Arc::new(MockProvider::text("Hello back!"));
        let config = AgentConfig {
            deliver: true,
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let queued = state
            .message_pipeline()
            .next_for_channel("signal")
            .expect("expected queued signal delivery");
        assert_eq!(
            queued.message.metadata.recipient_id.as_deref(),
            Some("fresh-peer")
        );
    }

    /// Helper: register a run with a specific cancel token so tests can control it.
    fn setup_session_and_run_with_token(
        state: &WsServerState,
        session_key: &str,
        run_id: &str,
        cancel_token: CancellationToken,
    ) -> sessions::Session {
        let session = state
            .session_store()
            .get_or_create_session(session_key, sessions::SessionMetadata::default())
            .unwrap();
        state
            .session_store()
            .append_message(sessions::ChatMessage::user(&session.id, "Hello"))
            .unwrap();
        {
            use crate::server::ws::{AgentRun, AgentRunStatus};
            let now = std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_millis() as u64;
            let mut registry = state.agent_run_registry.lock();
            registry.register(AgentRun {
                run_id: run_id.to_string(),
                session_key: session_key.to_string(),
                delivery_recipient_id: None,
                typing_context: None,
                status: AgentRunStatus::Queued,
                message: "Hello".to_string(),
                response: String::new(),
                error: None,
                created_at: now,
                started_at: None,
                completed_at: None,
                cancel_token,
                waiters: Vec::new(),
            });
        }
        session
    }

    /// Mock provider that pauses before sending events, giving time to cancel.
    struct SlowMockProvider {
        delay_ms: u64,
        events: parking_lot::Mutex<Vec<Vec<StreamEvent>>>,
    }

    impl SlowMockProvider {
        fn new(delay_ms: u64, events: Vec<Vec<StreamEvent>>) -> Self {
            let mut reversed = events;
            reversed.reverse();
            Self {
                delay_ms,
                events: parking_lot::Mutex::new(reversed),
            }
        }
    }

    #[async_trait]
    impl LlmProvider for SlowMockProvider {
        async fn complete(
            &self,
            _request: crate::agent::provider::CompletionRequest,
            _cancel_token: CancellationToken,
        ) -> Result<mpsc::Receiver<StreamEvent>, crate::agent::AgentError> {
            let events = {
                let mut responses = self.events.lock();
                responses.pop().unwrap_or_default()
            };
            let delay = self.delay_ms;
            let (tx, rx) = mpsc::channel(64);
            tokio::spawn(async move {
                for event in events {
                    tokio::time::sleep(std::time::Duration::from_millis(delay)).await;
                    let _ = tx.send(event).await;
                }
            });
            Ok(rx)
        }
    }

    #[tokio::test]
    async fn test_cancellation_before_start() {
        // Cancel the token before execute_run is called — mark_started should
        // return false, and execute_run should return Cancelled immediately.
        let (state, _tmp) = make_test_state();
        let run_id = "run-cancel-before";
        let session_key = "test-cancel-before";
        let token = CancellationToken::new();
        setup_session_and_run_with_token(&state, session_key, run_id, token.clone());

        // Cancel before running
        token.cancel();

        let provider = Arc::new(MockProvider::text("should not appear"));
        let config = AgentConfig {
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            token,
        )
        .await;

        assert!(
            matches!(result, Err(crate::agent::AgentError::Cancelled)),
            "expected Cancelled, got: {:?}",
            result
        );
    }

    #[tokio::test]
    async fn test_cancellation_mid_stream() {
        // Cancel the token while the stream is being consumed — the select!
        // branch should fire and return Cancelled.
        let (state, _tmp) = make_test_state();
        let run_id = "run-cancel-stream";
        let session_key = "test-cancel-stream";
        let token = CancellationToken::new();
        setup_session_and_run_with_token(&state, session_key, run_id, token.clone());

        // Slow provider: sends events with 200ms delay each
        let provider = Arc::new(SlowMockProvider::new(
            200,
            vec![vec![
                StreamEvent::TextDelta {
                    text: "chunk1".to_string(),
                    metadata: None,
                },
                StreamEvent::TextDelta {
                    text: "chunk2".to_string(),
                    metadata: None,
                },
                StreamEvent::TextDelta {
                    text: "chunk3".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ]],
        ));

        let cancel_token = token.clone();
        // Cancel after 300ms — should interrupt mid-stream
        tokio::spawn(async move {
            tokio::time::sleep(std::time::Duration::from_millis(300)).await;
            cancel_token.cancel();
        });

        let config = AgentConfig {
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            token,
        )
        .await;

        assert!(
            matches!(result, Err(crate::agent::AgentError::Cancelled)),
            "expected Cancelled, got: {:?}",
            result
        );
    }

    #[tokio::test]
    async fn test_cancellation_during_second_turn_stream() {
        // Cancel the token during the second turn's streaming.
        // Turn 1 completes fast (tool use), turn 2 streams slowly so the cancel fires mid-stream.
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-cancel-turn2";
        let session_key = "test-cancel-turn2";
        let token = CancellationToken::new();
        setup_session_and_run_with_token(&state, session_key, run_id, token.clone());

        // Turn 1 is fast (tool use), turn 2 is slow (300ms per event).
        // We cancel at 100ms into turn 2, which should catch it in the select! branch.
        let provider = Arc::new(SlowMockProvider::new(
            300,
            vec![
                // Turn 1
                vec![
                    StreamEvent::ToolUse {
                        id: "tool_1".to_string(),
                        name: "time".to_string(),
                        input: serde_json::json!({}),
                        metadata: None,
                    },
                    StreamEvent::Stop {
                        reason: StopReason::ToolUse,
                        usage: TokenUsage {
                            input_tokens: 10,
                            output_tokens: 5,
                        },
                    },
                ],
                // Turn 2 (slow — events arrive at 300ms intervals)
                vec![
                    StreamEvent::TextDelta {
                        text: "unreachable".to_string(),
                        metadata: None,
                    },
                    StreamEvent::Stop {
                        reason: StopReason::EndTurn,
                        usage: TokenUsage {
                            input_tokens: 10,
                            output_tokens: 5,
                        },
                    },
                ],
            ],
        ));

        let cancel_token = token.clone();
        // Cancel at 800ms — after turn 1 completes (~600ms) but before turn 2 finishes (~1200ms)
        tokio::spawn(async move {
            tokio::time::sleep(std::time::Duration::from_millis(800)).await;
            cancel_token.cancel();
        });

        let config = AgentConfig {
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            token,
        )
        .await;

        assert!(
            matches!(result, Err(crate::agent::AgentError::Cancelled)),
            "expected Cancelled, got: {:?}",
            result
        );
    }

    // ============== sanitize_provider_error Tests ==============

    #[test]
    fn test_sanitize_short_ascii() {
        let result = sanitize_provider_error("connection refused");
        assert_eq!(result, "connection refused");
    }

    #[test]
    fn test_sanitize_strips_api_key() {
        let msg = "auth error with key sk-ant-api03-abcdefghij1234567890";
        let result = sanitize_provider_error(msg);
        assert!(result.contains("[REDACTED]"));
        assert!(!result.contains("abcdefghij"));
    }

    #[test]
    fn test_sanitize_strips_authorization_header() {
        let msg = "proxy error: Authorization: Bearer sk-ant-secret1234567890";
        let result = sanitize_provider_error(msg);
        assert!(!result.contains("secret1234567890"));
    }

    #[test]
    fn test_sanitize_strips_x_api_key_header() {
        let msg = "gateway echo: x-api-key: sk-ant-api03-abcdefghij1234567890";
        let result = sanitize_provider_error(msg);
        assert!(!result.contains("abcdefghij"));
        assert!(result.contains("x-api-key: [REDACTED]"));
    }

    #[test]
    fn test_sanitize_truncates_long_ascii() {
        let msg = "x".repeat(600);
        let result = sanitize_provider_error(&msg);
        assert!(result.ends_with("... (truncated)"));
        // Should be 500 x's + "... (truncated)"
        assert_eq!(&result[..500], "x".repeat(500));
    }

    #[test]
    fn test_sanitize_no_truncation_at_exactly_500() {
        let msg = "y".repeat(500);
        let result = sanitize_provider_error(&msg);
        assert_eq!(result, msg);
    }

    #[test]
    fn test_sanitize_truncation_at_501() {
        let msg = "z".repeat(501);
        let result = sanitize_provider_error(&msg);
        assert!(result.ends_with("... (truncated)"));
    }

    #[test]
    fn test_sanitize_multibyte_truncation() {
        // Each '€' is 3 bytes. 167 * 3 = 501 bytes > 500 threshold.
        let msg = "€".repeat(167);
        let result = sanitize_provider_error(&msg);
        assert!(result.ends_with("... (truncated)"));
        // Must not panic — the old code would panic slicing mid-codepoint.
        // Boundary should be at 166 * 3 = 498 bytes (last char boundary <= 500).
        let prefix = &result[..result.len() - "... (truncated)".len()];
        assert_eq!(prefix, "€".repeat(166));
    }

    #[test]
    fn test_sanitize_empty_string() {
        let result = sanitize_provider_error("");
        assert_eq!(result, "");
    }

    #[tokio::test]
    async fn test_premature_stream_end_returns_stream_error() {
        // Mock provider that sends a TextDelta but never sends Stop —
        // simulates a premature stream termination (network drop).
        let (state, _tmp) = make_test_state();
        let run_id = "run-premature-end";
        let session_key = "test-premature-end";
        setup_session_and_run(&state, session_key, run_id);

        let provider = Arc::new(MockProvider::new(vec![vec![
            // TextDelta but no Stop event — channel will close
            StreamEvent::TextDelta {
                text: "partial output".to_string(),
                metadata: None,
            },
        ]]));

        let config = AgentConfig {
            max_turns: 5,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;

        assert!(
            matches!(&result, Err(AgentError::Stream(msg)) if msg.contains("stop")),
            "expected Stream error about missing stop event, got: {:?}",
            result,
        );
    }

    // ============== Tool Policy Enforcement Tests ==============

    #[tokio::test]
    async fn test_tool_policy_deny_list_blocks_tool_call() {
        // Configure a deny-list that blocks the "time" tool.
        // The LLM requests "time" anyway — the executor should return an error
        // result for that tool call and continue.
        use crate::agent::tool_policy::ToolPolicy;
        use std::collections::HashSet;

        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-policy-deny";
        let session_key = "test-policy-deny";
        setup_session_and_run(&state, session_key, run_id);

        // Turn 1: LLM requests "time" tool, Turn 2: text response
        let provider = Arc::new(MockProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "Done.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 5,
                    },
                },
            ],
        ]));

        let config = AgentConfig {
            max_turns: 5,
            tool_policy: ToolPolicy::DenyList(
                ["time"]
                    .iter()
                    .map(|s| s.to_string())
                    .collect::<HashSet<_>>(),
            ),
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;

        // The run should complete successfully — the denied tool returns an
        // error to the LLM which then produces a text response.
        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        // Verify the tool result was an error by checking session history
        let session = state
            .session_store()
            .get_session_by_key(session_key)
            .unwrap();
        let history = state
            .session_store()
            .get_history(&session.id, None, None)
            .unwrap();

        // Find the tool result message
        let tool_msg = history
            .iter()
            .find(|m| m.role == sessions::MessageRole::Tool)
            .expect("should have a tool result message");
        assert!(
            tool_msg.content.contains("not available"),
            "tool result should indicate tool is not available, got: {}",
            tool_msg.content
        );
    }

    #[tokio::test]
    async fn test_tool_policy_allow_list_blocks_unlisted_tool() {
        // Configure an allow-list with only "search" — the "time" tool is NOT listed.
        // The LLM requests "time" anyway — should be blocked.
        use crate::agent::tool_policy::ToolPolicy;
        use std::collections::HashSet;

        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-policy-allow";
        let session_key = "test-policy-allow";
        setup_session_and_run(&state, session_key, run_id);

        let provider = Arc::new(MockProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "OK".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 15,
                        output_tokens: 3,
                    },
                },
            ],
        ]));

        let config = AgentConfig {
            max_turns: 5,
            tool_policy: ToolPolicy::AllowList(
                ["search"]
                    .iter()
                    .map(|s| s.to_string())
                    .collect::<HashSet<_>>(),
            ),
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;

        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let session = state
            .session_store()
            .get_session_by_key(session_key)
            .unwrap();
        let history = state
            .session_store()
            .get_history(&session.id, None, None)
            .unwrap();

        let tool_msg = history
            .iter()
            .find(|m| m.role == sessions::MessageRole::Tool)
            .expect("should have a tool result message");
        assert!(
            tool_msg.content.contains("not available"),
            "tool result should indicate tool is not available, got: {}",
            tool_msg.content
        );
    }

    #[tokio::test]
    async fn test_tool_policy_allow_all_permits_tool() {
        // With AllowAll policy, the "time" tool should execute normally.
        use crate::agent::tool_policy::ToolPolicy;

        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-policy-all";
        let session_key = "test-policy-all";
        setup_session_and_run(&state, session_key, run_id);

        let provider = Arc::new(MockProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "The time is now.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 10,
                    },
                },
            ],
        ]));

        let config = AgentConfig {
            max_turns: 5,
            tool_policy: ToolPolicy::AllowAll,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;

        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let session = state
            .session_store()
            .get_session_by_key(session_key)
            .unwrap();
        let history = state
            .session_store()
            .get_history(&session.id, None, None)
            .unwrap();

        // The tool result should be a successful time response, not an error
        let tool_msg = history
            .iter()
            .find(|m| m.role == sessions::MessageRole::Tool)
            .expect("should have a tool result message");
        assert!(
            tool_msg.content.contains("timestamp"),
            "tool result should contain timestamp from successful execution, got: {}",
            tool_msg.content
        );
    }

    // ============== Exfiltration Guard Tests ==============

    #[tokio::test]
    async fn test_exfiltration_guard_blocks_sensitive_tool() {
        // When exfiltration_guard is enabled, a tool call to an
        // exfiltration-sensitive tool should be blocked at dispatch and
        // produce an error result containing "exfiltration guard".
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-exfil-guard";
        let session_key = "test-exfil-guard";
        setup_session_and_run(&state, session_key, run_id);

        // Turn 1: LLM requests "web_fetch", Turn 2: text response
        let provider = Arc::new(MockProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_exfil".to_string(),
                    name: "web_fetch".to_string(),
                    input: serde_json::json!({"url": "https://evil.com"}),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "Blocked.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 5,
                    },
                },
            ],
        ]));

        let config = AgentConfig {
            max_turns: 5,
            exfiltration_guard: true,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;

        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        // Verify the tool result was an exfiltration guard error
        let session = state
            .session_store()
            .get_session_by_key(session_key)
            .unwrap();
        let history = state
            .session_store()
            .get_history(&session.id, None, None)
            .unwrap();

        let tool_msg = history
            .iter()
            .find(|m| m.role == sessions::MessageRole::Tool)
            .expect("should have a tool result message");
        assert!(
            tool_msg.content.contains("exfiltration guard"),
            "tool result should mention exfiltration guard, got: {}",
            tool_msg.content
        );
    }

    #[tokio::test]
    async fn test_exfiltration_guard_disabled_allows_sensitive_tool() {
        // When exfiltration_guard is false (default), exfiltration-sensitive
        // tools should work normally — they should NOT be blocked.
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-exfil-off";
        let session_key = "test-exfil-off";
        setup_session_and_run(&state, session_key, run_id);

        // Turn 1: LLM requests "time" (non-sensitive), Turn 2: text
        let provider = Arc::new(MockProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "Done.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 5,
                    },
                },
            ],
        ]));

        let config = AgentConfig {
            max_turns: 5,
            exfiltration_guard: false,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;

        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        // Tool result should be a successful time response, not an error
        let session = state
            .session_store()
            .get_session_by_key(session_key)
            .unwrap();
        let history = state
            .session_store()
            .get_history(&session.id, None, None)
            .unwrap();

        let tool_msg = history
            .iter()
            .find(|m| m.role == sessions::MessageRole::Tool)
            .expect("should have a tool result message");
        assert!(
            !tool_msg.content.contains("exfiltration guard"),
            "tool result should NOT mention exfiltration guard when disabled, got: {}",
            tool_msg.content
        );
    }

    #[tokio::test]
    async fn test_exfiltration_guard_allows_non_sensitive_tool() {
        // Even with exfiltration_guard enabled, non-sensitive tools should
        // execute normally.
        let (state, _tmp) = make_test_state_with_tools();
        let run_id = "run-exfil-nonsens";
        let session_key = "test-exfil-nonsens";
        setup_session_and_run(&state, session_key, run_id);

        let provider = Arc::new(MockProvider::new(vec![
            vec![
                StreamEvent::ToolUse {
                    id: "tool_1".to_string(),
                    name: "time".to_string(),
                    input: serde_json::json!({}),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::ToolUse,
                    usage: TokenUsage {
                        input_tokens: 10,
                        output_tokens: 5,
                    },
                },
            ],
            vec![
                StreamEvent::TextDelta {
                    text: "The time is now.".to_string(),
                    metadata: None,
                },
                StreamEvent::Stop {
                    reason: StopReason::EndTurn,
                    usage: TokenUsage {
                        input_tokens: 20,
                        output_tokens: 10,
                    },
                },
            ],
        ]));

        let config = AgentConfig {
            max_turns: 5,
            exfiltration_guard: true,
            ..Default::default()
        };

        let result = execute_run(
            run_id.to_string(),
            session_key.to_string(),
            config,
            state.clone(),
            provider,
            CancellationToken::new(),
        )
        .await;

        assert!(result.is_ok(), "execute_run failed: {:?}", result.err());

        let session = state
            .session_store()
            .get_session_by_key(session_key)
            .unwrap();
        let history = state
            .session_store()
            .get_history(&session.id, None, None)
            .unwrap();

        let tool_msg = history
            .iter()
            .find(|m| m.role == sessions::MessageRole::Tool)
            .expect("should have a tool result message");
        // "time" tool should execute successfully (returns timestamp)
        assert!(
            tool_msg.content.contains("timestamp"),
            "non-sensitive tool should execute normally with guard enabled, got: {}",
            tool_msg.content
        );
    }

    // ============== sanitize_provider_error Additional Coverage ==============

    #[test]
    fn test_sanitize_provider_error_strips_api_key() {
        let msg = "Error: invalid key sk-ant-api03-abcdefghij1234567890";
        let result = sanitize_provider_error(msg);
        assert!(!result.contains("sk-ant-api03"));
        assert!(result.contains("[REDACTED]"));
    }

    #[test]
    fn test_usage_provider_for_nearai_model() {
        assert_eq!(
            usage_provider_for_model("nearai:google/gemma-4-31B-it"),
            Some("nearai")
        );
    }

    #[test]
    fn test_sanitize_provider_error_strips_x_api_key_header() {
        let msg = "Request failed: x-api-key: sk-ant-abcdefghij1234567890";
        let result = sanitize_provider_error(msg);
        assert!(!result.contains("sk-ant-"));
        assert!(result.contains("[REDACTED]"));
    }

    #[test]
    fn test_sanitize_provider_error_strips_authorization_header() {
        let msg = "HTTP error: Authorization: Bearer tok_abcdefghij1234567890";
        let result = sanitize_provider_error(msg);
        assert!(!result.contains("tok_abcdefghij"));
        assert!(result.contains("[REDACTED]"));
    }

    #[test]
    fn test_sanitize_provider_error_truncates_long_messages() {
        let long_msg = "x".repeat(1000);
        let result = sanitize_provider_error(&long_msg);
        assert!(result.len() < 600); // 500 + "... (truncated)"
        assert!(result.ends_with("... (truncated)"));
    }

    #[test]
    fn test_sanitize_provider_error_utf8_boundary() {
        // Each emoji is 4 bytes. 126 emojis = 504 bytes > 500
        let msg = "\u{1f525}".repeat(126);
        let result = sanitize_provider_error(&msg);
        // Should not panic and should truncate safely
        assert!(result.ends_with("... (truncated)"));
        // Verify no partial UTF-8
        assert!(result.is_char_boundary(result.len() - "... (truncated)".len()));
    }

    #[test]
    fn test_sanitize_provider_error_short_message_unchanged() {
        let msg = "Connection refused";
        let result = sanitize_provider_error(msg);
        assert_eq!(result, "Connection refused");
    }

    #[test]
    fn test_sanitize_provider_error_empty_string() {
        let result = sanitize_provider_error("");
        assert_eq!(result, "");
    }

    #[test]
    fn test_sanitize_provider_error_exactly_500_bytes() {
        let msg = "a".repeat(500);
        let result = sanitize_provider_error(&msg);
        assert_eq!(result, msg); // Should NOT be truncated at exactly 500
    }

    fn fake_tool_call(id: &str, name: &str) -> (String, String, serde_json::Value) {
        (
            id.to_string(),
            name.to_string(),
            serde_json::json!({"city": "SF"}),
        )
    }

    fn empty_stream_result() -> StreamResult {
        StreamResult {
            turn_text: String::new(),
            assistant_blocks: Vec::new(),
            pending_tool_calls: Vec::new(),
            stop_reason: StopReason::EndTurn,
            turn_usage: TokenUsage::default(),
        }
    }

    #[test]
    fn test_pending_tool_cancellation_payloads_empty_input_produces_nothing() {
        let payloads = pending_tool_cancellation_payloads(&[], None);
        assert!(payloads.is_empty());
    }

    #[test]
    fn test_pending_tool_cancellation_payloads_omits_reason_when_none() {
        let payloads =
            pending_tool_cancellation_payloads(&[fake_tool_call("tu-1", "get_weather")], None);
        assert_eq!(payloads.len(), 1);
        assert_eq!(payloads[0]["toolUseId"], "tu-1");
        assert_eq!(payloads[0]["name"], "get_weather");
        assert!(
            payloads[0].get("reason").is_none(),
            "reason must be absent when not provided: {:?}",
            payloads[0],
        );
        assert!(
            payloads[0].get("input").is_none(),
            "input must NOT be included — it was already sent in the original tool_use event",
        );
    }

    #[test]
    fn test_pending_tool_cancellation_payloads_includes_reason_when_some() {
        let payloads = pending_tool_cancellation_payloads(
            &[fake_tool_call("tu-2", "delete_thing")],
            Some("cancelled"),
        );
        assert_eq!(payloads.len(), 1);
        assert_eq!(payloads[0]["reason"], "cancelled");
    }

    #[test]
    fn test_pending_tool_cancellation_payloads_one_per_tool_call_in_order() {
        let payloads = pending_tool_cancellation_payloads(
            &[
                fake_tool_call("tu-1", "a"),
                fake_tool_call("tu-2", "b"),
                fake_tool_call("tu-3", "c"),
            ],
            Some("stream stalled"),
        );
        assert_eq!(payloads.len(), 3);
        assert_eq!(payloads[0]["toolUseId"], "tu-1");
        assert_eq!(payloads[0]["name"], "a");
        assert_eq!(payloads[1]["toolUseId"], "tu-2");
        assert_eq!(payloads[2]["toolUseId"], "tu-3");
        for payload in &payloads {
            assert_eq!(payload["reason"], "stream stalled");
        }
    }

    #[tokio::test]
    async fn test_handle_stream_event_non_tool_stop_cancels_pending_tool_history() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-stop-cancel";
        let session_key = "session-stop-cancel";
        setup_session_and_run(&state, session_key, run_id);

        let mut result = empty_stream_result();
        let seq = AtomicU64::new(0);

        handle_stream_event(
            StreamEvent::TextDelta {
                text: "before".to_string(),
                metadata: None,
            },
            &mut result,
            &state,
            run_id,
            session_key,
            &seq,
        )
        .unwrap();
        handle_stream_event(
            StreamEvent::ToolUse {
                id: "tu-1".to_string(),
                name: "get_weather".to_string(),
                input: serde_json::json!({"city": "SF"}),
                metadata: None,
            },
            &mut result,
            &state,
            run_id,
            session_key,
            &seq,
        )
        .unwrap();

        assert_eq!(result.pending_tool_calls.len(), 1);
        assert!(result
            .assistant_blocks
            .iter()
            .any(|block| matches!(block, ContentBlock::ToolUse { id, .. } if id == "tu-1")));

        let outcome = handle_stream_event(
            StreamEvent::Stop {
                reason: StopReason::EndTurn,
                usage: TokenUsage::default(),
            },
            &mut result,
            &state,
            run_id,
            session_key,
            &seq,
        );

        assert!(matches!(outcome, Ok(true)));
        assert!(result.pending_tool_calls.is_empty());
        assert!(!result
            .assistant_blocks
            .iter()
            .any(|block| matches!(block, ContentBlock::ToolUse { .. })));
        assert!(result
            .assistant_blocks
            .iter()
            .any(|block| matches!(block, ContentBlock::Text { text, .. } if text == "before")));
    }

    #[tokio::test]
    async fn test_cancel_pending_tool_calls_is_idempotent() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-cancel-idempotent";
        let session_key = "session-cancel-idempotent";
        setup_session_and_run(&state, session_key, run_id);

        let mut result = empty_stream_result();
        result.assistant_blocks.push(ContentBlock::ToolUse {
            id: "tu-1".to_string(),
            name: "get_weather".to_string(),
            input: serde_json::json!({"city": "SF"}),
            metadata: None,
        });
        result
            .pending_tool_calls
            .push(fake_tool_call("tu-1", "get_weather"));
        let seq = AtomicU64::new(0);

        assert_eq!(
            cancel_pending_tool_calls(&state, run_id, &seq, &mut result, Some("boom")),
            1
        );
        assert_eq!(
            cancel_pending_tool_calls(&state, run_id, &seq, &mut result, Some("boom")),
            0
        );
        assert_eq!(seq.load(Ordering::Relaxed), 1);
        assert!(result.pending_tool_calls.is_empty());
        assert!(!result
            .assistant_blocks
            .iter()
            .any(|block| matches!(block, ContentBlock::ToolUse { .. })));
    }

    #[tokio::test]
    async fn test_process_llm_stream_provider_error_after_tool_use_clears_pending_tool_once() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-error-after-tool";
        let session_key = "session-error-after-tool";
        setup_session_and_run(&state, session_key, run_id);

        let (tx, mut rx) = mpsc::channel(4);
        tx.send(StreamEvent::ToolUse {
            id: "tu-1".to_string(),
            name: "get_weather".to_string(),
            input: serde_json::json!({"city": "SF"}),
            metadata: None,
        })
        .await
        .unwrap();
        tx.send(StreamEvent::Error {
            message: "provider failed".to_string(),
            usage: Some(TokenUsage {
                input_tokens: 9,
                output_tokens: 3,
            }),
        })
        .await
        .unwrap();
        drop(tx);

        let seq = AtomicU64::new(0);
        let (result, outcome) = process_llm_stream(
            &mut rx,
            &state,
            run_id,
            session_key,
            &seq,
            &CancellationToken::new(),
        )
        .await;

        assert!(
            matches!(outcome, Err(AgentError::Provider(message)) if message == "provider failed")
        );
        assert!(result.pending_tool_calls.is_empty());
        assert!(!result
            .assistant_blocks
            .iter()
            .any(|block| matches!(block, ContentBlock::ToolUse { .. })));
        assert_eq!(result.turn_usage.input_tokens, 9);
        assert_eq!(result.turn_usage.output_tokens, 3);
        assert_eq!(
            seq.load(Ordering::Relaxed),
            4,
            "tool_use, tool_use_cancelled, agent error, and chat error should be the only broadcasts"
        );
    }

    #[tokio::test]
    async fn test_handle_stream_event_error_preserves_usage_from_event() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-error-usage";
        let session_key = "session-error-usage";
        setup_session_and_run(&state, session_key, run_id);

        let mut result = StreamResult {
            turn_text: String::new(),
            assistant_blocks: Vec::new(),
            pending_tool_calls: Vec::new(),
            stop_reason: StopReason::EndTurn,
            turn_usage: TokenUsage::default(),
        };
        let seq = AtomicU64::new(0);
        let usage = TokenUsage {
            input_tokens: 123,
            output_tokens: 7,
        };
        let outcome = handle_stream_event(
            StreamEvent::Error {
                message: "boom".to_string(),
                usage: Some(usage),
            },
            &mut result,
            &state,
            run_id,
            session_key,
            &seq,
        );
        assert!(outcome.is_err(), "Error event must surface as Err");
        assert_eq!(
            result.turn_usage.input_tokens, 123,
            "Error.usage must populate result.turn_usage so the caller can bill partial turns"
        );
        assert_eq!(result.turn_usage.output_tokens, 7);
    }

    #[tokio::test]
    async fn test_handle_stream_event_error_leaves_prior_usage_when_none() {
        let (state, _tmp) = make_test_state();
        let run_id = "run-error-no-usage";
        let session_key = "session-error-no-usage";
        setup_session_and_run(&state, session_key, run_id);

        let mut result = StreamResult {
            turn_text: String::new(),
            assistant_blocks: Vec::new(),
            pending_tool_calls: Vec::new(),
            stop_reason: StopReason::EndTurn,
            turn_usage: TokenUsage {
                input_tokens: 42,
                output_tokens: 17,
            },
        };
        let seq = AtomicU64::new(0);
        let outcome = handle_stream_event(
            StreamEvent::Error {
                message: "transport".to_string(),
                usage: None,
            },
            &mut result,
            &state,
            run_id,
            session_key,
            &seq,
        );
        assert!(outcome.is_err());
        assert_eq!(
            result.turn_usage.input_tokens, 42,
            "Error.usage=None must NOT clobber prior accumulated usage"
        );
        assert_eq!(result.turn_usage.output_tokens, 17);
    }
}
