a.tolmachev
2088 lines
65 KiB
Rust
2088 lines
65 KiB
Rust
use std::{
|
|
convert::Infallible,
|
|
sync::Arc,
|
|
time::{Duration, Instant},
|
|
};
|
|
|
|
use axum::{
|
|
Json, Router,
|
|
extract::{Path, State},
|
|
http::{
|
|
HeaderMap, HeaderValue, StatusCode,
|
|
header::{self, ACCEPT, AUTHORIZATION},
|
|
},
|
|
response::{
|
|
IntoResponse, Response,
|
|
sse::{Event, KeepAlive, Sse},
|
|
},
|
|
routing::get,
|
|
};
|
|
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
|
use crank_core::{
|
|
AsyncJobHandle, AsyncJobId, InvocationLevel, InvocationLog, InvocationLogId, InvocationSource,
|
|
InvocationStatus, JobStatus, PlatformApiKeyScope, StreamSession, StreamSessionId, StreamStatus,
|
|
};
|
|
use crank_registry::{
|
|
CreateAsyncJobRequest, CreateInvocationLogRequest, CreateStreamSessionRequest,
|
|
PostgresRegistry, PublishedAgentTool, UpdateAsyncJobStatusRequest,
|
|
UpdateStreamSessionStateRequest,
|
|
};
|
|
use crank_runtime::{RuntimeError, RuntimeExecutor, RuntimeOperation};
|
|
use futures_util::stream;
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::{Value, json};
|
|
use sha2::{Digest, Sha256};
|
|
use time::{OffsetDateTime, format_description::well_known::Rfc3339};
|
|
|
|
use crate::{
|
|
catalog::PublishedToolCatalog,
|
|
jsonrpc::{
|
|
CURRENT_PROTOCOL_VERSION, DEFAULT_PROTOCOL_VERSION, is_notification, is_request,
|
|
is_response, jsonrpc_error, jsonrpc_result, method_name, negotiated_protocol_version,
|
|
params, request_id,
|
|
},
|
|
session::{SessionState, SessionStore},
|
|
};
|
|
|
|
const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
|
|
const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";
|
|
|
|
#[derive(Clone, Copy)]
|
|
enum ResponseMode {
|
|
Json,
|
|
Sse,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct AppState {
|
|
registry: PostgresRegistry,
|
|
catalog: PublishedToolCatalog,
|
|
runtime: RuntimeExecutor,
|
|
sessions: SessionStore,
|
|
allowed_origins: AllowedOrigins,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct AllowedOrigins {
|
|
public_origin: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
struct InitializeParams {
|
|
#[serde(rename = "protocolVersion")]
|
|
protocol_version: String,
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
struct ToolCallParams {
|
|
name: String,
|
|
#[serde(default)]
|
|
arguments: Value,
|
|
}
|
|
|
|
#[derive(Clone, Copy)]
|
|
enum GeneratedToolKind {
|
|
Base,
|
|
SessionStart,
|
|
SessionPoll,
|
|
SessionStop,
|
|
AsyncJobStart,
|
|
AsyncJobStatus,
|
|
AsyncJobResult,
|
|
AsyncJobCancel,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
struct ResolvedToolCall {
|
|
tool: PublishedAgentTool,
|
|
kind: GeneratedToolKind,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct SessionControlArgs {
|
|
session_id: String,
|
|
}
|
|
|
|
#[derive(Debug, Deserialize)]
|
|
struct AsyncJobControlArgs {
|
|
job_id: String,
|
|
}
|
|
|
|
#[derive(Clone, Debug, Deserialize)]
|
|
struct AgentRoutePath {
|
|
workspace_slug: String,
|
|
agent_slug: String,
|
|
}
|
|
|
|
pub fn build_app(
|
|
registry: PostgresRegistry,
|
|
refresh_interval: Duration,
|
|
public_base_url: Option<String>,
|
|
) -> Router {
|
|
let state = Arc::new(AppState {
|
|
registry: registry.clone(),
|
|
catalog: PublishedToolCatalog::new(registry, refresh_interval),
|
|
runtime: RuntimeExecutor::new(),
|
|
sessions: SessionStore::new(),
|
|
allowed_origins: AllowedOrigins::new(public_base_url),
|
|
});
|
|
|
|
Router::new()
|
|
.route("/health", get(health))
|
|
.route(
|
|
"/v1/{workspace_slug}/{agent_slug}",
|
|
get(mcp_get).post(mcp_post).delete(mcp_delete),
|
|
)
|
|
.with_state(state)
|
|
}
|
|
|
|
async fn health() -> Json<Value> {
|
|
Json(json!({
|
|
"service": "mcp-server",
|
|
"status": "ok"
|
|
}))
|
|
}
|
|
|
|
async fn mcp_get(
|
|
Path(path): Path<AgentRoutePath>,
|
|
State(state): State<Arc<AppState>>,
|
|
headers: HeaderMap,
|
|
) -> Response {
|
|
if let Err(status) = validate_origin(&state.allowed_origins, &headers) {
|
|
return status.into_response();
|
|
}
|
|
|
|
if let Err(status) = validate_get_accept_header(&headers) {
|
|
return status.into_response();
|
|
}
|
|
|
|
if let Err(status) =
|
|
require_platform_api_key(&state, &path, &headers, PlatformApiKeyScope::Read).await
|
|
{
|
|
return status.into_response();
|
|
}
|
|
|
|
let session_id = match session_id_from_headers(&headers) {
|
|
Ok(Some(session_id)) => session_id,
|
|
Ok(None) => return StatusCode::BAD_REQUEST.into_response(),
|
|
Err(status) => return status.into_response(),
|
|
};
|
|
|
|
let Some(session) = state.sessions.get(&session_id).await else {
|
|
return StatusCode::NOT_FOUND.into_response();
|
|
};
|
|
|
|
if session.workspace_slug != path.workspace_slug || session.agent_slug != path.agent_slug {
|
|
return StatusCode::NOT_FOUND.into_response();
|
|
}
|
|
|
|
if !session.initialized {
|
|
return StatusCode::BAD_REQUEST.into_response();
|
|
}
|
|
|
|
if let Err(status) = validate_session_protocol_version(&headers, &session.protocol_version) {
|
|
return status.into_response();
|
|
}
|
|
|
|
sse_response(
|
|
StatusCode::OK,
|
|
stream::pending::<Result<Event, Infallible>>(),
|
|
Some(&session_id),
|
|
Some(&session.protocol_version),
|
|
)
|
|
}
|
|
|
|
async fn mcp_delete(
|
|
Path(path): Path<AgentRoutePath>,
|
|
State(state): State<Arc<AppState>>,
|
|
headers: HeaderMap,
|
|
) -> Response {
|
|
if let Err(status) =
|
|
require_platform_api_key(&state, &path, &headers, PlatformApiKeyScope::Read).await
|
|
{
|
|
return status.into_response();
|
|
}
|
|
|
|
match session_id_from_headers(&headers) {
|
|
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>,
|
|
) -> Response {
|
|
if let Err(status) = validate_origin(&state.allowed_origins, &headers) {
|
|
return status.into_response();
|
|
}
|
|
|
|
let response_mode = match negotiate_post_response_mode(&headers) {
|
|
Ok(mode) => mode,
|
|
Err(status) => return status.into_response(),
|
|
};
|
|
|
|
if is_response(&message) || is_notification(&message) && method_name(&message).is_none() {
|
|
return StatusCode::ACCEPTED.into_response();
|
|
}
|
|
|
|
let protocol_version = match protocol_version_from_headers(&headers) {
|
|
Ok(value) => value,
|
|
Err(status) => return status.into_response(),
|
|
};
|
|
|
|
if let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) {
|
|
if let Ok(session_id) = session_id.to_str() {
|
|
if let Some(session) = state.sessions.get(session_id).await {
|
|
if let Err(status) =
|
|
validate_session_protocol_version(&headers, &session.protocol_version)
|
|
{
|
|
return status.into_response();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
let required_scope = match method_name(&message) {
|
|
Some("tools/call") => PlatformApiKeyScope::Write,
|
|
_ => PlatformApiKeyScope::Read,
|
|
};
|
|
if let Err(status) = require_platform_api_key(&state, &path, &headers, required_scope).await {
|
|
return status.into_response();
|
|
}
|
|
|
|
match method_name(&message) {
|
|
Some("initialize") if is_request(&message) => {
|
|
handle_initialize(state, &path, &message, response_mode).await
|
|
}
|
|
Some("notifications/initialized") if is_notification(&message) => {
|
|
handle_initialized_notification(state, &path, &headers).await
|
|
}
|
|
Some("ping") if is_request(&message) => {
|
|
let session = match require_initialized_session(&state, &path, &headers, &message).await
|
|
{
|
|
Ok(session) => session,
|
|
Err(response) => return response,
|
|
};
|
|
|
|
transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_result(
|
|
request_id(&message),
|
|
json!({ "protocolVersion": session.protocol_version }),
|
|
),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
)
|
|
}
|
|
Some("tools/list") if is_request(&message) => {
|
|
let session = match require_initialized_session(&state, &path, &headers, &message).await
|
|
{
|
|
Ok(session) => session,
|
|
Err(response) => return response,
|
|
};
|
|
|
|
match state
|
|
.catalog
|
|
.list_tools(&session.workspace_slug, &session.agent_slug)
|
|
.await
|
|
{
|
|
Ok(tools) => {
|
|
let definitions = tools.iter().flat_map(tool_definitions).collect::<Vec<_>>();
|
|
|
|
transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_result(request_id(&message), json!({ "tools": definitions })),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
)
|
|
}
|
|
Err(error) => internal_jsonrpc_error(&message, error),
|
|
}
|
|
}
|
|
Some("tools/call") if is_request(&message) => {
|
|
let session = match require_initialized_session(&state, &path, &headers, &message).await
|
|
{
|
|
Ok(session) => session,
|
|
Err(response) => return response,
|
|
};
|
|
let tool_call_params: ToolCallParams = match serde_json::from_value(params(&message)) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(request_id(&message), -32602, error.to_string()),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
);
|
|
}
|
|
};
|
|
|
|
let arguments = if tool_call_params.arguments.is_null() {
|
|
json!({})
|
|
} else {
|
|
tool_call_params.arguments
|
|
};
|
|
|
|
match state
|
|
.catalog
|
|
.list_tools(&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,
|
|
resolved,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
None => transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(
|
|
request_id(&message),
|
|
-32602,
|
|
format!("tool {} was not found", tool_call_params.name),
|
|
),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
),
|
|
},
|
|
Err(error) => internal_jsonrpc_error(&message, error),
|
|
}
|
|
}
|
|
Some(method) if is_notification(&message) => {
|
|
let _ = method;
|
|
StatusCode::ACCEPTED.into_response()
|
|
}
|
|
Some(method) => transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(
|
|
request_id(&message),
|
|
-32601,
|
|
format!("method {method} is not supported"),
|
|
),
|
|
response_mode,
|
|
None,
|
|
Some(&protocol_version),
|
|
),
|
|
None => transport_response(
|
|
StatusCode::BAD_REQUEST,
|
|
jsonrpc_error(Value::Null, -32600, "invalid JSON-RPC message"),
|
|
response_mode,
|
|
None,
|
|
Some(&protocol_version),
|
|
),
|
|
}
|
|
}
|
|
|
|
async fn handle_tool_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
resolved: ResolvedToolCall,
|
|
arguments: Value,
|
|
) -> Response {
|
|
match resolved.kind {
|
|
GeneratedToolKind::Base => {
|
|
handle_base_tool_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
GeneratedToolKind::SessionStart => {
|
|
handle_session_start_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
GeneratedToolKind::SessionPoll => {
|
|
handle_session_poll_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
GeneratedToolKind::SessionStop => {
|
|
handle_session_stop_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
GeneratedToolKind::AsyncJobStart => {
|
|
handle_async_job_start_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
GeneratedToolKind::AsyncJobStatus => {
|
|
handle_async_job_status_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
GeneratedToolKind::AsyncJobResult => {
|
|
handle_async_job_result_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
GeneratedToolKind::AsyncJobCancel => {
|
|
handle_async_job_cancel_call(
|
|
state,
|
|
session,
|
|
message,
|
|
response_mode,
|
|
resolved.tool,
|
|
arguments,
|
|
)
|
|
.await
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Serialize, Deserialize)]
|
|
struct StoredSessionState {
|
|
input: Value,
|
|
summary: Value,
|
|
items: Vec<Value>,
|
|
next_index: usize,
|
|
batch_size: usize,
|
|
}
|
|
|
|
async fn handle_base_tool_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let operation = runtime_operation(&tool);
|
|
let request_preview = build_request_preview(&state.runtime, &operation, &arguments);
|
|
let started_at = Instant::now();
|
|
let is_window_mode = matches!(
|
|
operation
|
|
.execution_config
|
|
.streaming
|
|
.as_ref()
|
|
.map(|streaming| streaming.mode),
|
|
Some(crank_core::ExecutionMode::Window)
|
|
);
|
|
|
|
let result = if is_window_mode {
|
|
state
|
|
.runtime
|
|
.execute_window(&operation, &arguments)
|
|
.await
|
|
.map(|output| {
|
|
json!({
|
|
"summary": output.summary,
|
|
"items": output.items,
|
|
"cursor": output.cursor,
|
|
"window_complete": output.window_complete,
|
|
"truncated": output.truncated,
|
|
"has_more": output.has_more,
|
|
})
|
|
})
|
|
} else {
|
|
state.runtime.execute(&operation, &arguments).await
|
|
};
|
|
|
|
match result {
|
|
Ok(output) => {
|
|
let _ = persist_invocation(
|
|
&state,
|
|
&tool,
|
|
InvocationRecord {
|
|
tool_name: &tool.tool_name,
|
|
status: InvocationStatus::Ok,
|
|
level: InvocationLevel::Info,
|
|
message: "agent tool call completed",
|
|
status_code: None,
|
|
error_kind: None,
|
|
duration: started_at.elapsed(),
|
|
request_preview,
|
|
response_preview: output.clone(),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
success_tool_response(message, response_mode, &session.protocol_version, output)
|
|
}
|
|
Err(error) => {
|
|
let _ = persist_invocation(
|
|
&state,
|
|
&tool,
|
|
InvocationRecord {
|
|
tool_name: &tool.tool_name,
|
|
status: InvocationStatus::Error,
|
|
level: InvocationLevel::Error,
|
|
message: &error.to_string(),
|
|
status_code: None,
|
|
error_kind: Some(runtime_error_code(&error)),
|
|
duration: started_at.elapsed(),
|
|
request_preview,
|
|
response_preview: Value::Null,
|
|
},
|
|
)
|
|
.await;
|
|
|
|
tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
runtime_error_code(&error),
|
|
error.to_string(),
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_session_start_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let runtime_operation = runtime_operation(&tool);
|
|
let request_preview = build_request_preview(&state.runtime, &runtime_operation, &arguments);
|
|
let started_at = Instant::now();
|
|
let Some(streaming) = runtime_operation.execution_config.streaming.as_ref() else {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"streaming_config_error",
|
|
"streaming config is required for session tools".to_owned(),
|
|
);
|
|
};
|
|
|
|
match state
|
|
.runtime
|
|
.execute_session_seed(&runtime_operation, &arguments)
|
|
.await
|
|
{
|
|
Ok(seed) => {
|
|
let batch_size = streaming.max_items.unwrap_or(10).max(1) as usize;
|
|
let preview_count = seed.items.len().min(batch_size);
|
|
let preview_items = seed.items[..preview_count].to_vec();
|
|
let next_index = preview_count;
|
|
let session_id =
|
|
StreamSessionId::new(format!("sess_{}", uuid::Uuid::now_v7().simple()));
|
|
let now = now_rfc3339();
|
|
let expires_at = add_millis(&now, streaming.max_session_lifetime_ms.unwrap_or(60_000));
|
|
let session_record = StreamSession {
|
|
id: session_id.clone(),
|
|
workspace_id: tool.workspace_id.clone(),
|
|
agent_id: Some(tool.agent_id.clone()),
|
|
operation_id: tool.operation.id.clone(),
|
|
protocol: tool.operation.protocol,
|
|
mode: crank_core::ExecutionMode::Session,
|
|
status: StreamStatus::Running,
|
|
cursor: (next_index < seed.items.len()).then(|| json!(next_index)),
|
|
state: json!(StoredSessionState {
|
|
input: arguments.clone(),
|
|
summary: seed.summary.clone(),
|
|
items: seed.items.clone(),
|
|
next_index,
|
|
batch_size,
|
|
}),
|
|
expires_at: expires_at.clone(),
|
|
last_poll_at: Some(now.clone()),
|
|
created_at: now.clone(),
|
|
closed_at: None,
|
|
};
|
|
|
|
if let Err(error) = state
|
|
.registry
|
|
.create_stream_session(CreateStreamSessionRequest {
|
|
session: &session_record,
|
|
})
|
|
.await
|
|
{
|
|
return internal_jsonrpc_error(message, error);
|
|
}
|
|
|
|
let output = json!({
|
|
"session_id": session_id.as_str(),
|
|
"status": "running",
|
|
"expires_at": expires_at,
|
|
"poll_after_ms": streaming.poll_interval_ms.unwrap_or(1000),
|
|
"preview": {
|
|
"summary": seed.summary,
|
|
"items": preview_items,
|
|
}
|
|
});
|
|
|
|
let _ = persist_invocation(
|
|
&state,
|
|
&tool,
|
|
InvocationRecord {
|
|
tool_name: streaming
|
|
.tool_family
|
|
.start_tool_name
|
|
.as_deref()
|
|
.unwrap_or(&tool.tool_name),
|
|
status: InvocationStatus::Ok,
|
|
level: InvocationLevel::Info,
|
|
message: "stream session started",
|
|
status_code: None,
|
|
error_kind: None,
|
|
duration: started_at.elapsed(),
|
|
request_preview,
|
|
response_preview: output.clone(),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
success_tool_response(message, response_mode, &session.protocol_version, output)
|
|
}
|
|
Err(error) => {
|
|
let _ = persist_invocation(
|
|
&state,
|
|
&tool,
|
|
InvocationRecord {
|
|
tool_name: streaming
|
|
.tool_family
|
|
.start_tool_name
|
|
.as_deref()
|
|
.unwrap_or(&tool.tool_name),
|
|
status: InvocationStatus::Error,
|
|
level: InvocationLevel::Error,
|
|
message: &error.to_string(),
|
|
status_code: None,
|
|
error_kind: Some(runtime_error_code(&error)),
|
|
duration: started_at.elapsed(),
|
|
request_preview,
|
|
response_preview: Value::Null,
|
|
},
|
|
)
|
|
.await;
|
|
|
|
tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
runtime_error_code(&error),
|
|
error.to_string(),
|
|
)
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn handle_session_poll_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let control: SessionControlArgs = match serde_json::from_value(arguments.clone()) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(request_id(message), -32602, error.to_string()),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
);
|
|
}
|
|
};
|
|
|
|
let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"streaming_config_error",
|
|
"streaming config is required for session tools".to_owned(),
|
|
);
|
|
};
|
|
|
|
let now = now_rfc3339();
|
|
let loaded = match state
|
|
.registry
|
|
.get_stream_session(&StreamSessionId::new(control.session_id.clone()))
|
|
.await
|
|
{
|
|
Ok(Some(session_record)) => session_record,
|
|
Ok(None) => {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"stream_session_not_found",
|
|
format!("stream session {} was not found", control.session_id),
|
|
);
|
|
}
|
|
Err(error) => return internal_jsonrpc_error(message, error),
|
|
};
|
|
|
|
if loaded.is_expired(&now) {
|
|
let _ = state
|
|
.registry
|
|
.update_stream_session_state(UpdateStreamSessionStateRequest {
|
|
session_id: &loaded.id,
|
|
current_status: loaded.status,
|
|
next_status: StreamStatus::Expired,
|
|
cursor: loaded.cursor.as_ref(),
|
|
state: &loaded.state,
|
|
expires_at: None,
|
|
last_poll_at: Some(&now),
|
|
closed_at: Some(&now),
|
|
})
|
|
.await;
|
|
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"stream_session_expired",
|
|
format!("stream session {} has expired", control.session_id),
|
|
);
|
|
}
|
|
|
|
if !loaded.can_poll(&now) {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"stream_session_not_found",
|
|
format!("stream session {} is not running", control.session_id),
|
|
);
|
|
}
|
|
|
|
let mut state_payload: StoredSessionState = match serde_json::from_value(loaded.state.clone()) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"streaming_payload_error",
|
|
error.to_string(),
|
|
);
|
|
}
|
|
};
|
|
|
|
let next_end =
|
|
(state_payload.next_index + state_payload.batch_size).min(state_payload.items.len());
|
|
let items = state_payload.items[state_payload.next_index..next_end].to_vec();
|
|
state_payload.next_index = next_end;
|
|
let has_more = state_payload.next_index < state_payload.items.len();
|
|
let next_status = if has_more {
|
|
StreamStatus::Running
|
|
} else {
|
|
StreamStatus::Stopped
|
|
};
|
|
let cursor = has_more.then(|| json!(state_payload.next_index));
|
|
|
|
let updated = match state
|
|
.registry
|
|
.update_stream_session_state(UpdateStreamSessionStateRequest {
|
|
session_id: &loaded.id,
|
|
current_status: loaded.status,
|
|
next_status,
|
|
cursor: cursor.as_ref(),
|
|
state: &json!(state_payload),
|
|
expires_at: Some(&add_millis(
|
|
&now,
|
|
streaming.max_session_lifetime_ms.unwrap_or(60_000),
|
|
)),
|
|
last_poll_at: Some(&now),
|
|
closed_at: matches!(next_status, StreamStatus::Stopped).then_some(now.as_str()),
|
|
})
|
|
.await
|
|
{
|
|
Ok(session_record) => session_record,
|
|
Err(error) => return internal_jsonrpc_error(message, error),
|
|
};
|
|
|
|
let output = json!({
|
|
"session_id": updated.id.as_str(),
|
|
"status": serialize_stream_status(updated.status),
|
|
"expires_at": updated.expires_at,
|
|
"summary": serde_json::from_value::<StoredSessionState>(updated.state.clone()).map(|value| value.summary).unwrap_or(Value::Null),
|
|
"items": items,
|
|
"cursor": updated.cursor,
|
|
"has_more": updated.status == StreamStatus::Running,
|
|
});
|
|
|
|
let _ = persist_invocation(
|
|
&state,
|
|
&tool,
|
|
InvocationRecord {
|
|
tool_name: streaming
|
|
.tool_family
|
|
.poll_tool_name
|
|
.as_deref()
|
|
.unwrap_or(&tool.tool_name),
|
|
status: InvocationStatus::Ok,
|
|
level: InvocationLevel::Info,
|
|
message: "stream session polled",
|
|
status_code: None,
|
|
error_kind: None,
|
|
duration: Duration::ZERO,
|
|
request_preview: json!({ "session_id": control.session_id }),
|
|
response_preview: output.clone(),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
success_tool_response(message, response_mode, &session.protocol_version, output)
|
|
}
|
|
|
|
async fn handle_session_stop_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
_tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let control: SessionControlArgs = match serde_json::from_value(arguments.clone()) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(request_id(message), -32602, error.to_string()),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
);
|
|
}
|
|
};
|
|
|
|
match state
|
|
.registry
|
|
.close_stream_session(
|
|
&StreamSessionId::new(control.session_id.clone()),
|
|
&now_rfc3339(),
|
|
)
|
|
.await
|
|
{
|
|
Ok(()) => success_tool_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
json!({
|
|
"session_id": control.session_id,
|
|
"status": "stopped"
|
|
}),
|
|
),
|
|
Err(error) => internal_jsonrpc_error(message, error),
|
|
}
|
|
}
|
|
|
|
async fn handle_async_job_start_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let operation = runtime_operation(&tool);
|
|
let request_preview = build_request_preview(&state.runtime, &operation, &arguments);
|
|
let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"streaming_config_error",
|
|
"streaming config is required for async job tools".to_owned(),
|
|
);
|
|
};
|
|
|
|
let now = now_rfc3339();
|
|
let job = AsyncJobHandle {
|
|
id: AsyncJobId::new(format!("job_{}", uuid::Uuid::now_v7().simple())),
|
|
workspace_id: tool.workspace_id.clone(),
|
|
agent_id: Some(tool.agent_id.clone()),
|
|
operation_id: tool.operation.id.clone(),
|
|
status: JobStatus::Running,
|
|
progress: json!({ "pct": 0 }),
|
|
result: None,
|
|
error: None,
|
|
expires_at: Some(add_millis(&now, 300_000)),
|
|
created_at: now.clone(),
|
|
updated_at: now.clone(),
|
|
finished_at: None,
|
|
};
|
|
|
|
if let Err(error) = state
|
|
.registry
|
|
.create_async_job(CreateAsyncJobRequest { job: &job })
|
|
.await
|
|
{
|
|
return internal_jsonrpc_error(message, error);
|
|
}
|
|
|
|
let registry = state.registry.clone();
|
|
let tool_for_task = tool.clone();
|
|
let arguments_for_task = arguments.clone();
|
|
let job_id = job.id.clone();
|
|
tokio::spawn(async move {
|
|
let runtime = RuntimeExecutor::new();
|
|
let task_operation = runtime_operation(&tool_for_task);
|
|
let result = runtime.execute(&task_operation, &arguments_for_task).await;
|
|
let finished_at = now_rfc3339();
|
|
|
|
let update_result = match result {
|
|
Ok(output) => {
|
|
registry
|
|
.update_async_job_status(UpdateAsyncJobStatusRequest {
|
|
job_id: &job_id,
|
|
current_status: JobStatus::Running,
|
|
next_status: JobStatus::Completed,
|
|
progress: &json!({ "pct": 100 }),
|
|
result: Some(&output),
|
|
error: None,
|
|
expires_at: None,
|
|
updated_at: &finished_at,
|
|
finished_at: Some(&finished_at),
|
|
})
|
|
.await
|
|
}
|
|
Err(error) => {
|
|
registry
|
|
.update_async_job_status(UpdateAsyncJobStatusRequest {
|
|
job_id: &job_id,
|
|
current_status: JobStatus::Running,
|
|
next_status: JobStatus::Failed,
|
|
progress: &json!({ "pct": 100 }),
|
|
result: None,
|
|
error: Some(&json!({
|
|
"code": runtime_error_code(&error),
|
|
"message": error.to_string()
|
|
})),
|
|
expires_at: None,
|
|
updated_at: &finished_at,
|
|
finished_at: Some(&finished_at),
|
|
})
|
|
.await
|
|
}
|
|
};
|
|
|
|
let _ = update_result;
|
|
});
|
|
|
|
let output = json!({
|
|
"job_id": job.id.as_str(),
|
|
"status": "running",
|
|
"progress": job.progress,
|
|
});
|
|
|
|
let _ = persist_invocation(
|
|
&state,
|
|
&tool,
|
|
InvocationRecord {
|
|
tool_name: streaming
|
|
.tool_family
|
|
.start_tool_name
|
|
.as_deref()
|
|
.unwrap_or(&tool.tool_name),
|
|
status: InvocationStatus::Ok,
|
|
level: InvocationLevel::Info,
|
|
message: "async job started",
|
|
status_code: None,
|
|
error_kind: None,
|
|
duration: Duration::ZERO,
|
|
request_preview,
|
|
response_preview: output.clone(),
|
|
},
|
|
)
|
|
.await;
|
|
|
|
success_tool_response(message, response_mode, &session.protocol_version, output)
|
|
}
|
|
|
|
async fn handle_async_job_status_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
_tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(request_id(message), -32602, error.to_string()),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
);
|
|
}
|
|
};
|
|
|
|
let job = match state
|
|
.registry
|
|
.get_async_job(&AsyncJobId::new(control.job_id.clone()))
|
|
.await
|
|
{
|
|
Ok(Some(job)) => job,
|
|
Ok(None) => {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"async_job_not_found",
|
|
format!("async job {} was not found", control.job_id),
|
|
);
|
|
}
|
|
Err(error) => return internal_jsonrpc_error(message, error),
|
|
};
|
|
|
|
success_tool_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
json!({
|
|
"job_id": job.id.as_str(),
|
|
"status": serialize_job_status(job.status),
|
|
"progress": job.progress,
|
|
"updated_at": job.updated_at,
|
|
"finished_at": job.finished_at,
|
|
}),
|
|
)
|
|
}
|
|
|
|
async fn handle_async_job_result_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
_tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(request_id(message), -32602, error.to_string()),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
);
|
|
}
|
|
};
|
|
|
|
let job = match state
|
|
.registry
|
|
.get_async_job(&AsyncJobId::new(control.job_id.clone()))
|
|
.await
|
|
{
|
|
Ok(Some(job)) => job,
|
|
Ok(None) => {
|
|
return tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"async_job_not_found",
|
|
format!("async job {} was not found", control.job_id),
|
|
);
|
|
}
|
|
Err(error) => return internal_jsonrpc_error(message, error),
|
|
};
|
|
|
|
match job.status {
|
|
JobStatus::Completed => success_tool_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
job.result.unwrap_or(Value::Null),
|
|
),
|
|
JobStatus::Failed => tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"async_job_failed",
|
|
job.error
|
|
.and_then(|value| value.get("message").cloned())
|
|
.and_then(|value| value.as_str().map(ToOwned::to_owned))
|
|
.unwrap_or_else(|| "async job failed".to_owned()),
|
|
),
|
|
JobStatus::Cancelled => tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"async_job_cancelled",
|
|
"async job was cancelled".to_owned(),
|
|
),
|
|
_ => tool_error_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
"async_job_not_ready",
|
|
"async job result is not ready".to_owned(),
|
|
),
|
|
}
|
|
}
|
|
|
|
async fn handle_async_job_cancel_call(
|
|
state: Arc<AppState>,
|
|
session: &SessionState,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
_tool: PublishedAgentTool,
|
|
arguments: Value,
|
|
) -> Response {
|
|
let control: AsyncJobControlArgs = match serde_json::from_value(arguments.clone()) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(request_id(message), -32602, error.to_string()),
|
|
response_mode,
|
|
None,
|
|
Some(&session.protocol_version),
|
|
);
|
|
}
|
|
};
|
|
|
|
match state
|
|
.registry
|
|
.cancel_async_job(&AsyncJobId::new(control.job_id.clone()), &now_rfc3339())
|
|
.await
|
|
{
|
|
Ok(()) => success_tool_response(
|
|
message,
|
|
response_mode,
|
|
&session.protocol_version,
|
|
json!({
|
|
"job_id": control.job_id,
|
|
"status": "cancelled"
|
|
}),
|
|
),
|
|
Err(error) => internal_jsonrpc_error(message, error),
|
|
}
|
|
}
|
|
|
|
async fn handle_initialize(
|
|
state: Arc<AppState>,
|
|
path: &AgentRoutePath,
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
) -> Response {
|
|
let initialize_params: InitializeParams = match serde_json::from_value(params(message)) {
|
|
Ok(value) => value,
|
|
Err(error) => {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(request_id(message), -32602, error.to_string()),
|
|
response_mode,
|
|
None,
|
|
Some(DEFAULT_PROTOCOL_VERSION),
|
|
);
|
|
}
|
|
};
|
|
let Some(protocol_version) = negotiated_protocol_version(&initialize_params.protocol_version)
|
|
else {
|
|
return transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_error(
|
|
request_id(message),
|
|
-32602,
|
|
format!(
|
|
"unsupported protocol version {}",
|
|
initialize_params.protocol_version
|
|
),
|
|
),
|
|
response_mode,
|
|
None,
|
|
Some(DEFAULT_PROTOCOL_VERSION),
|
|
);
|
|
};
|
|
let session_id = state
|
|
.sessions
|
|
.create(protocol_version, &path.workspace_slug, &path.agent_slug)
|
|
.await;
|
|
|
|
transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_result(
|
|
request_id(message),
|
|
json!({
|
|
"protocolVersion": protocol_version,
|
|
"capabilities": {
|
|
"tools": {
|
|
"listChanged": false
|
|
}
|
|
},
|
|
"serverInfo": {
|
|
"name": "crank-mcp-server",
|
|
"version": env!("CARGO_PKG_VERSION")
|
|
}
|
|
}),
|
|
),
|
|
response_mode,
|
|
Some(&session_id),
|
|
Some(protocol_version),
|
|
)
|
|
}
|
|
|
|
async fn handle_initialized_notification(
|
|
state: Arc<AppState>,
|
|
path: &AgentRoutePath,
|
|
headers: &HeaderMap,
|
|
) -> Response {
|
|
match session_id_from_headers(headers) {
|
|
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(),
|
|
}
|
|
}
|
|
|
|
async fn require_initialized_session(
|
|
state: &Arc<AppState>,
|
|
path: &AgentRoutePath,
|
|
headers: &HeaderMap,
|
|
message: &Value,
|
|
) -> Result<SessionState, Response> {
|
|
let session_id = match session_id_from_headers(headers) {
|
|
Ok(Some(session_id)) => session_id,
|
|
Ok(None) => return Err(StatusCode::BAD_REQUEST.into_response()),
|
|
Err(status) => return Err(status.into_response()),
|
|
};
|
|
|
|
let Some(session) = state.sessions.get(&session_id).await else {
|
|
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,
|
|
jsonrpc_error(request_id(message), -32002, "session is not initialized"),
|
|
None,
|
|
Some(&session.protocol_version),
|
|
));
|
|
}
|
|
|
|
Ok(session)
|
|
}
|
|
|
|
async fn require_platform_api_key(
|
|
state: &Arc<AppState>,
|
|
path: &AgentRoutePath,
|
|
headers: &HeaderMap,
|
|
required_scope: PlatformApiKeyScope,
|
|
) -> Result<(), StatusCode> {
|
|
let secret = bearer_token(headers).ok_or(StatusCode::UNAUTHORIZED)?;
|
|
let secret_hash = hash_access_secret(secret);
|
|
let Some(api_key) = state
|
|
.registry
|
|
.get_platform_api_key_by_secret_for_workspace_slug(&path.workspace_slug, &secret_hash)
|
|
.await
|
|
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
|
|
else {
|
|
return Err(StatusCode::UNAUTHORIZED);
|
|
};
|
|
|
|
if !allows_scope(&api_key.api_key.scopes, required_scope) {
|
|
return Err(StatusCode::FORBIDDEN);
|
|
}
|
|
|
|
let used_at = OffsetDateTime::now_utc()
|
|
.format(&Rfc3339)
|
|
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
|
state
|
|
.registry
|
|
.touch_platform_api_key(&api_key.api_key.workspace_id, &api_key.api_key.id, &used_at)
|
|
.await
|
|
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn bearer_token(headers: &HeaderMap) -> Option<&str> {
|
|
let value = headers.get(AUTHORIZATION)?.to_str().ok()?;
|
|
let (scheme, token) = value.split_once(' ')?;
|
|
if !scheme.eq_ignore_ascii_case("Bearer") || token.is_empty() {
|
|
return None;
|
|
}
|
|
|
|
Some(token)
|
|
}
|
|
|
|
fn allows_scope(scopes: &[PlatformApiKeyScope], required_scope: PlatformApiKeyScope) -> bool {
|
|
match required_scope {
|
|
PlatformApiKeyScope::Read => scopes.iter().any(|scope| {
|
|
matches!(
|
|
scope,
|
|
PlatformApiKeyScope::Read
|
|
| PlatformApiKeyScope::Write
|
|
| PlatformApiKeyScope::Deploy
|
|
)
|
|
}),
|
|
PlatformApiKeyScope::Write => scopes.iter().any(|scope| {
|
|
matches!(
|
|
scope,
|
|
PlatformApiKeyScope::Write | PlatformApiKeyScope::Deploy
|
|
)
|
|
}),
|
|
PlatformApiKeyScope::Deploy => scopes
|
|
.iter()
|
|
.any(|scope| matches!(scope, PlatformApiKeyScope::Deploy)),
|
|
}
|
|
}
|
|
|
|
fn validate_origin(
|
|
allowed_origins: &AllowedOrigins,
|
|
headers: &HeaderMap,
|
|
) -> Result<(), StatusCode> {
|
|
let Some(origin) = headers.get(header::ORIGIN) else {
|
|
return Ok(());
|
|
};
|
|
let Ok(origin) = origin.to_str() else {
|
|
return Err(StatusCode::FORBIDDEN);
|
|
};
|
|
|
|
if allowed_origins.is_allowed(origin) {
|
|
return Ok(());
|
|
}
|
|
|
|
Err(StatusCode::FORBIDDEN)
|
|
}
|
|
|
|
fn negotiate_post_response_mode(headers: &HeaderMap) -> Result<ResponseMode, StatusCode> {
|
|
let Some(accept) = headers.get(ACCEPT) else {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
};
|
|
let Ok(accept) = accept.to_str() else {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
};
|
|
|
|
let mut saw_json = false;
|
|
let mut saw_sse = false;
|
|
let mut preferred = None;
|
|
|
|
for part in accept.split(',') {
|
|
let media_type = part
|
|
.split(';')
|
|
.next()
|
|
.unwrap_or_default()
|
|
.trim()
|
|
.to_ascii_lowercase();
|
|
|
|
match media_type.as_str() {
|
|
"application/json" => {
|
|
saw_json = true;
|
|
if preferred.is_none() {
|
|
preferred = Some(ResponseMode::Json);
|
|
}
|
|
}
|
|
"text/event-stream" => {
|
|
saw_sse = true;
|
|
if preferred.is_none() {
|
|
preferred = Some(ResponseMode::Sse);
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
if saw_json && saw_sse {
|
|
return preferred.ok_or(StatusCode::NOT_ACCEPTABLE);
|
|
}
|
|
|
|
Err(StatusCode::NOT_ACCEPTABLE)
|
|
}
|
|
|
|
fn validate_get_accept_header(headers: &HeaderMap) -> Result<(), StatusCode> {
|
|
let Some(accept) = headers.get(ACCEPT) else {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
};
|
|
let Ok(accept) = accept.to_str() else {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
};
|
|
|
|
if accept
|
|
.split(',')
|
|
.map(|part| {
|
|
part.split(';')
|
|
.next()
|
|
.unwrap_or_default()
|
|
.trim()
|
|
.to_ascii_lowercase()
|
|
})
|
|
.any(|media_type| media_type == "text/event-stream")
|
|
{
|
|
return Ok(());
|
|
}
|
|
|
|
Err(StatusCode::NOT_ACCEPTABLE)
|
|
}
|
|
|
|
fn protocol_version_from_headers(headers: &HeaderMap) -> Result<String, StatusCode> {
|
|
let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else {
|
|
return Ok(DEFAULT_PROTOCOL_VERSION.to_owned());
|
|
};
|
|
let Ok(version) = version.to_str() else {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
};
|
|
|
|
if negotiated_protocol_version(version).is_some() {
|
|
return Ok(version.to_owned());
|
|
}
|
|
|
|
Err(StatusCode::BAD_REQUEST)
|
|
}
|
|
|
|
fn validate_session_protocol_version(
|
|
headers: &HeaderMap,
|
|
negotiated_session_version: &str,
|
|
) -> Result<(), StatusCode> {
|
|
let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else {
|
|
return Ok(());
|
|
};
|
|
let Ok(version) = version.to_str() else {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
};
|
|
|
|
if negotiated_protocol_version(version).is_none() {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
if version != negotiated_session_version {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn session_id_from_headers(headers: &HeaderMap) -> Result<Option<String>, StatusCode> {
|
|
let Some(session_id) = headers.get(HEADER_MCP_SESSION_ID) else {
|
|
return Ok(None);
|
|
};
|
|
let Ok(session_id) = session_id.to_str() else {
|
|
return Err(StatusCode::BAD_REQUEST);
|
|
};
|
|
|
|
Ok(Some(session_id.to_owned()))
|
|
}
|
|
|
|
fn internal_jsonrpc_error(message: &Value, error: impl std::fmt::Display) -> Response {
|
|
transport_response(
|
|
StatusCode::INTERNAL_SERVER_ERROR,
|
|
jsonrpc_error(request_id(message), -32603, error.to_string()),
|
|
ResponseMode::Json,
|
|
None,
|
|
Some(DEFAULT_PROTOCOL_VERSION),
|
|
)
|
|
}
|
|
|
|
fn build_request_preview(
|
|
runtime: &RuntimeExecutor,
|
|
operation: &RuntimeOperation,
|
|
arguments: &Value,
|
|
) -> Value {
|
|
match runtime.prepare_request(operation, arguments) {
|
|
Ok(prepared) => json!({
|
|
"path": prepared.path_params,
|
|
"query": prepared.query_params,
|
|
"headers": prepared.headers,
|
|
"grpc": prepared.grpc.unwrap_or(Value::Null),
|
|
"variables": prepared.variables.unwrap_or(Value::Null),
|
|
"body": prepared.body.unwrap_or(Value::Null)
|
|
}),
|
|
Err(_) => Value::Null,
|
|
}
|
|
}
|
|
|
|
struct InvocationRecord<'a> {
|
|
tool_name: &'a str,
|
|
status: InvocationStatus,
|
|
level: InvocationLevel,
|
|
message: &'a str,
|
|
status_code: Option<u16>,
|
|
error_kind: Option<&'a str>,
|
|
duration: Duration,
|
|
request_preview: Value,
|
|
response_preview: Value,
|
|
}
|
|
|
|
async fn persist_invocation(
|
|
state: &Arc<AppState>,
|
|
tool: &PublishedAgentTool,
|
|
record: InvocationRecord<'_>,
|
|
) -> Result<(), crank_registry::RegistryError> {
|
|
let created_at = OffsetDateTime::now_utc()
|
|
.format(&Rfc3339)
|
|
.unwrap_or_else(|_| "1970-01-01T00:00:00Z".to_owned());
|
|
let duration_ms = u64::try_from(record.duration.as_millis()).unwrap_or(u64::MAX);
|
|
let log = InvocationLog {
|
|
id: InvocationLogId::new(format!("log_{}", uuid::Uuid::now_v7().simple())),
|
|
workspace_id: tool.workspace_id.clone(),
|
|
agent_id: Some(tool.agent_id.clone()),
|
|
operation_id: tool.operation.id.clone(),
|
|
source: InvocationSource::AgentToolCall,
|
|
level: record.level,
|
|
status: record.status,
|
|
tool_name: record.tool_name.to_owned(),
|
|
message: record.message.to_owned(),
|
|
request_id: None,
|
|
status_code: record.status_code,
|
|
duration_ms,
|
|
error_kind: record.error_kind.map(ToOwned::to_owned),
|
|
request_preview: record.request_preview,
|
|
response_preview: record.response_preview,
|
|
created_at,
|
|
};
|
|
|
|
state
|
|
.registry
|
|
.create_invocation_log(CreateInvocationLogRequest { log: &log })
|
|
.await
|
|
}
|
|
|
|
fn success_tool_response(
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
protocol_version: &str,
|
|
output: Value,
|
|
) -> Response {
|
|
transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_result(
|
|
request_id(message),
|
|
json!({
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": serde_json::to_string_pretty(&output).unwrap_or_else(|_| "{}".to_owned())
|
|
}
|
|
],
|
|
"structuredContent": output,
|
|
"isError": false
|
|
}),
|
|
),
|
|
response_mode,
|
|
None,
|
|
Some(protocol_version),
|
|
)
|
|
}
|
|
|
|
fn tool_error_response(
|
|
message: &Value,
|
|
response_mode: ResponseMode,
|
|
protocol_version: &str,
|
|
code: &str,
|
|
error_message: String,
|
|
) -> Response {
|
|
transport_response(
|
|
StatusCode::OK,
|
|
jsonrpc_result(
|
|
request_id(message),
|
|
json!({
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": error_message
|
|
}
|
|
],
|
|
"structuredContent": {
|
|
"error": {
|
|
"code": code,
|
|
"message": error_message
|
|
}
|
|
},
|
|
"isError": true
|
|
}),
|
|
),
|
|
response_mode,
|
|
None,
|
|
Some(protocol_version),
|
|
)
|
|
}
|
|
|
|
fn now_rfc3339() -> String {
|
|
OffsetDateTime::now_utc()
|
|
.format(&Rfc3339)
|
|
.unwrap_or_else(|_| "1970-01-01T00:00:00Z".to_owned())
|
|
}
|
|
|
|
fn add_millis(timestamp: &str, millis: u64) -> String {
|
|
let Ok(parsed) = OffsetDateTime::parse(timestamp, &Rfc3339) else {
|
|
return timestamp.to_owned();
|
|
};
|
|
let delta = time::Duration::milliseconds(i64::try_from(millis).unwrap_or(i64::MAX));
|
|
|
|
(parsed + delta)
|
|
.format(&Rfc3339)
|
|
.unwrap_or_else(|_| timestamp.to_owned())
|
|
}
|
|
|
|
fn serialize_stream_status(status: StreamStatus) -> &'static str {
|
|
match status {
|
|
StreamStatus::Created => "created",
|
|
StreamStatus::Running => "running",
|
|
StreamStatus::Stopped => "stopped",
|
|
StreamStatus::Failed => "failed",
|
|
StreamStatus::Expired => "expired",
|
|
}
|
|
}
|
|
|
|
fn serialize_job_status(status: JobStatus) -> &'static str {
|
|
match status {
|
|
JobStatus::Created => "created",
|
|
JobStatus::Running => "running",
|
|
JobStatus::Completed => "completed",
|
|
JobStatus::Failed => "failed",
|
|
JobStatus::Cancelled => "cancelled",
|
|
JobStatus::Expired => "expired",
|
|
}
|
|
}
|
|
|
|
fn runtime_error_code(error: &RuntimeError) -> &'static str {
|
|
match error {
|
|
RuntimeError::Schema(_) => "schema_validation_error",
|
|
RuntimeError::Mapping(_) => "mapping_error",
|
|
RuntimeError::GraphqlAdapter(_) => "adapter_execution_error",
|
|
RuntimeError::GrpcAdapter(_) => "adapter_execution_error",
|
|
RuntimeError::RestAdapter(_) => "adapter_execution_error",
|
|
RuntimeError::UnsupportedProtocol { .. } => "unsupported_protocol",
|
|
RuntimeError::MissingStreamingConfig { .. } => "streaming_config_error",
|
|
RuntimeError::UnsupportedExecutionMode { .. } => "streaming_mode_error",
|
|
RuntimeError::InvalidPreparedRequest { .. } => "runtime_error",
|
|
RuntimeError::InvalidStreamingPayload { .. } => "streaming_payload_error",
|
|
}
|
|
}
|
|
|
|
fn hash_access_secret(secret: &str) -> String {
|
|
let digest = Sha256::digest(secret.as_bytes());
|
|
URL_SAFE_NO_PAD.encode(digest)
|
|
}
|
|
|
|
fn json_response(
|
|
status: StatusCode,
|
|
payload: Value,
|
|
session_id: Option<&str>,
|
|
protocol_version: Option<&str>,
|
|
) -> Response {
|
|
let mut response = (status, Json(payload)).into_response();
|
|
|
|
if let Some(session_id) = session_id {
|
|
response.headers_mut().insert(
|
|
HEADER_MCP_SESSION_ID,
|
|
HeaderValue::from_str(session_id)
|
|
.unwrap_or_else(|_| HeaderValue::from_static("invalid")),
|
|
);
|
|
}
|
|
|
|
if let Some(protocol_version) = protocol_version {
|
|
response.headers_mut().insert(
|
|
HEADER_MCP_PROTOCOL_VERSION,
|
|
HeaderValue::from_str(protocol_version)
|
|
.unwrap_or_else(|_| HeaderValue::from_static(CURRENT_PROTOCOL_VERSION)),
|
|
);
|
|
}
|
|
|
|
response
|
|
}
|
|
|
|
fn transport_response(
|
|
status: StatusCode,
|
|
payload: Value,
|
|
response_mode: ResponseMode,
|
|
session_id: Option<&str>,
|
|
protocol_version: Option<&str>,
|
|
) -> Response {
|
|
if status == StatusCode::OK && matches!(response_mode, ResponseMode::Sse) {
|
|
let payload = payload.to_string();
|
|
let stream = stream::once(async move { Ok(Event::default().data(payload)) });
|
|
|
|
return sse_response(status, stream, session_id, protocol_version);
|
|
}
|
|
|
|
json_response(status, payload, session_id, protocol_version)
|
|
}
|
|
|
|
fn sse_response<S>(
|
|
status: StatusCode,
|
|
stream: S,
|
|
session_id: Option<&str>,
|
|
protocol_version: Option<&str>,
|
|
) -> Response
|
|
where
|
|
S: futures_util::stream::Stream<Item = Result<Event, Infallible>> + Send + 'static,
|
|
{
|
|
let mut response = (
|
|
status,
|
|
Sse::new(stream).keep_alive(KeepAlive::new().interval(Duration::from_secs(15))),
|
|
)
|
|
.into_response();
|
|
|
|
if let Some(session_id) = session_id {
|
|
response.headers_mut().insert(
|
|
HEADER_MCP_SESSION_ID,
|
|
HeaderValue::from_str(session_id)
|
|
.unwrap_or_else(|_| HeaderValue::from_static("invalid")),
|
|
);
|
|
}
|
|
|
|
if let Some(protocol_version) = protocol_version {
|
|
response.headers_mut().insert(
|
|
HEADER_MCP_PROTOCOL_VERSION,
|
|
HeaderValue::from_str(protocol_version)
|
|
.unwrap_or_else(|_| HeaderValue::from_static(CURRENT_PROTOCOL_VERSION)),
|
|
);
|
|
}
|
|
|
|
response
|
|
}
|
|
|
|
fn tool_definitions(tool: &PublishedAgentTool) -> Vec<Value> {
|
|
let mut definitions = Vec::new();
|
|
let operation = &tool.operation;
|
|
|
|
match operation
|
|
.execution_config
|
|
.streaming
|
|
.as_ref()
|
|
.map(|value| value.mode)
|
|
{
|
|
Some(crank_core::ExecutionMode::Session) => {
|
|
let streaming = operation
|
|
.execution_config
|
|
.streaming
|
|
.as_ref()
|
|
.expect("streaming");
|
|
if let (Some(start), Some(poll), Some(stop)) = (
|
|
streaming.tool_family.start_tool_name.as_ref(),
|
|
streaming.tool_family.poll_tool_name.as_ref(),
|
|
streaming.tool_family.stop_tool_name.as_ref(),
|
|
) {
|
|
definitions.push(tool_definition(
|
|
start,
|
|
&format!("{} Start", tool.tool_title),
|
|
&format!("Start session for {}", tool.tool_description),
|
|
schema_to_json_schema(&operation.input_schema),
|
|
));
|
|
definitions.push(tool_definition(
|
|
poll,
|
|
&format!("{} Poll", tool.tool_title),
|
|
&format!("Poll session for {}", tool.tool_description),
|
|
id_input_schema("session_id"),
|
|
));
|
|
definitions.push(tool_definition(
|
|
stop,
|
|
&format!("{} Stop", tool.tool_title),
|
|
&format!("Stop session for {}", tool.tool_description),
|
|
id_input_schema("session_id"),
|
|
));
|
|
}
|
|
}
|
|
Some(crank_core::ExecutionMode::AsyncJob) => {
|
|
let streaming = operation
|
|
.execution_config
|
|
.streaming
|
|
.as_ref()
|
|
.expect("streaming");
|
|
if let (Some(start), Some(status), Some(result), Some(cancel)) = (
|
|
streaming.tool_family.start_tool_name.as_ref(),
|
|
streaming.tool_family.status_tool_name.as_ref(),
|
|
streaming.tool_family.result_tool_name.as_ref(),
|
|
streaming.tool_family.cancel_tool_name.as_ref(),
|
|
) {
|
|
definitions.push(tool_definition(
|
|
start,
|
|
&format!("{} Start", tool.tool_title),
|
|
&format!("Start async job for {}", tool.tool_description),
|
|
schema_to_json_schema(&operation.input_schema),
|
|
));
|
|
definitions.push(tool_definition(
|
|
status,
|
|
&format!("{} Status", tool.tool_title),
|
|
&format!("Get job status for {}", tool.tool_description),
|
|
id_input_schema("job_id"),
|
|
));
|
|
definitions.push(tool_definition(
|
|
result,
|
|
&format!("{} Result", tool.tool_title),
|
|
&format!("Get async job result for {}", tool.tool_description),
|
|
id_input_schema("job_id"),
|
|
));
|
|
definitions.push(tool_definition(
|
|
cancel,
|
|
&format!("{} Cancel", tool.tool_title),
|
|
&format!("Cancel async job for {}", tool.tool_description),
|
|
id_input_schema("job_id"),
|
|
));
|
|
}
|
|
}
|
|
_ => {
|
|
definitions.push(tool_definition(
|
|
&tool.tool_name,
|
|
&tool.tool_title,
|
|
&tool.tool_description,
|
|
schema_to_json_schema(&operation.input_schema),
|
|
));
|
|
}
|
|
}
|
|
|
|
definitions
|
|
}
|
|
|
|
fn tool_definition(name: &str, title: &str, description: &str, input_schema: Value) -> Value {
|
|
json!({
|
|
"name": name,
|
|
"title": title,
|
|
"description": description,
|
|
"inputSchema": input_schema
|
|
})
|
|
}
|
|
|
|
fn id_input_schema(field_name: &str) -> Value {
|
|
json!({
|
|
"type": "object",
|
|
"properties": {
|
|
field_name: {
|
|
"type": "string"
|
|
}
|
|
},
|
|
"required": [field_name]
|
|
})
|
|
}
|
|
|
|
fn resolve_generated_tool(
|
|
tools: &[PublishedAgentTool],
|
|
tool_name: &str,
|
|
) -> Option<ResolvedToolCall> {
|
|
for tool in tools {
|
|
if tool.tool_name == tool_name {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::Base,
|
|
});
|
|
}
|
|
|
|
let Some(streaming) = tool.operation.execution_config.streaming.as_ref() else {
|
|
continue;
|
|
};
|
|
|
|
match streaming.mode {
|
|
crank_core::ExecutionMode::Session => {
|
|
if streaming.tool_family.start_tool_name.as_deref() == Some(tool_name) {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::SessionStart,
|
|
});
|
|
}
|
|
if streaming.tool_family.poll_tool_name.as_deref() == Some(tool_name) {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::SessionPoll,
|
|
});
|
|
}
|
|
if streaming.tool_family.stop_tool_name.as_deref() == Some(tool_name) {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::SessionStop,
|
|
});
|
|
}
|
|
}
|
|
crank_core::ExecutionMode::AsyncJob => {
|
|
if streaming.tool_family.start_tool_name.as_deref() == Some(tool_name) {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::AsyncJobStart,
|
|
});
|
|
}
|
|
if streaming.tool_family.status_tool_name.as_deref() == Some(tool_name) {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::AsyncJobStatus,
|
|
});
|
|
}
|
|
if streaming.tool_family.result_tool_name.as_deref() == Some(tool_name) {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::AsyncJobResult,
|
|
});
|
|
}
|
|
if streaming.tool_family.cancel_tool_name.as_deref() == Some(tool_name) {
|
|
return Some(ResolvedToolCall {
|
|
tool: tool.clone(),
|
|
kind: GeneratedToolKind::AsyncJobCancel,
|
|
});
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
fn runtime_operation(tool: &PublishedAgentTool) -> RuntimeOperation {
|
|
let mut operation = RuntimeOperation::from(tool.operation.clone());
|
|
operation.tool_name = tool.tool_name.clone();
|
|
operation.tool_description.title = tool.tool_title.clone();
|
|
operation.tool_description.description = tool.tool_description.clone();
|
|
operation
|
|
}
|
|
|
|
fn schema_to_json_schema(schema: &crank_schema::Schema) -> Value {
|
|
match schema.kind {
|
|
crank_schema::SchemaKind::Object => {
|
|
let mut properties = serde_json::Map::new();
|
|
let mut required = Vec::new();
|
|
|
|
for (field_name, field_schema) in &schema.fields {
|
|
properties.insert(field_name.clone(), schema_to_json_schema(field_schema));
|
|
|
|
if field_schema.required {
|
|
required.push(Value::String(field_name.clone()));
|
|
}
|
|
}
|
|
|
|
json!({
|
|
"type": "object",
|
|
"properties": properties,
|
|
"required": required
|
|
})
|
|
}
|
|
crank_schema::SchemaKind::Array => json!({
|
|
"type": "array",
|
|
"items": schema
|
|
.items
|
|
.as_deref()
|
|
.map(schema_to_json_schema)
|
|
.unwrap_or_else(|| json!({}))
|
|
}),
|
|
crank_schema::SchemaKind::String => json!({ "type": "string" }),
|
|
crank_schema::SchemaKind::Integer => json!({ "type": "integer" }),
|
|
crank_schema::SchemaKind::Number => json!({ "type": "number" }),
|
|
crank_schema::SchemaKind::Boolean => json!({ "type": "boolean" }),
|
|
crank_schema::SchemaKind::Enum => json!({
|
|
"type": "string",
|
|
"enum": schema.enum_values
|
|
}),
|
|
crank_schema::SchemaKind::Null => json!({ "type": "null" }),
|
|
crank_schema::SchemaKind::Oneof => json!({
|
|
"anyOf": schema.variants.iter().map(schema_to_json_schema).collect::<Vec<_>>()
|
|
}),
|
|
}
|
|
}
|
|
|
|
impl AllowedOrigins {
|
|
fn new(public_base_url: Option<String>) -> Self {
|
|
Self {
|
|
public_origin: public_base_url.and_then(|value| extract_origin(&value)),
|
|
}
|
|
}
|
|
|
|
fn is_allowed(&self, origin: &str) -> bool {
|
|
if origin.starts_with("http://localhost")
|
|
|| origin.starts_with("http://127.0.0.1")
|
|
|| origin.starts_with("https://localhost")
|
|
|| origin.starts_with("https://127.0.0.1")
|
|
{
|
|
return true;
|
|
}
|
|
|
|
match &self.public_origin {
|
|
Some(public_origin) => public_origin == origin,
|
|
None => false,
|
|
}
|
|
}
|
|
}
|
|
|
|
fn extract_origin(url: &str) -> Option<String> {
|
|
let mut parts = url.split('/');
|
|
let scheme = parts.next()?;
|
|
let empty = parts.next()?;
|
|
let authority = parts.next()?;
|
|
|
|
if empty.is_empty() {
|
|
return Some(format!("{scheme}//{authority}"));
|
|
}
|
|
|
|
None
|
|
}
|