aion-worker-mcp 0.31.0

//! From tools to activities: one Aion worker per MCP server, serving the
//! task queue named after the server, with one activity per tool named by
//! the tool's identifier form (`get-sum` → `get_sum`). A workflow document
//! declares exactly that: `worker everything` / `action get_sum(a: Int, b:
//! Int) -> Sum`. Input is the tool's arguments object; output is the tool's
//! result as the server returned it.

use std::sync::Arc;
use std::time::Duration;

use aion_worker::activity::encode_payload;
use aion_worker::{ActivityContext, ActivityFailure, HandlerFuture, Worker, WorkerConfig};
use anyhow::Context;
use serde_json::{Value, json};

use crate::client::{ClientError, McpClient, Tool};
use crate::config::activity_name;

/// Everything a worker needs besides its queue and identity, which come from
/// the server name.
#[derive(Debug, Clone, Copy)]
pub struct WorkerShape {
    /// Maximum concurrent activities per server.
    pub concurrency: usize,
    /// First reconnect delay.
    pub reconnect_initial_backoff: Duration,
    /// Longest reconnect delay.
    pub reconnect_max_backoff: Duration,
    /// Reconnect attempts before giving up (at least 1).
    pub reconnect_max_attempts: usize,
    /// Periodic heartbeat while a call is in flight, if any.
    pub heartbeat_every: Option<Duration>,
}

/// Lists one server's tools and builds the worker that serves them on the
/// queue named after the server, identified as `<identity>/<server>`.
///
/// # Errors
///
/// A server that cannot list its tools, lists none, or lists two tools whose
/// identifier forms collide.
pub async fn build_worker(
    endpoint: &str,
    identity: &str,
    shape: WorkerShape,
    client: &Arc<McpClient>,
) -> anyhow::Result<Worker> {
    let config = WorkerConfig::builder()
        .endpoint(endpoint.to_owned())
        .task_queue(client.name().to_owned())
        .identity(format!("{identity}/{}", client.name()))
        .max_concurrency(shape.concurrency)
        .reconnect_initial_backoff(shape.reconnect_initial_backoff)
        .reconnect_max_backoff(shape.reconnect_max_backoff)
        .reconnect_max_attempts(shape.reconnect_max_attempts)
        .build()
        .with_context(|| format!("worker configuration for `{}`", client.name()))?;
    let mut builder = Worker::builder(config);
    let tools = client
        .list_tools()
        .await
        .with_context(|| format!("listing tools of `{}`", client.name()))?;
    anyhow::ensure!(
        !tools.is_empty(),
        "server `{}` lists no tools",
        client.name()
    );
    let mut served: std::collections::BTreeMap<String, String> = std::collections::BTreeMap::new();
    for tool in tools {
        let activity = activity_name(&tool.name);
        if let Some(other) = served.get(&activity) {
            anyhow::bail!(
                "server `{}`: tools `{other}` and `{}` would both be served as `{activity}`; rename one",
                client.name(),
                tool.name
            );
        }
        served.insert(activity.clone(), tool.name.clone());
        let handler = make_handler(Arc::clone(client), tool.clone(), shape.heartbeat_every);
        builder = builder
            .register_activity::<Value, Value, _>(activity.clone(), handler)
            .with_context(|| format!("registering `{}`.`{activity}`", client.name()))?;
        let output_schema = tool.output_schema.clone().unwrap_or(Value::Null);
        tracing::info!(
            queue = client.name(),
            activity,
            tool = %tool.name,
            description = tool.description.as_deref().unwrap_or(""),
            input_schema = %tool.input_schema,
            %output_schema,
            "registered"
        );
    }
    builder
        .build()
        .with_context(|| format!("building the worker for `{}`", client.name()))
}

fn make_handler(
    client: Arc<McpClient>,
    tool: Tool,
    heartbeat_every: Option<Duration>,
) -> impl for<'c> Fn(Value, &'c ActivityContext) -> HandlerFuture<'c, Value> + Send + Sync + 'static
{
    move |input: Value, context: &ActivityContext| {
        let client = Arc::clone(&client);
        let tool_name = tool.name.clone();
        Box::pin(async move {
            let arguments = match input {
                Value::Null => json!({}),
                Value::Object(map) => Value::Object(map),
                other => {
                    return Err(ActivityFailure::terminal(format!(
                        "`{}.{tool_name}` takes a JSON object of arguments, got {}",
                        client.name(),
                        kind_of(&other)
                    )));
                }
            };
            let heartbeat_on_progress = |params: &Value| {
                let detail = encode_payload(params).ok();
                if let Err(error) = context.heartbeat(detail) {
                    tracing::warn!(%error, "heartbeat on progress failed");
                }
            };
            let call = client.call_tool(
                &tool_name,
                arguments,
                context.cancelled(),
                heartbeat_on_progress,
            );
            let outcome = match heartbeat_every {
                Some(every) => {
                    tokio::pin!(call);
                    let mut ticks = tokio::time::interval(every);
                    ticks.tick().await; // the first tick is immediate; the call has just begun
                    loop {
                        tokio::select! {
                            result = &mut call => break result,
                            _ = ticks.tick() => {
                                if let Err(error) = context.heartbeat(None) {
                                    tracing::warn!(%error, "periodic heartbeat failed");
                                }
                            }
                        }
                    }
                }
                None => call.await,
            };
            match outcome {
                Ok(result) => {
                    if result.get("isError").and_then(Value::as_bool) == Some(true) {
                        let mut failure = ActivityFailure::terminal(format!(
                            "`{}.{tool_name}` reported an error: {}",
                            client.name(),
                            content_text(&result)
                        ));
                        if let Ok(detail) = encode_payload(&result) {
                            failure = failure.with_detail(detail);
                        }
                        return Err(failure);
                    }
                    Ok(result)
                }
                Err(ClientError::ServerGone { .. } | ClientError::Io { .. }) => Err(
                    ActivityFailure::retryable(outcome_message(&tool_name, &client)),
                ),
                Err(error) => Err(ActivityFailure::terminal(error.to_string())),
            }
        })
    }
}

fn outcome_message(tool: &str, client: &McpClient) -> String {
    format!(
        "`{}.{tool}`: the MCP server is gone; the call may be retried once it is back",
        client.name()
    )
}

/// The text blocks of a result, joined, for a failure message.
fn content_text(result: &Value) -> String {
    let texts: Vec<&str> = result
        .get("content")
        .and_then(Value::as_array)
        .map(|blocks| {
            blocks
                .iter()
                .filter_map(|b| b.get("text").and_then(Value::as_str))
                .collect()
        })
        .unwrap_or_default();
    if texts.is_empty() {
        result.to_string()
    } else {
        texts.join("\n")
    }
}

fn kind_of(value: &Value) -> &'static str {
    match value {
        Value::Null => "null",
        Value::Bool(_) => "a boolean",
        Value::Number(_) => "a number",
        Value::String(_) => "a string",
        Value::Array(_) => "an array",
        Value::Object(_) => "an object",
    }
}

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

    #[test]
    fn content_text_joins_text_blocks_and_falls_back_to_the_raw_result() {
        let r = json!({"content":[{"type":"text","text":"a"},{"type":"image","data":"x"},{"type":"text","text":"b"}]});
        assert_eq!(content_text(&r), "a\nb");
        let r = json!({"structuredContent":{"k":1}});
        assert_eq!(content_text(&r), r.to_string());
    }
}