effectfence 0.2.0

//! Wrap mode — `effectfence wrap -- <command that starts any MCP server>`
//!
//! The 10× surface: instead of asking agents to *politely* call
//! fence_prepare/fence_commit, effectfence stands in FRONT of an existing
//! MCP server and fences every tool call automatically:
//!
//! ```text
//!   agent/client ──MCP──> effectfence wrap ──MCP──> real tool server
//! ```
//!
//! - Tool list is mirrored 1:1 from the child (names, schemas, docs).
//! - Every `tools/call` derives an intent from `hash(tool + args)`:
//!   identical duplicate calls — same tool, same arguments — execute the
//!   child ONCE; later duplicates get the recorded result replayed.
//! - Concurrent identical calls: one forwards; the rest WAIT for it and are
//!   handed its recorded result (up to `IN_FLIGHT_WAIT`), rather than being
//!   refused. A duplicate that outlasts the wait is told the original is
//!   still running and was not executed twice.
//! - A forward that outlives the intent lease keeps its claim by beating
//!   while it waits, so a slow child tool cannot have its call taken over
//!   and run a second time.
//! - Non-identical calls pass straight through, unfenced.
//!
//! Honest v1 scope (documented, not hidden): tools only (no resource or
//! prompt passthrough yet); intent derivation treats *byte-identical
//! canonical arguments* as "the same action" — an agent that varies a
//! timestamp argument defeats dedup (that direction fails SAFE: the call
//! runs, nothing corrupts); state is in-memory per wrap process.

use std::sync::Arc;

use anyhow::{Context as _, Result};
use rmcp::{
    ErrorData as McpError, RoleClient, RoleServer, ServerHandler, ServiceExt,
    model::{
        CallToolRequestParams, CallToolResponse, CallToolResult, ContentBlock, Implementation,
        ListToolsResult, PaginatedRequestParams, ServerCapabilities, ServerInfo, Tool,
    },
    service::{RequestContext, RunningService},
    transport::{ConfigureCommandExt, TokioChildProcess, stdio},
};
use sha2::{Digest, Sha256};

use crate::fence::{
    Admission, EffectFence, EffectRequest, FenceError, VectorClock, prepare_effect_fence,
};

/// Derive the intent for a tool call: SHA-256 over the tool name and the
/// JSON arguments serialized with sorted keys (via BTreeMap round-trip),
/// so key order differences don't split identical actions.
fn derive_intent(tool: &str, args: &serde_json::Value) -> String {
    fn canonicalize(v: &serde_json::Value) -> serde_json::Value {
        match v {
            serde_json::Value::Object(m) => {
                let sorted: std::collections::BTreeMap<String, serde_json::Value> = m
                    .iter()
                    .map(|(k, val)| (k.clone(), canonicalize(val)))
                    .collect();
                serde_json::to_value(sorted).expect("BTreeMap serialization cannot fail")
            }
            serde_json::Value::Array(a) => {
                serde_json::Value::Array(a.iter().map(canonicalize).collect())
            }
            other => other.clone(),
        }
    }
    let canon = serde_json::to_string(&canonicalize(args)).expect("Value serialization");
    let mut h = Sha256::new();
    h.update(tool.as_bytes());
    h.update([0u8]);
    h.update(canon.as_bytes());
    format!("wrap:{}:{}", tool, hex::encode(h.finalize()))
}

mod hex {
    pub fn encode(bytes: impl AsRef<[u8]>) -> String {
        use std::fmt::Write;
        let mut out = String::new();
        for b in bytes.as_ref() {
            write!(out, "{b:02x}").expect("write to String");
        }
        out
    }
}

#[derive(Clone)]
struct WrapServer {
    child: Arc<RunningService<RoleClient, ()>>,
    fence: EffectFence,
    tools: Arc<Vec<Tool>>,
    child_name: Arc<str>,
}

/// How long a duplicate waits for the in-flight original to commit before
/// giving up. Longer than most tool calls, shorter than a stuck one.
const IN_FLIGHT_WAIT: std::time::Duration = std::time::Duration::from_secs(30);
/// Poll interval while waiting. Short enough to feel immediate, long enough
/// not to spin a core.
const IN_FLIGHT_POLL: std::time::Duration = std::time::Duration::from_millis(25);

/// Fraction of the lease to wait between heartbeats while a child call is
/// still running. A third leaves room for two missed beats before the lease
/// a genuinely dead holder is supposed to lose actually lapses.
const HEARTBEAT_DIVISOR: u32 = 3;

impl WrapServer {
    /// Wait for the caller holding `intent` to finish, then return its
    /// recorded result. Returns `None` if it has not committed in time --
    /// which never means "run it anyway", only "we still do not know".
    async fn await_holder(
        &self,
        intent: &str,
        args: &serde_json::Value,
        tool: &str,
    ) -> Option<CallToolResult> {
        let deadline = tokio::time::Instant::now() + IN_FLIGHT_WAIT;
        while tokio::time::Instant::now() < deadline {
            tokio::time::sleep(IN_FLIGHT_POLL).await;
            // Re-admitting is how we observe the outcome: once the holder
            // commits, the same intent answers Replay with its cert.
            match prepare_effect_fence(
                &self.fence,
                EffectRequest {
                    intent: intent.to_string(),
                    parent: None,
                    domain: format!("tool:{tool}"),
                    tool: tool.to_string(),
                    args: args.clone(),
                    read_set: vec![],
                    agent: "effectfence-wrap".to_string(),
                    known_clock: VectorClock::new(),
                },
            ) {
                Ok(Admission::Replay(cert)) => {
                    return serde_json::from_value::<CallToolResult>(cert.result.clone()).ok();
                }
                // Still running -- keep waiting.
                Err(FenceError::IntentInFlight { .. }) => continue,
                // The holder released the intent (tool-level error) or its
                // lease lapsed, so re-admitting just handed US the lease.
                // Drop it: we are not going to silently re-run the holder's
                // effect, and keeping it would leave the intent in flight,
                // blocking the caller's own retry until the lease expires.
                Ok(Admission::Fresh(prepared)) => {
                    self.fence.clear_intent(&prepared.intent);
                    return None;
                }
                // The holder aborted or failed; we are not going to silently
                // re-run its effect, so report rather than execute.
                _ => return None,
            }
        }
        None
    }

    async fn forward(&self, req: CallToolRequestParams) -> Result<CallToolResult, McpError> {
        self.child
            .call_tool(req)
            .await
            .map_err(|e| McpError::internal_error(format!("child server error: {e}"), None))
    }
}

impl ServerHandler for WrapServer {
    fn get_info(&self) -> ServerInfo {
        // ServerInfo (= InitializeResult) is #[non_exhaustive]: construct via
        // its constructor, then set public fields.
        let mut info = ServerInfo::new(ServerCapabilities::builder().enable_tools().build());
        info.server_info = Implementation::new(
            format!("effectfence-wrap({})", self.child_name),
            env!("CARGO_PKG_VERSION"),
        );
        info.instructions = Some(format!(
            "Every tool of the wrapped server `{}` is fenced: an identical duplicate call \
             (same tool, same arguments) will NOT re-execute -- it returns the recorded \
             result of the first execution. A concurrent identical call waits for the \
             one in flight and is handed its result, so the action runs exactly once.",
            self.child_name
        ));
        info
    }

    async fn list_tools(
        &self,
        _request: Option<PaginatedRequestParams>,
        _context: RequestContext<RoleServer>,
    ) -> Result<ListToolsResult, McpError> {
        Ok(ListToolsResult {
            tools: (*self.tools).clone(),
            ..Default::default()
        })
    }

    async fn call_tool(
        &self,
        request: CallToolRequestParams,
        _context: RequestContext<RoleServer>,
    ) -> Result<CallToolResponse, McpError> {
        let args_value = serde_json::Value::Object(request.arguments.clone().unwrap_or_default());
        let intent = derive_intent(&request.name, &args_value);
        let tool_name = request.name.to_string();
        // Kept for the in-flight wait path below: the originals are moved into
        // the EffectRequest.
        let intent_for_wait = intent.clone();
        let args_for_wait = args_value.clone();

        let admission = prepare_effect_fence(
            &self.fence,
            EffectRequest {
                intent,
                parent: None,
                domain: format!("tool:{}", request.name),
                tool: request.name.to_string(),
                args: args_value,
                read_set: vec![],
                agent: "effectfence-wrap".to_string(),
                known_clock: VectorClock::new(),
            },
        );

        match admission {
            Ok(Admission::Fresh(prepared)) => {
                // A child tool slower than the lease would have its claim
                // taken over mid-call and the SAME action run a second time --
                // the exact double-execution this proxy exists to stop. Say
                // "still alive" until the child answers; stop the moment it
                // does, so a wrap process that dies still releases on schedule.
                let beater = {
                    let fence = self.fence.clone();
                    let intent = prepared.intent.clone();
                    let every = fence.lease_ttl() / HEARTBEAT_DIVISOR;
                    tokio::spawn(async move {
                        loop {
                            tokio::time::sleep(every).await;
                            if !fence.heartbeat(&intent) {
                                break; // no longer in flight -- nothing to hold
                            }
                        }
                    })
                };
                let forwarded = self.forward(request).await;
                beater.abort();
                let result = match forwarded {
                    Ok(r) => r,
                    Err(e) => {
                        // The transport died before the child answered, so
                        // whether the effect fired is UNKNOWN. Fence the intent
                        // (this is the double-charge case the crate exists for);
                        // freeing it here would be the fail-dangerous direction.
                        crate::fence::abort_effect(&self.fence, prepared, e.to_string());
                        return Err(e);
                    }
                };
                let stored = serde_json::to_value(&result)
                    .map_err(|e| McpError::internal_error(e.to_string(), None))?;
                // A tool-level error is the child ANSWERING: it ran and reported
                // failure. That is neither a terminal success nor an unknown
                // outcome, so release the lease outright and let an identical
                // retry execute. Fencing here (what abort_effect does) refused
                // the retry for the whole result_ttl -- 24h by default -- and
                // wrap exposes no clear_intent, so one transient error from any
                // wrapped tool bricked that exact call for the rest of the day.
                if result.is_error.unwrap_or(false) {
                    self.fence.clear_intent(&prepared.intent);
                    return Ok(result.into());
                }
                match crate::fence::commit_effect_cert(&self.fence, prepared, stored) {
                    Ok(_cert) => Ok(result.into()),
                    Err(e) => {
                        // Bookkeeping failed; the call already ran. Serve the
                        // real result -- never turn success into an error.
                        tracing_warn(&format!("wrap: commit failed: {e}"));
                        Ok(result.into())
                    }
                }
            }
            Ok(Admission::Replay(cert)) => {
                let replay: CallToolResult =
                    serde_json::from_value(cert.result.clone()).map_err(|e| {
                        McpError::internal_error(
                            format!("stored result no longer deserializes: {e}"),
                            None,
                        )
                    })?;
                Ok(replay.into())
            }
            // An identical call that arrives while the first is still running
            // is the ordinary twin-caller race, not an error: the caller wants
            // the ONE result, and it is moments away. Returning a failure here
            // made the common case look broken -- an agent sees a tool error
            // and reports failure for work that is about to succeed. Wait for
            // the holder to commit, then hand back its recorded result.
            Err(FenceError::IntentInFlight { .. }) => {
                match self
                    .await_holder(&intent_for_wait, &args_for_wait, &tool_name)
                    .await
                {
                    Some(replay) => Ok(replay.into()),
                    None => Ok(CallToolResponse::Complete(CallToolResult::error(vec![
                        ContentBlock::text(format!(
                            "effectfence: an identical call to `{tool_name}` is still running and \
                             did not finish within {}s. It was NOT executed a second time. Retry \
                             to receive its recorded result.",
                            IN_FLIGHT_WAIT.as_secs()
                        )),
                    ]))),
                }
            }
            // A fenced intent is NOT retryable by re-reading and trying again:
            // an earlier attempt died with its outcome unknown, and the fence
            // holds until a human reconciles. Saying "try again" here sent
            // agents into a retry loop that could never succeed.
            Err(err @ FenceError::IntentFailed { .. }) => {
                Ok(CallToolResponse::Complete(CallToolResult::error(vec![
                    ContentBlock::text(format!(
                        "effectfence: call fenced ({err}). An earlier identical call to \
                         `{tool_name}` failed with its outcome UNKNOWN, so this one was NOT \
                         executed and retrying it cannot clear the fence. Check the downstream \
                         system for whether the first attempt took effect; if it did not, this \
                         action needs a new intent (vary the arguments) or an operator reset."
                    )),
                ])))
            }
            Err(err) => Ok(CallToolResponse::Complete(CallToolResult::error(vec![
                ContentBlock::text(format!(
                    "effectfence: call refused ({err}). This action was NOT executed. If it was \
                 intentionally a NEW action, vary the arguments; if a dependency moved, re-read \
                 and try again."
                )),
            ]))),
        }
    }
}

fn tracing_warn(msg: &str) {
    eprintln!("[effectfence-wrap] {msg}");
}

/// Lease and result windows for wrap, from the environment.
///
/// `FenceConfig::lease_ttl`'s contract is that it must exceed the longest
/// effect you run, and wrap previously pinned it to the 60 s default with no
/// way to change it — an operator whose tool takes longer had no remedy at
/// all. (A forward now heartbeats, so the lease no longer has to bound the
/// slowest call; this is for operators who want the window explicit.)
///
/// `EFFECTFENCE_LEASE_SECS` — in-flight lease, default 60.
/// `EFFECTFENCE_RESULT_TTL_SECS` — how long a finished outcome is replayable,
/// default 86400. Both ignore values that do not parse as a positive integer,
/// so a typo keeps the safe default rather than silently disabling fencing.
fn fence_config_from_env() -> crate::fence::FenceConfig {
    fn secs(key: &str, default: u64) -> std::time::Duration {
        let parsed = std::env::var(key)
            .ok()
            .and_then(|v| v.trim().parse::<u64>().ok())
            .filter(|n| *n > 0);
        if std::env::var(key).is_ok() && parsed.is_none() {
            eprintln!(
                "[effectfence-wrap] ignoring {key}: not a positive integer; using {default}s"
            );
        }
        std::time::Duration::from_secs(parsed.unwrap_or(default))
    }
    crate::fence::FenceConfig {
        lease_ttl: secs("EFFECTFENCE_LEASE_SECS", 60),
        result_ttl: secs("EFFECTFENCE_RESULT_TTL_SECS", 24 * 60 * 60),
    }
}

/// Entry point: spawn the child MCP server, mirror its tools, serve fenced.
pub async fn run_wrap(child_cmd: Vec<String>) -> Result<()> {
    let (program, rest) = child_cmd
        .split_first()
        .context("wrap: no child command given. Usage: effectfence wrap -- <command> [args...]")?;

    let transport = TokioChildProcess::new(tokio::process::Command::new(program).configure(|c| {
        c.args(rest);
    }))
    .context("wrap: failed to spawn child MCP server")?;

    let child = ().serve(transport).await.context("wrap: MCP handshake with child failed")?;

    let tools = child
        .list_all_tools()
        .await
        .context("wrap: could not list child tools")?;

    eprintln!(
        "[effectfence-wrap] fencing {} tool(s) from `{}`",
        tools.len(),
        program
    );

    let server = WrapServer {
        child: Arc::new(child),
        fence: EffectFence::with_config(fence_config_from_env()),
        tools: Arc::new(tools),
        child_name: Arc::from(program.as_str()),
    };

    let service = server.serve(stdio()).await?;
    service.waiting().await?;
    Ok(())
}