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 => {
|
||||
|
||||
Reference in New Issue
Block a user