IdeaSDK/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** replies.
//!
//! 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::str::FromStr;
use std::sync::Arc;
use std::time::Instant;

use application::OrchestratorService;
use domain::{AgentToolPolicy, DomainEvent, IssueRef, 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::policy::ToolPolicyRegistry;
use super::tickets::TicketToolProvider;
use super::tools::{self, ToolMapError};

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

/// 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>>,
    /// Optional public ticket provider. The MCP surface says `ticket`; the
    /// provider maps those calls to application/domain `Issue` use cases.
    ticket_tools: Option<Arc<dyn TicketToolProvider>>,
    /// Optional per-requester MCP tool policy registry for constrained sessions.
    tool_policies: Option<Arc<ToolPolicyRegistry>>,
}

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,
            ticket_tools: None,
            tool_policies: None,
        }
    }

    /// 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
    }

    /// Attaches the public ticket tool provider.
    #[must_use]
    pub fn with_ticket_tools(mut self, ticket_tools: Arc<dyn TicketToolProvider>) -> Self {
        self.ticket_tools = Some(ticket_tools);
        self
    }

    /// Attaches the requester-scoped MCP tool policy registry.
    #[must_use]
    pub fn with_tool_policies(mut self, tool_policies: Arc<ToolPolicyRegistry>) -> Self {
        self.tool_policies = Some(tool_policies);
        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(),
            ticket_tools: self.ticket_tools.clone(),
            tool_policies: self.tool_policies.clone(),
        }
    }

    /// 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.** A slow request must **not** stall the read
    /// loop: otherwise a single long-running tool call parks the whole connection and
    /// every later call 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 policy = self
            .tool_policies
            .as_ref()
            .and_then(|registry| registry.get(&self.requester));
        let tools: Vec<Value> = tools::catalogue()
            .into_iter()
            .filter(|t| {
                policy
                    .as_ref()
                    .map_or(true, |policy| policy.permits(t.name))
            })
            .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.
    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!({}));

        if let Some(policy) = self
            .tool_policies
            .as_ref()
            .and_then(|registry| registry.get(&self.requester))
        {
            self.enforce_tool_policy(&policy, &name, &arguments)?;
        }

        // Diagnostics begin beacon (best-effort, jamais le corps task/result) : trace
        // l'entrée d'un `tools/call`, son tool, le peer demandeur et la cible/longueurs.
        let started = Instant::now();
        let requester_label = if self.requester.is_empty() {
            "mcp".to_owned()
        } else {
            self.requester.clone()
        };
        let arg_target = arguments
            .get("target")
            .and_then(Value::as_str)
            .unwrap_or("-")
            .to_owned();
        let task_len = arguments
            .get("task")
            .and_then(Value::as_str)
            .map_or(0, str::len);
        application::diag!(
            "[mcp] tools_call begin tool={name} requester={requester_label} target={arg_target} \
             task_len={task_len}",
        );

        if tools::is_ticket_tool(&name) {
            let result = match &self.ticket_tools {
                Some(provider) => {
                    provider
                        .handle_ticket_tool(&self.project, &self.requester, &name, arguments)
                        .await
                }
                None => Err(super::tickets::TicketToolError::new(
                    "notConfigured",
                    "ticket tools are not configured",
                )),
            };
            self.publish_processed(&name, result.is_ok());
            return match result {
                Ok(value) => {
                    let text = serde_json::to_string(&value).unwrap_or_else(|_| "null".to_owned());
                    application::diag!(
                        "[mcp] tools_call end tool={name} requester={requester_label} \
                         target={arg_target} ok=true is_error=false result_len={} elapsed_ms={}",
                        text.len(),
                        started.elapsed().as_millis(),
                    );
                    Ok(tool_result_text(&text, false))
                }
                Err(err) => {
                    let text =
                        serde_json::to_string(&err.to_value()).unwrap_or_else(|_| err.to_string());
                    application::diag!(
                        "[mcp] tools_call end tool={name} requester={requester_label} \
                         target={arg_target} ok=false is_error=true result_len={} elapsed_ms={}",
                        text.len(),
                        started.elapsed().as_millis(),
                    );
                    Ok(tool_result_text(&text, true))
                }
            };
        }

        // The handshake-provided requester is still passed to the mapper for tools that
        // need peer identity.
        let command = match tools::map_tool_call(&name, &arguments, &self.requester) {
            Ok(command) => command,
            Err(e) => {
                application::diag!(
                    "[mcp] tools_call end tool={name} requester={requester_label} mapped=err \
                     err={e} elapsed_ms={}",
                    started.elapsed().as_millis(),
                );
                return Err(map_err_to_jsonrpc(e));
            }
        };

        let result = self.service.dispatch(&self.project, command).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());
        // End beacon : longueurs uniquement (jamais le corps), ok/is_error et elapsed —
        // ferme la trace ouverte par « begin ».
        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());
                application::diag!(
                    "[mcp] tools_call end tool={name} requester={requester_label} \
                     target={arg_target} ok=true is_error=false result_len={} elapsed_ms={}",
                    text.len(),
                    started.elapsed().as_millis(),
                );
                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) => {
                let detail = e.to_string();
                application::diag!(
                    "[mcp] tools_call end tool={name} requester={requester_label} \
                     target={arg_target} ok=false is_error=true result_len={} elapsed_ms={}",
                    detail.len(),
                    started.elapsed().as_millis(),
                );
                Ok(tool_result_text(&detail, true))
            }
        }
    }

    fn enforce_tool_policy(
        &self,
        policy: &AgentToolPolicy,
        name: &str,
        arguments: &Value,
    ) -> Result<(), JsonRpcError> {
        if !policy.permits(name) {
            return Err(JsonRpcError::new(
                error_codes::INVALID_PARAMS,
                format!(
                    "tool `{name}` is not permitted for requester {}",
                    self.requester
                ),
            ));
        }
        if matches!(
            name,
            "idea_ticket_update"
                | "idea_ticket_update_status"
                | "idea_ticket_update_priority"
                | "idea_ticket_update_carnet"
                | "idea_ticket_link"
                | "idea_ticket_unlink"
        ) {
            let raw_ref = arguments
                .get("ref")
                .and_then(Value::as_str)
                .ok_or_else(|| {
                    JsonRpcError::new(
                        error_codes::INVALID_PARAMS,
                        format!("tool `{name}` requires a ticket ref under the active policy"),
                    )
                })?;
            let issue_ref = IssueRef::from_str(raw_ref).map_err(|e| {
                JsonRpcError::new(
                    error_codes::INVALID_PARAMS,
                    format!("invalid ticket ref: {e}"),
                )
            })?;
            if !policy.permits_ticket_mutation(name, issue_ref) {
                return Err(JsonRpcError::new(
                    error_codes::INVALID_PARAMS,
                    format!("tool `{name}` is not permitted for ticket {issue_ref}"),
                ));
            }
        }
        Ok(())
    }

    /// 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())
        }
    }
}