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 => {
+74 -21
View File
@@ -4,7 +4,7 @@ use std::{
time::{Duration, Instant},
};
use crank_registry::{PostgresRegistry, RegistryError, RegistryOperation};
use crank_registry::{PostgresRegistry, PublishedAgentTool, RegistryError};
use tokio::sync::RwLock;
use tracing::info;
@@ -12,14 +12,20 @@ use tracing::info;
pub struct PublishedToolCatalog {
registry: PostgresRegistry,
refresh_interval: Duration,
cached: Arc<RwLock<CachedCatalog>>,
cached: Arc<RwLock<HashMap<CatalogKey, CachedCatalog>>>,
}
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct CatalogKey {
workspace_slug: String,
agent_slug: String,
}
#[derive(Default)]
struct CachedCatalog {
loaded_at: Option<Instant>,
tools: Vec<RegistryOperation>,
tools_by_name: HashMap<String, RegistryOperation>,
tools: Vec<PublishedAgentTool>,
tools_by_name: HashMap<String, PublishedAgentTool>,
}
impl PublishedToolCatalog {
@@ -27,30 +33,47 @@ impl PublishedToolCatalog {
Self {
registry,
refresh_interval,
cached: Arc::new(RwLock::new(CachedCatalog::default())),
cached: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn list_tools(&self) -> Result<Vec<RegistryOperation>, RegistryError> {
self.refresh_if_stale().await?;
pub async fn list_tools(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<Vec<PublishedAgentTool>, RegistryError> {
self.refresh_if_stale(workspace_slug, agent_slug).await?;
let guard = self.cached.read().await;
Ok(guard.tools.clone())
Ok(guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.clone())
.unwrap_or_default())
}
pub async fn get_tool(
&self,
workspace_slug: &str,
agent_slug: &str,
tool_name: &str,
) -> Result<Option<RegistryOperation>, RegistryError> {
self.refresh_if_stale().await?;
) -> Result<Option<PublishedAgentTool>, RegistryError> {
self.refresh_if_stale(workspace_slug, agent_slug).await?;
let guard = self.cached.read().await;
Ok(guard.tools_by_name.get(tool_name).cloned())
Ok(guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.and_then(|entry| entry.tools_by_name.get(tool_name))
.cloned())
}
async fn refresh_if_stale(&self) -> Result<(), RegistryError> {
async fn refresh_if_stale(
&self,
workspace_slug: &str,
agent_slug: &str,
) -> Result<(), RegistryError> {
let key = CatalogKey::new(workspace_slug, agent_slug);
let should_refresh = {
let guard = self.cached.read().await;
match guard.loaded_at {
match guard.get(&key).and_then(|entry| entry.loaded_at) {
Some(loaded_at) => loaded_at.elapsed() >= self.refresh_interval,
None => true,
}
@@ -60,25 +83,55 @@ impl PublishedToolCatalog {
return Ok(());
}
let tools = self.registry.list_published_operations().await?;
let tools = match self
.registry
.get_published_agent_tools_by_slug(workspace_slug, agent_slug)
.await
{
Ok(tools) => tools,
Err(RegistryError::PublishedAgentNotFound { .. }) => Vec::new(),
Err(error) => return Err(error),
};
let tools_by_name = tools
.iter()
.cloned()
.map(|operation| (operation.name.clone(), operation))
.map(|tool| (tool.tool_name.clone(), tool))
.collect::<HashMap<_, _>>();
let mut guard = self.cached.write().await;
let previous_count = guard.tools.len();
let previous_count = guard
.get(&key)
.map(|entry| entry.tools.len())
.unwrap_or_default();
guard.loaded_at = Some(Instant::now());
guard.tools = tools;
guard.tools_by_name = tools_by_name;
guard.insert(
key,
CachedCatalog {
loaded_at: Some(Instant::now()),
tools,
tools_by_name,
},
);
info!(
published_tool_count = guard.tools.len(),
workspace_slug,
agent_slug,
published_tool_count = guard
.get(&CatalogKey::new(workspace_slug, agent_slug))
.map(|entry| entry.tools.len())
.unwrap_or_default(),
previous_published_tool_count = previous_count,
"published tool catalog refreshed"
"published agent catalog refreshed"
);
Ok(())
}
}
impl CatalogKey {
fn new(workspace_slug: &str, agent_slug: &str) -> Self {
Self {
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
}
}
}
+186 -21
View File
@@ -51,12 +51,14 @@ mod tests {
use axum::{Json, Router, http::header, routing::post};
use crank_adapter_grpc::test_support as grpc_test_support;
use crank_core::{
DescriptorId, ExecutionConfig, GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod,
Operation, OperationId, OperationStatus, Protocol, RestTarget, Target, ToolDescription,
WorkspaceId,
Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, DescriptorId,
ExecutionConfig, GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod, Operation,
OperationId, OperationStatus, Protocol, RestTarget, Target, ToolDescription, WorkspaceId,
};
use crank_mapping::{MappingRule, MappingSet};
use crank_registry::{PostgresRegistry, PublishRequest};
use crank_registry::{
CreateAgentRequest, PostgresRegistry, PublishAgentRequest, PublishRequest,
};
use crank_schema::{Schema, SchemaKind};
use serde_json::{Value, json};
use sqlx::{Executor, postgres::PgPoolOptions};
@@ -68,6 +70,10 @@ mod tests {
WorkspaceId::new("ws_default")
}
fn test_workspace_slug() -> &'static str {
"default"
}
#[tokio::test]
async fn initializes_lists_and_calls_published_tool_via_mcp() {
let registry = test_registry().await;
@@ -88,6 +94,7 @@ mod tests {
})
.await
.unwrap();
publish_agent_for_operation(&registry, &operation, "sales-rest").await;
let base_url = spawn_mcp_server(build_app(
registry,
@@ -96,11 +103,12 @@ mod tests {
))
.await;
let client = reqwest::Client::new();
let initialized_session = initialize_session(&client, &base_url).await;
let mcp_url = agent_mcp_url(&base_url, "sales-rest");
let initialized_session = initialize_session(&client, &mcp_url).await;
let tools = post_jsonrpc(
&client,
&base_url,
&mcp_url,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
@@ -112,7 +120,7 @@ mod tests {
.await;
let call_result = post_jsonrpc(
&client,
&base_url,
&mcp_url,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
@@ -156,6 +164,7 @@ mod tests {
})
.await
.unwrap();
publish_agent_for_operation(&registry, &operation, "sales-graphql").await;
let base_url = spawn_mcp_server(build_app(
registry,
@@ -164,11 +173,12 @@ mod tests {
))
.await;
let client = reqwest::Client::new();
let initialized_session = initialize_session(&client, &base_url).await;
let mcp_url = agent_mcp_url(&base_url, "sales-graphql");
let initialized_session = initialize_session(&client, &mcp_url).await;
let call_result = post_jsonrpc(
&client,
&base_url,
&mcp_url,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
@@ -211,6 +221,7 @@ mod tests {
})
.await
.unwrap();
publish_agent_for_operation(&registry, &operation, "sales-grpc").await;
let base_url = spawn_mcp_server(build_app(
registry,
@@ -219,11 +230,12 @@ mod tests {
))
.await;
let client = reqwest::Client::new();
let initialized_session = initialize_session(&client, &base_url).await;
let mcp_url = agent_mcp_url(&base_url, "sales-grpc");
let initialized_session = initialize_session(&client, &mcp_url).await;
let call_result = post_jsonrpc(
&client,
&base_url,
&mcp_url,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
@@ -256,8 +268,9 @@ mod tests {
))
.await;
let client = reqwest::Client::new();
let mcp_url = agent_mcp_url(&base_url, "sales-init");
let initialize_response = client
.post(format!("{base_url}/mcp"))
.post(&mcp_url)
.header(header::ACCEPT, "application/json, text/event-stream")
.json(&json!({
"jsonrpc": "2.0",
@@ -278,7 +291,7 @@ mod tests {
.unwrap()
.to_owned();
let tools_list = client
.post(format!("{base_url}/mcp"))
.post(&mcp_url)
.header(header::ACCEPT, "application/json, text/event-stream")
.header("MCP-Session-Id", session_id)
.header("MCP-Protocol-Version", "2025-11-25")
@@ -309,11 +322,12 @@ mod tests {
))
.await;
let client = reqwest::Client::new();
let initialized_session = initialize_session(&client, &base_url).await;
let mcp_url = agent_mcp_url(&base_url, "sales-refresh");
let initialized_session = initialize_session(&client, &mcp_url).await;
let before_publish = post_jsonrpc(
&client,
&base_url,
&mcp_url,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
@@ -339,10 +353,11 @@ mod tests {
})
.await
.unwrap();
publish_agent_for_operation(&registry, &operation, "sales-refresh").await;
let after_publish = post_jsonrpc(
&client,
&base_url,
&mcp_url,
Some(&initialized_session),
json!({
"jsonrpc": "2.0",
@@ -360,9 +375,87 @@ mod tests {
);
}
async fn initialize_session(client: &reqwest::Client, base_url: &str) -> String {
#[tokio::test]
async fn lists_only_bound_tools_for_agent_context() {
let registry = test_registry().await;
let upstream_base_url = spawn_upstream_server().await;
let operation_a = test_operation(&upstream_base_url, "crm_create_lead");
let operation_b = test_operation(&upstream_base_url, "crm_update_lead");
for operation in [&operation_a, &operation_b] {
registry
.create_operation(&test_workspace_id(), operation, Some("alice"))
.await
.unwrap();
registry
.publish_operation(PublishRequest {
workspace_id: &test_workspace_id(),
operation_id: &operation.id,
version: 1,
published_at: "2026-03-26T10:00:00Z",
published_by: Some("alice"),
})
.await
.unwrap();
}
publish_agent_with_bindings(
&registry,
"sales-a",
vec![binding_for_operation(&operation_a)],
)
.await;
publish_agent_with_bindings(
&registry,
"sales-b",
vec![binding_for_operation(&operation_b)],
)
.await;
let base_url = spawn_mcp_server(build_app(
registry,
Duration::from_millis(0),
Some("https://crank.example.com".to_owned()),
))
.await;
let client = reqwest::Client::new();
let agent_a_url = agent_mcp_url(&base_url, "sales-a");
let agent_b_url = agent_mcp_url(&base_url, "sales-b");
let session_a = initialize_session(&client, &agent_a_url).await;
let session_b = initialize_session(&client, &agent_b_url).await;
let tools_a = post_jsonrpc(
&client,
&agent_a_url,
Some(&session_a),
json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/list",
"params": {}
}),
)
.await;
let tools_b = post_jsonrpc(
&client,
&agent_b_url,
Some(&session_b),
json!({
"jsonrpc": "2.0",
"id": 2,
"method": "tools/list",
"params": {}
}),
)
.await;
assert_eq!(tools_a["result"]["tools"][0]["name"], "crm_create_lead");
assert_eq!(tools_b["result"]["tools"][0]["name"], "crm_update_lead");
}
async fn initialize_session(client: &reqwest::Client, mcp_url: &str) -> String {
let initialize_response = client
.post(format!("{base_url}/mcp"))
.post(mcp_url)
.header(header::ACCEPT, "application/json, text/event-stream")
.json(&json!({
"jsonrpc": "2.0",
@@ -384,7 +477,7 @@ mod tests {
.to_owned();
let initialized_response = client
.post(format!("{base_url}/mcp"))
.post(mcp_url)
.header(header::ACCEPT, "application/json, text/event-stream")
.header("MCP-Session-Id", &session_id)
.header("MCP-Protocol-Version", "2025-11-25")
@@ -404,12 +497,12 @@ mod tests {
async fn post_jsonrpc(
client: &reqwest::Client,
base_url: &str,
mcp_url: &str,
session_id: Option<&str>,
payload: Value,
) -> Value {
let mut request = client
.post(format!("{base_url}/mcp"))
.post(mcp_url)
.header(header::ACCEPT, "application/json, text/event-stream")
.header("MCP-Protocol-Version", "2025-11-25");
@@ -462,6 +555,78 @@ mod tests {
format!("http://{}", address)
}
fn agent_mcp_url(base_url: &str, agent_slug: &str) -> String {
format!("{base_url}/v1/{}/{}", test_workspace_slug(), agent_slug)
}
async fn publish_agent_for_operation(
registry: &PostgresRegistry,
operation: &Operation<Schema, MappingSet>,
agent_slug: &str,
) {
publish_agent_with_bindings(registry, agent_slug, vec![binding_for_operation(operation)])
.await;
}
async fn publish_agent_with_bindings(
registry: &PostgresRegistry,
agent_slug: &str,
bindings: Vec<AgentOperationBinding>,
) {
let agent_id = AgentId::new(format!("agent_{agent_slug}"));
let agent = Agent {
id: agent_id.clone(),
workspace_id: test_workspace_id(),
slug: agent_slug.to_owned(),
display_name: format!("Agent {agent_slug}"),
description: "Curated MCP toolset".to_owned(),
status: AgentStatus::Draft,
current_draft_version: 1,
latest_published_version: None,
created_at: "2026-03-26T10:00:00Z".to_owned(),
updated_at: "2026-03-26T10:00:00Z".to_owned(),
published_at: None,
};
let version = AgentVersion {
agent_id: agent_id.clone(),
version: 1,
status: AgentStatus::Draft,
instructions: json!({}),
tool_selection_policy: json!({}),
created_at: "2026-03-26T10:00:00Z".to_owned(),
};
registry
.create_agent(CreateAgentRequest {
agent: &agent,
version: &version,
bindings: &bindings,
})
.await
.unwrap();
registry
.publish_agent(PublishAgentRequest {
workspace_id: &test_workspace_id(),
agent_id: &agent_id,
version: 1,
published_at: "2026-03-26T10:00:00Z",
published_by: Some("alice"),
})
.await
.unwrap();
}
fn binding_for_operation(operation: &Operation<Schema, MappingSet>) -> AgentOperationBinding {
AgentOperationBinding {
operation_id: operation.id.clone(),
operation_version: 1,
tool_name: operation.name.clone(),
tool_title: operation.tool_description.title.clone(),
tool_description_override: None,
enabled: true,
}
}
async fn create_lead(Json(payload): Json<Value>) -> Json<Value> {
Json(json!({
"id": "lead_123",
+10 -1
View File
@@ -12,6 +12,8 @@ pub struct SessionStore {
pub struct SessionState {
pub protocol_version: String,
pub initialized: bool,
pub workspace_slug: String,
pub agent_slug: String,
}
impl SessionStore {
@@ -21,7 +23,12 @@ impl SessionStore {
}
}
pub async fn create(&self, protocol_version: &str) -> String {
pub async fn create(
&self,
protocol_version: &str,
workspace_slug: &str,
agent_slug: &str,
) -> String {
let session_id = Uuid::now_v7().to_string();
let mut guard = self.inner.write().await;
@@ -30,6 +37,8 @@ impl SessionStore {
SessionState {
protocol_version: protocol_version.to_owned(),
initialized: false,
workspace_slug: workspace_slug.to_owned(),
agent_slug: agent_slug.to_owned(),
},
);