mnemo-mcp-server 0.5.23

use std::io::IsTerminal;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};

use clap::{Parser, Subcommand};
use rmcp::{ServiceExt, transport::stdio};
use tokio::sync::Notify;

mod attest;
mod commands;
mod manifest;
mod safe_spawn;

use attest::CatalogAttestor;

use mnemo_core::anomaly::outlier::train_baseline;
use mnemo_core::embedding::openai::OpenAiEmbedding;
use mnemo_core::embedding::{EmbeddingProvider, NoopEmbedding};
use mnemo_core::encryption::ContentEncryption;
use mnemo_core::index::VectorIndex;
use mnemo_core::index::usearch::UsearchIndex;
use mnemo_core::query::MnemoEngine;
use mnemo_core::search::FullTextIndex;
use mnemo_core::search::tantivy_index::TantivyFullTextIndex;
use mnemo_core::storage::StorageBackend;
use mnemo_core::storage::duckdb::DuckDbStorage;
use mnemo_mcp::server::MnemoServer;

#[derive(Parser)]
#[command(name = "mnemo", about = "MCP-native memory database for AI agents")]
struct Cli {
    /// Path to the database file
    #[arg(long, default_value = "mnemo.db", env = "MNEMO_DB_PATH")]
    db_path: PathBuf,

    /// OpenAI API key for embeddings
    #[arg(long, env = "OPENAI_API_KEY")]
    openai_api_key: Option<String>,

    /// Embedding model name
    #[arg(
        long,
        default_value = "text-embedding-3-small",
        env = "MNEMO_EMBEDDING_MODEL"
    )]
    embedding_model: String,

    /// Embedding dimensions
    #[arg(long, default_value = "1536", env = "MNEMO_DIMENSIONS")]
    dimensions: usize,

    /// Default agent ID
    #[arg(long, default_value = "default", env = "MNEMO_AGENT_ID")]
    agent_id: String,

    /// Default organization ID
    #[arg(long, env = "MNEMO_ORG_ID")]
    org_id: Option<String>,

    /// Path to ONNX embedding model (uses local inference instead of OpenAI)
    #[arg(long, env = "MNEMO_ONNX_MODEL_PATH")]
    onnx_model_path: Option<String>,

    /// PostgreSQL connection URL (enables PostgreSQL backend instead of DuckDB)
    #[arg(long, env = "MNEMO_POSTGRES_URL")]
    postgres_url: Option<String>,

    /// REST API port (starts an HTTP server alongside MCP stdio)
    #[arg(long, env = "MNEMO_REST_PORT")]
    rest_port: Option<u16>,

    /// Idle timeout in seconds — auto-shutdown after no requests (0 = disabled)
    #[arg(long, default_value = "0", env = "MNEMO_IDLE_TIMEOUT")]
    idle_timeout_seconds: u64,

    /// AES-256-GCM encryption key (64-char hex string) for at-rest content encryption
    #[arg(long, env = "MNEMO_ENCRYPTION_KEY")]
    encryption_key: Option<String>,

    /// Interval in seconds between TTL sweeps (0 = disabled). A sweep hard-deletes
    /// every memory whose `expires_at` is in the past and emits MemoryExpired
    /// audit events.
    #[arg(long, default_value = "0", env = "MNEMO_TTL_SWEEP_INTERVAL")]
    ttl_sweep_interval_seconds: u64,

    #[command(subcommand)]
    command: Option<Command>,
}

#[derive(Subcommand)]
enum Command {
    /// Manage the per-agent embedding-space baseline used by the z-score
    /// outlier detector (v0.3.3, Task A).
    Baseline(BaselineArgs),
    /// Start the MCP STDIO server in hardened mode using a TOML manifest
    /// (v0.4.0-rc3 Task B2).
    ///
    /// Defends against the OX-MCP "exfiltrate-then-act" disclosure
    /// (2026-04-24): refuses inherited secrets, JSON-injection argv, and
    /// untrusted parent processes before any engine state is touched. All
    /// privileged knobs come from the TOML manifest — env vars and
    /// command-line flags cannot grant capabilities.
    McpServer(McpServerArgs),
    /// Replay a JSONL dataset of `{query, expected}` rows against an
    /// in-memory engine and emit a per-row latency / top-k report
    /// (v0.4.0-rc3 Task B6).
    ///
    /// The bundled dataset at `crates/mnemo-core/benches/data/longmemeval_m.jsonl`
    /// is the default when `--dataset` is omitted. Used to compare
    /// configuration sweeps (provenance on/off, recency half-life,
    /// hybrid weights) against a fixed prompt set.
    Eval(EvalArgs),
    /// Run a measurement-only benchmark (v0.4.9).
    ///
    /// Currently exposes one subcommand: `embeddings`, which runs the
    /// embedding-backend selection bench plus an SLA-aware recommender
    /// anchored on arXiv:2605.23618 (GE2 vs local encoders — quality
    /// and latency). The bench is measurement and recommendation
    /// only: no retrieval defaults change, no RRF-weights change,
    /// no managed-cloud default.
    #[command(subcommand)]
    Bench(BenchCommand),
    /// Compliance primitives (mnemo-compliance).
    ///
    /// Currently exposes `retention`, which prints a processing-log
    /// retention-conformance profile (DPDP Rules 2025 / EU AI Act Art.19 /
    /// HIPAA §164.312(b)) and checks that the active storage backend can
    /// honour its append-only retention floor — failing loud with a typed
    /// error if it cannot.
    #[command(subcommand)]
    Compliance(ComplianceCommand),
}

#[derive(Subcommand)]
enum ComplianceCommand {
    /// Print a retention-conformance profile and gate it against the active
    /// storage backend's append-only guarantee.
    Retention(RetentionArgs),
}

#[derive(clap::Args)]
struct RetentionArgs {
    /// Which obligation profile to load.
    #[arg(long, value_parser = ["dpdp", "eu-ai-act-art19", "hipaa"], default_value = "dpdp")]
    profile: String,
    /// Override the retention floor in days (defaults to the obligation's
    /// legal minimum). An operator may set a *longer* floor for its policy.
    #[arg(long)]
    floor_days: Option<u32>,
}

#[derive(Subcommand)]
enum BenchCommand {
    /// Embedding-backend selection bench. For every configured backend
    /// (Noop + bench-local hashing baseline always; OpenAI if
    /// `OPENAI_API_KEY` is set; ONNX if `MNEMO_ONNX_MODEL_PATH` is set
    /// and `mnemo-core` was built with the `onnx` feature), measure
    /// nDCG@10, recall@10, p50/p95 single-vector embed latency, and
    /// throughput at batch sizes 1/8/32, then recommend the
    /// highest-nDCG backend whose p95 ≤ the SLO.
    Embeddings(BenchEmbeddingsArgs),
}

#[derive(clap::Args)]
struct BenchEmbeddingsArgs {
    /// p95 latency SLO in milliseconds. The recommender picks the
    /// highest-nDCG backend whose measured p95 ≤ this value.
    #[arg(long, default_value_t = 50.0)]
    slo_ms: f64,
    /// Embedding dimensions to construct each backend with. The
    /// fixture is small, so the default 384 is fine for both
    /// MiniLM-class local models and `text-embedding-3-small` at
    /// reduced dim.
    #[arg(long, default_value_t = 384)]
    dimensions: usize,
    /// Number of single-vector embed calls timed per backend for
    /// p50/p95. Default 32; raise for tighter percentile estimates.
    #[arg(long, default_value_t = 32)]
    latency_samples: usize,
}

#[derive(clap::Args)]
struct BaselineArgs {
    /// Train and persist a baseline from every non-deleted memory for this agent.
    #[arg(long)]
    train: bool,

    /// Agent ID to train or inspect the baseline for. Falls back to
    /// `--agent-id` / `MNEMO_AGENT_ID` when omitted.
    #[arg(long)]
    agent_id: Option<String>,

    /// Maximum records to load when training (defaults to `MAX_BATCH_QUERY_LIMIT`).
    #[arg(long, default_value = "10000")]
    limit: usize,
}

#[derive(clap::Args)]
struct McpServerArgs {
    /// Path to the TOML manifest carrying every privileged knob.
    #[arg(long)]
    manifest: PathBuf,
    /// Print a ready-to-paste `[tool_catalog_pin]` block for the tools this
    /// exact binary advertises (after applying any manifest `[role_filter]`),
    /// then exit WITHOUT serving. Pipe into a file and set
    /// `tool_catalog_pin_path` in the manifest to enable serve-time
    /// tool-catalog attestation (arXiv 2604.20994).
    #[arg(long)]
    print_catalog_pin: bool,
}

#[derive(clap::Args)]
struct EvalArgs {
    /// Path to a JSONL dataset of `{id, content, query, expected}` rows.
    /// Defaults to the bundled LongMemEval_M sample.
    #[arg(long)]
    dataset: Option<PathBuf>,
    /// Where to write per-row results as JSONL. Defaults to stdout.
    #[arg(long)]
    output: Option<PathBuf>,
    /// Recall limit per query.
    #[arg(long, default_value = "5")]
    limit: usize,
    /// Request a provenance receipt on every recall.
    #[arg(long)]
    with_provenance: bool,
    /// HMAC key (hex, >=32 bytes) for the provenance signer. Required
    /// when `--with-provenance` is set.
    #[arg(long)]
    provenance_key_hex: Option<String>,
    /// Recall strategy ("semantic", "hybrid", "lexical").
    #[arg(long, default_value = "hybrid")]
    strategy: String,
}

/// DocTrace experience-memory tier gate. Off unless
/// `MNEMO_EXPERIENCE_MEMORY` is `1` / `true` (case-insensitive), so the
/// default server behaviour is unchanged.
fn experience_memory_enabled() -> bool {
    std::env::var("MNEMO_EXPERIENCE_MEMORY")
        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
        .unwrap_or(false)
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::from_default_env().add_directive("mnemo=info".parse()?),
        )
        .with_writer(std::io::stderr)
        .init();

    let cli = Cli::parse();

    // Dispatch one-shot subcommands before any server setup.
    match &cli.command {
        Some(Command::Baseline(args)) => return run_baseline(&cli, args).await,
        Some(Command::McpServer(args)) => return run_mcp_server(&cli, args).await,
        Some(Command::Eval(args)) => return run_eval(&cli, args).await,
        Some(Command::Bench(sub)) => return run_bench(sub).await,
        Some(Command::Compliance(sub)) => return run_compliance(sub).await,
        None => {}
    }

    // Initialize embedding provider (ONNX > OpenAI > Noop)
    let embedding: Arc<dyn EmbeddingProvider> = if let Some(ref onnx_path) = cli.onnx_model_path {
        tracing::info!("Using ONNX local embeddings from {}", onnx_path);
        Arc::new(mnemo_core::embedding::onnx::OnnxEmbedding::new(
            onnx_path,
            cli.dimensions,
        )?)
    } else if let Some(api_key) = cli.openai_api_key {
        tracing::info!("Using OpenAI embeddings ({})", cli.embedding_model);
        Arc::new(OpenAiEmbedding::new(
            api_key,
            cli.embedding_model,
            cli.dimensions,
        ))
    } else {
        tracing::warn!(
            "No OPENAI_API_KEY set, using noop embeddings (semantic search will not work)"
        );
        Arc::new(NoopEmbedding::new(cli.dimensions))
    };

    // Build engine based on backend selection
    // Keep a reference to the DuckDB vector index for shutdown save
    #[allow(unused_assignments)]
    let mut duckdb_index: Option<Arc<UsearchIndex>> = None;
    let engine = if let Some(_pg_url) = &cli.postgres_url {
        #[cfg(feature = "postgres")]
        {
            let pg_storage =
                Arc::new(mnemo_postgres::PgStorage::connect(_pg_url, cli.dimensions).await?);
            // Share the storage pool so pgvector ANN search (semantic / auto /
            // graph / domain_scoped recall) runs against the HNSW index (#99).
            let pg_index = Arc::new(mnemo_postgres::PgVectorIndex::with_pool(
                pg_storage.pool(),
                cli.dimensions,
            ));
            tracing::info!("Using PostgreSQL backend");
            let mut eng = MnemoEngine::new(
                pg_storage,
                pg_index,
                embedding,
                cli.agent_id.clone(),
                cli.org_id.clone(),
            );
            if let Some(ref key_hex) = cli.encryption_key {
                let enc = ContentEncryption::from_hex(key_hex)?;
                eng = eng.with_encryption(Arc::new(enc));
                tracing::info!("At-rest encryption enabled");
            }
            if experience_memory_enabled() {
                eng = eng.with_experience_memory();
                tracing::info!("Experience-memory tier (DocTrace) enabled");
            }
            Arc::new(eng)
        }
        #[cfg(not(feature = "postgres"))]
        {
            return Err("PostgreSQL support not enabled. Rebuild with --features postgres".into());
        }
    } else {
        // DuckDB backend (default)
        let storage = Arc::new(DuckDbStorage::open(&cli.db_path)?);
        tracing::info!("Database opened at {:?}", cli.db_path);

        let index = Arc::new(UsearchIndex::new(cli.dimensions)?);

        // Load existing index if available
        let index_path = cli.db_path.with_extension("usearch");
        if index_path.exists() {
            index.load(&index_path)?;
            tracing::info!("Loaded vector index ({} vectors)", index.len());
        }

        // Initialize full-text index
        let ft_path = cli.db_path.with_extension("tantivy");
        let full_text = Arc::new(TantivyFullTextIndex::new(&ft_path)?);
        tracing::info!(
            "Full-text index ready at {:?} ({} docs)",
            ft_path,
            full_text.len()
        );

        // Keep a clone of the actual index for shutdown save
        duckdb_index = Some(index.clone());

        let mut eng = MnemoEngine::new(
            storage,
            index.clone(),
            embedding,
            cli.agent_id.clone(),
            cli.org_id.clone(),
        )
        .with_full_text(full_text.clone());
        if let Some(ref key_hex) = cli.encryption_key {
            let enc = ContentEncryption::from_hex(key_hex)?;
            eng = eng.with_encryption(Arc::new(enc));
            tracing::info!("At-rest encryption enabled");
        }
        if experience_memory_enabled() {
            eng = eng.with_experience_memory();
            tracing::info!("Experience-memory tier (DocTrace) enabled");
        }
        Arc::new(eng)
    };

    // Optionally start REST API server
    #[cfg(feature = "rest")]
    if let Some(port) = cli.rest_port {
        let rest_engine = engine.clone();
        tokio::spawn(async move {
            let app = mnemo_rest::router(rest_engine);
            match tokio::net::TcpListener::bind(format!("0.0.0.0:{port}")).await {
                Ok(listener) => {
                    tracing::info!("REST API listening on 0.0.0.0:{port}");
                    if let Err(e) = axum::serve(listener, app).await {
                        tracing::error!("REST server failed: {e}");
                    }
                }
                Err(e) => {
                    tracing::error!("Failed to bind REST port {port}: {e}");
                }
            }
        });
    }

    // Shared activity tracker for idle timeout
    let activity_tracker = if cli.idle_timeout_seconds > 0 {
        Some(Arc::new(AtomicU64::new(
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_secs(),
        )))
    } else {
        None
    };

    // Shared shutdown signal
    let shutdown_notify = Arc::new(Notify::new());

    // Start idle timeout watchdog (for scale-to-zero)
    if let Some(ref tracker) = activity_tracker {
        let timeout = cli.idle_timeout_seconds;
        let watchdog_tracker = tracker.clone();
        let watchdog_engine = engine.clone();
        let watchdog_shutdown = shutdown_notify.clone();
        tokio::spawn(async move {
            loop {
                tokio::time::sleep(std::time::Duration::from_secs(10)).await;
                let last = watchdog_tracker.load(Ordering::Relaxed);
                let now = std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap()
                    .as_secs();
                if now - last > timeout {
                    tracing::info!(
                        "Idle timeout reached ({timeout}s), shutting down for scale-to-zero"
                    );
                    // Checkpoint before exit so state can be restored on next start
                    match watchdog_engine
                        .checkpoint(mnemo_core::query::checkpoint::CheckpointRequest {
                            thread_id: "__shutdown__".to_string(),
                            agent_id: None,
                            branch_name: Some("main".to_string()),
                            state_snapshot: serde_json::json!({"reason": "idle_timeout"}),
                            label: Some("auto-shutdown".to_string()),
                            metadata: None,
                        })
                        .await
                    {
                        Ok(resp) => tracing::info!("Shutdown checkpoint created: {}", resp.id),
                        Err(e) => tracing::warn!("Failed to create shutdown checkpoint: {e}"),
                    }
                    watchdog_shutdown.notify_one();
                    return;
                }
            }
        });

        tracing::info!("Idle timeout watchdog enabled: {timeout}s");
    }

    // Signal handler for graceful shutdown (Ctrl+C / SIGTERM)
    let signal_shutdown = shutdown_notify.clone();
    tokio::spawn(async move {
        if let Err(e) = tokio::signal::ctrl_c().await {
            tracing::error!("Failed to listen for Ctrl+C: {e}");
            return;
        }
        tracing::info!("Received shutdown signal");
        signal_shutdown.notify_one();
    });

    // Start TTL sweeper that hard-deletes expired memories on a fixed cadence.
    // Disabled when ttl_sweep_interval_seconds == 0.
    if cli.ttl_sweep_interval_seconds > 0 {
        let ttl_interval = cli.ttl_sweep_interval_seconds;
        let ttl_engine = engine.clone();
        let ttl_shutdown = shutdown_notify.clone();
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(std::time::Duration::from_secs(ttl_interval));
            // Skip the immediate first tick so startup isn't surprised by a sweep.
            interval.tick().await;
            loop {
                tokio::select! {
                    _ = interval.tick() => {
                        match ttl_engine.run_ttl_sweep().await {
                            Ok(report) if report.swept_count > 0 || !report.errors.is_empty() => {
                                tracing::info!(
                                    swept = report.swept_count,
                                    errors = report.errors.len(),
                                    "TTL sweep complete"
                                );
                            }
                            Ok(_) => {}
                            Err(e) => tracing::warn!("TTL sweep failed: {e}"),
                        }
                    }
                    _ = ttl_shutdown.notified() => return,
                }
            }
        });
        tracing::info!("TTL sweeper enabled (every {ttl_interval}s)");
    }

    // Create and start MCP server
    let mut server = MnemoServer::new(engine);
    if let Some(ref tracker) = activity_tracker {
        server = server.with_activity_tracker(tracker.clone());
    }
    tracing::info!("Starting Mnemo MCP server on stdio");

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

    // Wait for either MCP service to end or a shutdown signal
    tokio::select! {
        result = service.waiting() => {
            if let Err(e) = result {
                tracing::error!("MCP service error: {e}");
            }
        }
        _ = shutdown_notify.notified() => {
            tracing::info!("Shutdown initiated, saving state...");
        }
    }

    // Save DuckDB vector index on shutdown (using the actual populated index)
    if let Some(ref index) = duckdb_index {
        let index_path = cli.db_path.with_extension("usearch");
        tracing::info!("Saving vector index ({} vectors)...", index.len());
        if let Err(e) = index.save(&index_path) {
            tracing::error!("Failed to save vector index: {}", e);
        }
    }

    Ok(())
}

/// Handle `mnemo baseline --train --agent-id <id>` (v0.3.3 Task A).
///
/// Loads every non-deleted memory for the agent from DuckDB, computes
/// per-dimension mean + diagonal variance over the records that carry an
/// embedding, and persists the result to the `embedding_baseline` table.
/// Subsequent `remember` calls with
/// `PoisoningPolicy::with_outlier_threshold(z)` set will be scored
/// against this baseline.
async fn run_baseline(cli: &Cli, args: &BaselineArgs) -> Result<(), Box<dyn std::error::Error>> {
    if !args.train {
        return Err(
            "baseline: nothing to do — pass `--train` to train and persist a baseline".into(),
        );
    }
    let agent_id = args
        .agent_id
        .clone()
        .unwrap_or_else(|| cli.agent_id.clone());
    if agent_id.is_empty() {
        return Err("baseline: --agent-id is required (or set MNEMO_AGENT_ID)".into());
    }

    tracing::info!(
        agent = %agent_id,
        db = ?cli.db_path,
        "training embedding baseline"
    );

    let storage = Arc::new(DuckDbStorage::open(&cli.db_path)?);
    let filter = mnemo_core::storage::MemoryFilter {
        agent_id: Some(agent_id.clone()),
        ..Default::default()
    };
    let records = storage.list_memories(&filter, args.limit, 0).await?;
    let with_emb = records.iter().filter(|r| r.embedding.is_some()).count();
    tracing::info!(
        total = records.len(),
        with_embedding = with_emb,
        "loaded records"
    );

    let Some(baseline) = train_baseline(&agent_id, &records) else {
        return Err(format!(
            "baseline: not enough embedded records to train for agent {agent_id} (found {with_emb})"
        )
        .into());
    };

    storage
        .insert_or_update_embedding_baseline(&baseline)
        .await?;
    println!(
        "baseline trained for agent '{}' — n={} d={} updated_at={}",
        baseline.agent_id,
        baseline.n,
        baseline.mu.len(),
        baseline.updated_at
    );
    Ok(())
}

/// Handle `mnemo mcp-server --manifest <path>` (v0.4.0-rc3 Task B2).
///
/// Runs the safe-spawn gauntlet against the OX-MCP threat model
/// (2026-04-24) BEFORE constructing any engine state, then starts the
/// existing MCP STDIO server.
async fn run_mcp_server(cli: &Cli, args: &McpServerArgs) -> Result<(), Box<dyn std::error::Error>> {
    let manifest = manifest::Manifest::load(&args.manifest)?;
    tracing::info!(
        manifest = ?args.manifest,
        allowed_tools = ?manifest.allowed_tools,
        allowed_parents = ?manifest.allowed_parents,
        "manifest loaded"
    );

    // Gauntlet step 1: refuse inherited secrets (override with
    // `MNEMO_REJECT_INHERITED_SECRETS=0` for opt-out testing only).
    let reject_secrets = std::env::var("MNEMO_REJECT_INHERITED_SECRETS").as_deref() != Ok("0");
    safe_spawn::check_inherited_secrets(std::env::vars(), reject_secrets)?;

    // Gauntlet step 2: refuse JSON-injection-style argv. Operators must
    // express config via the manifest, not via `--config`/`-c`.
    let argv: Vec<String> = std::env::args().collect();
    safe_spawn::check_args_pattern(&argv)?;

    // Gauntlet step 3: refuse untrusted parent processes when stdin is
    // not a TTY. Parent basename comes from `MNEMO_PARENT_BASENAME` so
    // the operator's launcher controls the trust assertion (we avoid
    // pulling in libc / sysctl just to read /proc).
    let parent_basename = std::env::var("MNEMO_PARENT_BASENAME").ok();
    let has_tty = std::io::stdin().is_terminal();
    safe_spawn::check_parent_process(
        parent_basename.as_deref(),
        has_tty,
        &manifest.allowed_parents,
    )?;
    tracing::info!(
        has_tty,
        parent = parent_basename.as_deref().unwrap_or("<unknown>"),
        "safe-spawn gauntlet passed"
    );

    // The manifest pins the agent set the operator has approved for
    // this binary. An empty set means "any agent", matching the
    // permissive default the test suite exercises.
    if !manifest.allowed_agents.is_empty() && !manifest.allowed_agents.contains(&cli.agent_id) {
        return Err(format!(
            "refused to start: agent_id {:?} is not in manifest.allowed_agents (got {:?})",
            cli.agent_id, manifest.allowed_agents
        )
        .into());
    }

    // The manifest's `audit_log_path` is the destination for future
    // append-only audit exports (see B4). It is logged here so an
    // operator running the binary can see exactly what the manifest
    // committed them to before any traffic flows.
    tracing::info!(
        audit_log_path = ?manifest.audit_log_path,
        "audit log destination configured"
    );

    // Load the HMAC keystore the manifest points at and attach a
    // `ProvenanceSigner` to the engine. With the signer attached, every
    // `recall(..., with_provenance=true)` returns a verifiable receipt
    // (B1) — and crucially, the key material reaches the engine via a
    // chmod-restricted file, never via env or argv.
    let keystore = manifest::Keystore::load(&manifest.keystore_path)?;
    let key_bytes = keystore.key_bytes()?;
    let signer = mnemo_core::provenance::ProvenanceSigner::new(&keystore.key_id, &key_bytes);
    tracing::info!(
        key_id = %signer.key_id(),
        "provenance signer attached"
    );

    // v0.4.0 (P0-1) — load the optional tool-catalog pin and build an attestor.
    // A malformed pin refuses startup here; the actual attestation against the
    // server's advertised catalog runs once below, after the server (and any
    // role filter) is built and before it serves stdio.
    let tool_attestor: Option<attest::PinnedAttestor> =
        if let Some(pin_path) = manifest.tool_catalog_pin_path.as_ref() {
            let pin = attest::catalog_pin::load(pin_path)?;
            tracing::info!(
                pin_signer = %pin.signer,
                pin_tool_count = pin.tools.len(),
                pin_catalog_sha = %hex::encode(pin.catalog_sha256()),
                "MCP tool-catalog pin loaded"
            );
            Some(attest::PinnedAttestor::new(pin))
        } else {
            tracing::warn!(
                "no tool_catalog_pin_path in manifest — running without \
                 catalog-poisoning defense (arXiv 2604.20994). Set \
                 `tool_catalog_pin_path` to enable."
            );
            None
        };
    // v0.4.2 (A1) — load the optional `[role_filter]` block. Building the
    // `ManifestRoleFilter` here means a malformed manifest refuses startup rather
    // than silently accepting an unenforceable filter, and the filter is attached to
    // the server below, so a denied tool is hidden from `tools/list` and rejected by
    // `tools/call` with `-32601`.
    //
    // Transport limitation (the thing an auditor asks about): the stdio transport
    // carries no per-call caller identity, so dispatch builds a `CallerContext` from
    // the engine's default agent id with no roles. A `deny` list therefore acts as a
    // server-wide tool denylist on the stdio transport, not as per-caller RBAC.
    let role_filter = manifest.role_filter.as_ref().map(|cfg| {
        let filter = mnemo_mcp::role_filter::ManifestRoleFilter::new(cfg.clone());
        if filter.is_noop() {
            // A `[role_filter]` block that denies nothing is a real footgun: the
            // operator likely intended a restriction that is not expressed here.
            tracing::warn!(
                default_policy = ?cfg.default,
                caller_role_count = cfg.caller_roles.len(),
                allow_entries = cfg.allow.len(),
                deny_entries = cfg.deny.len(),
                is_noop = true,
                "MCP [role_filter] block present but is a no-op (no roles, no allow, no \
                 deny, default allow_all) — every tool stays reachable; it enforces nothing"
            );
        } else {
            tracing::info!(
                default_policy = ?cfg.default,
                caller_role_count = cfg.caller_roles.len(),
                allow_entries = cfg.allow.len(),
                deny_entries = cfg.deny.len(),
                is_noop = false,
                "MCP [role_filter] enforcement active — a denied tool is hidden from \
                 tools/list and rejected by tools/call with -32601. Note: on stdio there \
                 is no per-call caller identity, so this is a server-wide tool denylist, \
                 not per-caller RBAC"
            );
        }
        Arc::new(filter)
    });
    if role_filter.is_none() {
        tracing::info!(
            "no [role_filter] block in manifest — every advertised tool is reachable \
             (pre-v0.4.2 behaviour preserved). See \
             https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization"
        );
    }

    // Embedding provider: mirror the default startup path. ONNX > OpenAI > Noop.
    let embedding: Arc<dyn EmbeddingProvider> = if let Some(ref onnx_path) = cli.onnx_model_path {
        tracing::info!("Using ONNX local embeddings from {}", onnx_path);
        Arc::new(mnemo_core::embedding::onnx::OnnxEmbedding::new(
            onnx_path,
            cli.dimensions,
        )?)
    } else if let Some(ref api_key) = cli.openai_api_key {
        tracing::info!("Using OpenAI embeddings ({})", cli.embedding_model);
        Arc::new(OpenAiEmbedding::new(
            api_key.clone(),
            cli.embedding_model.clone(),
            cli.dimensions,
        ))
    } else {
        tracing::warn!(
            "No OPENAI_API_KEY set, using noop embeddings (semantic search will not work)"
        );
        Arc::new(NoopEmbedding::new(cli.dimensions))
    };

    // Storage: hardened mode is DuckDB-only (PostgreSQL connection
    // strings are exactly the kind of capability the manifest is meant
    // to keep out of env). `cli.db_path` still applies — it is the
    // path-only knob in the CLI.
    let storage = Arc::new(DuckDbStorage::open(&cli.db_path)?);
    let index = Arc::new(UsearchIndex::new(cli.dimensions)?);
    let index_path = cli.db_path.with_extension("usearch");
    if index_path.exists() {
        index.load(&index_path)?;
        tracing::info!("Loaded vector index ({} vectors)", index.len());
    }
    let ft_path = cli.db_path.with_extension("tantivy");
    let full_text = Arc::new(TantivyFullTextIndex::new(&ft_path)?);
    tracing::info!(
        "Full-text index ready at {:?} ({} docs)",
        ft_path,
        full_text.len()
    );
    let mut eng = MnemoEngine::new(
        storage,
        index.clone(),
        embedding,
        cli.agent_id.clone(),
        cli.org_id.clone(),
    )
    .with_full_text(full_text)
    .with_provenance_signer(Arc::new(signer));
    if let Some(ref key_hex) = cli.encryption_key {
        let enc = ContentEncryption::from_hex(key_hex)?;
        eng = eng.with_encryption(Arc::new(enc));
        tracing::info!("At-rest encryption enabled");
    }
    let engine = Arc::new(eng);

    let shutdown_notify = Arc::new(Notify::new());
    let signal_shutdown = shutdown_notify.clone();
    tokio::spawn(async move {
        if let Err(e) = tokio::signal::ctrl_c().await {
            tracing::error!("Failed to listen for Ctrl+C: {e}");
            return;
        }
        tracing::info!("Received shutdown signal");
        signal_shutdown.notify_one();
    });

    let mut server = MnemoServer::new(engine.clone());
    if let Some(filter) = role_filter.clone() {
        server = server.with_role_filter(filter);
    }

    // `--print-catalog-pin`: emit a ready-to-paste pin for exactly the tools
    // this binary advertises (post role-filter), then exit WITHOUT serving.
    if args.print_catalog_pin {
        let toml = render_catalog_pin_toml(
            &server.advertised_tool_catalog(),
            "REPLACE-ME:catalog-pin",
            &chrono::Utc::now().to_rfc3339(),
        );
        print!("{toml}");
        return Ok(());
    }

    // v0.4.0 (P0-1) — serve-time tool-catalog attestation (arXiv 2604.20994).
    // Runs ONCE, over the compiled `#[tool]` catalog the server will advertise
    // (post role-filter), before it serves stdio. On stdio the catalog is static
    // after boot, so a boot-time check is complete for this transport: it defends
    // against a substituted/tampered binary, a hostile dependency that injects or
    // renames a tool, and pin drift after a version bump. It is NOT per-request
    // and is only as strong as the manifest file's permissions.
    if let Some(attestor) = tool_attestor.as_ref() {
        let fingerprints: Vec<attest::ToolFingerprint> = server
            .advertised_tool_catalog()
            .iter()
            .map(|(name, desc, schema)| attest::fingerprint_tool(name, desc, schema))
            .collect();
        let verdict = attestor.attest(&fingerprints)?;
        // Audit EVERY verdict (before any early Err) so even a refused startup
        // leaves an `mcp_tool_catalog_drift` trail. The module doc promises this.
        record_catalog_drift_event(&engine, &cli.agent_id, &verdict).await;

        let removed_only = verdict.is_removed_only_drift();
        match &verdict {
            attest::AttestationVerdict::Match => {
                tracing::info!(
                    pin_signer = %attestor.baseline().signer,
                    catalog_sha = %hex::encode(attestor.baseline().catalog_sha256()),
                    tool_count = fingerprints.len(),
                    "MCP tool-catalog attestation PASSED — advertised catalog matches the pin"
                );
            }
            attest::AttestationVerdict::Drift { removed, .. }
                if removed_only && manifest.allow_removed_drift =>
            {
                let names: Vec<&str> = removed.iter().map(|t| t.name.as_str()).collect();
                tracing::warn!(
                    removed = ?names,
                    "MCP tool-catalog removed-only drift ACCEPTED via allow_removed_drift — \
                     the advertised catalog is a subset of the pin"
                );
            }
            attest::AttestationVerdict::Drift {
                added,
                removed,
                mutated,
            } => {
                let names = |v: &[attest::ToolFingerprint]| {
                    v.iter().map(|t| t.name.clone()).collect::<Vec<_>>()
                };
                return Err(format!(
                    "MCP tool-catalog attestation FAILED (arXiv 2604.20994): advertised catalog \
                     drifted from the pin — added={:?} mutated={:?} removed={:?}. Refusing to \
                     serve. Regenerate the pin with `--print-catalog-pin` if this change is \
                     intended{}.",
                    names(added.as_slice()),
                    names(mutated.as_slice()),
                    names(removed.as_slice()),
                    if removed_only {
                        ", or set allow_removed_drift = true for a removed-only downgrade"
                    } else {
                        ""
                    }
                )
                .into());
            }
            attest::AttestationVerdict::Reject { reason } => {
                return Err(format!(
                    "MCP tool-catalog attestation REJECTED (arXiv 2604.20994): {reason}. \
                     Refusing to serve."
                )
                .into());
            }
        }
    }

    tracing::info!("Starting Mnemo MCP server on stdio (hardened mode)");
    let service = server.serve(stdio()).await?;
    tokio::select! {
        result = service.waiting() => {
            if let Err(e) = result {
                tracing::error!("MCP service error: {e}");
            }
        }
        _ = shutdown_notify.notified() => {
            tracing::info!("Shutdown initiated, saving state...");
        }
    }
    let index_path = cli.db_path.with_extension("usearch");
    tracing::info!("Saving vector index ({} vectors)...", index.len());
    if let Err(e) = index.save(&index_path) {
        tracing::error!("Failed to save vector index: {}", e);
    }
    Ok(())
}

/// Wrap a value as a TOML basic string, escaping `\` and `"`. Our values
/// (tool names, an RFC3339 timestamp, a signer id) carry no control chars.
fn toml_basic_string(s: &str) -> String {
    format!("\"{}\"", s.replace('\\', "\\\\").replace('"', "\\\""))
}

/// Render a ready-to-paste `[tool_catalog_pin]` block from an advertised
/// catalog (`(name, description, input_schema_json)` triples). The emitted TOML
/// round-trips through [`attest::catalog_pin::load`] (asserted in tests), and
/// each `schema_sha256` is the same fingerprint the attestor computes, so a pin
/// generated from a binary Matches that binary's own advertised catalog.
fn render_catalog_pin_toml(
    catalog: &[(String, String, String)],
    signer: &str,
    signed_at: &str,
) -> String {
    let mut out = String::new();
    out.push_str(
        "# Generated by `mnemo mcp-server --manifest <m> --print-catalog-pin`.\n\
         # Point the manifest's `tool_catalog_pin_path` at this file to enable\n\
         # serve-time tool-catalog attestation (arXiv 2604.20994). Replace the\n\
         # signer placeholder with a stable `<host>:<key_id>` before committing.\n",
    );
    out.push_str("[tool_catalog_pin]\n");
    out.push_str(&format!("signer = {}\n", toml_basic_string(signer)));
    out.push_str(&format!("signed_at = {}\n", toml_basic_string(signed_at)));
    for (name, desc, schema) in catalog {
        let fp = attest::fingerprint_tool(name, desc, schema);
        out.push_str("\n[[tool_catalog_pin.tools]]\n");
        out.push_str(&format!("name = {}\n", toml_basic_string(name)));
        out.push_str(&format!("schema_sha256 = \"{}\"\n", fp.schema_hex()));
    }
    out
}

/// Append one `McpToolCatalogDrift` audit event recording an attestation
/// verdict (the module doc promises "Every verdict is recorded"). Best-effort:
/// a storage error is logged, never propagated, so it can run before a refused
/// startup returns `Err` without masking the real reason.
async fn record_catalog_drift_event(
    engine: &mnemo_core::query::MnemoEngine,
    agent_id: &str,
    verdict: &attest::AttestationVerdict,
) {
    let names =
        |v: &[attest::ToolFingerprint]| v.iter().map(|t| t.name.clone()).collect::<Vec<String>>();
    let (label, added, removed, mutated, reason) = match verdict {
        attest::AttestationVerdict::Match => ("match", vec![], vec![], vec![], None),
        attest::AttestationVerdict::Drift {
            added,
            removed,
            mutated,
        } => ("drift", names(added), names(removed), names(mutated), None),
        attest::AttestationVerdict::Reject { reason } => {
            ("reject", vec![], vec![], vec![], Some(reason.clone()))
        }
    };
    let payload = serde_json::json!({
        "verdict": label,
        "added": added,
        "removed": removed,
        "mutated": mutated,
        "reason": reason,
    });
    let event = mnemo_core::query::event_builder::build_event(
        engine,
        agent_id,
        mnemo_core::model::event::EventType::McpToolCatalogDrift,
        payload,
        "mcp_tool_catalog_attestation",
        None,
    )
    .await;
    if let Err(e) = engine.storage.insert_event(&event).await {
        tracing::error!(event_id = %event.id, error = %e, "failed to record McpToolCatalogDrift audit event");
    }
}

/// Handle `mnemo eval` (v0.4.0-rc3 Task B6).
///
/// Replays a JSONL dataset of `{id, content, query, expected}` rows
/// against an in-memory engine and emits a per-row JSONL report
/// (latency_ms, top_k, hit). Used to compare config sweeps
/// (provenance on/off, hybrid weights, recency half-life) against a
/// fixed prompt set without spinning up a full deployment.
async fn run_eval(cli: &Cli, args: &EvalArgs) -> Result<(), Box<dyn std::error::Error>> {
    use std::io::{BufWriter, Write};
    use std::time::Instant;

    use mnemo_core::query::recall::RecallRequest;
    use mnemo_core::query::remember::RememberRequest;

    #[derive(serde::Deserialize)]
    struct Row {
        id: String,
        content: String,
        query: String,
        expected: String,
    }

    let dataset_path = args.dataset.clone().unwrap_or_else(|| {
        // The bundled LongMemEval_M lives in mnemo-core's bench data
        // dir relative to the workspace root. We resolve via
        // CARGO_MANIFEST_DIR for robustness across the
        // `cargo install` install path.
        let here = PathBuf::from(env!("CARGO_MANIFEST_DIR"));
        here.join("..")
            .join("mnemo-core")
            .join("benches")
            .join("data")
            .join("longmemeval_m.jsonl")
    });

    let text = std::fs::read_to_string(&dataset_path)
        .map_err(|e| format!("eval: failed to read dataset {dataset_path:?}: {e}"))?;
    let rows: Vec<Row> = text
        .lines()
        .filter(|l| !l.trim().is_empty())
        .map(|l| serde_json::from_str::<Row>(l).map_err(|e| format!("eval: bad row '{l}': {e}")))
        .collect::<Result<_, _>>()?;
    if rows.is_empty() {
        return Err(format!("eval: dataset {dataset_path:?} is empty").into());
    }

    // Build engine. Eval is always in-memory so a config sweep does
    // not pollute the operator's persisted DB.
    let storage = Arc::new(DuckDbStorage::open_in_memory()?);
    let index = Arc::new(UsearchIndex::new(cli.dimensions)?);
    let embedding: Arc<dyn EmbeddingProvider> = Arc::new(NoopEmbedding::new(cli.dimensions));
    let mut eng = MnemoEngine::new(
        storage,
        index,
        embedding,
        cli.agent_id.clone(),
        cli.org_id.clone(),
    );
    if args.with_provenance {
        let key_hex = args.provenance_key_hex.as_ref().ok_or(
            "eval: --with-provenance requires --provenance-key-hex (>=32 raw bytes hex-encoded)",
        )?;
        let key_bytes = hex::decode(key_hex)
            .map_err(|e| format!("eval: --provenance-key-hex not valid hex: {e}"))?;
        if key_bytes.len() < 32 {
            return Err(format!(
                "eval: --provenance-key-hex must decode to >= 32 bytes (got {})",
                key_bytes.len()
            )
            .into());
        }
        let signer = mnemo_core::provenance::ProvenanceSigner::new("eval-key", &key_bytes);
        eng = eng.with_provenance_signer(Arc::new(signer));
    }
    let engine = Arc::new(eng);

    // Seed the engine with each row's content. Eval queries hit the
    // same engine so we can measure end-to-end recall latency.
    for r in &rows {
        let mut req = RememberRequest::new(r.content.clone());
        req.tags = Some(vec![format!("eval-id:{}", r.id)]);
        engine.remember(req).await?;
    }

    // Open the output sink. None means stdout.
    let mut out: Box<dyn Write> = match &args.output {
        Some(path) => Box::new(BufWriter::new(std::fs::File::create(path)?)),
        None => Box::new(BufWriter::new(std::io::stdout().lock())),
    };

    let mut hits = 0usize;
    let mut total_latency_us: u128 = 0;
    for r in &rows {
        let recall = RecallRequest {
            query: r.query.clone(),
            agent_id: None,
            limit: Some(args.limit),
            memory_type: None,
            memory_types: None,
            scope: None,
            min_importance: None,
            tags: None,
            org_id: None,
            strategy: Some(args.strategy.clone()),
            temporal_range: None,
            recency_half_life_hours: None,
            hybrid_weights: None,
            rrf_k: None,
            as_of: None,
            explain: None,
            with_provenance: if args.with_provenance {
                Some(true)
            } else {
                None
            },
            mode: None,
            current_fact_resolver: None,
            orientation_cache: None,
            evidence_budget: None,
            retained_token_budget: None,
            domain_scope: None,
            reasoning_trust: None,
        };
        let t0 = Instant::now();
        let resp = engine.recall(recall).await?;
        let elapsed_us = t0.elapsed().as_micros();
        total_latency_us += elapsed_us;

        let recalled_contents: Vec<String> =
            resp.memories.iter().map(|m| m.content.clone()).collect();
        let hit = recalled_contents
            .iter()
            .any(|c| c.to_lowercase().contains(&r.expected.to_lowercase()));
        if hit {
            hits += 1;
        }

        let row = serde_json::json!({
            "id": r.id,
            "query": r.query,
            "expected": r.expected,
            "recalled_count": resp.memories.len(),
            "recalled_top1": resp.memories.first().map(|m| m.content.clone()),
            "hit": hit,
            "latency_us": elapsed_us,
            "provenance_present": resp.provenance.is_some(),
        });
        writeln!(out, "{}", serde_json::to_string(&row)?)?;
    }

    let n = rows.len() as f64;
    let avg_latency_us = total_latency_us as f64 / n;
    let summary = serde_json::json!({
        "summary": true,
        "rows": rows.len(),
        "hits": hits,
        "hit_rate": hits as f64 / n,
        "avg_latency_us": avg_latency_us,
        "strategy": args.strategy,
        "with_provenance": args.with_provenance,
    });
    writeln!(out, "{}", serde_json::to_string(&summary)?)?;
    out.flush()?;
    Ok(())
}

/// v0.4.9 — `mnemo bench <subcommand>` dispatch.
async fn run_bench(sub: &BenchCommand) -> Result<(), Box<dyn std::error::Error>> {
    match sub {
        BenchCommand::Embeddings(args) => run_bench_embeddings(args).await,
    }
}

async fn run_compliance(sub: &ComplianceCommand) -> Result<(), Box<dyn std::error::Error>> {
    match sub {
        ComplianceCommand::Retention(args) => run_compliance_retention(args).await,
    }
}

async fn run_compliance_retention(args: &RetentionArgs) -> Result<(), Box<dyn std::error::Error>> {
    use mnemo_compliance::RetentionProfile;

    let mut profile = match args.profile.as_str() {
        "dpdp" => RetentionProfile::dpdp_rules(),
        "eu-ai-act-art19" => RetentionProfile::eu_ai_act_art19(),
        "hipaa" => RetentionProfile::hipaa_164_312b(),
        other => return Err(format!("unknown retention profile '{other}'").into()),
    };
    if let Some(days) = args.floor_days {
        profile = profile.with_floor_days(days);
    }

    // Gate the floor against the active backend's append-only guarantee. The
    // default embedded backend is DuckDB, whose `agent_events` log is
    // append-only, so this passes; a backend that could not honour the floor
    // would surface `ComplianceError::RetentionFloorUnsupported` naming itself.
    let storage = DuckDbStorage::open_in_memory()?;
    let backend = storage.backend_name();
    let backend_ok = profile
        .assert_backend_can_retain(backend, storage.events_are_append_only())
        .is_ok();

    let out = serde_json::json!({
        "profile": profile.name,
        "obligation": profile.obligation,
        "floor_days": profile.floor_days,
        "commencement": profile.commencement,
        "source_url": profile.source_url,
        "backend": backend,
        "backend_can_honour_floor": backend_ok,
    });
    println!("{}", serde_json::to_string_pretty(&out)?);

    // Fail loud (non-zero exit) if the backend cannot honour the floor.
    profile.assert_backend_can_retain(backend, storage.events_are_append_only())?;
    Ok(())
}

async fn run_bench_embeddings(
    args: &BenchEmbeddingsArgs,
) -> Result<(), Box<dyn std::error::Error>> {
    let opts = mnemo_embeddings_bench::RunOptions {
        dimensions: args.dimensions,
        latency_samples: args.latency_samples,
    };
    let results = mnemo_embeddings_bench::run_all(&opts).await;
    let rows: Vec<mnemo_embeddings_bench::BackendRow> = results
        .iter()
        .filter_map(|(_, row, _)| row.clone())
        .collect();
    let rec = mnemo_embeddings_bench::recommend(&rows, args.slo_ms);
    let table = mnemo_embeddings_bench::render_table(&results, &rec);
    print!("{table}");
    Ok(())
}

#[cfg(test)]
mod catalog_pin_tests {
    use super::*;

    /// The `--print-catalog-pin` output must round-trip through the loader, and
    /// the loaded fingerprints must equal what the attestor computes from the
    /// same catalog — i.e. a pin generated from a binary Matches that binary.
    #[test]
    fn rendered_pin_round_trips_and_matches() {
        let catalog = vec![
            (
                "mnemo.recall".to_string(),
                "Search memories".to_string(),
                r#"{"type":"object"}"#.to_string(),
            ),
            (
                "mnemo.verify".to_string(),
                "Verify hash chain".to_string(),
                r#"{"type":"object","properties":{}}"#.to_string(),
            ),
        ];
        let toml = render_catalog_pin_toml(&catalog, "test:signer", "2026-07-30T00:00:00Z");

        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("pin.toml");
        std::fs::write(&path, &toml).unwrap();
        let pin = attest::catalog_pin::load(&path).expect("emitted pin must load");
        assert_eq!(pin.signer, "test:signer");
        assert_eq!(pin.tools.len(), 2);

        // The loaded pin must MATCH the same catalog run through the attestor.
        let fingerprints: Vec<attest::ToolFingerprint> = catalog
            .iter()
            .map(|(n, d, s)| attest::fingerprint_tool(n, d, s))
            .collect();
        let attestor = attest::PinnedAttestor::new(pin);
        assert_eq!(
            attestor.attest(&fingerprints).unwrap(),
            attest::AttestationVerdict::Match
        );
    }
}