агенты: добавить поиск инструментов по каталогу

This commit is contained in:
2026-07-21 13:12:46 +03:00
parent 63f8ee333f
commit 99bd05c145
35 changed files with 2088 additions and 155 deletions
+25 -38
View File
@@ -43,7 +43,7 @@ use crate::{
DEFAULT_PROTOCOL_VERSION, is_notification, is_request, is_response, jsonrpc_error,
jsonrpc_result, method_name, negotiated_protocol_version, params, request_id,
},
manifest::tool_definitions,
manifest::catalog_tool_definitions,
rate_limit::{
enforce_post_rate_limit, enforce_transport_rate_limit, rate_limited_jsonrpc_response,
rate_limited_status_response,
@@ -53,6 +53,7 @@ use crate::{
ToolErrorContract, generic_tool_error_contract, runtime_error_code,
tool_error_contract_from_runtime, tool_error_text, tool_error_value,
},
tool_search::handle_catalog_tool_call,
transport::{
AllowedOrigins, HEADER_MCP_SESSION_ID, ResponseMode, json_response,
negotiate_post_response_mode, protocol_version_from_headers, resolve_request_id,
@@ -98,7 +99,7 @@ struct ApprovalDecisionPayload {
}
#[derive(Clone)]
struct ResolvedToolCall {
pub(super) struct ResolvedToolCall {
tool: PublishedAgentTool,
}
@@ -679,11 +680,11 @@ async fn mcp_post(
match state
.catalog
.list_tools(&session.workspace_slug, &session.agent_slug)
.get_catalog(&session.workspace_slug, &session.agent_slug)
.await
{
Ok(tools) => {
let definitions = tools.iter().flat_map(tool_definitions).collect::<Vec<_>>();
Ok(catalog) => {
let definitions = catalog_tool_definitions(&catalog);
transport_response(
StatusCode::OK,
@@ -715,45 +716,31 @@ async fn mcp_post(
}
};
let mut arguments = if tool_call_params.arguments.is_null() {
let arguments = if tool_call_params.arguments.is_null() {
json!({})
} else {
tool_call_params.arguments
};
let confirmation_token = take_confirmation_token(&mut arguments);
match state
.catalog
.list_tools(&session.workspace_slug, &session.agent_slug)
.get_catalog(&session.workspace_slug, &session.agent_slug)
.await
{
Ok(tools) => match resolve_generated_tool(&tools, &tool_call_params.name) {
Some(resolved) => {
handle_tool_call(
state.clone(),
&session,
&message,
response_mode,
&credential,
resolved,
arguments,
confirmation_token,
&transport_request_id,
)
.await
}
None => transport_response(
StatusCode::OK,
jsonrpc_error(
request_id(&message),
-32602,
format!("tool {} was not found", tool_call_params.name),
),
Ok(catalog) => {
handle_catalog_tool_call(
state.clone(),
&session,
&message,
response_mode,
None,
Some(&session.protocol_version),
),
},
&credential,
&catalog,
&tool_call_params.name,
arguments,
&transport_request_id,
)
.await
}
Err(error) => internal_jsonrpc_error(&message, error),
}
}
@@ -785,7 +772,7 @@ async fn mcp_post(
}
#[allow(clippy::too_many_arguments)]
async fn handle_tool_call(
pub(super) async fn handle_tool_call(
state: Arc<AppState>,
session: &SessionState,
message: &Value,
@@ -1381,7 +1368,7 @@ fn internal_jsonrpc_error(message: &Value, error: impl std::fmt::Display) -> Res
)
}
fn take_confirmation_token(arguments: &mut Value) -> Option<String> {
pub(super) fn take_confirmation_token(arguments: &mut Value) -> Option<String> {
let Value::Object(object) = arguments else {
return None;
};
@@ -1479,7 +1466,7 @@ fn success_tool_response(
)
}
fn tool_error_response(
pub(super) fn tool_error_response(
message: &Value,
response_mode: ResponseMode,
protocol_version: &str,
@@ -1517,7 +1504,7 @@ fn add_millis(timestamp: OffsetDateTime, millis: u64) -> OffsetDateTime {
timestamp + delta
}
fn resolve_generated_tool(
pub(super) fn resolve_generated_tool(
tools: &[PublishedAgentTool],
tool_name: &str,
) -> Option<ResolvedToolCall> {
+32 -23
View File
@@ -5,7 +5,7 @@ use std::{
};
use crank_core::{CacheScope, CoordinationStateStore, CoordinationStateValue};
use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
use crank_registry::{PostgresRegistry, PublishedAgentCatalog, PublishedAgentTool, RegistryError};
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, RwLock};
use tracing::{info, warn};
@@ -27,15 +27,14 @@ struct CatalogKey {
agent_slug: String,
}
#[derive(Default)]
struct CachedCatalog {
loaded_at: Option<Instant>,
tools: Vec<PublishedAgentTool>,
catalog: PublishedAgentCatalog,
}
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
struct CatalogSnapshot {
tools: Vec<PublishedAgentTool>,
catalog: PublishedAgentCatalog,
generated_at_ms: u64,
}
@@ -59,12 +58,23 @@ impl PublishedToolCatalog {
workspace_slug: &str,
agent_slug: &str,
) -> Result<Vec<PublishedAgentTool>, RegistryError> {
Ok(self.get_catalog(workspace_slug, agent_slug).await?.tools)
}
pub async fn get_catalog(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<PublishedAgentCatalog, RegistryError> {
self.refresh_if_stale(workspace_slug, agent_slug).await?;
let guard = self.cached.read().await;
Ok(guard
guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.clone())
.unwrap_or_default())
.map(|entry| entry.catalog.clone())
.ok_or_else(|| RegistryError::PublishedAgentNotFound {
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
})
}
async fn refresh_if_stale(
@@ -106,42 +116,41 @@ impl PublishedToolCatalog {
return Ok(());
}
if let Some((tools, age)) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
log_catalog_analysis(workspace_slug, agent_slug, "shared_cache", &tools);
if let Some((catalog, age)) = self.load_shared_snapshot(workspace_slug, agent_slug).await {
log_catalog_analysis(workspace_slug, agent_slug, "shared_cache", &catalog.tools);
let mut guard = self.cached.write().await;
guard.insert(
key,
CachedCatalog {
loaded_at: Instant::now().checked_sub(age),
tools,
catalog,
},
);
return Ok(());
}
let tools = match self
let catalog = match self
.registry
.get_published_agent_tools_by_slug(workspace_slug, agent_slug)
.get_published_agent_catalog_by_slug(workspace_slug, agent_slug)
.await
{
Ok(tools) => tools,
Err(RegistryError::PublishedAgentNotFound { .. }) => Vec::new(),
Ok(catalog) => catalog,
Err(error) => return Err(error),
};
log_catalog_analysis(workspace_slug, agent_slug, "postgres", &tools);
self.store_shared_snapshot(workspace_slug, agent_slug, &tools)
log_catalog_analysis(workspace_slug, agent_slug, "postgres", &catalog.tools);
self.store_shared_snapshot(workspace_slug, agent_slug, &catalog)
.await;
let mut guard = self.cached.write().await;
let previous_count = guard
.get(&key)
.map(|entry| entry.tools.len())
.map(|entry| entry.catalog.tools.len())
.unwrap_or_default();
guard.insert(
key,
CachedCatalog {
loaded_at: Some(Instant::now()),
tools,
catalog,
},
);
@@ -150,7 +159,7 @@ impl PublishedToolCatalog {
agent_slug,
published_tool_count = guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.len())
.map(|entry| entry.catalog.tools.len())
.unwrap_or_default(),
previous_published_tool_count = previous_count,
"published agent catalog refreshed"
@@ -163,7 +172,7 @@ impl PublishedToolCatalog {
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Option<(Vec<PublishedAgentTool>, Duration)> {
) -> Option<(PublishedAgentCatalog, Duration)> {
if self.refresh_interval.is_zero() {
return None;
}
@@ -179,21 +188,21 @@ impl PublishedToolCatalog {
};
let snapshot = serde_json::from_value::<CatalogSnapshot>(value.payload).ok()?;
let age = Duration::from_millis(now_unix_ms().saturating_sub(snapshot.generated_at_ms));
(age < self.refresh_interval).then_some((snapshot.tools, age))
(age < self.refresh_interval).then_some((snapshot.catalog, age))
}
async fn store_shared_snapshot(
&self,
workspace_slug: &str,
agent_slug: &str,
tools: &[PublishedAgentTool],
catalog: &PublishedAgentCatalog,
) {
let Some(ttl) = catalog_snapshot_ttl(self.refresh_interval) else {
return;
};
let payload = match serde_json::to_value(CatalogSnapshot {
tools: tools.to_vec(),
catalog: catalog.clone(),
generated_at_ms: now_unix_ms(),
}) {
Ok(payload) => payload,
+1
View File
@@ -8,6 +8,7 @@ pub mod manifest;
mod rate_limit;
pub mod session;
pub mod tool_error;
mod tool_search;
mod transport;
pub use app::{build_app, build_app_with_background_workers};
+103 -3
View File
@@ -1,10 +1,110 @@
use crank_core::{
HttpMethod, OperationSafetyClass, OperationSafetyPolicy, Target, ToolCatalogAnalysis,
ToolCatalogAnalysisError, ToolQualityCatalogTool, analyze_tool_catalog,
HttpMethod, OperationSafetyClass, OperationSafetyPolicy, SearchableTool, Target,
ToolAccessMode, ToolCatalogAnalysis, ToolCatalogAnalysisError, ToolQualityCatalogTool,
ToolSelectionPolicy, analyze_tool_catalog,
};
use crank_registry::PublishedAgentTool;
use crank_registry::{PublishedAgentCatalog, PublishedAgentTool};
use serde_json::{Value, json};
pub const SEARCH_TOOLS_NAME: &str = "search_tools";
pub const CALL_TOOL_NAME: &str = "call_tool";
pub fn catalog_tool_definitions(catalog: &PublishedAgentCatalog) -> Vec<Value> {
match catalog.tool_selection_policy.mode {
ToolAccessMode::Direct => catalog.tools.iter().flat_map(tool_definitions).collect(),
ToolAccessMode::Search => search_mode_tool_definitions(&catalog.tool_selection_policy),
}
}
pub fn searchable_tools(catalog: &PublishedAgentCatalog) -> Vec<SearchableTool> {
catalog
.tools
.iter()
.map(|tool| {
let groups = catalog
.tool_selection_policy
.groups
.iter()
.filter(|group| group.tool_names.iter().any(|name| name == &tool.tool_name))
.collect::<Vec<_>>();
SearchableTool {
name: tool.tool_name.clone(),
title: tool.tool_title.clone(),
description: tool.tool_description.clone(),
input_schema: tool_definitions(tool)
.into_iter()
.next()
.and_then(|definition| definition.get("inputSchema").cloned())
.unwrap_or_else(|| json!({"type": "object"})),
group_ids: groups.iter().map(|group| group.id.clone()).collect(),
group_context: groups
.iter()
.map(|group| format!("{} {}", group.name, group.description))
.collect::<Vec<_>>()
.join(" "),
}
})
.collect()
}
fn search_mode_tool_definitions(policy: &ToolSelectionPolicy) -> Vec<Value> {
let group_catalog = if policy.groups.is_empty() {
"Разделы каталога не заданы; выполняйте поиск по всему каталогу.".to_owned()
} else {
policy
.groups
.iter()
.map(|group| format!("{}{}: {}", group.id, group.name, group.description))
.collect::<Vec<_>>()
.join("\n")
};
vec![
tool_definition(
SEARCH_TOOLS_NAME,
"Подобрать инструменты",
&format!(
"Находит подходящие инструменты агента и возвращает их полные входные схемы. Доступные разделы:\n{group_catalog}"
),
json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Краткое описание требуемого действия"
},
"group_ids": {
"type": "array",
"items": {"type": "string"},
"description": "Необязательный список разделов каталога"
},
"max_results": {
"type": "integer",
"minimum": 1,
"maximum": 20,
"default": policy.search.max_results
}
},
"required": ["query"]
}),
),
tool_definition(
CALL_TOOL_NAME,
"Вызвать найденный инструмент",
"Вызывает инструмент, ранее найденный через search_tools. Передайте имя, аргументы по полученной схеме и версию каталога из результата поиска.",
json!({
"type": "object",
"properties": {
"name": {"type": "string"},
"arguments": {"type": "object"},
"catalog_revision": {"type": "string"}
},
"required": ["name", "arguments", "catalog_revision"]
}),
),
]
}
pub fn tool_definitions(tool: &PublishedAgentTool) -> Vec<Value> {
let safety = effective_safety_policy(tool);
let requires_confirmation = safety.class.requires_confirmation();
@@ -0,0 +1,276 @@
use std::{collections::BTreeSet, sync::Arc};
use axum::{http::StatusCode, response::Response};
use crank_core::{ToolAccessMode, search_tool_catalog};
use crank_registry::PublishedAgentCatalog;
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 Some(resolved) = resolve_generated_tool(&catalog.tools, tool_name) else {
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),
)
}