mermaid_cli/app/

run_non_interactive.rs

//! Headless driver for `mermaid run <prompt>`.
//!
//! Same reducer + same effect runner + same providers + same tools
//! as the interactive path. Differences: no `TerminalGuard`, no
//! crossterm events, no tick timer, no render. One synthetic
//! `Msg::SubmitPrompt` seeds the reducer; the loop spins until
//! `state.turn == Idle` and the queue is empty.

use std::path::PathBuf;
use std::time::Duration;

use anyhow::Result;

use crate::app::lifecycle::RuntimeLifecycle;
use crate::cli::OutputFormat;
use crate::effect::EffectRunner;
use crate::engine::{
    DriveExit, DrivePolicy, Engine, EngineHandle, Inbox, Observation, OnCancel, StepObserver,
    StopWhen,
};
use crate::providers::ToolRegistry;
use mermaid_domain::Config;
use mermaid_domain::{Msg, RUN_EVENT_PROTOCOL_VERSION, RunEvent, State};
use mermaid_model::models::MessageRole;

/// Output shape the CLI prints.
#[derive(Debug, Default)]
pub struct RunResult {
    pub response: String,
    pub reasoning: Option<String>,
    pub total_tokens: usize,
    pub errors: Vec<String>,
    /// Conversation/session id that owns this run — resumable with
    /// `mermaid run --resume <id>`.
    pub session_id: String,
    /// `--output-schema` runs: the response parsed as JSON, present only when
    /// it parsed AND validated against the schema.
    pub structured_output: Option<serde_json::Value>,
}

/// Per-invocation options for `run_non_interactive`.
///
/// Added as a struct so new flags can land without reshuffling the
/// function's positional args. All fields default to "no change".
#[derive(Debug, Default, Clone)]
pub struct RunOptions {
    /// When true, register an empty `ToolRegistry` — the model sees no
    /// tools and can't take actions. Dry-run mode for
    /// `mermaid run --no-execute`.
    pub no_execute: bool,
    /// Durable runtime task that owns this run, when launched through
    /// `mermaidd` or `mermaid run` task creation.
    pub task_id: Option<String>,
    /// External cancellation. When it fires, the driver injects
    /// `Msg::CancelTurn` — the same message the TUI's Esc sends — so the
    /// reducer unwinds the turn gracefully (tool process tree killed, turn
    /// `JoinSet` drained). If the reducer hasn't reached `Idle` within a grace
    /// window after that, the drive loop hard-stops.
    pub cancel: Option<tokio_util::sync::CancellationToken>,
    /// Wall-clock budget override. `None` keeps the built-in 20-minute
    /// deadline.
    pub deadline: Option<Duration>,
    /// When true, the driver streams the run lifecycle to stdout as
    /// newline-delimited `RunEvent` JSON (`mermaid run --format ndjson`). Off
    /// for the daemon scheduler and every other caller, which own their own
    /// output.
    pub stream_ndjson: bool,
    /// Saved conversation to seed the session with (`--resume <id>` /
    /// `--continue`). The run appends to the SAME session id, so repeated
    /// `--resume <id>` invocations chain naturally.
    pub seed: Option<mermaid_domain::ConversationHistory>,
    /// `--output-schema`: JSON Schema the final answer must conform to. The
    /// agentic loop runs normally; one extra FORMATTING turn (no tools,
    /// native constrained output where supported) reshapes the final answer,
    /// validated client-side. See `run_formatting_turn`.
    pub output_schema: Option<serde_json::Value>,
    /// Live `RunEvent` tee for daemon task subscriptions (`subscribe_task`).
    /// Every event that would print on an NDJSON stream is also broadcast
    /// here (send is sync + non-blocking; no-receiver errors are ignored).
    pub event_tx: Option<tokio::sync::broadcast::Sender<mermaid_domain::RunEvent>>,
    /// Where to publish this run's [`EngineHandle`], once there is one.
    ///
    /// The handle is the run's mailbox: it is what lets the daemon put a
    /// message into a task that is already happening, rather than only watching
    /// it and killing it. Sent exactly once, as soon as the engine exists —
    /// which is before the first model call, so a caller that awaits it can
    /// reach the whole run.
    ///
    /// An `mpsc` rather than a `oneshot` only because `RunOptions` is `Clone`
    /// and a `oneshot::Sender` is not.
    pub handle_tx: Option<tokio::sync::mpsc::Sender<EngineHandle<RunEvent>>>,
    /// `mermaid run --plan`: enter plan mode before the prompt seeds, so the
    /// run explores read-only and delivers a plan file.
    pub plan: bool,
    /// `--plan-autoaccept`: the headless approval starts implementation
    /// immediately instead of ending the run at the plan.
    pub plan_autoaccept: bool,
}

/// Drive one prompt to completion with explicit per-call options. Bounded by a
/// generous 20-minute wall-clock so a runaway model doesn't hang a script.
///
/// # Errors
///
/// Resolving the provider for `model_id`, the run exceeding the 20-minute
/// wall-clock, and a failure in the driving loop. A model that answers with a
/// refusal, a tool that fails, and a turn that ends without finishing the task
/// are all `Ok` — that is what the returned [`RunResult`] describes, and the
/// caller turns it into an exit code.
#[expect(
    clippy::too_many_lines,
    reason = "a boot sequence whose steps each touch the same config, state, runner and options, \
     in an order the comments defend (project context before the seed, scratchpad before the \
     first request, the handle before the engine takes the runner); a helper per step would need \
     most of those values and would hide the ordering that the races require"
)]
pub async fn run_non_interactive_with(
    mut config: Config,
    cwd: PathBuf,
    model_id: String,
    prompt: String,
    opts: RunOptions,
) -> Result<RunResult> {
    // `--plan-autoaccept`: the headless exit_plan_mode path consults these —
    // auto-approve with post_approve=start flows the run straight from the
    // approved plan into implementation.
    if opts.plan_autoaccept {
        config.plan.auto_approve = true;
        config.plan.post_approve = Some(mermaid_domain::PlanPostApprove::Start);
    }

    // Fold enabled plugins' MCP servers + agent types into the merged
    // config before anything consumes it (same policy as the interactive
    // path; warnings go to stderr — there is no transcript here yet).
    let plugin_assets = crate::app::plugin_assets::load();
    for warning in crate::app::plugin_assets::apply(&mut config, &plugin_assets) {
        eprintln!("mermaid: {warning}");
    }
    let providers = std::sync::Arc::new(crate::providers::ProviderFactory::new(config.clone()));
    // F6 `--no-execute`: build an empty tool registry so the model can
    // plan but never act. MCP init below is also skipped to match.
    let tools = if opts.no_execute {
        std::sync::Arc::new(ToolRegistry::new())
    } else {
        ToolRegistry::build(&config, providers.clone())
    };
    let (mut runner, mut msg_rx) =
        EffectRunner::pair_from_with_task(cwd.clone(), providers, tools, opts.task_id.clone());
    runner = runner.without_terminal_title();

    // Captured before `model_id` is moved into `State`, for the NDJSON stream.
    let stream_ndjson = opts.stream_ndjson;
    let event_model = model_id.clone();

    // One event bus for the whole run. The caller's when it supplied one — the
    // daemon subscribes before the run starts, to catch the line that names the
    // session — and otherwise one of ours, but only when a handle was asked
    // for: an unwatched run must not pay to project events nobody reads. The
    // handle and the observer share it, or `subscribe()` would be silent.
    let event_tx = match (opts.event_tx.clone(), opts.handle_tx.is_some()) {
        (Some(tx), _) => Some(tx),
        (None, true) => Some(tokio::sync::broadcast::channel(RUN_EVENT_BUS_CAPACITY).0),
        (None, false) => None,
    };

    let mut state = State::new(
        config.clone(),
        cwd.clone(),
        model_id,
        chrono::Local::now(),
        std::env::temp_dir(),
    );
    // `--resume <id>` / `--continue`: seed the session from the saved
    // conversation (same machinery as the interactive path — meters restored,
    // orphan tool pairs repaired via normalize_history), then backfill
    // provenance blanks. The seeded id survives, so the run appends to the
    // same `.mermaid/conversations/<id>.json`.
    if let Some(history) = opts.seed.clone() {
        state.seed_conversation(history);
    }
    state
        .ui
        .pending_msgs
        .push_back(Msg::SessionProvenanceResolved(
            crate::session::probe_session_provenance(&cwd),
        ));
    let session_id = state.session.conversation.id.clone();
    let mut lifecycle = RuntimeLifecycle::new();

    // Load project instructions + the memory index synchronously. The
    // interactive TUI gets these from the config watcher's first poll
    // (`run.rs`), which the headless driver never spawns — so without this the
    // model call would go out with no MERMAID.md/AGENTS.md and no memory, while
    // `mermaid doctor` reports them loaded. `build_chat_request` reads them off
    // `state`, so they must be in place before the seed below.
    let (instructions, memory, skills) =
        crate::app::instructions::load_project_context(&cwd, &config.memory);
    state.instructions = instructions;
    state.memory = memory;
    state.skills = skills;
    state.plugin_commands = plugin_assets.commands;

    // Bootstrap effects (MCP init) before the first prompt.
    //
    // Skip MCP init when `--no-execute` — MCP tools would advertise
    // through the registry we just emptied, so spinning up their
    // processes is wasted work.
    if !config.mcp_servers.is_empty() && !opts.no_execute {
        runner.dispatch(mermaid_domain::Cmd::InitMcpServers(
            config.mcp_servers.clone(),
        ));
    }

    // A resumed session may carry an in-flight checklist; hand it to the
    // TaskBroker so headless task tool calls continue the restored list.
    if !state.session.conversation.tasks.tasks.is_empty() {
        runner.dispatch(mermaid_domain::Cmd::SyncTaskStore(
            state.session.conversation.tasks.clone(),
        ));
    }

    // Materialize the per-session scratch dir SYNCHRONOUSLY and stamp it
    // before the seed. The interactive path goes through
    // `Cmd::EnsureScratchpad` -> `Msg::ScratchpadReady`, but here the first
    // (often only) request is built immediately below — the async round
    // trip loses that race and the prompt ships without the scratchpad
    // path, so the model can't address it until a tool turn has elapsed.
    // Same rationale as the synchronous project-context load above.
    // `session_id` was captured after any seed, so a resumed run adopts
    // the dir keyed by its restored conversation id.
    match crate::session::scratchpad::ensure(&cwd, &session_id) {
        Ok(path) => state.session.scratchpad = Some(path),
        Err(err) => tracing::warn!(%err, "scratchpad unavailable for this run"),
    }

    // First line of the NDJSON stream: protocol + run identity. Broadcast
    // to any daemon subscriber regardless of stdout streaming.
    let started = RunEvent::SessionStarted {
        protocol_version: RUN_EVENT_PROTOCOL_VERSION,
        cli_version: env!("CARGO_PKG_VERSION").to_string(),
        model: event_model,
        task_id: opts.task_id.clone(),
        session_id: session_id.clone(),
    };
    if stream_ndjson {
        emit_run_event(&started);
    }
    if let Some(tx) = &event_tx {
        let _ = tx.send(started);
    }

    // Publish this run's mailbox before the engine takes the runner, so a
    // caller awaiting the handle can reach the run from its first model call
    // onward. The bus is the caller's own when it supplied one — the daemon
    // subscribes before the run starts, to catch the line that names the
    // session — and otherwise one of ours, which keeps `subscribe()` total.
    if let Some((want, events)) = opts.handle_tx.as_ref().zip(event_tx.clone()) {
        let _ = want.try_send(EngineHandle::new(runner.sender(), events));
    }

    // Everything the reducer needs is in place: hand the state and the runner
    // to the engine, which owns the drive from here.
    let mut engine = Engine::new(state, runner).with_observer(RunStream {
        stream_ndjson,
        event_tx: event_tx.clone(),
    });

    // `--plan`: flip into plan mode BEFORE the prompt seeds, through the
    // same reducer path as the interactive `/plan` (path allocation, model
    // swap, prompt injection all included).
    //
    // Both this and the seed below go through `reduce`, not `step`: they are
    // synthetic inputs the driver manufactures, and projecting them onto the
    // public stream would announce events no client sent.
    if opts.plan {
        engine.reduce(
            chrono::Local::now(),
            Msg::Slash(mermaid_domain::SlashCmd::Plan(None)),
        );
    }

    // Seed the turn. The clock is injected as data so the reducer stays pure.
    engine.reduce(
        chrono::Local::now(),
        Msg::SubmitPrompt {
            text: prompt,
            attachment_ids: vec![],
        },
    );

    let deadline = opts.deadline.unwrap_or(Duration::from_secs(20 * 60));
    let cancel = opts.cancel.clone();
    let policy = DrivePolicy {
        stop: StopWhen::Settled,
        cancel: cancel.clone(),
        // Cancellation unwinds the turn through the reducer — the same
        // `Msg::CancelTurn` the TUI's Esc sends — rather than dropping it.
        on_cancel: OnCancel::Unwind {
            grace: CANCEL_GRACE,
        },
        deadline: Some(deadline),
    };
    let mut inbox = Inbox::new(&mut msg_rx).with_lifecycle(&mut lifecycle);

    if engine.drive(&mut inbox, &policy).await == DriveExit::TimedOut {
        let (_, runner, _) = engine.into_parts();
        return timed_out(runner, "non-interactive run", deadline).await;
    }

    let mut result = build_result(engine.state());

    // `--output-schema`: one extra formatting turn on the completed run.
    if let Some(schema) = opts.output_schema.clone() {
        let cancelled = cancel.as_ref().is_some_and(|t| t.is_cancelled());
        if cancelled || engine.state().should_exit || result.response.is_empty() {
            result
                .errors
                .push("output_schema: skipped (run ended without a final answer)".to_string());
        } else {
            let exit = run_formatting_turn(&mut engine, schema.clone(), &mut inbox, &policy).await;
            if exit == DriveExit::TimedOut {
                let (_, runner, _) = engine.into_parts();
                return timed_out(runner, "output-schema formatting turn", deadline).await;
            }
            apply_schema_outcome(&mut result, engine.state(), &schema);
        }
    }

    let (_, runner, _) = engine.into_parts();
    runner.shutdown().await;
    // Terminal line of the stream: the aggregated result — sent to daemon
    // subscribers even when stdout NDJSON is off (that's how a
    // `subscribe_task` stream knows to close).
    let terminal = RunEvent::Result {
        response: result.response.clone(),
        reasoning: result.reasoning.clone(),
        total_tokens: result.total_tokens as u64,
        errors: result.errors.clone(),
        session_id: result.session_id.clone(),
        structured_output: result.structured_output.clone(),
    };
    if stream_ndjson {
        emit_run_event(&terminal);
    }
    if let Some(tx) = &event_tx {
        let _ = tx.send(terminal);
    }
    Ok(result)
}

/// Slots on the run's event bus. Matches the daemon's own: a lagged
/// subscriber gets a `RecvError::Lagged` marker rather than backpressure into
/// the run, which must never stall on someone watching it.
const RUN_EVENT_BUS_CAPACITY: usize = 1024;

/// How long a cancelled run may keep unwinding before the drive hard-stops.
/// Generous next to the turn scope's own ~2s teardown bound.
const CANCEL_GRACE: Duration = Duration::from_secs(15);

/// Everything a headless run says out loud, as an engine observer.
///
/// Runs before `update` consumes the message, so the projection sees fields the
/// reducer strips, and each message is projected ONCE — printed and/or
/// broadcast. Most messages have no projection at all.
struct RunStream {
    stream_ndjson: bool,
    event_tx: Option<tokio::sync::broadcast::Sender<RunEvent>>,
}

impl StepObserver for RunStream {
    fn observe(&mut self, obs: Observation<'_>) {
        // Plumbing notices ("Starting the local Ollama server…") have no
        // renderer here — mirror them to stderr live so the user isn't staring
        // at silence during an up-to-15s server start. stderr, not stdout: the
        // response payload must stay clean for scripts.
        if let Msg::TransientStatus { text } = obs.msg {
            eprintln!("{text}");
        }
        if !self.stream_ndjson && self.event_tx.is_none() {
            return;
        }
        let Some(event) = RunEvent::from_msg(obs.msg) else {
            return;
        };
        if self.stream_ndjson {
            emit_run_event(&event);
        }
        if let Some(tx) = &self.event_tx {
            let _ = tx.send(event);
        }
    }
}

/// Shut the runner down, then report the timeout.
///
/// The deadline is a drive-policy arm rather than a `timeout()` wrapper
/// precisely so this path can exist: the wrapper returned `Err` through `?`
/// while the runner was still owned by the caller's stack frame, and dropping a
/// runner mid-flight leaks its MCP children (#76 — fixed in `drive_child` at
/// the time, and only there).
/// Takes the runner rather than the engine: this future is live at two await
/// points inside the driver's own, and moving the engine in would put a
/// `State`-sized slot in both — which is how `run_non_interactive_with`'s
/// future crosses `clippy::large_futures`.
async fn timed_out(runner: EffectRunner, what: &str, deadline: Duration) -> Result<RunResult> {
    runner.shutdown().await;
    Err(anyhow::anyhow!(
        "{what} exceeded {} seconds",
        deadline.as_secs()
    ))
}

/// The synthetic prompt that drives the `--output-schema` formatting turn.
const FORMAT_PROMPT: &str = "Convert your final answer into a single JSON object that \
conforms to the provided schema. Respond with only the JSON object - no prose, no code fences.";

/// Run the dedicated `--output-schema` formatting turn: set the schema rider
/// on state (`build_chat_request` drops all tools + adapters engage native
/// constrained output), seed the format prompt, and drive to settled.
async fn run_formatting_turn<O: StepObserver>(
    engine: &mut Engine<EffectRunner, O>,
    schema: serde_json::Value,
    inbox: &mut Inbox<'_>,
    policy: &DrivePolicy,
) -> DriveExit {
    engine.state_mut().output_schema = Some(schema);
    engine.reduce(
        chrono::Local::now(),
        Msg::SubmitPrompt {
            text: FORMAT_PROMPT.to_string(),
            attachment_ids: vec![],
        },
    );
    engine.drive(inbox, policy).await
}

/// Fold the formatting turn's outcome into the run result: the reshaped text
/// replaces the response (code fences stripped), and `structured_output` is
/// set only when the text parses AND validates. Every failure keeps the best
/// available text and records an `output_schema:` error — a run that produced
/// an answer never returns empty output.
fn apply_schema_outcome(result: &mut RunResult, state: &State, schema: &serde_json::Value) {
    // The formatting reply is the last assistant message; total tokens grew.
    let formatted = build_result(state);
    result.total_tokens = formatted.total_tokens;
    if formatted.response.is_empty() || formatted.response == result.response {
        result
            .errors
            .push("output_schema: formatting turn produced no output".to_string());
        return;
    }
    let text = strip_code_fences(&formatted.response);
    result.response = text.to_string();
    let parsed: serde_json::Value = match serde_json::from_str(text) {
        Ok(v) => v,
        Err(e) => {
            result
                .errors
                .push(format!("output_schema: response is not valid JSON: {e}"));
            return;
        },
    };
    let validator = match jsonschema::validator_for(schema) {
        Ok(v) => v,
        Err(e) => {
            result
                .errors
                .push(format!("output_schema: schema did not compile: {e}"));
            return;
        },
    };
    if let Some(err) = validator.iter_errors(&parsed).next() {
        result
            .errors
            .push(format!("output_schema: response does not conform: {err}"));
        return;
    }
    result.structured_output = Some(parsed);
}

/// Trim a single wrapping markdown code fence (with optional info string) —
/// models wrap JSON in ` ```json ` fences despite instructions.
fn strip_code_fences(text: &str) -> &str {
    let t = text.trim();
    let Some(rest) = t.strip_prefix("```") else {
        return t;
    };
    let Some(rest) = rest.split_once('\n').map(|(_, r)| r) else {
        return t;
    };
    match rest.strip_suffix("```") {
        Some(inner) => inner.trim(),
        None => t,
    }
}

/// Write one `RunEvent` as a JSON line to stdout (the NDJSON SDK stream).
fn emit_run_event(event: &RunEvent) {
    println!("{}", serde_json::to_string(event).unwrap_or_default());
}

/// Walk the committed message history and pull out the last
/// assistant response + any errors encountered.
fn build_result(state: &State) -> RunResult {
    let mut out = RunResult {
        total_tokens: state.session.cumulative_token_usage.total_tokens(),
        session_id: state.session.conversation.id.clone(),
        ..RunResult::default()
    };

    for msg in state.session.messages() {
        for action in &msg.actions {
            if let mermaid_domain::ActionResult::Error { error } = &action.result {
                out.errors
                    .push(format!("{}: {}", action.action_type, error));
            }
        }
    }

    // The final reply may span an auto-continue chain: the last assistant
    // message plus any `Continuation`-kind messages leading up to it. Walk
    // back to the chain's head, then join the segments in order with the
    // same conservative resume-echo trim the transcript stitch uses —
    // otherwise `--output text/json` silently dropped everything before the
    // last continuation.
    let messages = state.session.messages();
    if let Some(last_idx) = messages
        .iter()
        .rposition(|m| m.role == MessageRole::Assistant)
    {
        let mut head_idx = last_idx;
        while head_idx > 0
            && messages[head_idx].kind == mermaid_model::models::ChatMessageKind::Continuation
            && let Some(prev_idx) = messages[..head_idx]
                .iter()
                .rposition(|m| m.role == MessageRole::Assistant)
            // Only step onto a real bubble: a compaction checkpoint or the
            // empty error-carrier is never part of the reply chain.
            && matches!(
                messages[prev_idx].kind,
                mermaid_model::models::ChatMessageKind::Normal
                    | mermaid_model::models::ChatMessageKind::Continuation
            )
            && messages[prev_idx].tool_calls.is_none()
        {
            head_idx = prev_idx;
        }
        let mut response = String::new();
        let mut reasoning: Option<String> = None;
        for msg in messages[head_idx..=last_idx]
            .iter()
            .filter(|m| m.role == MessageRole::Assistant)
        {
            let skip = mermaid_model::utils::continuation_overlap(&response, &msg.content);
            response.push_str(&msg.content[skip..]);
            if let Some(t) = &msg.thinking {
                match &mut reasoning {
                    Some(r) => {
                        r.push_str("\n\n");
                        r.push_str(t);
                    },
                    None => reasoning = Some(t.clone()),
                }
            }
        }
        out.response = response;
        out.reasoning = reasoning;
    }

    out
}

/// Render a `RunResult` in the requested output format.
#[must_use]
pub fn format_result(result: &RunResult, format: OutputFormat) -> String {
    match format {
        OutputFormat::Text => {
            if result.response.is_empty() && !result.errors.is_empty() {
                result.errors.join("\n")
            } else {
                result.response.clone()
            }
        },
        OutputFormat::Markdown => {
            let mut out = result.response.clone();
            if !result.errors.is_empty() {
                out.push_str("\n\n---\n\n## Errors\n\n");
                for e in &result.errors {
                    out.push_str(&format!("- {e}\n"));
                }
            }
            out
        },
        OutputFormat::Json => {
            // Typed single-object form — the same shape as the streamed terminal
            // `RunEvent::Result`, so the golden test pins this output too.
            let event = RunEvent::Result {
                response: result.response.clone(),
                reasoning: result.reasoning.clone(),
                total_tokens: result.total_tokens as u64,
                errors: result.errors.clone(),
                session_id: result.session_id.clone(),
                structured_output: result.structured_output.clone(),
            };
            serde_json::to_string_pretty(&event).unwrap_or_default()
        },
        OutputFormat::Ndjson => {
            // Events were streamed live during the run; nothing to print here.
            String::new()
        },
    }
}

#[cfg(test)]
mod tests {
    use super::build_result;

    #[test]
    fn build_result_joins_an_auto_continued_reply() {
        use mermaid_model::models::{ChatMessage, ChatMessageKind};
        let mut state = mermaid_domain::State::new(
            mermaid_domain::Config::default(),
            std::path::PathBuf::from("/tmp/p"),
            "ollama/test".to_string(),
            chrono::Local::now(),
            std::path::PathBuf::from("/tmp"),
        );
        state
            .session
            .append(ChatMessage::user("audit the widget"), state.now);
        state.session.append(
            ChatMessage::assistant("part one covers the resolver internals"),
            state.now,
        );
        // The continuation echoes the tail of part one — joined output must
        // carry it exactly once.
        let mut cont = ChatMessage::assistant("the resolver internals, part two the adapters.");
        cont.kind = ChatMessageKind::Continuation;
        state.session.append(cont, state.now);

        let result = build_result(&state);
        assert_eq!(
            result.response, "part one covers the resolver internals, part two the adapters.",
            "headless output joins the whole chain, echo trimmed"
        );
    }

    #[test]
    fn build_result_without_chain_takes_the_last_reply() {
        use mermaid_model::models::ChatMessage;
        let mut state = mermaid_domain::State::new(
            mermaid_domain::Config::default(),
            std::path::PathBuf::from("/tmp/p"),
            "ollama/test".to_string(),
            chrono::Local::now(),
            std::path::PathBuf::from("/tmp"),
        );
        state.session.append(ChatMessage::user("first"), state.now);
        state
            .session
            .append(ChatMessage::assistant("earlier reply"), state.now);
        state.session.append(ChatMessage::user("second"), state.now);
        state
            .session
            .append(ChatMessage::assistant("final reply"), state.now);
        let result = build_result(&state);
        assert_eq!(result.response, "final reply");
    }

    #[test]
    fn strip_code_fences_unwraps_single_fence_only() {
        use super::strip_code_fences;
        assert_eq!(strip_code_fences("{\"a\":1}"), "{\"a\":1}");
        assert_eq!(strip_code_fences("```json\n{\"a\":1}\n```"), "{\"a\":1}");
        assert_eq!(strip_code_fences("```\n{\"a\":1}\n```"), "{\"a\":1}");
        // Unterminated fence -> left alone (don't mangle).
        assert_eq!(
            strip_code_fences("```json\n{\"a\":1}"),
            "```json\n{\"a\":1}"
        );
        assert_eq!(strip_code_fences("  {\"a\":1}  "), "{\"a\":1}");
    }

    fn schema_state(reply: &str) -> mermaid_domain::State {
        use mermaid_model::models::ChatMessage;
        let mut state = mermaid_domain::State::new(
            mermaid_domain::Config::default(),
            std::path::PathBuf::from("/tmp/p"),
            "ollama/test".to_string(),
            chrono::Local::now(),
            std::path::PathBuf::from("/tmp"),
        );
        state.session.append(ChatMessage::user("q"), state.now);
        state
            .session
            .append(ChatMessage::assistant("the plain answer"), state.now);
        state
            .session
            .append(ChatMessage::user("format it"), state.now);
        state
            .session
            .append(ChatMessage::assistant(reply), state.now);
        state
    }

    fn base_result() -> super::RunResult {
        super::RunResult {
            response: "the plain answer".to_string(),
            ..super::RunResult::default()
        }
    }

    #[test]
    fn schema_outcome_valid_json_sets_structured_output() {
        let schema = serde_json::json!({
            "type": "object",
            "properties": {"answer": {"type": "integer"}},
            "required": ["answer"],
        });
        let state = schema_state("```json\n{\"answer\": 42}\n```");
        let mut result = base_result();
        super::apply_schema_outcome(&mut result, &state, &schema);
        assert_eq!(result.response, "{\"answer\": 42}");
        assert_eq!(
            result.structured_output,
            Some(serde_json::json!({"answer": 42}))
        );
        assert!(result.errors.is_empty(), "{:?}", result.errors);
    }

    #[test]
    fn schema_outcome_invalid_json_keeps_text_and_records() {
        let schema = serde_json::json!({"type": "object"});
        let state = schema_state("not json at all");
        let mut result = base_result();
        super::apply_schema_outcome(&mut result, &state, &schema);
        assert_eq!(result.response, "not json at all");
        assert!(result.structured_output.is_none());
        assert!(
            result.errors.iter().any(|e| e.contains("not valid JSON")),
            "{:?}",
            result.errors
        );
    }

    #[test]
    fn schema_outcome_nonconforming_json_records_reason() {
        let schema = serde_json::json!({
            "type": "object",
            "properties": {"answer": {"type": "integer"}},
            "required": ["answer"],
        });
        let state = schema_state("{\"wrong\": true}");
        let mut result = base_result();
        super::apply_schema_outcome(&mut result, &state, &schema);
        assert!(result.structured_output.is_none());
        assert!(
            result.errors.iter().any(|e| e.contains("does not conform")),
            "{:?}",
            result.errors
        );
    }

    #[test]
    fn schema_outcome_no_new_reply_keeps_original() {
        // The formatting turn produced nothing new (same last assistant
        // message) -> keep the agent's answer, record the failure.
        let schema = serde_json::json!({"type": "object"});
        use mermaid_model::models::ChatMessage;
        let mut state = mermaid_domain::State::new(
            mermaid_domain::Config::default(),
            std::path::PathBuf::from("/tmp/p"),
            "ollama/test".to_string(),
            chrono::Local::now(),
            std::path::PathBuf::from("/tmp"),
        );
        state.session.append(ChatMessage::user("q"), state.now);
        state
            .session
            .append(ChatMessage::assistant("the plain answer"), state.now);
        let mut result = base_result();
        super::apply_schema_outcome(&mut result, &state, &schema);
        assert_eq!(result.response, "the plain answer");
        assert!(result.structured_output.is_none());
        assert!(
            result
                .errors
                .iter()
                .any(|e| e.contains("produced no output")),
            "{:?}",
            result.errors
        );
    }

    #[test]
    fn result_event_carries_structured_output_to_subscribers() {
        // The broadcast tee is exercised end-to-end by the daemon integration
        // test; here pin the terminal event SHAPE subscribers rely on (a
        // `result` type ends the stream).
        let event = mermaid_domain::RunEvent::Result {
            response: "done".to_string(),
            reasoning: None,
            total_tokens: 3,
            errors: vec![],
            session_id: "s".to_string(),
            structured_output: None,
        };
        let wire = serde_json::to_string(&event).unwrap();
        assert!(wire.contains("\"type\":\"result\""), "{wire}");
        let (tx, mut rx) = tokio::sync::broadcast::channel::<mermaid_domain::RunEvent>(4);
        tx.send(event.clone()).unwrap();
        assert_eq!(rx.try_recv().unwrap(), event);
    }

    // When to stop driving now belongs to `StopWhen::Settled` in
    // `crate::engine`, tested there against the real loop rather than against
    // a boolean predicate standing in for it.
}