aion-server 0.31.0

//! The assistant tool service: the editing loop, one session, no tasks.

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

use crate::ServerState;

use super::caller::AssistantSessionCaller;
use super::catalog::{
    ASSISTANT_CONTEXT_TOOL, ASSISTANT_DOCUMENT_CHECK_TOOL, ASSISTANT_DOCUMENT_EDIT_TOOL,
    assistant_tool_catalog,
};
use super::context_tool;
use super::document_check_tool;
use super::document_edit_tool;

/// The tool implementation behind `/assistant/mcp`.
pub(crate) struct AssistantToolService {
    state: ServerState,
    catalog: ToolCatalog,
    instructions: String,
}

impl AssistantToolService {
    /// Build the service over the server's shared state.
    ///
    /// # Errors
    ///
    /// [`CatalogError`] when the published schema fails to compile.
    pub(crate) fn new(state: ServerState) -> Result<Self, CatalogError> {
        Ok(Self {
            state,
            catalog: assistant_tool_catalog()?,
            instructions: INSTRUCTIONS.to_owned(),
        })
    }
}

#[async_trait::async_trait]
impl ToolService for AssistantToolService {
    type Caller = AssistantSessionCaller;

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

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

    /// Nothing here is task-shaped. Each tool answers within its own request —
    /// a read, a validate-and-append, a synchronous check; a task would be
    /// durable work outliving the request, and there is none.
    fn task_decision(&self, _call: &ToolCall) -> ToolTaskDecision {
        ToolTaskDecision::Inline
    }

    async fn call(
        &self,
        caller: &Self::Caller,
        call: ToolCall,
    ) -> Result<ToolOutcome, ToolFailure> {
        match call.name.as_str() {
            ASSISTANT_CONTEXT_TOOL => {
                context_tool::assistant_context(self.state.assistant_sessions(), caller).await
            }
            ASSISTANT_DOCUMENT_EDIT_TOOL => {
                document_edit_tool::assistant_document_edit(
                    self.state.assistant_sessions(),
                    caller,
                    serde_json::Value::Object(call.arguments),
                )
                .await
            }
            ASSISTANT_DOCUMENT_CHECK_TOOL => {
                document_check_tool::assistant_document_check(&self.state, caller).await
            }
            // Unreachable through the dispatcher, which refuses an unpublished
            // name before this runs. Handled anyway, and handled by NAMING the
            // catalogue: an agent that tried a general-surface tool here must be
            // told it is on the wrong route rather than that the tool does not
            // exist, which would be false.
            other => Err(ToolFailure::new(
                format!(
                    "`{other}` is not a tool on the assistant route. This route publishes \
                     `{ASSISTANT_CONTEXT_TOOL}`, `{ASSISTANT_DOCUMENT_EDIT_TOOL}` and \
                     `{ASSISTANT_DOCUMENT_CHECK_TOOL}`; the workflow tools are on this server's \
                     general MCP endpoint, under a different credential."
                ),
                json!({ "code": "unknown_tool", "tool": other }),
            )),
        }
    }

    /// No call here is task-shaped, so none has a handle to mint.
    fn task_handle(&self, _call: &ToolCall, _outcome: &ToolOutcome) -> Option<String> {
        None
    }

    /// This route issues no task handles, so every handle presented to it was
    /// issued by something else and is refused as unknown rather than looked up.
    async fn resolve_task(
        &self,
        _caller: &Self::Caller,
        task_id: &str,
    ) -> Result<TaskProjection, TaskResolveError> {
        Err(TaskResolveError::Unknown {
            task_id: task_id.to_owned(),
        })
    }
}

/// What `server/discover` teaches an agent about this route.
const INSTRUCTIONS: &str = "This endpoint belongs to ONE assistant conversation — the one whose bearer you called with. \
     It publishes the editing loop, three tools about that conversation and nothing else.\n\n\
     `assistant_context` answers what the operator has on screen in the Aion console right now: \
     the page, the concepts that page explains, and the document their editor is showing with \
     any selection in it. Call it before asking the operator where anything is. It returns the \
     document's path, so you never need to ask which directory a repository is in or which file \
     is open — asking for either is asking for something you were already given.\n\n\
     `assistant_document_edit` edits that document in place, as replace-exactly-once operations \
     quoted from what you read. The edits land in the operator's editor as you make them, and \
     the operator keeps or reverts them — so make small, named changes one concept at a time, \
     and never read or save files to reach this document: it is the operator's buffer, and \
     saving stays the operator's act.\n\n\
     `assistant_document_check` runs the editor's own AWL check over that document as it now \
     stands, your edits included. Check after you edit and fix what it reports before telling \
     the operator what you changed.\n\n\
     None of these name a session or a file: your credential already names the conversation, and \
     there is no way through this endpoint to reach another one. This server's workflow tools \
     are NOT here; they are on its general MCP endpoint under a different credential.";

#[cfg(test)]
mod tests {
    use super::super::catalog::ASSISTANT_CONTEXT_TOOL;
    use super::INSTRUCTIONS;

    /// The instructions have to carry the one behaviour Tom named: the agent
    /// must never ask which directory the repository is in. If the sentence goes,
    /// the behaviour goes with it, because prose is the whole mechanism here.
    #[test]
    fn the_instructions_tell_the_agent_not_to_ask_where_the_repository_is() {
        assert!(INSTRUCTIONS.contains(ASSISTANT_CONTEXT_TOOL));
        assert!(
            INSTRUCTIONS.contains("which directory a repository is in"),
            "the instruction that stops the agent asking for the repo path is load-bearing"
        );
    }
}