IdeA/crates/infrastructure/src/orchestrator/mcp/server.rs

//! [`McpServer`] — the IdeA MCP **driving adapter** (cadrage Décision 4).
//!
//! This is the **pair of [`FsOrchestratorWatcher`](super::super::FsOrchestratorWatcher)**:
//! another entry door onto the *same* [`OrchestratorService::dispatch`]. Where the
//! watcher reads a JSON file, this server reads a JSON-RPC `tools/call`; both build
//! an [`OrchestratorCommand`] and return its [`OrchestratorOutcome`] unchanged. It
//! invents **no semantics** and **never re-routes** the reply — for `idea_ask_agent`
//! it returns `outcome.reply` inline, exactly what `dispatch` produced.
//!
//! It implements the strict MCP minimum over JSON-RPC 2.0:
//! - `initialize` → advertise protocol version + tool capability,
//! - `tools/list` → the [`tools::catalogue`],
//! - `tools/call` → map → `dispatch` → MCP tool result,
//! - notifications (e.g. `notifications/initialized`) → accepted, no reply.
//!
//! It receives `Arc<OrchestratorService>` and the target [`Project`] by injection
//! at the composition root (M3), exactly like the watcher — no application logic is
//! duplicated here.

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

use application::OrchestratorService;
use domain::{DomainEvent, OrchestrationSource, Project};
use serde_json::{json, Value};
use tokio::sync::mpsc;

use super::jsonrpc::{
    error_codes, JsonRpcError, JsonRpcRequest, JsonRpcResponse, Transport, TransportError,
    JSONRPC_VERSION,
};
use super::tools::{self, ToolMapError};

/// The MCP protocol version this server speaks (advertised on `initialize`).
const MCP_PROTOCOL_VERSION: &str = "2024-11-05";

/// **Server-side safety net** bounding the synchronous `idea_ask_agent` rendezvous.
///
/// `idea_ask_agent` is the only tool whose `tools/call` *blocks* awaiting another
/// agent's `idea_reply` (a synchronous rendezvous, resolved deep in the application
/// layer). The application layer already bounds the rendezvous itself, but this
/// adapter adds its own **finite** outer bound so a wedged target can never park a
/// `tools/call` task forever: on expiry the call returns a clean JSON-RPC error
/// instead of hanging. Generous on purpose (delegated turns can be very long), but
/// never infinite. Every other tool (`idea_reply`, `idea_list_agents`, …) is left
/// untouched — they don't rendezvous.
const ASK_RENDEZVOUS_TIMEOUT: Duration = Duration::from_secs(24 * 60 * 60);

/// The IdeA MCP server: an entry adapter over [`OrchestratorService::dispatch`].
///
/// Cheap to clone the dependencies it holds; one instance serves one project's
/// agents over one transport (M3 owns one server per open project, beside the
/// watcher).
pub struct McpServer {
    service: Arc<OrchestratorService>,
    project: Project,
    /// Optional event sink (a domain [`EventBus`](domain::ports::EventBus) publish
    /// closure), the twin of the file watcher's. When set, a processed `tools/call`
    /// publishes [`DomainEvent::OrchestratorRequestProcessed`] tagged
    /// [`OrchestrationSource::Mcp`] so the presentation layer can surface that this
    /// delegation arrived through the MCP door. `None` ⇒ no publication (the server
    /// still works), so existing call sites stay valid.
    events: Option<Arc<dyn Fn(DomainEvent) + Send + Sync>>,
    /// Identity of the connected peer — the real agent id carried in the loopback
    /// handshake (cadrage v5 §1.4). When non-empty it becomes
    /// [`DomainEvent::OrchestratorRequestProcessed`]'s `requester_id` instead of the
    /// frozen `"mcp"` placeholder; empty ⇒ the legacy `"mcp"` label (back-compat for
    /// M2 callers and connections that arrive without a requester).
    requester: String,
    /// Sink optionnel de **readiness de démarrage** : appelé avec le `requester` brut
    /// (handshake) la première fois que le peer émet `initialize` (son CLI est up et
    /// parle MCP). La composition root y branche
    /// `OrchestratorService::release_agent_cold_start` pour livrer un 1er tour différé
    /// (fix race cold-launch, signal MCP). L'infra ne connaît pas `AgentId` : la
    /// composition root parse le `requester`. `None` ⇒ no-op.
    ready_sink: Option<Arc<dyn Fn(&str) + Send + Sync>>,
    /// Finite outer bound applied **only** to the `idea_ask_agent` rendezvous (see
    /// [`ASK_RENDEZVOUS_TIMEOUT`]). Carried per instance so tests can shrink it to a
    /// few milliseconds without polluting the public API; production code always uses
    /// the default.
    ask_rendezvous_timeout: Duration,
}

impl McpServer {
    /// Builds the server from the injected application service and target project.
    /// No event sink — use [`with_events`](Self::with_events) to surface processed
    /// requests on the bus.
    #[must_use]
    pub fn new(service: Arc<OrchestratorService>, project: Project) -> Self {
        Self {
            service,
            project,
            events: None,
            requester: String::new(),
            ready_sink: None,
            ask_rendezvous_timeout: ASK_RENDEZVOUS_TIMEOUT,
        }
    }

    /// Attaches an event sink so each processed `tools/call` is republished on the
    /// bus tagged [`OrchestrationSource::Mcp`] — the MCP twin of the file watcher's
    /// `events` closure. Additive: callers that do not need observability keep using
    /// [`new`](Self::new).
    #[must_use]
    pub fn with_events(mut self, events: Arc<dyn Fn(DomainEvent) + Send + Sync>) -> Self {
        self.events = Some(events);
        self
    }

    /// Attache un sink de **readiness de démarrage** : appelé avec l'id (handshake
    /// `requester`) de l'agent connecté la première fois qu'il émet `initialize` (son CLI
    /// est up et parle MCP). La composition root y branche
    /// `OrchestratorService::release_agent_cold_start` pour livrer un éventuel 1er tour
    /// différé (fix race cold-launch, signal MCP). Additif : sans sink, no-op.
    #[must_use]
    pub fn with_ready_sink(mut self, ready_sink: Arc<dyn Fn(&str) + Send + Sync>) -> Self {
        self.ready_sink = Some(ready_sink);
        self
    }

    /// Returns a per-connection clone of this server tagged with the connected
    /// peer's `requester` id (the loopback handshake's `requester`, cadrage v5 §1.4).
    ///
    /// The base server (one per project, M3) is shared by every connection; each
    /// accepted peer derives its own contextualized server so concurrent peers never
    /// share a requester. An empty `requester` keeps the legacy `"mcp"` label.
    ///
    /// Cheap: only `Arc`s, a `Project` clone, and the requester `String` are copied.
    #[must_use]
    pub fn for_requester(&self, requester: impl Into<String>) -> Self {
        Self {
            service: Arc::clone(&self.service),
            project: self.project.clone(),
            events: self.events.clone(),
            requester: requester.into(),
            ready_sink: self.ready_sink.clone(),
            ask_rendezvous_timeout: self.ask_rendezvous_timeout,
        }
    }

    /// **Test seam** (doc-hidden): shrinks the `idea_ask_agent` rendezvous bound
    /// (see [`ASK_RENDEZVOUS_TIMEOUT`]) so a wedged-target test need not wait the
    /// full production timeout. Doc-hidden so it does not widen the documented public
    /// surface; production code never calls it.
    #[doc(hidden)]
    #[must_use]
    pub fn with_ask_rendezvous_timeout(mut self, timeout: Duration) -> Self {
        self.ask_rendezvous_timeout = timeout;
        self
    }

    /// Serves JSON-RPC messages from `transport`, tagging every processed
    /// `tools/call` with `requester` as the delegating agent (cadrage v5 §1.4).
    ///
    /// A thin wrapper over [`serve`](Self::serve) on a per-connection
    /// [`for_requester`](Self::for_requester) clone: the caller (the per-peer accept
    /// loop) need not mutate the shared project server. An empty `requester` yields
    /// the legacy `"mcp"` label.
    pub async fn serve_as<T: Transport>(&self, requester: impl Into<String>, transport: &mut T) {
        self.for_requester(requester).serve(transport).await;
    }

    /// Serves JSON-RPC messages from `transport` until it closes (clean EOF).
    ///
    /// Every inbound line is handled in isolation: a malformed line or an unknown
    /// method yields a JSON-RPC error response, **never a panic** and never a
    /// dropped connection. Notifications (no `id`) are processed without a reply.
    ///
    /// **Full-duplex / non-blocking (anti-wedge).** A request whose handling blocks —
    /// notably `idea_ask_agent`, which awaits another agent's `idea_reply` in a
    /// synchronous rendezvous — must **not** stall the read loop: otherwise a single
    /// in-flight ask parks the whole connection and every later call (even a
    /// rendezvous-free `idea_list_agents`) is never even read. So each inbound message
    /// is handled on its own spawned task that owns a cheap per-call clone of the
    /// server; finished responses are funnelled back through an `mpsc` channel and
    /// written by the same loop. Responses carry their JSON-RPC `id`, so out-of-order
    /// completion is fine (the full-duplex bridge correlates by id). Notifications
    /// (`handle_raw` ⇒ `None`) emit nothing.
    pub async fn serve<T: Transport>(&self, transport: &mut T) {
        let (tx, mut rx) = mpsc::unbounded_channel::<Option<Vec<u8>>>();
        // Number of spawned handler tasks not yet observed on `rx`. While `> 0`, an
        // EOF on the read side must NOT exit immediately: we keep draining `rx` so
        // already-computed responses still reach the transport (graceful shutdown).
        // A response-less notification decrements without ever sending — see below.
        let mut in_flight: usize = 0;
        // Once the peer closed its read side, stop accepting new inbound; only drain.
        let mut reading = true;
        loop {
            tokio::select! {
                // Drain ready responses first so a flood of inbound never starves
                // writes; correctness does not depend on the bias.
                biased;

                outbound = rx.recv() => {
                    // `tx` is held by the loop, so `recv` only yields `None` once it
                    // is dropped — which never happens before the loop returns.
                    let Some(slot) = outbound else { break };
                    in_flight -= 1;
                    if let Some(bytes) = slot {
                        if transport.send(&bytes).await.is_err() {
                            break;
                        }
                    }
                    // Peer gone and every spawned task accounted for ⇒ done.
                    if !reading && in_flight == 0 {
                        break;
                    }
                }

                incoming = transport.recv(), if reading => {
                    let raw = match incoming {
                        Ok(bytes) => bytes,
                        // Peer closed / I/O error: stop reading but keep draining the
                        // responses of tasks still in flight before returning.
                        Err(TransportError::Closed) | Err(TransportError::Io(_)) => {
                            reading = false;
                            if in_flight == 0 {
                                break;
                            }
                            continue;
                        }
                    };
                    // Own a cheap clone (Arc/String/Project) so the task is `'static`
                    // and never borrows the transport. Every spawned task sends
                    // exactly one slot back (`Some(bytes)` for a reply, `None` for a
                    // notification) so `in_flight` is always reconciled.
                    in_flight += 1;
                    let server = self.for_requester(self.requester.clone());
                    let tx = tx.clone();
                    tokio::spawn(async move {
                        let slot = match server.handle_raw(&raw).await {
                            Some(response) => serde_json::to_vec(&response).ok(),
                            None => None,
                        };
                        // Receiver dropped ⇒ the loop has stopped; discard.
                        let _ = tx.send(slot);
                    });
                }
            }
        }
    }

    /// Parses one raw message and produces the response to send back, or `None`
    /// for a notification (no `id`) that owes no reply.
    ///
    /// Kept standalone (no transport) so the whole request→response behaviour is
    /// unit-testable over plain bytes, with no I/O.
    pub async fn handle_raw(&self, raw: &[u8]) -> Option<JsonRpcResponse> {
        let request: JsonRpcRequest = match serde_json::from_slice(raw) {
            Ok(r) => r,
            Err(e) => {
                // Could not correlate (no id parsed) ⇒ null id per JSON-RPC.
                return Some(JsonRpcResponse::error(
                    Value::Null,
                    JsonRpcError::new(error_codes::PARSE_ERROR, format!("invalid json: {e}")),
                ));
            }
        };

        if request.jsonrpc != JSONRPC_VERSION {
            let id = request.id.unwrap_or(Value::Null);
            return Some(JsonRpcResponse::error(
                id,
                JsonRpcError::new(
                    error_codes::INVALID_REQUEST,
                    format!("unsupported jsonrpc version: {}", request.jsonrpc),
                ),
            ));
        }

        // Notification (no id): never reply (the `?` short-circuits to `None`).
        let id = request.id.clone()?;

        let result = self.dispatch_method(&request.method, request.params).await;
        Some(match result {
            Ok(value) => JsonRpcResponse::success(id, value),
            Err(error) => JsonRpcResponse::error(id, error),
        })
    }

    /// Routes a method name to its handler.
    async fn dispatch_method(
        &self,
        method: &str,
        params: Option<Value>,
    ) -> Result<Value, JsonRpcError> {
        match method {
            "initialize" => {
                self.notify_ready();
                Ok(self.initialize_result())
            }
            "tools/list" => Ok(self.tools_list_result()),
            "tools/call" => self.tools_call(params.unwrap_or(Value::Null)).await,
            other => Err(JsonRpcError::new(
                error_codes::METHOD_NOT_FOUND,
                format!("method not found: {other}"),
            )),
        }
    }

    /// Notifie le sink de readiness de démarrage avec l'identité du peer connecté
    /// (handshake `requester`). No-op si pas de sink ou requester vide (peer legacy/anonyme).
    fn notify_ready(&self) {
        if let Some(sink) = &self.ready_sink {
            if !self.requester.is_empty() {
                sink(&self.requester);
            }
        }
    }

    /// The `initialize` result: protocol version, server identity, and the fact
    /// that we expose tools.
    fn initialize_result(&self) -> Value {
        json!({
            "protocolVersion": MCP_PROTOCOL_VERSION,
            "capabilities": { "tools": {} },
            "serverInfo": { "name": "idea-orchestrator", "version": env!("CARGO_PKG_VERSION") }
        })
    }

    /// The `tools/list` result: the catalogue as MCP tool descriptors.
    fn tools_list_result(&self) -> Value {
        let tools: Vec<Value> = tools::catalogue()
            .into_iter()
            .map(|t| {
                json!({
                    "name": t.name,
                    "description": t.description,
                    "inputSchema": t.input_schema,
                })
            })
            .collect();
        json!({ "tools": tools })
    }

    /// `tools/call`: map the tool to an [`OrchestratorCommand`], `dispatch` it, and
    /// fold the [`OrchestratorOutcome`](application::OrchestratorOutcome) into an MCP
    /// tool result. The `idea_ask_agent` reply is returned **inline** — never
    /// re-routed.
    async fn tools_call(&self, params: Value) -> Result<Value, JsonRpcError> {
        let name = params
            .get("name")
            .and_then(Value::as_str)
            .ok_or_else(|| JsonRpcError::new(error_codes::INVALID_PARAMS, "missing tool `name`"))?
            .to_owned();
        let arguments = params.get("arguments").cloned().unwrap_or(json!({}));

        // `idea_reply` needs the connected peer's identity as `from`; every other tool
        // ignores `requester`. The handshake-provided requester is the source of truth.
        let command =
            tools::map_tool_call(&name, &arguments, &self.requester).map_err(map_err_to_jsonrpc)?;

        // `idea_ask_agent` is the only tool that blocks on a synchronous rendezvous
        // (the target's `idea_reply`). Bound *only* that path with a finite outer
        // timeout (server-side safety net, see [`ASK_RENDEZVOUS_TIMEOUT`]): on expiry
        // return a clean JSON-RPC error instead of parking the task forever. Every
        // other tool dispatches unbounded — they never rendezvous.
        let dispatch = self.service.dispatch(&self.project, command);
        let result = if name == "idea_ask_agent" {
            match tokio::time::timeout(self.ask_rendezvous_timeout, dispatch).await {
                Ok(result) => result,
                Err(_elapsed) => {
                    return Err(JsonRpcError::new(
                        error_codes::INTERNAL_ERROR,
                        "idea_ask_agent : aucune réponse de la cible avant le timeout du rendezvous",
                    ));
                }
            }
        } else {
            dispatch.await
        };
        // Surface the processed delegation on the bus, tagged as the MCP door — the
        // twin of the file watcher's publish. `ok` mirrors the dispatch outcome; the
        // action is the tool name. No-op when no sink is attached.
        self.publish_processed(&name, result.is_ok());
        match result {
            Ok(outcome) => {
                // The text the agent sees: the reply for an `ask`, else the detail.
                let text = outcome
                    .reply
                    .clone()
                    .unwrap_or_else(|| outcome.detail.clone());
                Ok(tool_result_text(&text, false))
            }
            // A failed IdeA command is reported as a tool execution error
            // (`isError: true`) rather than a protocol error: the agent can read
            // it and decide, and the connection stays healthy.
            Err(e) => Ok(tool_result_text(&e.to_string(), true)),
        }
    }

    /// Publishes [`DomainEvent::OrchestratorRequestProcessed`] for a handled
    /// `tools/call`, tagged [`OrchestrationSource::Mcp`]. The requester id is the
    /// connected peer's real agent id (carried in the loopback handshake, cadrage v5
    /// §1.4), falling back to the legacy `"mcp"` label when no identity was supplied.
    /// No-op without an event sink.
    fn publish_processed(&self, action: &str, ok: bool) {
        if let Some(publish) = &self.events {
            let requester_id = if self.requester.is_empty() {
                "mcp".to_owned()
            } else {
                self.requester.clone()
            };
            publish(DomainEvent::OrchestratorRequestProcessed {
                requester_id,
                action: action.to_owned(),
                ok,
                source: OrchestrationSource::Mcp,
            });
        }
    }
}

/// Builds an MCP `tools/call` result with a single text content block.
fn tool_result_text(text: &str, is_error: bool) -> Value {
    json!({
        "content": [ { "type": "text", "text": text } ],
        "isError": is_error
    })
}

/// Maps a [`ToolMapError`] to the right JSON-RPC error code.
fn map_err_to_jsonrpc(err: ToolMapError) -> JsonRpcError {
    match err {
        ToolMapError::UnknownTool(_) => {
            JsonRpcError::new(error_codes::METHOD_NOT_FOUND, err.to_string())
        }
        ToolMapError::BadArguments(_) | ToolMapError::Invalid(_) => {
            JsonRpcError::new(error_codes::INVALID_PARAMS, err.to_string())
        }
    }
}