crank/crates/crank-community-mcp/src/tool_search.rs

use std::{collections::BTreeSet, sync::Arc};

use axum::{http::StatusCode, response::Response};
use crank_core::{ToolAccessMode, search_tool_catalog};
use crank_registry::PublishedAgentCatalog;
use crank_trace::{ErrorCategory, Stage, StageOutcome};
use serde::Deserialize;
use serde_json::{Value, json};

use crate::{
    app::{
        AppState, handle_tool_call, resolve_generated_tool, take_confirmation_token,
        tool_error_response,
    },
    auth::VerifiedMachineCredential,
    jsonrpc::{jsonrpc_error, jsonrpc_result, request_id},
    manifest::{CALL_TOOL_NAME, SEARCH_TOOLS_NAME, searchable_tools},
    session::SessionState,
    tool_error::generic_tool_error_contract,
    transport::{ResponseMode, transport_response},
};

#[derive(Debug, Deserialize)]
struct SearchToolsArguments {
    query: String,
    #[serde(default)]
    group_ids: Vec<String>,
    max_results: Option<usize>,
}

#[derive(Debug, Deserialize)]
struct ProxyToolCallArguments {
    name: String,
    #[serde(default)]
    arguments: Value,
    catalog_revision: String,
}

#[allow(clippy::too_many_arguments)]
pub(super) async fn handle_catalog_tool_call(
    state: Arc<AppState>,
    session: &SessionState,
    message: &Value,
    response_mode: ResponseMode,
    credential: &VerifiedMachineCredential,
    catalog: &PublishedAgentCatalog,
    tool_name: &str,
    arguments: Value,
    transport_request_id: &str,
) -> Response {
    match catalog.tool_selection_policy.mode {
        ToolAccessMode::Direct => {
            execute_catalog_tool(
                state,
                session,
                message,
                response_mode,
                credential,
                catalog,
                tool_name,
                arguments,
                transport_request_id,
            )
            .await
        }
        ToolAccessMode::Search if tool_name == SEARCH_TOOLS_NAME => {
            handle_search_tools(message, response_mode, session, catalog, arguments)
        }
        ToolAccessMode::Search if tool_name == CALL_TOOL_NAME => {
            let proxy: ProxyToolCallArguments = match serde_json::from_value(arguments) {
                Ok(proxy) => proxy,
                Err(error) => {
                    return invalid_arguments_response(
                        message,
                        response_mode,
                        &session.protocol_version,
                        error.to_string(),
                    );
                }
            };
            if proxy.catalog_revision != catalog_revision(catalog) {
                return tool_error_response(
                    message,
                    response_mode,
                    &session.protocol_version,
                    generic_tool_error_contract(
                        "catalog_revision_changed",
                        format!(
                            "catalog revision {} is no longer current",
                            proxy.catalog_revision
                        ),
                        transport_request_id,
                        true,
                        Some(
                            "Повторите search_tools и вызовите инструмент с новой версией каталога.",
                        ),
                    ),
                );
            }
            execute_catalog_tool(
                state,
                session,
                message,
                response_mode,
                credential,
                catalog,
                &proxy.name,
                proxy.arguments,
                transport_request_id,
            )
            .await
        }
        ToolAccessMode::Search => {
            tool_not_found_response(message, response_mode, &session.protocol_version, tool_name)
        }
    }
}

#[allow(clippy::too_many_arguments)]
async fn execute_catalog_tool(
    state: Arc<AppState>,
    session: &SessionState,
    message: &Value,
    response_mode: ResponseMode,
    credential: &VerifiedMachineCredential,
    catalog: &PublishedAgentCatalog,
    tool_name: &str,
    mut arguments: Value,
    transport_request_id: &str,
) -> Response {
    let resolve_span = Stage::McpToolsResolve.span();
    let resolved = resolve_span.in_scope(|| resolve_generated_tool(&catalog.tools, tool_name));
    let resolved = match resolved {
        Some(resolved) => {
            StageOutcome::Success.record(&resolve_span);
            drop(resolve_span);
            resolved
        }
        None => {
            StageOutcome::Error.record(&resolve_span);
            ErrorCategory::Catalog.record(&resolve_span);
            drop(resolve_span);
            return tool_not_found_response(
                message,
                response_mode,
                &session.protocol_version,
                tool_name,
            );
        }
    };
    let confirmation_token = take_confirmation_token(&mut arguments);
    handle_tool_call(
        state,
        session,
        message,
        response_mode,
        credential,
        resolved,
        arguments,
        confirmation_token,
        transport_request_id,
    )
    .await
}

fn handle_search_tools(
    message: &Value,
    response_mode: ResponseMode,
    session: &SessionState,
    catalog: &PublishedAgentCatalog,
    arguments: Value,
) -> Response {
    let search: SearchToolsArguments = match serde_json::from_value(arguments) {
        Ok(search) => search,
        Err(error) => {
            return invalid_arguments_response(
                message,
                response_mode,
                &session.protocol_version,
                error.to_string(),
            );
        }
    };
    if search.query.trim().is_empty() {
        return invalid_arguments_response(
            message,
            response_mode,
            &session.protocol_version,
            "query must not be empty".to_owned(),
        );
    }
    let known_group_ids = catalog
        .tool_selection_policy
        .groups
        .iter()
        .map(|group| group.id.as_str())
        .collect::<BTreeSet<_>>();
    if let Some(group_id) = search
        .group_ids
        .iter()
        .find(|group_id| !known_group_ids.contains(group_id.as_str()))
    {
        return invalid_arguments_response(
            message,
            response_mode,
            &session.protocol_version,
            format!("unknown tool group {group_id}"),
        );
    }

    let configured_limit = catalog.tool_selection_policy.search.max_results;
    let requested_limit = search.max_results.unwrap_or(configured_limit).clamp(1, 20);
    let tools = search_tool_catalog(
        &searchable_tools(catalog),
        &search.query,
        &search.group_ids,
        requested_limit.min(configured_limit),
    )
    .into_iter()
    .map(|found| {
        json!({
            "name": found.tool.name,
            "title": found.tool.title,
            "description": found.tool.description,
            "inputSchema": found.tool.input_schema,
            "group_ids": found.tool.group_ids,
            "score": found.score,
        })
    })
    .collect::<Vec<_>>();
    let result = json!({
        "catalog_revision": catalog_revision(catalog),
        "tools": tools,
    });
    let text = serde_json::to_string_pretty(&result).unwrap_or_else(|_| result.to_string());

    transport_response(
        StatusCode::OK,
        jsonrpc_result(
            request_id(message),
            json!({
                "content": [{"type": "text", "text": text}],
                "structuredContent": result,
                "isError": false
            }),
        ),
        response_mode,
        None,
        Some(&session.protocol_version),
    )
}

fn catalog_revision(catalog: &PublishedAgentCatalog) -> String {
    format!("agent-version-{}", catalog.agent_version)
}

fn invalid_arguments_response(
    message: &Value,
    response_mode: ResponseMode,
    protocol_version: &str,
    detail: String,
) -> Response {
    transport_response(
        StatusCode::OK,
        jsonrpc_error(request_id(message), -32602, detail),
        response_mode,
        None,
        Some(protocol_version),
    )
}

fn tool_not_found_response(
    message: &Value,
    response_mode: ResponseMode,
    protocol_version: &str,
    tool_name: &str,
) -> Response {
    transport_response(
        StatusCode::OK,
        jsonrpc_error(
            request_id(message),
            -32602,
            format!("tool {tool_name} was not found"),
        ),
        response_mode,
        None,
        Some(protocol_version),
    )
}