feat: add agent publishing foundation

This commit is contained in:
a.tolmachev
2026-03-29 22:10:15 +03:00
parent d757adb192
commit 7b7699cf8b
18 changed files with 1520 additions and 105 deletions
+103 -35
View File
@@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration};
use axum::{
Json, Router,
extract::State,
extract::{Path, State},
http::{
HeaderMap, HeaderValue, StatusCode,
header::{self, ACCEPT},
@@ -10,7 +10,7 @@ use axum::{
response::{IntoResponse, Response},
routing::get,
};
use crank_registry::{PostgresRegistry, RegistryOperation};
use crank_registry::{PostgresRegistry, PublishedAgentTool};
use crank_runtime::{RuntimeError, RuntimeExecutor, RuntimeOperation};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
@@ -54,6 +54,12 @@ struct ToolCallParams {
arguments: Value,
}
#[derive(Clone, Debug, Deserialize)]
struct AgentRoutePath {
workspace_slug: String,
agent_slug: String,
}
pub fn build_app(
registry: PostgresRegistry,
refresh_interval: Duration,
@@ -68,8 +74,10 @@ pub fn build_app(
Router::new()
.route("/health", get(health))
.route("/", get(mcp_get).post(mcp_post).delete(mcp_delete))
.route("/mcp", get(mcp_get).post(mcp_post).delete(mcp_delete))
.route(
"/v1/{workspace_slug}/{agent_slug}",
get(mcp_get).post(mcp_post).delete(mcp_delete),
)
.with_state(state)
}
@@ -80,25 +88,37 @@ async fn health() -> Json<Value> {
}))
}
async fn mcp_get() -> Response {
async fn mcp_get(Path(_path): Path<AgentRoutePath>) -> Response {
StatusCode::METHOD_NOT_ALLOWED.into_response()
}
async fn mcp_delete(State(state): State<Arc<AppState>>, headers: HeaderMap) -> Response {
async fn mcp_delete(
Path(path): Path<AgentRoutePath>,
State(state): State<Arc<AppState>>,
headers: HeaderMap,
) -> Response {
match session_id_from_headers(&headers) {
Ok(Some(session_id)) => {
if state.sessions.delete(&session_id).await {
StatusCode::NO_CONTENT.into_response()
} else {
StatusCode::NOT_FOUND.into_response()
Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
Some(session)
if session.workspace_slug == path.workspace_slug
&& session.agent_slug == path.agent_slug =>
{
if state.sessions.delete(&session_id).await {
StatusCode::NO_CONTENT.into_response()
} else {
StatusCode::NOT_FOUND.into_response()
}
}
}
Some(_) => StatusCode::NOT_FOUND.into_response(),
None => StatusCode::NOT_FOUND.into_response(),
},
Ok(None) => StatusCode::BAD_REQUEST.into_response(),
Err(status) => status.into_response(),
}
}
async fn mcp_post(
Path(path): Path<AgentRoutePath>,
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(message): Json<Value>,
@@ -121,12 +141,15 @@ async fn mcp_post(
};
match method_name(&message) {
Some("initialize") if is_request(&message) => handle_initialize(state, &message).await,
Some("initialize") if is_request(&message) => {
handle_initialize(state, &path, &message).await
}
Some("notifications/initialized") if is_notification(&message) => {
handle_initialized_notification(state, &headers).await
handle_initialized_notification(state, &path, &headers).await
}
Some("ping") if is_request(&message) => {
let session = match require_initialized_session(&state, &headers, &message).await {
let session = match require_initialized_session(&state, &path, &headers, &message).await
{
Ok(session) => session,
Err(response) => return response,
};
@@ -142,12 +165,17 @@ async fn mcp_post(
)
}
Some("tools/list") if is_request(&message) => {
let session = match require_initialized_session(&state, &headers, &message).await {
let session = match require_initialized_session(&state, &path, &headers, &message).await
{
Ok(session) => session,
Err(response) => return response,
};
match state.catalog.list_tools().await {
match state
.catalog
.list_tools(&session.workspace_slug, &session.agent_slug)
.await
{
Ok(tools) => {
let definitions = tools.iter().map(tool_definition).collect::<Vec<_>>();
@@ -162,7 +190,8 @@ async fn mcp_post(
}
}
Some("tools/call") if is_request(&message) => {
let session = match require_initialized_session(&state, &headers, &message).await {
let session = match require_initialized_session(&state, &path, &headers, &message).await
{
Ok(session) => session,
Err(response) => return response,
};
@@ -184,11 +213,19 @@ async fn mcp_post(
tool_call_params.arguments
};
match state.catalog.get_tool(&tool_call_params.name).await {
Ok(Some(operation)) => {
match state
.catalog
.get_tool(
&session.workspace_slug,
&session.agent_slug,
&tool_call_params.name,
)
.await
{
Ok(Some(tool)) => {
match state
.runtime
.execute(&RuntimeOperation::from(operation), &arguments)
.execute(&runtime_operation(tool), &arguments)
.await
{
Ok(output) => json_response(
@@ -270,7 +307,11 @@ async fn mcp_post(
}
}
async fn handle_initialize(state: Arc<AppState>, message: &Value) -> Response {
async fn handle_initialize(
state: Arc<AppState>,
path: &AgentRoutePath,
message: &Value,
) -> Response {
let initialize_params: InitializeParams = match serde_json::from_value(params(message)) {
Ok(value) => value,
Err(error) => {
@@ -298,7 +339,10 @@ async fn handle_initialize(state: Arc<AppState>, message: &Value) -> Response {
Some(DEFAULT_PROTOCOL_VERSION),
);
};
let session_id = state.sessions.create(protocol_version).await;
let session_id = state
.sessions
.create(protocol_version, &path.workspace_slug, &path.agent_slug)
.await;
json_response(
StatusCode::OK,
@@ -322,15 +366,26 @@ async fn handle_initialize(state: Arc<AppState>, message: &Value) -> Response {
)
}
async fn handle_initialized_notification(state: Arc<AppState>, headers: &HeaderMap) -> Response {
async fn handle_initialized_notification(
state: Arc<AppState>,
path: &AgentRoutePath,
headers: &HeaderMap,
) -> Response {
match session_id_from_headers(headers) {
Ok(Some(session_id)) => {
if state.sessions.mark_initialized(&session_id).await {
StatusCode::ACCEPTED.into_response()
} else {
StatusCode::NOT_FOUND.into_response()
Ok(Some(session_id)) => match state.sessions.get(&session_id).await {
Some(session)
if session.workspace_slug == path.workspace_slug
&& session.agent_slug == path.agent_slug =>
{
if state.sessions.mark_initialized(&session_id).await {
StatusCode::ACCEPTED.into_response()
} else {
StatusCode::NOT_FOUND.into_response()
}
}
}
Some(_) => StatusCode::NOT_FOUND.into_response(),
None => StatusCode::NOT_FOUND.into_response(),
},
Ok(None) => StatusCode::BAD_REQUEST.into_response(),
Err(status) => status.into_response(),
}
@@ -338,6 +393,7 @@ async fn handle_initialized_notification(state: Arc<AppState>, headers: &HeaderM
async fn require_initialized_session(
state: &Arc<AppState>,
path: &AgentRoutePath,
headers: &HeaderMap,
message: &Value,
) -> Result<SessionState, Response> {
@@ -351,6 +407,10 @@ async fn require_initialized_session(
return Err(StatusCode::NOT_FOUND.into_response());
};
if session.workspace_slug != path.workspace_slug || session.agent_slug != path.agent_slug {
return Err(StatusCode::NOT_FOUND.into_response());
}
if !session.initialized {
return Err(json_response(
StatusCode::OK,
@@ -471,15 +531,23 @@ fn json_response(
response
}
fn tool_definition(operation: &RegistryOperation) -> Value {
fn tool_definition(tool: &PublishedAgentTool) -> Value {
json!({
"name": operation.name,
"title": operation.tool_description.title,
"description": operation.tool_description.description,
"inputSchema": schema_to_json_schema(&operation.input_schema)
"name": tool.tool_name,
"title": tool.tool_title,
"description": tool.tool_description,
"inputSchema": schema_to_json_schema(&tool.operation.input_schema)
})
}
fn runtime_operation(tool: PublishedAgentTool) -> RuntimeOperation {
let mut operation = RuntimeOperation::from(tool.operation);
operation.tool_name = tool.tool_name;
operation.tool_description.title = tool.tool_title;
operation.tool_description.description = tool.tool_description;
operation
}
fn schema_to_json_schema(schema: &crank_schema::Schema) -> Value {
match schema.kind {
crank_schema::SchemaKind::Object => {