aion-worker-mcp 0.31.0

//! `aion-mcp-worker` — serve the tools of Model Context Protocol servers as
//! Aion activities.
//!
//! Launch it with an `mcpServers` document (inline JSON or a file). It starts
//! every server in the document over stdio, performs the MCP initialize
//! handshake, lists each server's tools, and serves one Aion worker per
//! server: the task queue is the server's name, each tool is an activity
//! named by the tool's identifier form (`get-sum` → `get_sum`), so a
//! workflow document declares `worker everything` / `action get_sum(...)`.
//! The activity input is the tool's JSON arguments;
//! the activity output is the tool's result exactly as the server returned it
//! (`content`, `structuredContent`, `isError`). A tool that answers
//! `isError: true` fails the activity terminally with the tool's own words;
//! a server that dies mid-call fails it retryably; a cancelled activity sends
//! the server MCP's cancellation notice.
//!
//! Every operational value is a flag: nothing here assumes a queue, a
//! concurrency, a backoff or a heartbeat cadence.

mod client;
mod config;
mod serve;

use std::path::PathBuf;
use std::time::Duration;

use anyhow::{Context, bail};
use clap::Parser;
use tracing_subscriber::EnvFilter;

/// Serve MCP servers' tools as Aion activities on one task queue.
#[derive(Debug, Parser)]
#[command(name = "aion-mcp-worker", version, about)]
struct Args {
    /// gRPC endpoint of the Aion server the worker dials.
    #[arg(long)]
    endpoint: String,
    /// Worker identity prefix; each server's worker is `<identity>/<server>`.
    #[arg(long)]
    identity: String,
    /// Maximum concurrent activities per server.
    #[arg(long)]
    concurrency: usize,
    /// First reconnect delay, in seconds.
    #[arg(long)]
    reconnect_initial_backoff: u64,
    /// Longest reconnect delay, in seconds.
    #[arg(long)]
    reconnect_max_backoff: u64,
    /// Reconnect attempts before giving up (at least 1; the SDK refuses 0).
    #[arg(long)]
    reconnect_max_attempts: usize,
    /// The `mcpServers` document, inline. Exactly one of --config / --config-file.
    #[arg(long, conflicts_with = "config_file")]
    config: Option<String>,
    /// Path to the `mcpServers` document. Exactly one of --config / --config-file.
    #[arg(long, conflicts_with = "config")]
    config_file: Option<PathBuf>,
    /// Send an activity heartbeat this often (seconds) while a tool call is in
    /// flight. Omit to heartbeat only on the server's progress notifications.
    #[arg(long)]
    heartbeat_every: Option<u64>,
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")),
        )
        .with_writer(std::io::stderr)
        .init();
    let args = Args::parse();

    let document = match (&args.config, &args.config_file) {
        (Some(inline), None) => inline.clone(),
        (None, Some(path)) => std::fs::read_to_string(path)
            .with_context(|| format!("reading the mcpServers document at {}", path.display()))?,
        _ => bail!("give exactly one of --config <json> or --config-file <path>"),
    };
    let config = config::McpConfig::parse(&document).context("the mcpServers document")?;

    let shape = serve::WorkerShape {
        concurrency: args.concurrency,
        reconnect_initial_backoff: Duration::from_secs(args.reconnect_initial_backoff),
        reconnect_max_backoff: Duration::from_secs(args.reconnect_max_backoff),
        reconnect_max_attempts: args.reconnect_max_attempts,
        heartbeat_every: args.heartbeat_every.map(Duration::from_secs),
    };

    let clients = client::spawn_all(&config).await?;
    let mut workers = Vec::with_capacity(clients.len());
    for client in &clients {
        match serve::build_worker(&args.endpoint, &args.identity, shape, client).await {
            Ok(worker) => workers.push(worker),
            Err(error) => {
                for client in &clients {
                    client.shutdown().await;
                }
                return Err(error);
            }
        }
    }
    let activities: usize = workers.iter().map(|w| w.activity_types().len()).sum();
    tracing::info!(
        identity = %args.identity,
        servers = clients.len(),
        activities,
        queues = ?clients.iter().map(|c| c.name().to_owned()).collect::<Vec<_>>(),
        "aion-mcp-worker ready"
    );

    // One ctrl-c drains every worker; a worker that ends on its own (its
    // reconnects exhausted) ends the process with its error, so a dead queue
    // is never served in silence.
    let (stop_tx, stop_rx) = tokio::sync::watch::channel(false);
    tokio::spawn(async move {
        if let Err(error) = tokio::signal::ctrl_c().await {
            tracing::error!(%error, "ctrl-c handler failed; serving until the connections end");
            return;
        }
        tracing::info!("ctrl-c: draining");
        let _ = stop_tx.send(true);
    });
    let mut runs = tokio::task::JoinSet::new();
    for (worker, client) in workers.into_iter().zip(&clients) {
        let mut stop = stop_rx.clone();
        let queue = client.name().to_owned();
        runs.spawn(async move {
            let shutdown = async move {
                let _ = stop.wait_for(|stopped| *stopped).await;
            };
            worker
                .run_until(shutdown)
                .await
                .with_context(|| format!("worker serving `{queue}` ended"))
        });
    }
    let mut outcome = Ok(());
    while let Some(joined) = runs.join_next().await {
        match joined {
            Ok(Ok(())) => {}
            Ok(Err(error)) => {
                tracing::error!(error = %format!("{error:#}"), "a worker ended with an error; stopping the rest");
                outcome = Err(error);
                runs.abort_all();
            }
            Err(join) => {
                tracing::error!(%join, "a worker task panicked or was cancelled");
                if outcome.is_ok() {
                    outcome = Err(anyhow::anyhow!("worker task ended abnormally: {join}"));
                }
            }
        }
    }

    for client in &clients {
        client.shutdown().await;
    }
    outcome
}