Pricing over MCP

mod dispatcher;

use chio_core::{
    capability::scope::{ChioScope, Operation, ToolGrant},
    crypto::Keypair,
};
use chio_egress_contract::HttpEgressContract;
use chio_kernel::{
    ChioKernel, KernelConfig, DEFAULT_CHECKPOINT_BATCH_SIZE, DEFAULT_MAX_STREAM_DURATION_SECS,
    DEFAULT_MAX_STREAM_TOTAL_BYTES,
};
use chio_mcp_edge::{ChioMcpEdge, McpEdgeConfig};
use chio_openapi_mcp_bridge::{BridgeConfig, OpenApiMcpBridge, OwnedBridgeToolServer};
use std::{collections::BTreeSet, io::BufReader, path::PathBuf};

fn grant(tool: &str, max_invocations: Option<u32>) -> ToolGrant {
    ToolGrant {
        server_id: "pricing".into(),
        tool_name: tool.into(),
        operations: vec![Operation::Invoke],
        constraints: Vec::new(),
        max_invocations,
        max_cost_per_invocation: None,
        max_total_cost: None,
        dpop_required: None,
    }
}

#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct RolePolicy {
    allow_read: bool,
    allow_write: bool,
    max_writes: u32,
}

fn main() -> anyhow::Result<()> {
    let directory = PathBuf::from(std::env::args().nth(1).ok_or_else(|| {
        anyhow::anyhow!("Usage: pricing-mcp NEW_SESSION_DIRECTORY [basalt|corbel]")
    })?);
    let role = std::env::args().nth(2).unwrap_or_else(|| "basalt".into());
    let policy_path = std::env::var_os("CHIO_PRICING_AUTHORITY_FILE")
        .map(std::path::PathBuf::from)
        .unwrap_or_else(|| "authority.json".into());
    let policy_bytes = std::fs::read(policy_path)?;
    let policies: std::collections::BTreeMap<String, RolePolicy> =
        serde_json::from_slice(&policy_bytes)?;
    let policy = policies
        .get(&role)
        .ok_or_else(|| anyhow::anyhow!("Unknown role {role}; choose a role in authority.json"))?;
    anyhow::ensure!(
        policy.max_writes <= 50,
        "max_writes must be between 0 and 50"
    );
    let mut grants = Vec::new();
    if policy.allow_read {
        grants.extend([grant("listQuotes", None), grant("readQuote", None)]);
    }
    if policy.allow_write && policy.max_writes > 0 {
        grants.push(grant("quote", Some(policy.max_writes)));
    }
    // Refuse reuse: each explicitly started client session receives a new grant.
    std::fs::create_dir(&directory)?;
    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        std::fs::set_permissions(&directory, std::fs::Permissions::from_mode(0o700))?;
    }
    std::fs::write(directory.join("authority.json"), &policy_bytes)?;
    std::fs::write(directory.join("role.txt"), &role)?;
    let spec = std::fs::read_to_string("openapi.yaml")?;
    let policy_hash = chio_core::sha256_hex(&chio_core::canonical_json_bytes(
        &serde_json::json!({
            "openapi": spec, "authority": serde_json::from_slice::<serde_json::Value>(&policy_bytes)?, "role": role
        }),
    )?);
    let authority =
        chio_control_plane::DurableAdmissionRuntime::open(&directory.join("admission.db"))?;
    let keypair = authority.kernel_keypair();
    std::fs::write(
        directory.join("kernel-public-key.txt"),
        keypair.public_key().to_hex(),
    )?;
    let mut kernel = ChioKernel::new(KernelConfig {
        keypair,
        ca_public_keys: Vec::new(),
        max_delegation_depth: 3,
        policy_hash,
        allow_sampling: false,
        allow_sampling_tool_use: false,
        allow_elicitation: false,
        max_stream_duration_secs: DEFAULT_MAX_STREAM_DURATION_SECS,
        max_stream_total_bytes: DEFAULT_MAX_STREAM_TOTAL_BYTES,
        require_web3_evidence: false,
        // Install the receipt store and durable admission authority below.
        // Revocations are process-local in this single-session host.
        allow_ephemeral_receipt_log: true,
        allow_ephemeral_revocation_store: true,
        checkpoint_batch_size: DEFAULT_CHECKPOINT_BATCH_SIZE,
        retention_config: None,
        memory_budget: chio_kernel::MemoryBudgetConfig::defaults(),
        deadlines: chio_kernel::HotPathDeadlineConfig::default(),
    });
    kernel.set_receipt_store(Box::new(chio_store_sqlite::SqliteReceiptStore::open(
        directory.join("receipts.db"),
    )?))?;
    authority.attach(&mut kernel)?;
    let service_key = Keypair::generate();
    let mut bridge = OpenApiMcpBridge::from_spec(
        &spec,
        BridgeConfig {
            server_id: "pricing".into(),
            server_name: "Larkfield pricing".into(),
            server_version: "1.0.0".into(),
            public_key: service_key.public_key().to_hex(),
            base_url: "http://127.0.0.1:8087".into(),
            egress_contract: Some(HttpEgressContract {
                tenant_egress_namespace: "larkfield-pricing-local".into(),
                allowed_schemes: BTreeSet::from(["http".into()]),
                allowed_authority_set: BTreeSet::from(["127.0.0.1:8087".into()]),
                deny_loopback: false,
                deny_link_local: true,
                deny_ipv6_ula: true,
                max_redirect_chain: 0,
                max_response_bytes: dispatcher::RESPONSE_LIMIT,
            }),
        },
    )?;
    bridge.set_dispatcher(dispatcher::build()?);
    let manifest = bridge.manifest_clone();
    let signed = chio_manifest::sign_manifest(&manifest, &service_key)?;
    chio_manifest::verify_manifest(&signed, &service_key.public_key())?;
    std::fs::write(
        directory.join("manifest.json"),
        serde_json::to_vec_pretty(&signed)?,
    )?;
    kernel.register_tool_server(Box::new(OwnedBridgeToolServer::from_bridge(bridge)));
    let agent = Keypair::generate();
    let capability = kernel.issue_capability(
        &agent.public_key(),
        ChioScope {
            grants,
            ..ChioScope::default()
        },
        1800,
    )?;
    std::fs::write(
        directory.join("capability.json"),
        serde_json::to_vec_pretty(&capability)?,
    )?;
    let mut edge = ChioMcpEdge::new(
        McpEdgeConfig {
            server_name: "Larkfield pricing, governed".into(),
            ..McpEdgeConfig::default()
        },
        kernel,
        agent.public_key().to_hex(),
        vec![capability],
        vec![manifest],
    )?;
    eprintln!(
        "Pricing MCP: role {role}, quote limit {}; evidence: {}",
        if policy.allow_write {
            policy.max_writes
        } else {
            0
        },
        directory.display()
    );
    edge.serve_stdio(BufReader::new(std::io::stdin()), std::io::stdout())?;
    Ok(())
}