aion-server 0.21.0

Aion workflow server library: HTTP, gRPC, WebSocket, and worker endpoints. Run it with the `aion` binary from the aion-cli crate.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

//! The Aion tool service: the MCP surface's only door into the engine.
//!
//! Every tool here is an IN-PROCESS call onto the same shared handlers the HTTP
//! and gRPC transports use. Nothing loops back through HTTP: a self-call would
//! be a second authorization path, a second serialization, and a second place
//! for the answer to differ from the one the console shows.

use aion_mcp::tasks::resolve::{TaskProjection, TaskResolveError};
use aion_mcp::tools::service::{
    ToolCall, ToolCatalog, ToolFailure, ToolOutcome, ToolService, ToolTaskDecision,
};
use serde_json::json;

use crate::{CallerIdentity, ServerState};

use super::catalog::aion_tool_catalog;
use super::instructions::instructions;
use super::task_ref;
use super::tools::{authoring, control, describe, history, list, start, transcript};

/// The Aion tool implementation.
pub struct AionToolService {
    state: ServerState,
    catalog: ToolCatalog,
    instructions: String,
}

impl AionToolService {
    /// Build the service over the server's shared state.
    ///
    /// # Errors
    ///
    /// [`ServerError::Config`](crate::ServerError::Config) when a published
    /// tool schema fails to compile — a server defect, surfaced at startup
    /// rather than on a caller's first `tools/call`.
    pub fn new(state: ServerState) -> Result<Self, crate::ServerError> {
        let catalog = aion_tool_catalog().map_err(|error| crate::ServerError::Config {
            message: format!("the MCP tool catalog is invalid: {error}"),
        })?;
        Ok(Self {
            state,
            catalog,
            instructions: instructions(),
        })
    }
}

#[async_trait::async_trait]
impl ToolService for AionToolService {
    type Caller = CallerIdentity;

    fn catalog(&self) -> &ToolCatalog {
        &self.catalog
    }

    fn instructions(&self) -> &str {
        &self.instructions
    }

    fn task_decision(&self, call: &ToolCall) -> ToolTaskDecision {
        if start::awaits_completion(call) {
            ToolTaskDecision::Task
        } else {
            ToolTaskDecision::Inline
        }
    }

    async fn call(
        &self,
        caller: &Self::Caller,
        call: ToolCall,
    ) -> Result<ToolOutcome, ToolFailure> {
        match call.name.as_str() {
            "describe_run" => describe::describe_run(&self.state, caller, &call).await,
            "read_transcript" => transcript::read_transcript(&self.state, caller, &call).await,
            "read_history" => history::read_history(&self.state, caller, &call).await,
            "list_runs" => list::list_runs(&self.state, caller, &call).await,
            "query" => control::query(&self.state, caller, &call).await,
            "start_run" => start::start_run(&self.state, caller, &call).await,
            "signal" => control::signal(&self.state, caller, &call).await,
            "cancel" => control::cancel(&self.state, caller, &call).await,
            "list_documents" => authoring::list_documents(&self.state, caller, &call).await,
            "read_document" => authoring::read_document(&self.state, caller, &call).await,
            "check_document" => authoring::check_document(&self.state, caller, &call).await,
            "save_document" => authoring::save_document(&self.state, caller, &call).await,
            "deploy_document" => authoring::deploy_document(&self.state, caller, &call).await,
            // Unreachable through the dispatcher, which refuses an unpublished
            // name with `-32602` before this is called. Handled anyway: a
            // catch-all that cannot be reached is still the difference between
            // a legible refusal and a panic if the catalog and this match ever
            // drift apart.
            other => Err(ToolFailure::new(
                format!("`{other}` is not a tool this server implements"),
                json!({ "code": "unknown_tool", "tool": other }),
            )),
        }
    }

    fn task_handle(&self, call: &ToolCall, outcome: &ToolOutcome) -> Option<String> {
        task_ref::task_handle(call, outcome)
    }

    async fn resolve_task(
        &self,
        caller: &Self::Caller,
        task_id: &str,
    ) -> Result<TaskProjection, TaskResolveError> {
        task_ref::resolve_task(&self.state, caller, task_id).await
    }
}