clawmates/crates/cm-api/src/topology_exec.rs

//! A [`cm_orchestrator::TurnExecutor`] that runs each topology turn as a real
//! ZeroClaw role-agent inside a container, driven over the gateway WebSocket
//! "drive" recipe.
//!
//! The orchestrator owns the topology graph (it sequences edges, meters, and
//! records the journal); this executor only runs *one* turn: it pairs with the
//! gateway, opens `/ws/chat?agent=<alias>`, sends the role+task+context prompt,
//! and streams the turn's events back into a [`TurnOutcome`].
//!
//! **§15 by construction:** the agents are provisioned tool-free (every
//! sensitive capability is a gated Clawmates MCP tool — the "door"), so a turn
//! takes no sandbox-leaving action here. If the gateway nonetheless emits an
//! `approval_request`, we record it as a **blocked** `GatedAction` and end the
//! turn — we never auto-approve.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use cm_domain::GatedCategory;
use cm_orchestrator::{GatedAction, OrchestratorError, TurnExecutor, TurnOutcome, TurnRequest};
use futures::{SinkExt, StreamExt};
use tokio::sync::Mutex;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite::Message;

/// Overall wall-clock budget for draining one turn's event stream. Must
/// exceed the daemon's own claude_cli provider timeout (600s on gw-04
/// via ZEROCLAW_providers__models__claude_cli__default__timeout_ms) —
/// otherwise the executor kills the ws before the daemon can reply and
/// we see a phantom "turn timed out" while the daemon still logs a
/// successful llm response coming back. 700s gives 100s of headroom so
/// a daemon that just barely made it under its own limit doesn't lose
/// its answer here.
const TURN_TIMEOUT: Duration = Duration::from_secs(700);

/// Drives ZeroClaw role-agents (in one container) to execute topology turns.
pub struct ZeroClawDriveExecutor {
    /// Gateway base URL, e.g. `http://127.0.0.1:42617`.
    gateway_url: String,
    /// One-time pairing code (minted into a bearer token on first use).
    pairing_code: String,
    /// Topology `node.role` → ZeroClaw agent alias.
    role_aliases: HashMap<String, String>,
    /// Alias used when a role isn't mapped.
    default_alias: String,
    /// Bearer token, paired lazily and reused across turns.
    token: Arc<Mutex<Option<String>>>,
    http: reqwest::Client,
}

impl ZeroClawDriveExecutor {
    /// Build an executor explicitly (used in tests).
    pub fn new(
        gateway_url: String,
        pairing_code: String,
        role_aliases: HashMap<String, String>,
        default_alias: String,
    ) -> Self {
        ZeroClawDriveExecutor {
            gateway_url: gateway_url.trim_end_matches('/').to_string(),
            pairing_code,
            role_aliases,
            default_alias,
            token: Arc::new(Mutex::new(None)),
            http: reqwest::Client::new(),
        }
    }

    /// Build from the environment:
    /// - `ZEROCLAW_GATEWAY_URL` (required) e.g. `http://127.0.0.1:42617`
    /// - `ZEROCLAW_TOKEN` (preferred) a durable bearer token — pair once
    ///   out-of-band, set this, and runs are repeatable. If set, no pairing.
    /// - `ZEROCLAW_PAIRING_CODE` (fallback) a one-time pairing code — consumed
    ///   on first use, so only good for a single run. One of TOKEN/CODE required.
    /// - `ZEROCLAW_AGENT_MAP` (optional) `role=alias,role=alias`
    /// - `ZEROCLAW_DEFAULT_AGENT` (optional, default `scout`)
    pub fn from_env() -> Result<Self, String> {
        let gateway_url =
            std::env::var("ZEROCLAW_GATEWAY_URL").map_err(|_| "ZEROCLAW_GATEWAY_URL not set")?;
        Self::from_env_for_gateway(gateway_url)
    }

    /// Same as [`from_env`] but with a caller-supplied gateway URL. Used by
    /// the research pipeline to point the executor at the per-topic team
    /// container spawned in `research_container::spawn` instead of the
    /// workspace-wide gateway from `ZEROCLAW_GATEWAY_URL`. The auth token,
    /// role map, and default alias still come from the parent server's
    /// env — they're propagated into the team container by
    /// `research_container::inherited_env` so both endpoints use the same
    /// credentials.
    pub fn from_env_for_gateway(gateway_url: String) -> Result<Self, String> {
        let token = std::env::var("ZEROCLAW_TOKEN")
            .ok()
            .filter(|t| !t.is_empty());
        let pairing_code = std::env::var("ZEROCLAW_PAIRING_CODE").unwrap_or_default();
        if token.is_none() && pairing_code.is_empty() {
            return Err("set ZEROCLAW_TOKEN or ZEROCLAW_PAIRING_CODE".to_string());
        }
        let default_alias =
            std::env::var("ZEROCLAW_DEFAULT_AGENT").unwrap_or_else(|_| "scout".to_string());
        let role_aliases = std::env::var("ZEROCLAW_AGENT_MAP")
            .ok()
            .map(|s| parse_agent_map(&s))
            .unwrap_or_default();
        let mut exec = Self::new(gateway_url, pairing_code, role_aliases, default_alias);
        exec.token = Arc::new(Mutex::new(token));
        Ok(exec)
    }

    /// Like [`from_env_for_gateway`] but with a caller-supplied pairing
    /// code — used by per-mission runtimes whose fresh daemons mint a
    /// new one-time code at startup. The env-derived ZEROCLAW_TOKEN
    /// is ignored (belongs to the shared runtime) so the lazy pair
    /// path runs and issues a bearer for this specific gateway.
    pub fn from_env_for_gateway_with_code(
        gateway_url: String,
        pairing_code: String,
    ) -> Result<Self, String> {
        if pairing_code.is_empty() {
            return Err("empty pairing_code".to_string());
        }
        let default_alias =
            std::env::var("ZEROCLAW_DEFAULT_AGENT").unwrap_or_else(|_| "scout".to_string());
        let role_aliases = std::env::var("ZEROCLAW_AGENT_MAP")
            .ok()
            .map(|s| parse_agent_map(&s))
            .unwrap_or_default();
        Ok(Self::new(
            gateway_url,
            pairing_code,
            role_aliases,
            default_alias,
        ))
    }

    fn alias_for(&self, role: &str) -> String {
        self.role_aliases
            .get(role)
            .cloned()
            .unwrap_or_else(|| self.default_alias.clone())
    }

    /// `POST {base}/pair` with the pairing code header → bearer token (cached).
    async fn ensure_paired(&self) -> Result<String, OrchestratorError> {
        let mut guard = self.token.lock().await;
        if let Some(tok) = guard.as_ref() {
            return Ok(tok.clone());
        }
        let res = self
            .http
            .post(format!("{}/pair", self.gateway_url))
            .header("X-Pairing-Code", &self.pairing_code)
            .header("Content-Type", "application/json")
            .body("{}")
            .send()
            .await
            .map_err(|e| OrchestratorError::Executor(format!("pair request failed: {e}")))?;
        if !res.status().is_success() {
            return Err(OrchestratorError::Executor(format!(
                "pair failed: {}",
                res.status()
            )));
        }
        let body: serde_json::Value = res
            .json()
            .await
            .map_err(|e| OrchestratorError::Executor(format!("pair response not json: {e}")))?;
        let token = body
            .get("token")
            .or_else(|| body.get("bearer"))
            .or_else(|| body.get("access_token"))
            .and_then(|v| v.as_str())
            .ok_or_else(|| OrchestratorError::Executor("pair response had no token".into()))?
            .to_string();
        *guard = Some(token.clone());
        Ok(token)
    }

    /// Mirror of `ProviderExecutor`'s prompt, flattened to one `content` string
    /// (the gateway `message` envelope carries a single content field).
    fn build_prompt(req: &TurnRequest) -> String {
        let system = format!(
            "You are the \"{}\" agent in a multi-agent system. Do your part of the task \
             concisely and return only your result.",
            req.role
        );
        let mut user = format!("Task: {}", req.task);
        if !req.context.is_empty() {
            user.push_str("\n\nContext from upstream agents:\n");
            for (i, c) in req.context.iter().enumerate() {
                user.push_str(&format!("[{i}] {c}\n"));
            }
        }
        format!("{system}\n\n{user}")
    }

    pub async fn drive(&self, alias: &str, prompt: &str) -> Result<TurnOutcome, OrchestratorError> {
        let token = self.ensure_paired().await?;
        let ws_base = if let Some(rest) = self.gateway_url.strip_prefix("https") {
            format!("wss{rest}")
        } else if let Some(rest) = self.gateway_url.strip_prefix("http") {
            format!("ws{rest}")
        } else {
            self.gateway_url.clone()
        };
        let ws_url = format!(
            "{ws_base}/ws/chat?agent={}&name=clawmates&token={}",
            urlencoding::encode(alias),
            urlencoding::encode(&token),
        );

        let (mut ws, _resp) = connect_async(&ws_url)
            .await
            .map_err(|e| OrchestratorError::Executor(format!("ws connect failed: {e}")))?;

        let envelope = serde_json::json!({ "type": "message", "content": prompt }).to_string();
        ws.send(Message::Text(envelope.into()))
            .await
            .map_err(|e| OrchestratorError::Executor(format!("ws send failed: {e}")))?;

        let outcome = tokio::time::timeout(TURN_TIMEOUT, Self::drain(&mut ws))
            .await
            .map_err(|_| OrchestratorError::Executor("turn timed out".into()))??;
        let _ = ws.close(None).await;
        Ok(outcome)
    }

    /// Use a runtime agent as a governance judge: drive `alias` with the judge
    /// prompt and parse the verdict (`DENY` anywhere ⇒ deny, else allow). This
    /// lets a **subscription-only** model (e.g. Kimi via `kimi_cli`) be the judge
    /// with no platform API key — the registry/SDK path GLM and Kimi can't take.
    /// Fail-open (returns `(true, …)`) so a judge outage never halts agents.
    pub async fn judge(&self, alias: &str, system: &str, user: &str) -> (bool, String) {
        let prompt = format!("{system}\n\n{user}");
        match self.drive(alias, &prompt).await {
            Ok(outcome) => {
                let text = outcome.output.trim().to_string();
                let allow = !text.to_uppercase().contains("DENY");
                (allow, text)
            }
            Err(e) => (true, format!("governor unreachable (fail-open): {e}")),
        }
    }

    /// Drive `alias` with a judging prompt and return its **raw** reply.
    ///
    /// [`Self::judge`] collapses the reply to a bool by substring-matching
    /// `DENY`, which only suits the governor's ALLOW/DENY contract and is
    /// fail-open. Callers that need a structured verdict — the phase
    /// completion evaluator wants `{"met":bool,"reason":string}` and must fail
    /// **closed** — need the text, and need the error rather than a
    /// synthesized permissive answer.
    pub async fn judge_raw(&self, alias: &str, system: &str, user: &str) -> Result<String, String> {
        let prompt = format!("{system}\n\n{user}");
        self.drive(alias, &prompt)
            .await
            .map(|outcome| outcome.output.trim().to_string())
            .map_err(|e| e.to_string())
    }

    /// Drive agent `alias` as a delegated sub-task and return its result. Reuses
    /// the same gateway drive as topology turns + the governor, so a delegated
    /// turn carries the same blocked-action / token instrumentation in its
    /// [`TurnOutcome`]. The target is tool-free behind the MCP door, so a
    /// delegated turn adds no new egress (§15 holds — the bridge only causes an
    /// in-workspace agent to run a turn; anything it does is independently gated
    /// at the door).
    pub async fn delegate(
        &self,
        alias: &str,
        task: &str,
        context: &[String],
    ) -> Result<TurnOutcome, OrchestratorError> {
        let mut prompt = format!(
            "You are being delegated a sub-task by another agent on your team. \
             Do it concisely and return only your result.\n\nTask: {task}"
        );
        if !context.is_empty() {
            prompt.push_str("\n\nContext from the delegating agent:\n");
            for (i, c) in context.iter().enumerate() {
                prompt.push_str(&format!("[{i}] {c}\n"));
            }
        }
        self.drive(alias, &prompt).await
    }

    /// Read frames until a terminal (`done`/`error`/`approval_request`) event.
    async fn drain<S>(ws: &mut S) -> Result<TurnOutcome, OrchestratorError>
    where
        S: StreamExt<Item = Result<Message, tokio_tungstenite::tungstenite::Error>>
            + SinkExt<Message>
            + Unpin,
    {
        let mut output = String::new();
        let mut tokens: u64 = 0;
        let mut gated: Vec<GatedAction> = Vec::new();

        while let Some(frame) = ws.next().await {
            let msg = frame.map_err(|e| OrchestratorError::Executor(format!("ws recv: {e}")))?;
            match msg {
                Message::Text(txt) => {
                    let v: serde_json::Value = serde_json::from_str(txt.as_str())
                        .map_err(|e| OrchestratorError::Executor(format!("bad frame: {e}")))?;
                    match v.get("type").and_then(|t| t.as_str()).unwrap_or("") {
                        "chunk" => {
                            if let Some(c) = v.get("content").and_then(|c| c.as_str()) {
                                output.push_str(c);
                            }
                        }
                        "done" => {
                            let input = v.get("input_tokens").and_then(|n| n.as_u64()).unwrap_or(0);
                            let out = v.get("output_tokens").and_then(|n| n.as_u64()).unwrap_or(0);
                            tokens = input + out;
                            break;
                        }
                        "approval_request" => {
                            // §15: agents are tool-free behind the MCP door, so
                            // this is unexpected. Record it as blocked, never
                            // auto-approve, and end the turn.
                            let tool = v.get("tool").and_then(|t| t.as_str()).unwrap_or("unknown");
                            let summary = v
                                .get("arguments_summary")
                                .and_then(|s| s.as_str())
                                .unwrap_or("");
                            gated.push(GatedAction {
                                category: GatedCategory::OutboundMessage,
                                summary: format!("{tool}: {summary}").trim().to_string(),
                                approved: false,
                            });
                            break;
                        }
                        "error" => {
                            let m = v
                                .get("message")
                                .and_then(|m| m.as_str())
                                .unwrap_or("agent error");
                            return Err(OrchestratorError::Executor(m.to_string()));
                        }
                        "aborted" => {
                            return Err(OrchestratorError::Executor("turn aborted".into()));
                        }
                        // session_start, thinking, tool_call, tool_result, …
                        _ => {}
                    }
                }
                Message::Ping(p) => {
                    let _ = ws.send(Message::Pong(p)).await;
                }
                Message::Close(_) => break,
                _ => {}
            }
        }

        Ok(TurnOutcome {
            output: output.trim().to_string(),
            tokens,
            gated,
        })
    }
}

impl TurnExecutor for ZeroClawDriveExecutor {
    async fn run_turn(&self, req: TurnRequest) -> Result<TurnOutcome, OrchestratorError> {
        // An explicit per-node agent (graph `node.attrs["agent"]`) wins, so one
        // request can pin a different model per role; otherwise use the map.
        let alias = req
            .agent
            .as_deref()
            .map(str::trim)
            .filter(|a| !a.is_empty())
            .map(str::to_string)
            .unwrap_or_else(|| {
                // Falling back here means the graph node was never bound to a
                // claw, so the turn runs as the default agent with the DEFAULT
                // agent's workspace and tools — not the mission's. That silently
                // produced whole missions of unusable output, so say so loudly.
                let fallback = self.alias_for(&req.role);
                eprintln!(
                    "topology_exec: node={} role={} has no bound agent — falling back to `{fallback}` \
                     (its workspace/tools, NOT the mission's)",
                    req.node_id, req.role,
                );
                fallback
            });
        let prompt = Self::build_prompt(&req);
        self.drive(&alias, &prompt).await
    }
}

/// Parse `role=alias,role=alias` into a map (blank/malformed entries skipped).
fn parse_agent_map(s: &str) -> HashMap<String, String> {
    s.split(',')
        .filter_map(|pair| {
            let (role, alias) = pair.split_once('=')?;
            let role = role.trim();
            let alias = alias.trim();
            if role.is_empty() || alias.is_empty() {
                None
            } else {
                Some((role.to_string(), alias.to_string()))
            }
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::extract::ws::{Message as AxMsg, WebSocket, WebSocketUpgrade};
    use axum::response::Response;
    use axum::routing::{get, post};
    use axum::{Json, Router};
    use serde_json::{json, Value};

    async fn pair() -> Json<Value> {
        Json(json!({ "token": "test-token" }))
    }

    fn frames() -> Vec<Value> {
        vec![
            json!({"type": "session_start", "session_id": "s1", "resumed": false}),
            json!({"type": "chunk", "content": "hel"}),
            json!({"type": "chunk", "content": "lo"}),
            json!({"type": "done", "input_tokens": 5, "output_tokens": 7}),
        ]
    }

    async fn ok_ws(ws: WebSocketUpgrade) -> Response {
        ws.on_upgrade(|mut socket: WebSocket| async move {
            let _ = socket.recv().await; // the client's message
            for f in frames() {
                let _ = socket.send(AxMsg::Text(f.to_string().into())).await;
            }
        })
    }

    async fn approval_ws(ws: WebSocketUpgrade) -> Response {
        ws.on_upgrade(|mut socket: WebSocket| async move {
            let _ = socket.recv().await;
            for f in [
                json!({"type": "chunk", "content": "working"}),
                json!({"type": "approval_request", "tool": "email.send", "arguments_summary": "to: [email protected]"}),
            ] {
                let _ = socket.send(AxMsg::Text(f.to_string().into())).await;
            }
        })
    }

    async fn serve(router: Router) -> String {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        tokio::spawn(async move {
            axum::serve(listener, router).await.unwrap();
        });
        format!("http://{addr}")
    }

    fn req() -> TurnRequest {
        TurnRequest {
            node_id: "a".into(),
            role: "researcher".into(),
            agent: None,
            attrs: Default::default(),
            task: "say hi".into(),
            context: vec![],
        }
    }

    #[test]
    fn explicit_node_agent_overrides_role_map() {
        let mut map = HashMap::new();
        map.insert("researcher".to_string(), "worker_glm".to_string());
        let exec = ZeroClawDriveExecutor::new("http://x".into(), "c".into(), map, "scout".into());
        // role map → worker_glm
        assert_eq!(exec.alias_for("researcher"), "worker_glm");
        // but an explicit per-node agent should win in run_turn's selection
        let agent = Some("worker_kimi".to_string());
        let chosen = agent
            .as_deref()
            .map(str::trim)
            .filter(|a| !a.is_empty())
            .map(str::to_string)
            .unwrap_or_else(|| exec.alias_for("researcher"));
        assert_eq!(chosen, "worker_kimi");
    }

    #[tokio::test]
    async fn drives_a_turn_and_accumulates_output_and_tokens() {
        let router = Router::new()
            .route("/pair", post(pair))
            .route("/ws/chat", get(ok_ws));
        let base = serve(router).await;
        let exec = ZeroClawDriveExecutor::new(base, "code".into(), HashMap::new(), "scout".into());

        let out = exec.run_turn(req()).await.unwrap();
        assert_eq!(out.output, "hello");
        assert_eq!(out.tokens, 12);
        assert!(out.gated.is_empty());
    }

    #[tokio::test]
    async fn approval_request_is_recorded_as_blocked() {
        let router = Router::new()
            .route("/pair", post(pair))
            .route("/ws/chat", get(approval_ws));
        let base = serve(router).await;
        let exec = ZeroClawDriveExecutor::new(base, "code".into(), HashMap::new(), "scout".into());

        let out = exec.run_turn(req()).await.unwrap();
        assert_eq!(out.gated.len(), 1);
        assert!(!out.gated[0].approved);
        assert!(out.gated[0].summary.contains("email.send"));
    }

    #[test]
    fn agent_map_parses_pairs() {
        let m = parse_agent_map("researcher=scout, writer=quill ,bad=,=x,ok=y");
        assert_eq!(m.get("researcher").unwrap(), "scout");
        assert_eq!(m.get("writer").unwrap(), "quill");
        assert_eq!(m.get("ok").unwrap(), "y");
        assert_eq!(m.len(), 3);
    }
}