agentd/intel/

mock.rs

// SPDX-License-Identifier: AGPL-3.0-only
//! A minimal built-in mock LLM for the observe-to-validate E2E suite.
//! Hidden mode: `agentd --internal-mock-llm <addr-file> [script]`.
//!
//! Binds a **loopback TCP** listener on `127.0.0.1:0` and writes the bound
//! `host:port` into `<addr-file>` atomically, via tmp-plus-rename, so the
//! launching harness can discover the endpoint by waiting for the file to
//! appear and is never able to read a half-written address. The harness then
//! hands agentd `--intelligence http://<addr>`; loopback plaintext is the
//! dev/test carve-out, since production intelligence is HTTPS-only.
//!
//! Speaks just enough OpenAI-compatible `/chat/completions` over that listener
//! to drive a *real* agentic loop without a live model: it reads the request and
//! returns a scripted assistant turn — a final answer or a tool call — switching
//! to a final answer once a tool result appears in the transcript, so the ReAct
//! cycle closes instead of spinning. Scripts: `final` (answer at once), `read`
//! (call `resource.read` then answer), `schedule` (call the `schedule` self-tool
//! then answer), `subscribe` (call the `subscribe` self-tool then answer),
//! `spawn-churn` (call `subagent.spawn` on *every* turn, never converging, so a
//! run issues a rapid burst of spawns and trips the spawn-rate limiter);
//! `slow`/`hang` hold the response to exercise the stuck/deadline detectors.
//! Small enough to ship, and it makes the loop and the self-* tools observable
//! end to end.
//!
//! **Programmable scripts:** `file:<path>` loads a JSON playbook, so a test can
//! script any conversation without adding a new built-in:
//!
//! ```json
//! { "turns": [ {"tool_calls": [{"name": "memory.set", "arguments": {"key": "k", "value": 1}}]},
//!              {"content": "done", "usage": {"prompt_tokens": 100, "completion_tokens": 20}} ],
//!   "match": [ {"when_contains": "\"preflight\"", "content": "{\"intent\":\"task\"}"} ] }
//! ```
//!
//! `match[]` rules are tried first (the first whose `when_contains` substring
//! appears in the request body answers); otherwise `turns[i]` answers where `i`
//! is the number of `role: tool` messages already in the transcript (clamped to
//! the last turn). A turn carries `content` (final text) or `tool_calls`, an
//! optional `usage`, and an optional `delay_ms`.

use std::io::{BufRead, BufReader, Read, Write};
use std::net::{TcpListener, TcpStream};

/// Start the mock IN-PROCESS and return its loopback address. This backs the
/// `intelligence.endpoints: mock:<script>` convenience, so a config runs fully
/// offline in ONE process: no key, no network, no second terminal. At most one
/// server is started per distinct script string and the address is memoised, so
/// repeated resolves of the same script share a listener instead of leaking a
/// thread per call. Scripts are the same as the hidden mode's
/// (`final` | `read` | `schedule` | `file:<playbook.json>`).
pub fn inprocess(script: &str) -> Result<String, String> {
    use std::collections::HashMap;
    use std::sync::{Mutex, OnceLock};
    static SERVERS: OnceLock<Mutex<HashMap<String, String>>> = OnceLock::new();
    let map = SERVERS.get_or_init(|| Mutex::new(HashMap::new()));
    if let Some(addr) = map.lock().expect("mock servers").get(script) {
        return Ok(addr.clone());
    }
    let addr_file = std::env::temp_dir().join(format!(
        "agentd-mockintel-{}-{}.addr",
        std::process::id(),
        crate::sha::sha256_hex(script.as_bytes())
            .chars()
            .take(12)
            .collect::<String>()
    ));
    let _ = std::fs::remove_file(&addr_file);
    let (af, sc) = (addr_file.to_string_lossy().to_string(), script.to_string());
    std::thread::Builder::new()
        .name("mock-intel".into())
        .spawn(move || run(&af, &sc))
        .map_err(|e| format!("mock intelligence: spawn: {e}"))?;
    let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5);
    loop {
        if let Ok(addr) = std::fs::read_to_string(&addr_file) {
            let addr = addr.trim().to_string();
            if !addr.is_empty() {
                map.lock()
                    .expect("mock servers")
                    .insert(script.to_string(), addr.clone());
                return Ok(addr);
            }
        }
        if std::time::Instant::now() > deadline {
            return Err("mock intelligence: server never announced its address".into());
        }
        std::thread::sleep(std::time::Duration::from_millis(5));
    }
}

/// Serve the mock LLM until the process is killed, announcing the bound
/// loopback address through `addr_file`. Returns the exit code.
pub fn run(addr_file: &str, script: &str) -> i32 {
    let listener = match TcpListener::bind("127.0.0.1:0") {
        Ok(l) => l,
        Err(e) => {
            eprintln!("mock-llm: bind 127.0.0.1:0: {e}");
            return crate::exit::GENERIC;
        }
    };
    if let Err(e) = crate::announce_addr(addr_file, &listener) {
        eprintln!("mock-llm: write {addr_file}: {e}");
        return crate::exit::GENERIC;
    }
    // One request per connection, since the intel client sends
    // `Connection: close`. Each is handled on its OWN thread so that a
    // `slow`/`hang` script sleeps only its own request: a sequential accept loop
    // would serialize every caller behind the slowest in-flight one, coupling
    // concurrent tests' timing and queueing two runs that share one mock behind
    // a 12s hang.
    for stream in listener.incoming().flatten() {
        let script = script.to_string();
        std::thread::spawn(move || handle(stream, &script));
    }
    0
}

fn handle(mut stream: TcpStream, script: &str) {
    let Some(body) = read_request_body(&mut stream) else {
        return;
    };
    if let Some(path) = script.strip_prefix("file:") {
        let payload = match std::fs::read_to_string(path)
            .map_err(|e| e.to_string())
            .and_then(|t| serde_json::from_str::<serde_json::Value>(&t).map_err(|e| e.to_string()))
        {
            Ok(playbook) => playbook_response(&playbook, &body),
            Err(e) => final_answer(&format!("mock-llm: cannot load playbook {path}: {e}")),
        };
        let resp = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            payload.len(),
            payload
        );
        let _ = stream.write_all(resp.as_bytes());
        let _ = stream.flush();
        return;
    }
    // `echo-system`: answer with the SYSTEM message verbatim. The system prompt
    // is whatever an operator's `context.template` produced, so echoing it makes
    // that observable end to end — a test asserts on what a model actually
    // receives rather than on an internal function's return value.
    if script == "echo-system" {
        let sys = serde_json::from_str::<serde_json::Value>(&body)
            .ok()
            .and_then(|v| {
                v["messages"].as_array().and_then(|ms| {
                    ms.iter()
                        .find(|m| m["role"] == "system")
                        .and_then(|m| m["content"].as_str().map(str::to_string))
                })
            })
            .unwrap_or_else(|| "(no system message)".to_string());
        let payload = final_answer(&sys);
        let resp = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
            payload.len(),
            payload
        );
        let _ = stream.write_all(resp.as_bytes());
        let _ = stream.flush();
        return;
    }
    // A `role:tool` message means the model already called a tool, so the next
    // turn is a final answer. The `gate` script needs the COUNT (define → run →
    // final is a three-phase conversation).
    let tool_results =
        body.matches("\"role\":\"tool\"").count() + body.matches("\"role\": \"tool\"").count();
    let saw_tool_result = tool_results > 0;
    // `slow`/`hang`: hold the response so the calling subagent stays alive in the
    // model call — `slow` (5s) lets the chaos suite catch a live subagent before
    // collapsing the tree; `hang` (long) keeps a run alive so a cancel/drain test
    // proves it was the teardown (not natural completion) that ended it.
    let script = match script {
        "slow" => {
            std::thread::sleep(std::time::Duration::from_secs(5));
            "final"
        }
        "hang" => {
            // Long enough to outlive a cancel/drain (kill ladder ~7s) so a test
            // proves the teardown ended it — but bounded so a *leaked* (uncancelled)
            // hang run can't pin the process-wide supervise lock for long.
            std::thread::sleep(std::time::Duration::from_secs(12));
            "final"
        }
        other => other,
    };
    // `gate`: the model AUTHORS a workflow containing a `human` gate, runs it —
    // the run blocks while the gate awaits the A2A reply — and then answers.
    let payload = if script == "gate" {
        match tool_results {
            0 => tool_call(
                "workflow.define",
                r#"{"workflow":{"start":"gate","nodes":{
                    "gate":{"kind":"human","payload":{"question":"approve the deploy?"},
                            "timeout_ms":30000,"writes":"verdict",
                            "edges":{"replied":"done","timeout":"esc"}},
                    "done":{"kind":"halt","status":"completed","result_from":"verdict"},
                    "esc":{"kind":"halt","status":"refused"}}}}"#,
            ),
            1 => tool_call("workflow.run", r#"{"workflow_id":"w1"}"#),
            _ => final_answer("gate flow complete"),
        }
    } else {
        response_json(script, saw_tool_result)
    };
    let resp = format!(
        "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{}",
        payload.len(),
        payload
    );
    let _ = stream.write_all(resp.as_bytes());
    let _ = stream.flush();
}

/// Read an HTTP/1.1 request: headers up to the blank line, then `Content-Length`
/// body bytes. Returns the request body (the chat-completions JSON).
fn read_request_body(stream: &mut TcpStream) -> Option<String> {
    let mut reader = BufReader::new(stream);
    let mut content_length = 0usize;
    loop {
        let mut line = String::new();
        if reader.read_line(&mut line).ok()? == 0 {
            return None; // EOF mid-headers
        }
        let t = line.trim_end();
        if t.is_empty() {
            break; // end of headers
        }
        if let Some(v) = t
            .strip_prefix("Content-Length:")
            .or_else(|| t.strip_prefix("content-length:"))
        {
            content_length = v.trim().parse().unwrap_or(0);
        }
    }
    let mut body = vec![0u8; content_length];
    reader.read_exact(&mut body).ok()?;
    Some(String::from_utf8_lossy(&body).into_owned())
}

/// The scripted assistant turn as an OpenAI chat-completion body.
fn response_json(script: &str, saw_tool_result: bool) -> String {
    match (script, saw_tool_result) {
        ("read", false) => tool_call("resource.read", r#"{"uri":"file:///in.json"}"#),
        ("read", true) => final_answer("read complete"),
        // Call an MCP *server* tool by its unprefixed catalogue name. The bench
        // harness points agentd at a stub MCP server serving `bench_echo`, so
        // the eval rig exercises a real tools/call round-trip end to end while
        // staying offline. Turn 2 answers once the tool result is seen.
        ("mcp-call", false) => tool_call("bench_echo", r#"{"query":"ping"}"#),
        ("mcp-call", true) => final_answer("bench tool called"),
        // Call a sandboxed `bash` tool served by the bench shell-bridge — the
        // SWE-bench / Terminal-Bench environment shape — then answer.
        ("shell-call", false) => tool_call("bash", r#"{"command":"echo pong > out.txt"}"#),
        ("shell-call", true) => final_answer("ran the command"),
        ("schedule", false) => tool_call(
            "schedule",
            r#"{"after_seconds":1,"instruction":"follow up"}"#,
        ),
        ("schedule", true) => final_answer("scheduled a follow-up"),
        ("subscribe", false) => tool_call("subscribe", r#"{"uri":"file:///watch.json"}"#),
        ("subscribe", true) => final_answer("now watching the resource"),
        // Delegate an objective to a declared remote A2A peer named "peer" and
        // then answer once the distillate comes back as a tool result. Drives
        // the agentd-as-A2A-client path end to end.
        ("a2a-delegate", false) => tool_call(
            "a2a.delegate",
            r#"{"peer":"peer","objective":"summarize the mesh","output_contract":"one line"}"#,
        ),
        ("a2a-delegate", true) => final_answer("delegated over a2a"),
        // Unlike read/schedule, which answer once a tool result is seen,
        // spawn-churn ignores `saw_tool_result` and emits another
        // `subagent.spawn` on every turn, so an in-loop run fires many rapid
        // detached spawns and exercises the spawn-rate limiter end to end.
        // `detach` keeps each accepted spawn fire-and-forget, which is what
        // keeps the burst fast enough to reach the limiter.
        ("spawn-churn", _) => tool_call(
            "subagent.spawn",
            r#"{"instruction":"do a trivial subtask","detach":true}"#,
        ),
        // Starts one workflow per turn, then answers. Paired with a workflow
        // whose only step messages the agent back, this is the
        // message → turn → run → message cycle in its shortest form — the one
        // the hop cap has to stop. Exactly one run per turn keeps the chain
        // LINEAR: a script that called on every turn would fan out instead,
        // and would be testing the step limiter rather than the hop cap.
        // Call a workflow REGISTERED AS A TOOL by its tool name, proving the
        // registration is callable and not merely advertised.
        ("wf-tool", false) => tool_call("billing.refund", r#"{"order_id":"A1"}"#),
        ("wf-tool", true) => final_answer("refund started"),
        ("wf-once", false) => tool_call("workflow.run", r#"{"name":"loop"}"#),
        ("wf-once", true) => final_answer("started the workflow"),
        // A structured JSON answer for the workflow `infer` node tests: the
        // executor parses and schema-checks this object.
        ("json", _) => final_answer(r#"{"verdict":"approve","score":9}"#),
        _ => final_answer("mock-llm done"),
    }
}

/// Answer from a `file:` playbook (see the module doc): a `match` rule by
/// request-body substring first, else the turn indexed by the number of tool
/// results already in the transcript.
fn playbook_response(playbook: &serde_json::Value, body: &str) -> String {
    let tool_results =
        body.matches("\"role\":\"tool\"").count() + body.matches("\"role\": \"tool\"").count();
    let turn = playbook
        .get("match")
        .and_then(serde_json::Value::as_array)
        .and_then(|rules| {
            rules.iter().find(|r| {
                r.get("when_contains")
                    .and_then(serde_json::Value::as_str)
                    .is_some_and(|needle| body.contains(needle))
            })
        })
        .or_else(|| {
            let turns = playbook.get("turns")?.as_array()?;
            turns.get(tool_results.min(turns.len().saturating_sub(1)))
        });
    let Some(turn) = turn else {
        return final_answer("mock-llm: empty playbook");
    };
    if let Some(ms) = turn.get("delay_ms").and_then(serde_json::Value::as_u64) {
        std::thread::sleep(std::time::Duration::from_millis(ms));
    }
    let usage = turn.get("usage").cloned();
    let mut resp: serde_json::Value = if let Some(calls) =
        turn.get("tool_calls").and_then(serde_json::Value::as_array)
    {
        let tool_calls: Vec<serde_json::Value> = calls
            .iter()
            .enumerate()
            .map(|(i, c)| {
                let args = c.get("arguments").cloned().unwrap_or(serde_json::json!({}));
                let args = match args {
                    serde_json::Value::String(s) => s,
                    other => other.to_string(),
                };
                serde_json::json!({"id": format!("call_{}", i + 1), "type": "function",
                    "function": {"name": c.get("name").and_then(serde_json::Value::as_str).unwrap_or(""), "arguments": args}})
            })
            .collect();
        serde_json::json!({
            "choices": [{"message": {"role": "assistant", "content": serde_json::Value::Null, "tool_calls": tool_calls},
                         "finish_reason": "tool_calls"}],
            "usage": {"prompt_tokens": 11, "completion_tokens": 7}
        })
    } else {
        let content = match turn.get("content") {
            Some(serde_json::Value::String(s)) => s.clone(),
            Some(other) => other.to_string(),
            None => "mock-llm done".to_string(),
        };
        serde_json::json!({
            "choices": [{"message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
            "usage": {"prompt_tokens": 11, "completion_tokens": 5}
        })
    };
    if let Some(u) = usage {
        resp["usage"] = u;
    }
    resp.to_string()
}

fn final_answer(text: &str) -> String {
    serde_json::json!({
        "choices": [{"message": {"role": "assistant", "content": text}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 11, "completion_tokens": 5}
    })
    .to_string()
}

fn tool_call(name: &str, args: &str) -> String {
    serde_json::json!({
        "choices": [{
            "message": {
                "role": "assistant",
                "content": serde_json::Value::Null,
                "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": name, "arguments": args}}]
            },
            "finish_reason": "tool_calls"
        }],
        "usage": {"prompt_tokens": 11, "completion_tokens": 7}
    })
    .to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::intel::openai;

    #[test]
    fn final_script_parses_to_a_completed_answer() {
        let resp = openai::parse_response(response_json("final", false).as_bytes()).unwrap();
        assert_eq!(resp.text.as_deref(), Some("mock-llm done"));
        assert!(resp.tool_calls.is_empty());
    }

    #[test]
    fn read_script_calls_then_answers() {
        // turn 1: a resource.read tool call
        let turn1 = openai::parse_response(response_json("read", false).as_bytes()).unwrap();
        assert!(turn1.wants_tools());
        assert_eq!(turn1.tool_calls[0].name, "resource.read");
        assert_eq!(turn1.tool_calls[0].arguments["uri"], "file:///in.json");
        // turn 2 (a tool result was seen): the final answer
        let turn2 = openai::parse_response(response_json("read", true).as_bytes()).unwrap();
        assert!(!turn2.wants_tools());
        assert_eq!(turn2.text.as_deref(), Some("read complete"));
    }

    #[test]
    fn schedule_script_calls_the_schedule_tool() {
        let turn1 = openai::parse_response(response_json("schedule", false).as_bytes()).unwrap();
        assert_eq!(turn1.tool_calls[0].name, "schedule");
        assert_eq!(turn1.tool_calls[0].arguments["after_seconds"], 1);
    }

    #[test]
    fn spawn_churn_never_converges() {
        // Every turn — whether or not a tool result was seen — is another
        // subagent.spawn with a valid (non-empty) instruction, so the run keeps
        // hammering the chokepoint instead of answering.
        for saw_tool in [false, true] {
            let turn =
                openai::parse_response(response_json("spawn-churn", saw_tool).as_bytes()).unwrap();
            assert!(turn.wants_tools(), "spawn-churn must keep calling tools");
            assert_eq!(turn.tool_calls[0].name, "subagent.spawn");
            assert_eq!(
                turn.tool_calls[0].arguments["instruction"],
                "do a trivial subtask"
            );
        }
    }

    #[test]
    fn a_playbook_answers_by_match_rule_then_by_turn_index() {
        let pb: serde_json::Value = serde_json::json!({
            "turns": [
                {"tool_calls": [{"name": "memory.set", "arguments": {"key": "k", "value": 1}}]},
                {"content": "all done", "usage": {"prompt_tokens": 500, "completion_tokens": 50}}
            ],
            "match": [{"when_contains": "PREFLIGHT", "content": {"intent": "status"}}]
        });
        // Turn 0: the scripted tool call (arguments serialized as a JSON string).
        let t0 = openai::parse_response(
            playbook_response(&pb, r#"{"messages":[{"role":"user","content":"hi"}]}"#).as_bytes(),
        )
        .unwrap();
        assert!(t0.wants_tools());
        assert_eq!(t0.tool_calls[0].name, "memory.set");
        assert_eq!(t0.tool_calls[0].arguments["value"], 1);
        // Turn 1 (one tool result seen): the final answer with the scripted usage.
        let t1 = openai::parse_response(
            playbook_response(&pb, r#"{"messages":[{"role":"tool","content":"ok"}]}"#).as_bytes(),
        )
        .unwrap();
        assert!(!t1.wants_tools());
        assert_eq!(t1.text.as_deref(), Some("all done"));
        assert_eq!(t1.usage.input_tokens, 500);
        assert_eq!(t1.usage.output_tokens, 50);
        // Beyond the last turn: clamps to the last.
        let t9 = openai::parse_response(
            playbook_response(&pb, r#"[{"role":"tool"},{"role":"tool"},{"role":"tool"}]"#)
                .as_bytes(),
        )
        .unwrap();
        assert_eq!(t9.text.as_deref(), Some("all done"));
        // A match rule wins over the turn index; object content is serialized.
        let m = openai::parse_response(
            playbook_response(
                &pb,
                r#"{"messages":[{"role":"system","content":"PREFLIGHT"}]}"#,
            )
            .as_bytes(),
        )
        .unwrap();
        assert_eq!(m.text.as_deref(), Some(r#"{"intent":"status"}"#));
    }

    #[test]
    fn subscribe_script_calls_the_subscribe_tool() {
        let turn1 = openai::parse_response(response_json("subscribe", false).as_bytes()).unwrap();
        assert_eq!(turn1.tool_calls[0].name, "subscribe");
        assert_eq!(turn1.tool_calls[0].arguments["uri"], "file:///watch.json");
        let turn2 = openai::parse_response(response_json("subscribe", true).as_bytes()).unwrap();
        assert!(!turn2.wants_tools());
    }
}