diff --git a/TASKS.md b/TASKS.md index 40447f8..5038a98 100644 --- a/TASKS.md +++ b/TASKS.md @@ -2,20 +2,20 @@ ## Current -### `feat/grpc-server-streaming-adapter` +### `feat/session-and-job-tools` Status: completed DoD: -- gRPC adapter supports bounded server-stream collection for `window` mode -- streamed protobuf messages are decoded into normalized JSON items -- adapter respects item/window limits and upstream timeout -- runtime dispatches gRPC `window` mode through streaming adapter path -- local gRPC upstream tests cover success, timeout and malformed method kind handling +- session and async_job tool families are published in the MCP tool catalog +- generated `start/poll/stop` and `start/status/result/cancel` tool calls work through JSON-RPC +- stream sessions and async jobs are persisted with explicit status transitions +- session and async_job tool calls are logged through observability +- MCP integration tests cover `session start -> poll -> stop` and `async job start -> status -> result` ## Next -- `feat/grpc-server-streaming-adapter` +- `feat/streaming-ui-config` ## Backlog diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index 04233e1..21b4541 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -19,10 +19,14 @@ use axum::{ }; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use crank_core::{ - InvocationLevel, InvocationLog, InvocationLogId, InvocationSource, InvocationStatus, - PlatformApiKeyScope, + AsyncJobHandle, AsyncJobId, InvocationLevel, InvocationLog, InvocationLogId, InvocationSource, + InvocationStatus, JobStatus, PlatformApiKeyScope, StreamSession, StreamSessionId, StreamStatus, +}; +use crank_registry::{ + CreateAsyncJobRequest, CreateInvocationLogRequest, CreateStreamSessionRequest, + PostgresRegistry, PublishedAgentTool, UpdateAsyncJobStatusRequest, + UpdateStreamSessionStateRequest, }; -use crank_registry::{CreateInvocationLogRequest, PostgresRegistry, PublishedAgentTool}; use crank_runtime::{RuntimeError, RuntimeExecutor, RuntimeOperation}; use futures_util::stream; use serde::{Deserialize, Serialize}; @@ -76,6 +80,34 @@ struct ToolCallParams { 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, @@ -273,7 +305,7 @@ async fn mcp_post( .await { Ok(tools) => { - let definitions = tools.iter().map(tool_definition).collect::>(); + let definitions = tools.iter().flat_map(tool_definitions).collect::>(); transport_response( StatusCode::OK, @@ -313,110 +345,33 @@ async fn mcp_post( match state .catalog - .get_tool( - &session.workspace_slug, - &session.agent_slug, - &tool_call_params.name, - ) + .list_tools(&session.workspace_slug, &session.agent_slug) .await { - Ok(Some(tool)) => { - let runtime_operation = runtime_operation(&tool); - let request_preview = - build_request_preview(&state.runtime, &runtime_operation, &arguments); - let started_at = Instant::now(); - - match state.runtime.execute(&runtime_operation, &arguments).await { - Ok(output) => transport_response( - StatusCode::OK, - { - let _ = persist_invocation( - &state, - &tool, - InvocationRecord { - 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; - 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 - }), - ) - }, + Ok(tools) => match resolve_generated_tool(&tools, &tool_call_params.name) { + Some(resolved) => { + handle_tool_call( + state.clone(), + &session, + &message, response_mode, - None, - Some(&session.protocol_version), - ), - Err(error) => { - let _ = persist_invocation( - &state, - &tool, - InvocationRecord { - 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; - transport_response( - StatusCode::OK, - jsonrpc_result( - request_id(&message), - json!({ - "content": [ - { - "type": "text", - "text": error.to_string() - } - ], - "structuredContent": { - "error": { - "code": runtime_error_code(&error), - "message": error.to_string() - } - }, - "isError": true - }), - ), - response_mode, - None, - Some(&session.protocol_version), - ) - } + resolved, + arguments, + ) + .await } - } - Ok(None) => transport_response( - StatusCode::OK, - jsonrpc_error( - request_id(&message), - -32602, - format!("tool {} was not found", tool_call_params.name), + 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), ), - response_mode, - None, - Some(&session.protocol_version), - ), + }, Err(error) => internal_jsonrpc_error(&message, error), } } @@ -445,6 +400,805 @@ async fn mcp_post( } } +async fn handle_tool_call( + state: Arc, + 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, + next_index: usize, + batch_size: usize, +} + +async fn handle_base_tool_call( + state: Arc, + 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(); + + match state.runtime.execute(&operation, &arguments).await { + 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, + 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, + 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::(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, + 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, + 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, + 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, + 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, + 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, path: &AgentRoutePath, @@ -794,6 +1548,7 @@ fn build_request_preview( } struct InvocationRecord<'a> { + tool_name: &'a str, status: InvocationStatus, level: InvocationLevel, message: &'a str, @@ -821,7 +1576,7 @@ async fn persist_invocation( source: InvocationSource::AgentToolCall, level: record.level, status: record.status, - tool_name: tool.tool_name.clone(), + tool_name: record.tool_name.to_owned(), message: record.message.to_owned(), request_id: None, status_code: record.status_code, @@ -838,6 +1593,104 @@ async fn persist_invocation( .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", @@ -936,15 +1789,189 @@ where response } -fn tool_definition(tool: &PublishedAgentTool) -> Value { +fn tool_definitions(tool: &PublishedAgentTool) -> Vec { + 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": tool.tool_name, - "title": tool.tool_title, - "description": tool.tool_description, - "inputSchema": schema_to_json_schema(&tool.operation.input_schema) + "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 { + 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(); diff --git a/apps/mcp-server/src/catalog.rs b/apps/mcp-server/src/catalog.rs index 3a2df58..02521b9 100644 --- a/apps/mcp-server/src/catalog.rs +++ b/apps/mcp-server/src/catalog.rs @@ -25,7 +25,6 @@ struct CatalogKey { struct CachedCatalog { loaded_at: Option, tools: Vec, - tools_by_name: HashMap, } impl PublishedToolCatalog { @@ -50,20 +49,6 @@ impl PublishedToolCatalog { .unwrap_or_default()) } - pub async fn get_tool( - &self, - workspace_slug: &str, - agent_slug: &str, - tool_name: &str, - ) -> Result, RegistryError> { - self.refresh_if_stale(workspace_slug, agent_slug).await?; - let guard = self.cached.read().await; - 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, workspace_slug: &str, @@ -92,11 +77,6 @@ impl PublishedToolCatalog { Err(RegistryError::PublishedAgentNotFound { .. }) => Vec::new(), Err(error) => return Err(error), }; - let tools_by_name = tools - .iter() - .cloned() - .map(|tool| (tool.tool_name.clone(), tool)) - .collect::>(); let mut guard = self.cached.write().await; let previous_count = guard .get(&key) @@ -108,7 +88,6 @@ impl PublishedToolCatalog { CachedCatalog { loaded_at: Some(Instant::now()), tools, - tools_by_name, }, ); diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index cd034f2..8c2c0b3 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -52,10 +52,11 @@ mod tests { use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; use crank_adapter_grpc::test_support as grpc_test_support; use crank_core::{ - Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, DescriptorId, - ExecutionConfig, GraphqlOperationType, GraphqlTarget, GrpcTarget, HttpMethod, Operation, - OperationId, OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyScope, - PlatformApiKeyStatus, Protocol, RestTarget, Target, ToolDescription, WorkspaceId, + Agent, AgentId, AgentOperationBinding, AgentStatus, AgentVersion, AggregationMode, + DescriptorId, ExecutionConfig, ExecutionMode, GraphqlOperationType, GraphqlTarget, + GrpcTarget, HttpMethod, Operation, OperationId, OperationStatus, PlatformApiKey, + PlatformApiKeyId, PlatformApiKeyScope, PlatformApiKeyStatus, Protocol, RestTarget, + StreamingConfig, Target, ToolDescription, ToolFamilyConfig, TransportBehavior, WorkspaceId, }; use crank_mapping::{MappingRule, MappingSet}; use crank_registry::{ @@ -67,6 +68,7 @@ mod tests { use sha2::{Digest, Sha256}; use sqlx::{Executor, postgres::PgPoolOptions}; use tokio::net::TcpListener; + use tokio::time::sleep; use crate::app::build_app; @@ -821,6 +823,297 @@ mod tests { assert_eq!(response.status(), reqwest::StatusCode::FORBIDDEN); } + #[tokio::test] + async fn exposes_and_runs_session_tool_family_via_mcp() { + let registry = test_registry().await; + let server_addr = grpc_test_support::spawn_unary_echo_server().await; + let operation = test_grpc_session_operation(&server_addr, "echo_stream_session"); + + 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_for_operation(®istry, &operation, "sales-session").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-session", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_app( + registry.clone(), + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let mcp_url = agent_mcp_url(&base_url, "sales-session"); + let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; + + let tools = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": {} + }), + ) + .await; + let tool_names = tools["result"]["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["name"].as_str().unwrap().to_owned()) + .collect::>(); + assert!(!tool_names.contains(&operation.name)); + assert!(tool_names.contains(&"echo_stream_session_start".to_owned())); + assert!(tool_names.contains(&"echo_stream_session_poll".to_owned())); + assert!(tool_names.contains(&"echo_stream_session_stop".to_owned())); + + let start_response = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": { + "name": "echo_stream_session_start", + "arguments": { + "message": "hello" + } + } + }), + ) + .await; + let session_id = start_response["result"]["structuredContent"]["session_id"] + .as_str() + .unwrap() + .to_owned(); + assert_eq!( + start_response["result"]["structuredContent"]["status"], + json!("running") + ); + + let poll_response = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 4, + "method": "tools/call", + "params": { + "name": "echo_stream_session_poll", + "arguments": { + "session_id": session_id + } + } + }), + ) + .await; + assert_eq!( + poll_response["result"]["structuredContent"]["session_id"], + json!(session_id) + ); + assert!( + poll_response["result"]["structuredContent"]["items"] + .as_array() + .is_some_and(|items| !items.is_empty()) + ); + + let stop_response = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": { + "name": "echo_stream_session_stop", + "arguments": { + "session_id": session_id + } + } + }), + ) + .await; + assert_eq!( + stop_response["result"]["structuredContent"], + json!({ + "session_id": session_id, + "status": "stopped" + }) + ); + } + + #[tokio::test] + async fn exposes_and_runs_async_job_tool_family_via_mcp() { + let registry = test_registry().await; + let upstream_base_url = spawn_upstream_server().await; + let operation = test_rest_async_job_operation(&upstream_base_url, "crm_async_create_lead"); + + 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_for_operation(®istry, &operation, "sales-async").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-async", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_app( + registry.clone(), + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let mcp_url = agent_mcp_url(&base_url, "sales-async"); + let initialized_session = initialize_session(&client, &mcp_url, &api_key).await; + + let tools = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/list", + "params": {} + }), + ) + .await; + let tool_names = tools["result"]["tools"] + .as_array() + .unwrap() + .iter() + .map(|tool| tool["name"].as_str().unwrap().to_owned()) + .collect::>(); + assert!(!tool_names.contains(&operation.name)); + assert!(tool_names.contains(&"crm_async_create_lead_start".to_owned())); + assert!(tool_names.contains(&"crm_async_create_lead_status".to_owned())); + assert!(tool_names.contains(&"crm_async_create_lead_result".to_owned())); + assert!(tool_names.contains(&"crm_async_create_lead_cancel".to_owned())); + + let start_response = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": { + "name": "crm_async_create_lead_start", + "arguments": { + "email": "user@example.com" + } + } + }), + ) + .await; + let job_id = start_response["result"]["structuredContent"]["job_id"] + .as_str() + .unwrap() + .to_owned(); + assert_eq!( + start_response["result"]["structuredContent"]["status"], + json!("running") + ); + + let mut status_response = Value::Null; + for _ in 0..20 { + status_response = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 4, + "method": "tools/call", + "params": { + "name": "crm_async_create_lead_status", + "arguments": { + "job_id": job_id + } + } + }), + ) + .await; + + if status_response["result"]["structuredContent"]["status"] == json!("completed") { + break; + } + + sleep(Duration::from_millis(25)).await; + } + + assert_eq!( + status_response["result"]["structuredContent"]["status"], + json!("completed") + ); + + let result_response = post_jsonrpc( + &client, + &mcp_url, + &api_key, + Some(&initialized_session), + json!({ + "jsonrpc": "2.0", + "id": 5, + "method": "tools/call", + "params": { + "name": "crm_async_create_lead_result", + "arguments": { + "job_id": job_id + } + } + }), + ) + .await; + assert_eq!( + result_response["result"]["structuredContent"], + json!({ "id": "lead_123" }) + ); + } + async fn initialize_session(client: &reqwest::Client, mcp_url: &str, api_key: &str) -> String { let initialize_response = client .post(mcp_url) @@ -926,7 +1219,9 @@ mod tests { } async fn spawn_upstream_server() -> String { - let app = Router::new().route("/crm/leads", post(create_lead)); + let app = Router::new() + .route("/crm/leads", post(create_lead)) + .route("/crm/slow-leads", post(create_slow_lead)); let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); @@ -1039,6 +1334,15 @@ mod tests { })) } + async fn create_slow_lead(Json(payload): Json) -> Json { + sleep(Duration::from_millis(250)).await; + + Json(json!({ + "id": "lead_123", + "email": payload["email"] + })) + } + async fn graphql_handler(Json(payload): Json) -> Json { let email = payload .get("variables") @@ -1274,6 +1578,98 @@ mod tests { } } + fn test_grpc_session_operation(server_addr: &str, name: &str) -> Operation { + let mut operation = test_grpc_operation(server_addr, name); + operation.target = Target::Grpc(GrpcTarget { + server_addr: server_addr.to_owned(), + package: "echo".to_owned(), + service: "EchoService".to_owned(), + method: "ServerEcho".to_owned(), + descriptor_ref: DescriptorId::new("desc_echo"), + descriptor_set_b64: grpc_test_support::descriptor_set_b64(), + }); + operation.display_name = "Server Echo Session".to_owned(); + operation.tool_description = ToolDescription { + title: "Server Echo Session".to_owned(), + description: "Streams echo messages through a bounded session".to_owned(), + tags: vec!["grpc".to_owned(), "streaming".to_owned()], + examples: Vec::new(), + }; + operation.execution_config.streaming = Some(StreamingConfig { + mode: ExecutionMode::Session, + transport_behavior: TransportBehavior::ServerStream, + window_duration_ms: Some(1_000), + poll_interval_ms: Some(250), + upstream_timeout_ms: Some(1_000), + idle_timeout_ms: Some(5_000), + max_session_lifetime_ms: Some(60_000), + max_items: Some(1), + max_bytes: Some(16 * 1024), + aggregation_mode: AggregationMode::SummaryPlusSamples, + summary_path: None, + items_path: Some("$.items".to_owned()), + cursor_path: None, + status_path: None, + done_path: Some("$.done".to_owned()), + redacted_paths: Vec::new(), + truncate_item_fields: false, + max_field_length: None, + drop_duplicates: false, + sampling_rate: None, + tool_family: ToolFamilyConfig { + start_tool_name: Some(format!("{name}_start")), + poll_tool_name: Some(format!("{name}_poll")), + stop_tool_name: Some(format!("{name}_stop")), + status_tool_name: None, + result_tool_name: None, + cancel_tool_name: None, + }, + }); + operation + } + + fn test_rest_async_job_operation(base_url: &str, name: &str) -> Operation { + let mut operation = test_operation(base_url, name); + operation.display_name = "Create Lead Async".to_owned(); + operation.tool_description = ToolDescription { + title: "Create Lead Async".to_owned(), + description: "Creates a CRM lead through an async job tool family".to_owned(), + tags: vec!["rest".to_owned(), "async_job".to_owned()], + examples: Vec::new(), + }; + operation.execution_config.streaming = Some(StreamingConfig { + mode: ExecutionMode::AsyncJob, + transport_behavior: TransportBehavior::DeferredResult, + window_duration_ms: None, + poll_interval_ms: Some(250), + upstream_timeout_ms: Some(2_000), + idle_timeout_ms: None, + max_session_lifetime_ms: Some(300_000), + max_items: None, + max_bytes: Some(16 * 1024), + aggregation_mode: AggregationMode::SummaryOnly, + summary_path: None, + items_path: None, + cursor_path: None, + status_path: None, + done_path: None, + redacted_paths: Vec::new(), + truncate_item_fields: false, + max_field_length: None, + drop_duplicates: false, + sampling_rate: None, + tool_family: ToolFamilyConfig { + start_tool_name: Some(format!("{name}_start")), + poll_tool_name: None, + stop_tool_name: None, + status_tool_name: Some(format!("{name}_status")), + result_tool_name: Some(format!("{name}_result")), + cancel_tool_name: Some(format!("{name}_cancel")), + }, + }); + operation + } + fn object_schema(field_name: &str) -> Schema { Schema { kind: SchemaKind::Object, diff --git a/crates/crank-runtime/src/executor.rs b/crates/crank-runtime/src/executor.rs index c4b2748..d73ba6f 100644 --- a/crates/crank-runtime/src/executor.rs +++ b/crates/crank-runtime/src/executor.rs @@ -74,6 +74,40 @@ impl RuntimeExecutor { crate::aggregation::collect_window_result(&adapter_response.body, streaming) } + pub async fn execute_session_seed( + &self, + operation: &RuntimeOperation, + input: &Value, + ) -> Result { + let Some(streaming) = operation.execution_config.streaming.as_ref() else { + return Err(RuntimeError::MissingStreamingConfig { + operation_id: operation.operation_id.as_str().to_owned(), + }); + }; + + if streaming.mode != ExecutionMode::Session { + return Err(RuntimeError::UnsupportedExecutionMode { + operation_id: operation.operation_id.as_str().to_owned(), + mode: streaming.mode, + }); + } + + let batch_size = streaming.max_items.unwrap_or(10).max(1); + let seed_limit = batch_size.saturating_mul(4); + let mut seeded_operation = operation.clone(); + if let Some(config) = seeded_operation.execution_config.streaming.as_mut() { + config.mode = ExecutionMode::Window; + config.window_duration_ms = config + .window_duration_ms + .or(config.poll_interval_ms) + .or(config.upstream_timeout_ms) + .or(Some(operation.execution_config.timeout_ms)); + config.max_items = Some(seed_limit); + } + + self.execute_window(&seeded_operation, input).await + } + pub fn prepare_request( &self, operation: &RuntimeOperation,