feat: add agent publishing foundation
This commit is contained in:
+103
-35
@@ -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 => {
|
||||
|
||||
@@ -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
@@ -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(®istry, &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(®istry, &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(®istry, &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(®istry, &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(
|
||||
®istry,
|
||||
"sales-a",
|
||||
vec![binding_for_operation(&operation_a)],
|
||||
)
|
||||
.await;
|
||||
publish_agent_with_bindings(
|
||||
®istry,
|
||||
"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",
|
||||
|
||||
@@ -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(),
|
||||
},
|
||||
);
|
||||
|
||||
|
||||
Reference in New Issue
Block a user