async-llm 0.9.0

use axum::{
    extract::{Path, State},
    http::{HeaderMap, StatusCode},
    response::sse::{Event, Sse},
    routing::{get, post},
    Json, Router,
};
use futures::{stream::BoxStream, StreamExt};
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, convert::Infallible, net::SocketAddr, sync::Arc};
use tokio::{net::TcpListener, sync::Notify};

/// A queued response variant.
#[derive(Clone, Debug, Deserialize, Serialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum MockResponse {
    Text {
        content: String,
    },
    ToolCall {
        name: String,
        input: serde_json::Value,
    },
    /// Several tool calls in one assistant message — parallel tool use, which is
    /// how a model asks two questions at once.
    ToolCalls {
        calls: Vec<(String, serde_json::Value)>,
    },
    /// `status` is used as the HTTP status code when returned as a plain JSON response,
    /// and as an SSE StreamError when returned in streaming mode.
    Error {
        status: u16,
        message: String,
    },
    TextStream {
        chunks: Vec<String>,
    },
    ToolCallStream {
        name: String,
        id: String,
        input: serde_json::Value,
    },
    Thinking {
        text: String,
        signature: String,
    },
    /// A response cut off by the output-token ceiling: `stop_reason: max_tokens`
    /// on the Anthropic wire, `finish_reason: length` on the OpenAI wire.
    Truncated {
        content: String,
    },
    /// A reasoning-model turn: a reasoning trace, then the answer. On the OpenAI
    /// wire the trace streams as `delta.reasoning_content` (DeepSeek/vLLM shape)
    /// before the content; the Anthropic wire renders only the answer text
    /// (reasoning on that wire is the `Thinking` variant).
    Reasoning {
        reasoning: String,
        content: String,
    },
    /// A stream that ends after `after` events without its terminal frame — the
    /// connection dropped mid-response. Reproduces #61 item 1.
    CutStream {
        chunks: Vec<String>,
        after: usize,
    },
    /// A tool call whose arguments are cut off mid-JSON, with no terminal frame.
    /// `partial_input_json` is emitted verbatim and is expected not to parse.
    CutToolCallStream {
        name: String,
        id: String,
        partial_input_json: String,
    },
}

pub(crate) struct QueueEntry {
    pub(crate) response: MockResponse,
    pub(crate) reached: Option<Arc<Notify>>,
    pub(crate) gate: Option<Arc<Notify>>,
    /// Held for this long before answering — a slow peer. With no client-side
    /// timeout this is indistinguishable from a hang (#61 item 5).
    pub(crate) delay: Option<std::time::Duration>,
}

impl QueueEntry {
    fn immediate(response: MockResponse) -> Self {
        Self {
            response,
            reached: None,
            gate: None,
            delay: None,
        }
    }
}

pub struct BlockHandle {
    gate: Arc<Notify>,
    reached: Arc<Notify>,
}

impl BlockHandle {
    pub async fn wait_until_received(&self) {
        self.reached.notified().await;
    }
    pub fn release(&self) {
        self.gate.notify_one();
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Scenario {
    #[serde(default)]
    pub description: String,
    pub responses: Vec<MockResponse>,
}

struct ScenarioState {
    responses: Vec<MockResponse>,
    cursor: usize,
}

impl ScenarioState {
    fn from_scenario(s: &Scenario) -> Self {
        Self {
            responses: s.responses.clone(),
            cursor: 0,
        }
    }
    fn next_response(&mut self) -> Option<MockResponse> {
        let resp = self.responses.get(self.cursor)?.clone();
        self.cursor += 1;
        Some(resp)
    }
}

#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ScenarioConfig {
    pub scenarios: HashMap<String, Scenario>,
}

#[derive(Default)]
pub(crate) struct MockState {
    queue: Mutex<Vec<QueueEntry>>,
    scenarios: Mutex<HashMap<String, Scenario>>,
    session_bindings: Mutex<HashMap<String, String>>,
    session_states: Mutex<HashMap<String, ScenarioState>>,
    /// Every inbound request body (both wires), for tests that assert on what
    /// reached the agent (e.g. the composed system prompt). Cleared by `/reset`.
    captured: Mutex<Vec<serde_json::Value>>,
    /// A one-shot hold applied to the next answer out of the queue, however it
    /// was queued. Set by `/reset` so it covers a test's *first* answer without
    /// the test having to know which queued entry that turns out to be.
    ///
    /// It exists because an instant provider is the unrealistic case: a real
    /// one takes hundreds of milliseconds to a first token, and a session's
    /// first turn starts with the session, before any browser has finished
    /// navigating to it. Answering with zero latency lets a whole turn happen
    /// in a window no client could have been watching.
    hold_next: Mutex<Option<std::time::Duration>>,
}

impl MockState {
    pub(crate) fn dequeue_entry(&self) -> Option<QueueEntry> {
        let mut entry = {
            let mut q = self.queue.lock();
            (!q.is_empty()).then(|| q.remove(0))
        }?;
        // The hold is spent on whichever answer comes first, and only once —
        // later turns in the same test run at full speed.
        if let Some(hold) = self.hold_next.lock().take() {
            if entry.delay.is_none() {
                entry.delay = Some(hold);
            }
        }
        Some(entry)
    }
    /// Record an inbound request body so `GET /received` can return it.
    pub(crate) fn capture(&self, body: serde_json::Value) {
        self.captured.lock().push(body);
    }
}

#[derive(Default)]
pub struct MockLlmServerBuilder {
    responses: Vec<MockResponse>,
    scenarios: HashMap<String, Scenario>,
    bind_all: bool,
    port: Option<u16>,
}

impl MockLlmServerBuilder {
    #[must_use]
    pub fn response(mut self, text: impl Into<String>) -> Self {
        self.responses.push(MockResponse::Text {
            content: text.into(),
        });
        self
    }
    #[must_use]
    pub fn tool_call(mut self, name: impl Into<String>, input: serde_json::Value) -> Self {
        self.responses.push(MockResponse::ToolCall {
            name: name.into(),
            input,
        });
        self
    }
    #[must_use]
    pub fn error(mut self, status: u16, message: impl Into<String>) -> Self {
        self.responses.push(MockResponse::Error {
            status,
            message: message.into(),
        });
        self
    }
    #[must_use]
    pub fn response_stream(mut self, chunks: impl IntoIterator<Item = impl Into<String>>) -> Self {
        self.responses.push(MockResponse::TextStream {
            chunks: chunks.into_iter().map(Into::into).collect(),
        });
        self
    }
    #[must_use]
    pub fn tool_call_stream(mut self, name: impl Into<String>, input: serde_json::Value) -> Self {
        self.responses.push(MockResponse::ToolCallStream {
            name: name.into(),
            id: format!("toolu_{}", uuid::Uuid::new_v4()),
            input,
        });
        self
    }
    #[must_use]
    pub fn thinking(mut self, text: impl Into<String>, signature: impl Into<String>) -> Self {
        self.responses.push(MockResponse::Thinking {
            text: text.into(),
            signature: signature.into(),
        });
        self
    }
    #[must_use]
    pub fn with_scenarios(mut self, config: ScenarioConfig) -> Self {
        self.scenarios = config.scenarios;
        self
    }
    #[must_use]
    pub fn bind_all_interfaces(mut self) -> Self {
        self.bind_all = true;
        self
    }
    #[must_use]
    pub fn port(mut self, port: u16) -> Self {
        self.port = Some(port);
        self
    }
    pub async fn build(self) -> MockLlmServer {
        let queue = self
            .responses
            .into_iter()
            .map(QueueEntry::immediate)
            .collect();
        let state = Arc::new(MockState {
            queue: Mutex::new(queue),
            scenarios: Mutex::new(self.scenarios),
            session_bindings: Mutex::new(HashMap::new()),
            session_states: Mutex::new(HashMap::new()),
            captured: Mutex::new(Vec::new()),
            hold_next: Mutex::new(None),
        });
        let app = Router::new()
            .route(
                "/v1/messages",
                post(crate::mock::anthropic::handle_messages),
            )
            .route(
                "/v1/chat/completions",
                post(crate::mock::openai::handle_chat_completions),
            )
            // Both spellings: a provider configured with a bare host reaches
            // `/responses`, one configured api.openai.com-style reaches
            // `/v1/responses`.
            .route("/responses", post(crate::mock::responses::handle_responses))
            .route(
                "/v1/responses",
                post(crate::mock::responses::handle_responses),
            )
            .route("/queue", post(handle_queue))
            .route("/received", get(handle_received))
            .route("/reset", post(handle_reset))
            .route("/scenarios/load", post(handle_load_scenarios))
            .route("/scenarios", get(handle_list_scenarios))
            .route(
                "/scenarios/{name}/register/{session_id}",
                post(handle_register_session),
            )
            .with_state(state.clone());
        let port = self.port.unwrap_or(0);
        let bind = if self.bind_all {
            format!("0.0.0.0:{port}")
        } else {
            format!("127.0.0.1:{port}")
        };
        let listener = TcpListener::bind(&bind).await.unwrap();
        let addr = listener.local_addr().unwrap();
        let handle = tokio::spawn(async move { axum::serve(listener, app).await.unwrap() });
        MockLlmServer {
            addr,
            _handle: handle,
            state,
        }
    }
}

pub struct MockLlmServer {
    addr: SocketAddr,
    _handle: tokio::task::JoinHandle<()>,
    state: Arc<MockState>,
}

impl MockLlmServer {
    #[must_use]
    pub fn builder() -> MockLlmServerBuilder {
        MockLlmServerBuilder::default()
    }
    #[must_use]
    pub fn url(&self) -> String {
        format!("http://{}", self.addr)
    }
    #[must_use]
    pub fn port(&self) -> u16 {
        self.addr.port()
    }
    pub fn queued_count(&self) -> usize {
        self.state.queue.lock().len()
    }
    pub fn queue_response(&self, text: impl Into<String>) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::Text {
                content: text.into(),
            }));
    }
    /// Queue one assistant message that makes several tool calls at once.
    pub fn queue_tool_calls(&self, calls: Vec<(String, serde_json::Value)>) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::ToolCalls { calls }));
    }

    pub fn queue_tool_call(&self, name: impl Into<String>, input: serde_json::Value) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::ToolCall {
                name: name.into(),
                input,
            }));
    }
    /// Queue an error response. On the Anthropic wire this becomes an SSE error
    /// event whose type is derived from `status`; on the OpenAI wire it becomes
    /// a real HTTP error status, which is how those backends signal failure.
    pub fn queue_error(&self, status: u16, message: impl Into<String>) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::Error {
                status,
                message: message.into(),
            }));
    }
    /// Queue a response the backend cut off at its output-token ceiling.
    pub fn queue_truncated(&self, content: impl Into<String>) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::Truncated {
                content: content.into(),
            }));
    }
    /// Queue a text response the server holds for `delay` before answering.
    pub fn queue_delayed(&self, text: impl Into<String>, delay: std::time::Duration) {
        self.state.queue.lock().push(QueueEntry {
            response: MockResponse::Text {
                content: text.into(),
            },
            reached: None,
            gate: None,
            delay: Some(delay),
        });
    }

    /// Queue a text stream cut off after `after` SSE events, with no terminal
    /// frame — a connection dropped mid-response.
    pub fn queue_cut_stream(
        &self,
        chunks: impl IntoIterator<Item = impl Into<String>>,
        after: usize,
    ) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::CutStream {
                chunks: chunks.into_iter().map(Into::into).collect(),
                after,
            }));
    }

    /// Queue a tool call whose arguments are cut off mid-JSON.
    pub fn queue_cut_tool_call(
        &self,
        name: impl Into<String>,
        id: impl Into<String>,
        partial_input_json: impl Into<String>,
    ) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::CutToolCallStream {
                name: name.into(),
                id: id.into(),
                partial_input_json: partial_input_json.into(),
            }));
    }

    /// Queue a reasoning-model turn: a reasoning trace, then the answer.
    pub fn queue_reasoning(&self, reasoning: impl Into<String>, content: impl Into<String>) {
        self.state
            .queue
            .lock()
            .push(QueueEntry::immediate(MockResponse::Reasoning {
                reasoning: reasoning.into(),
                content: content.into(),
            }));
    }
    pub fn blocking_response(&self, text: impl Into<String>) -> BlockHandle {
        let gate = Arc::new(Notify::new());
        let reached = Arc::new(Notify::new());
        self.state.queue.lock().push(QueueEntry {
            response: MockResponse::Text {
                content: text.into(),
            },
            reached: Some(Arc::clone(&reached)),
            gate: Some(Arc::clone(&gate)),
            delay: None,
        });
        BlockHandle { gate, reached }
    }
    pub fn load_scenarios(&self, config: ScenarioConfig) {
        self.state.scenarios.lock().extend(config.scenarios);
    }
    pub fn register_session(
        &self,
        session_id: impl Into<String>,
        scenario_name: impl Into<String>,
    ) {
        let session_id = session_id.into();
        let scenario_name = scenario_name.into();
        let scenarios = self.state.scenarios.lock();
        if let Some(scenario) = scenarios.get(&scenario_name) {
            let state = ScenarioState::from_scenario(scenario);
            drop(scenarios);
            self.state
                .session_states
                .lock()
                .insert(session_id.clone(), state);
            self.state
                .session_bindings
                .lock()
                .insert(session_id, scenario_name);
        }
    }
}

// ── internal request/response types ──────────────────────────────────────────

#[derive(Serialize)]
struct StatusResponse {
    status: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    message: Option<String>,
}

#[derive(Serialize)]
struct ScenariosListResponse {
    scenarios: Vec<String>,
}

pub(crate) enum ResponseKind {
    Json(axum::Json<serde_json::Value>),
    Sse(Sse<BoxStream<'static, Result<Event, Infallible>>>),
    HttpError(StatusCode, axum::Json<serde_json::Value>),
}

impl axum::response::IntoResponse for ResponseKind {
    fn into_response(self) -> axum::response::Response {
        match self {
            ResponseKind::Json(j) => j.into_response(),
            ResponseKind::Sse(s) => s.into_response(),
            ResponseKind::HttpError(status, body) => (status, body).into_response(),
        }
    }
}

// ── handlers ─────────────────────────────────────────────────────────────────

pub(crate) async fn handle_messages(
    State(state): State<Arc<MockState>>,
    headers: HeaderMap,
    Json(req): Json<serde_json::Value>,
) -> ResponseKind {
    state.capture(req.clone());
    let entry = if let Some(sid) = headers.get("X-Session-Id").and_then(|v| v.to_str().ok()) {
        let mut ss = state.session_states.lock();
        if let Some(scenario) = ss.get_mut(sid) {
            scenario.next_response().map(QueueEntry::immediate)
        } else {
            state.dequeue_entry()
        }
    } else {
        state.dequeue_entry()
    };

    if let Some(e) = &entry {
        if let Some(r) = &e.reached {
            r.notify_one();
        }
        if let Some(g) = &e.gate {
            g.notified().await;
        }
        if let Some(d) = e.delay {
            tokio::time::sleep(d).await;
        }
    }

    let response = entry.map(|e| e.response);
    let is_stream = req
        .get("stream")
        .and_then(serde_json::Value::as_bool)
        .unwrap_or(false);

    match response {
        Some(MockResponse::TextStream { chunks }) => sse_from_pairs(text_stream_sse(&chunks)),
        Some(MockResponse::ToolCallStream { name, id, input }) => {
            sse_from_pairs(tool_call_stream_sse(&name, &id, &input))
        }
        Some(MockResponse::CutStream { chunks, after }) => {
            sse_from_pairs(cut_text_stream_sse(&chunks, after))
        }
        Some(MockResponse::CutToolCallStream {
            name,
            id,
            partial_input_json,
        }) => sse_from_pairs(cut_tool_call_stream_sse(&name, &id, &partial_input_json)),
        other => {
            let resp = other;
            if is_stream {
                let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
                let tool_id = format!("toolu_{}", uuid::Uuid::new_v4());
                let pairs = match resp {
                    Some(MockResponse::Text { content }) => text_sse(&msg_id, &content),
                    Some(MockResponse::ToolCall { name, input }) => {
                        tool_sse(&msg_id, &tool_id, &name, &input)
                    }
                    Some(MockResponse::ToolCalls { calls }) => tools_sse(&msg_id, &calls),
                    Some(MockResponse::Thinking { text, signature }) => {
                        thinking_sse(&msg_id, &text, &signature)
                    }
                    Some(MockResponse::Truncated { content }) => truncated_sse(&msg_id, &content),
                    // Reasoning is an OpenAI-wire concept; here render just the answer.
                    Some(MockResponse::Reasoning { content, .. }) => text_sse(&msg_id, &content),
                    // Error in stream mode: emit a StreamError-compatible SSE event so
                    // async-anthropic parses it as AnthropicError::StreamError. The
                    // error `type` is derived from `status` so tests can choose a
                    // retryable (`overloaded_error`/`rate_limit_error`) or a
                    // non-retryable (`invalid_request_error`) failure deterministically.
                    Some(MockResponse::Error { status, message }) => {
                        let etype = match status {
                            429 => "rate_limit_error",
                            529 => "overloaded_error",
                            _ => "invalid_request_error",
                        };
                        vec![(
                            "error".into(),
                            // Top-level "type" + "message" matches async-anthropic's StreamError shape.
                            serde_json::json!({ "type": etype, "message": message }).to_string(),
                        )]
                    }
                    None => text_sse(&msg_id, "No mock response queued"),
                    Some(
                        MockResponse::TextStream { .. }
                        | MockResponse::ToolCallStream { .. }
                        | MockResponse::CutStream { .. }
                        | MockResponse::CutToolCallStream { .. },
                    ) => {
                        unreachable!()
                    }
                };
                sse_from_pairs(pairs)
            } else {
                match resp {
                    Some(MockResponse::Text { content }) => {
                        ResponseKind::Json(axum::Json(text_json(&content)))
                    }
                    Some(MockResponse::ToolCall { name, input }) => {
                        ResponseKind::Json(axum::Json(tool_json(&name, &input)))
                    }
                    Some(MockResponse::ToolCalls { calls }) => {
                        ResponseKind::Json(axum::Json(tools_json(&calls)))
                    }
                    Some(MockResponse::Thinking { text, signature }) => {
                        ResponseKind::Json(axum::Json(thinking_json(&text, &signature)))
                    }
                    Some(MockResponse::Truncated { content }) => {
                        ResponseKind::Json(axum::Json(text_json(&content)))
                    }
                    Some(MockResponse::Reasoning { content, .. }) => {
                        ResponseKind::Json(axum::Json(text_json(&content)))
                    }
                    Some(MockResponse::Error { status, message }) => {
                        let code = StatusCode::from_u16(status)
                            .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
                        ResponseKind::HttpError(code, axum::Json(error_json(&message)))
                    }
                    None => ResponseKind::Json(axum::Json(text_json("No mock response queued"))),
                    Some(
                        MockResponse::TextStream { .. }
                        | MockResponse::ToolCallStream { .. }
                        | MockResponse::CutStream { .. }
                        | MockResponse::CutToolCallStream { .. },
                    ) => {
                        unreachable!()
                    }
                }
            }
        }
    }
}

/// Append one response to the FIFO queue. The body is a full [`MockResponse`]
/// (tagged `type`, snake_case), so every variant is expressible over HTTP:
/// `{"type":"text","content":…}`, `{"type":"text_stream","chunks":[…]}`,
/// `{"type":"tool_call","name":…,"input":…}`, `{"type":"error","status":…,
/// "message":…}`, `{"type":"thinking","text":…,"signature":…}`.
///
/// An optional `"delayMs"` alongside those fields holds the answer back before
/// sending it — the HTTP face of the `delay` every queue entry already has, and
/// the only way an out-of-process test can make a turn *observably* in flight.
/// A browser test that means to watch a turn happen needs the turn to still be
/// happening when it looks.
async fn handle_queue(
    State(state): State<Arc<MockState>>,
    Json(body): Json<serde_json::Value>,
) -> Json<StatusResponse> {
    let delay = body
        .get("delayMs")
        .and_then(serde_json::Value::as_u64)
        .map(std::time::Duration::from_millis);
    let response: MockResponse = match serde_json::from_value(body) {
        Ok(r) => r,
        Err(e) => {
            return Json(StatusResponse {
                status: "error".into(),
                message: Some(e.to_string()),
            });
        }
    };
    state.queue.lock().push(QueueEntry {
        response,
        reached: None,
        gate: None,
        delay,
    });
    Json(StatusResponse {
        status: "queued".into(),
        message: None,
    })
}

/// Clear all per-test state: the FIFO queue and any per-session scenario
/// cursors/bindings. Loaded scenario definitions are left intact. Tests call
/// this before each case so nothing leaks across the serially-run suite.
async fn handle_reset(
    State(state): State<Arc<MockState>>,
    body: Option<Json<serde_json::Value>>,
) -> Json<StatusResponse> {
    state.queue.lock().clear();
    state.session_states.lock().clear();
    state.session_bindings.lock().clear();
    state.captured.lock().clear();
    // `holdFirstMs` arms a one-shot delay on this test's first answer. See
    // `MockState::hold_next` for why a zero-latency provider is the case worth
    // avoiding rather than the case worth defaulting to.
    *state.hold_next.lock() = body
        .and_then(|Json(b)| b.get("holdFirstMs").and_then(serde_json::Value::as_u64))
        .map(std::time::Duration::from_millis);
    Json(StatusResponse {
        status: "reset".into(),
        message: None,
    })
}

/// Return every captured request body, most-recent-first, so a test can assert
/// on what reached the agent (e.g. the composed system prompt).
async fn handle_received(State(state): State<Arc<MockState>>) -> Json<Vec<serde_json::Value>> {
    let mut bodies = state.captured.lock().clone();
    bodies.reverse();
    Json(bodies)
}

async fn handle_load_scenarios(
    State(state): State<Arc<MockState>>,
    Json(config): Json<ScenarioConfig>,
) -> Json<StatusResponse> {
    let count = config.scenarios.len();
    state.scenarios.lock().extend(config.scenarios);
    Json(StatusResponse {
        status: "loaded".into(),
        message: Some(format!("{count} scenarios loaded")),
    })
}

async fn handle_list_scenarios(State(state): State<Arc<MockState>>) -> Json<ScenariosListResponse> {
    let scenarios = state.scenarios.lock();
    Json(ScenariosListResponse {
        scenarios: scenarios.keys().cloned().collect(),
    })
}

async fn handle_register_session(
    State(state): State<Arc<MockState>>,
    Path((scenario_name, session_id)): Path<(String, String)>,
) -> Result<Json<StatusResponse>, (StatusCode, Json<StatusResponse>)> {
    let scenarios = state.scenarios.lock();
    if let Some(scenario) = scenarios.get(&scenario_name) {
        let scenario_state = ScenarioState::from_scenario(scenario);
        drop(scenarios);
        state
            .session_states
            .lock()
            .insert(session_id.clone(), scenario_state);
        state
            .session_bindings
            .lock()
            .insert(session_id.clone(), scenario_name.clone());
        Ok(Json(StatusResponse {
            status: "registered".into(),
            message: Some(format!(
                "Session {session_id} bound to scenario {scenario_name}"
            )),
        }))
    } else {
        Err((
            StatusCode::NOT_FOUND,
            Json(StatusResponse {
                status: "error".into(),
                message: Some(format!("Scenario '{scenario_name}' not found")),
            }),
        ))
    }
}

// ── SSE helpers ───────────────────────────────────────────────────────────────

pub(crate) fn sse_from_pairs(pairs: Vec<(String, String)>) -> ResponseKind {
    let events: Vec<Result<Event, Infallible>> = pairs
        .into_iter()
        .map(|(t, d)| Ok(Event::default().event(t).data(d)))
        .collect();
    ResponseKind::Sse(Sse::new(futures::stream::iter(events).boxed()))
}

fn text_sse(msg_id: &str, text: &str) -> Vec<(String, String)> {
    let tokens = u32::try_from(text.len() / 4).unwrap_or(u32::MAX);
    vec![
        (
            "message_start".into(),
            serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
        ),
        (
            "content_block_start".into(),
            serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}).to_string(),
        ),
        (
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":text}}).to_string(),
        ),
        (
            "content_block_stop".into(),
            serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
        ),
        (
            "message_delta".into(),
            serde_json::json!({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":tokens}}).to_string(),
        ),
        (
            "message_stop".into(),
            serde_json::json!({"type":"message_stop"}).to_string(),
        ),
    ]
}

/// Like [`text_sse`], but the turn ends with `stop_reason: max_tokens` — the
/// backend hit its output ceiling mid-answer.
fn truncated_sse(msg_id: &str, text: &str) -> Vec<(String, String)> {
    let mut pairs = text_sse(msg_id, text);
    for (event, data) in &mut pairs {
        if event == "message_delta" {
            *data = serde_json::json!({
                "type": "message_delta",
                "delta": {"stop_reason": "max_tokens", "stop_sequence": null},
                "usage": {"output_tokens": 5}
            })
            .to_string();
        }
    }
    pairs
}

fn tool_sse(
    msg_id: &str,
    tool_id: &str,
    name: &str,
    input: &serde_json::Value,
) -> Vec<(String, String)> {
    let input_str = serde_json::to_string(input).unwrap_or_default();
    vec![
        (
            "message_start".into(),
            serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
        ),
        (
            "content_block_start".into(),
            serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":tool_id,"name":name,"input":{}}}).to_string(),
        ),
        (
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":input_str}}).to_string(),
        ),
        (
            "content_block_stop".into(),
            serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
        ),
        (
            "message_delta".into(),
            serde_json::json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":20}}).to_string(),
        ),
        (
            "message_stop".into(),
            serde_json::json!({"type":"message_stop"}).to_string(),
        ),
    ]
}

/// One assistant message carrying several `tool_use` blocks, each its own
/// content-block index — the shape a provider streams for parallel tool use.
fn tools_sse(msg_id: &str, calls: &[(String, serde_json::Value)]) -> Vec<(String, String)> {
    let mut out = vec![(
        "message_start".to_string(),
        serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
    )];
    for (index, (name, input)) in calls.iter().enumerate() {
        let tool_id = format!("toolu_{}", uuid::Uuid::new_v4());
        let input_str = serde_json::to_string(input).unwrap_or_default();
        out.push((
            "content_block_start".to_string(),
            serde_json::json!({"type":"content_block_start","index":index,"content_block":{"type":"tool_use","id":tool_id,"name":name,"input":{}}}).to_string(),
        ));
        out.push((
            "content_block_delta".to_string(),
            serde_json::json!({"type":"content_block_delta","index":index,"delta":{"type":"input_json_delta","partial_json":input_str}}).to_string(),
        ));
        out.push((
            "content_block_stop".to_string(),
            serde_json::json!({"type":"content_block_stop","index":index}).to_string(),
        ));
    }
    out.push((
        "message_delta".to_string(),
        serde_json::json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":20}}).to_string(),
    ));
    out.push((
        "message_stop".to_string(),
        serde_json::json!({"type":"message_stop"}).to_string(),
    ));
    out
}

fn thinking_sse(msg_id: &str, text: &str, signature: &str) -> Vec<(String, String)> {
    vec![
        (
            "message_start".into(),
            serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":1}}}).to_string(),
        ),
        (
            "content_block_start".into(),
            serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":"","signature":""}}).to_string(),
        ),
        (
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":text}}).to_string(),
        ),
        (
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":signature}}).to_string(),
        ),
        (
            "content_block_stop".into(),
            serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
        ),
        (
            "message_delta".into(),
            serde_json::json!({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":5}}).to_string(),
        ),
        (
            "message_stop".into(),
            serde_json::json!({"type":"message_stop"}).to_string(),
        ),
    ]
}

fn text_stream_sse(chunks: &[String]) -> Vec<(String, String)> {
    let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
    let mut events = vec![
        (
            "message_start".into(),
            serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
        ),
        (
            "content_block_start".into(),
            serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}).to_string(),
        ),
    ];
    for chunk in chunks {
        events.push((
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":chunk}}).to_string(),
        ));
    }
    events.push((
        "content_block_stop".into(),
        serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
    ));
    events.push((
        "message_delta".into(),
        serde_json::json!({"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":chunks.len()}}).to_string(),
    ));
    events.push((
        "message_stop".into(),
        serde_json::json!({"type":"message_stop"}).to_string(),
    ));
    events
}

/// The opening frames plus `after` text deltas, and nothing else: no
/// `content_block_stop`, no `message_delta`, no `message_stop`.
fn cut_text_stream_sse(chunks: &[String], after: usize) -> Vec<(String, String)> {
    let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
    let mut events = vec![
        (
            "message_start".into(),
            serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
        ),
        (
            "content_block_start".into(),
            serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}).to_string(),
        ),
    ];
    for chunk in chunks.iter().take(after) {
        events.push((
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":chunk}}).to_string(),
        ));
    }
    events
}

/// A tool_use block whose `input_json_delta` is truncated mid-JSON, with no
/// `content_block_stop`, no `message_delta` and no `message_stop`.
fn cut_tool_call_stream_sse(name: &str, id: &str, partial: &str) -> Vec<(String, String)> {
    let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
    vec![
        (
            "message_start".into(),
            serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
        ),
        (
            "content_block_start".into(),
            serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":id,"name":name,"input":{}}}).to_string(),
        ),
        (
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":partial}}).to_string(),
        ),
    ]
}

fn tool_call_stream_sse(name: &str, id: &str, input: &serde_json::Value) -> Vec<(String, String)> {
    let msg_id = format!("msg_{}", uuid::Uuid::new_v4());
    let input_str = input.to_string();
    let fragments: Vec<String> = input_str
        .as_bytes()
        .chunks(10)
        .map(|c| String::from_utf8_lossy(c).to_string())
        .collect();
    let mut events = vec![
        (
            "message_start".into(),
            serde_json::json!({"type":"message_start","message":{"id":msg_id,"type":"message","role":"assistant","content":[],"model":"mock-model","stop_reason":null,"usage":{"input_tokens":10,"output_tokens":0}}}).to_string(),
        ),
        (
            "content_block_start".into(),
            serde_json::json!({"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":id,"name":name,"input":{}}}).to_string(),
        ),
    ];
    for frag in &fragments {
        events.push((
            "content_block_delta".into(),
            serde_json::json!({"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":frag}}).to_string(),
        ));
    }
    events.push((
        "content_block_stop".into(),
        serde_json::json!({"type":"content_block_stop","index":0}).to_string(),
    ));
    events.push((
        "message_delta".into(),
        serde_json::json!({"type":"message_delta","delta":{"stop_reason":"tool_use","stop_sequence":null},"usage":{"output_tokens":20}}).to_string(),
    ));
    events.push((
        "message_stop".into(),
        serde_json::json!({"type":"message_stop"}).to_string(),
    ));
    events
}

// ── JSON response helpers ─────────────────────────────────────────────────────

fn text_json(text: &str) -> serde_json::Value {
    serde_json::json!({
        "type": "message",
        "id": format!("msg_{}", uuid::Uuid::new_v4()),
        "role": "assistant",
        "content": [{"type": "text", "text": text}],
        "model": "mock-model",
        "stop_reason": "end_turn",
        "usage": {"input_tokens": 10, "output_tokens": text.len() / 4}
    })
}

fn tool_json(name: &str, input: &serde_json::Value) -> serde_json::Value {
    serde_json::json!({
        "type": "message",
        "id": format!("msg_{}", uuid::Uuid::new_v4()),
        "role": "assistant",
        "content": [{"type": "tool_use", "id": format!("toolu_{}", uuid::Uuid::new_v4()), "name": name, "input": input}],
        "model": "mock-model",
        "stop_reason": "tool_use",
        "usage": {"input_tokens": 10, "output_tokens": 20}
    })
}

fn tools_json(calls: &[(String, serde_json::Value)]) -> serde_json::Value {
    let content: Vec<serde_json::Value> = calls
        .iter()
        .map(|(name, input)| {
            serde_json::json!({"type": "tool_use", "id": format!("toolu_{}", uuid::Uuid::new_v4()), "name": name, "input": input})
        })
        .collect();
    serde_json::json!({
        "type": "message",
        "id": format!("msg_{}", uuid::Uuid::new_v4()),
        "role": "assistant",
        "content": content,
        "model": "mock-model",
        "stop_reason": "tool_use",
        "usage": {"input_tokens": 10, "output_tokens": 20}
    })
}

fn thinking_json(text: &str, signature: &str) -> serde_json::Value {
    serde_json::json!({
        "type": "message",
        "id": format!("msg_{}", uuid::Uuid::new_v4()),
        "role": "assistant",
        "content": [{"type": "thinking", "thinking": text, "signature": signature}],
        "model": "mock-model",
        "stop_reason": "end_turn",
        "usage": {"input_tokens": 10, "output_tokens": 5}
    })
}

fn error_json(message: &str) -> serde_json::Value {
    serde_json::json!({"type": "error", "error": {"type": "api_error", "message": message}})
}