diff --git a/apps/admin-api/src/service.rs b/apps/admin-api/src/service.rs index f69bad8..c9ba37c 100644 --- a/apps/admin-api/src/service.rs +++ b/apps/admin-api/src/service.rs @@ -2246,10 +2246,18 @@ impl AdminService { let created_at = now_string().map_err(|error| RuntimeError::SecretCrypto { details: error.to_string(), })?; - let expires_at = - add_millis(&created_at, 300_000).map_err(|error| RuntimeError::SecretCrypto { - details: error.to_string(), - })?; + let streaming = runtime.execution_config.streaming.as_ref().ok_or_else(|| { + RuntimeError::MissingStreamingConfig { + operation_id: runtime.operation_id.as_str().to_owned(), + } + })?; + let expires_at = add_millis( + &created_at, + streaming.max_session_lifetime_ms.unwrap_or(300_000), + ) + .map_err(|error| RuntimeError::SecretCrypto { + details: error.to_string(), + })?; let job = AsyncJobHandle { id: crank_core::AsyncJobId::new(new_prefixed_id("job")), workspace_id: workspace_id.clone(), diff --git a/apps/mcp-server/src/app.rs b/apps/mcp-server/src/app.rs index e2a8e57..321004b 100644 --- a/apps/mcp-server/src/app.rs +++ b/apps/mcp-server/src/app.rs @@ -889,6 +889,16 @@ async fn handle_session_poll_call( Err(error) => return internal_jsonrpc_error(message, error), }; + if !stream_session_belongs_to_tool(&loaded, &tool) { + return tool_error_response( + message, + response_mode, + &session.protocol_version, + "stream_session_not_found", + format!("stream session {} was not found", control.session_id), + ); + } + if loaded.is_expired(&now) { let _ = state .registry @@ -1008,7 +1018,7 @@ async fn handle_session_stop_call( session: &SessionState, message: &Value, response_mode: ResponseMode, - _tool: PublishedAgentTool, + tool: PublishedAgentTool, arguments: Value, ) -> Response { let control: SessionControlArgs = match serde_json::from_value(arguments.clone()) { @@ -1024,6 +1034,34 @@ async fn handle_session_stop_call( } }; + 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 !stream_session_belongs_to_tool(&loaded, &tool) { + return tool_error_response( + message, + response_mode, + &session.protocol_version, + "stream_session_not_found", + format!("stream session {} was not found", control.session_id), + ); + } + match state .registry .close_stream_session( @@ -1075,7 +1113,10 @@ async fn handle_async_job_start_call( progress: json!({ "pct": 0 }), result: None, error: None, - expires_at: Some(add_millis(&now, 300_000)), + expires_at: Some(add_millis( + &now, + streaming.max_session_lifetime_ms.unwrap_or(300_000), + )), created_at: now.clone(), updated_at: now.clone(), finished_at: None, @@ -1222,6 +1263,16 @@ async fn handle_async_job_status_call( Err(error) => return internal_jsonrpc_error(message, error), }; + if !async_job_belongs_to_tool(&job, &_tool) { + return tool_error_response( + message, + response_mode, + &session.protocol_version, + "async_job_not_found", + format!("async job {} was not found", control.job_id), + ); + } + success_tool_response( message, response_mode, @@ -1275,6 +1326,16 @@ async fn handle_async_job_result_call( Err(error) => return internal_jsonrpc_error(message, error), }; + if !async_job_belongs_to_tool(&job, &_tool) { + return tool_error_response( + message, + response_mode, + &session.protocol_version, + "async_job_not_found", + format!("async job {} was not found", control.job_id), + ); + } + match job.status { JobStatus::Completed => success_tool_response( message, @@ -1330,6 +1391,34 @@ async fn handle_async_job_cancel_call( } }; + match state + .registry + .get_async_job(&AsyncJobId::new(control.job_id.clone())) + .await + { + Ok(Some(job)) => { + if !async_job_belongs_to_tool(&job, &_tool) { + return tool_error_response( + message, + response_mode, + &session.protocol_version, + "async_job_not_found", + format!("async job {} was not found", control.job_id), + ); + } + } + 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 state .registry .cancel_async_job(&AsyncJobId::new(control.job_id.clone()), &now_rfc3339()) @@ -1588,7 +1677,7 @@ fn negotiate_post_response_mode(headers: &HeaderMap) -> Result Result<(), StatusCode> { Err(StatusCode::NOT_ACCEPTABLE) } +fn stream_session_belongs_to_tool( + session_record: &StreamSession, + tool: &PublishedAgentTool, +) -> bool { + session_record.workspace_id == tool.workspace_id + && session_record.agent_id.as_ref() == Some(&tool.agent_id) + && session_record.operation_id == tool.operation.id +} + +fn async_job_belongs_to_tool(job: &AsyncJobHandle, tool: &PublishedAgentTool) -> bool { + job.workspace_id == tool.workspace_id + && job.agent_id.as_ref() == Some(&tool.agent_id) + && job.operation_id == tool.operation.id +} + fn protocol_version_from_headers(headers: &HeaderMap) -> Result { let Some(version) = headers.get(HEADER_MCP_PROTOCOL_VERSION) else { return Ok(DEFAULT_PROTOCOL_VERSION.to_owned()); diff --git a/apps/mcp-server/src/main.rs b/apps/mcp-server/src/main.rs index 75b8048..330419c 100644 --- a/apps/mcp-server/src/main.rs +++ b/apps/mcp-server/src/main.rs @@ -568,6 +568,47 @@ mod tests { assert_eq!(after_delete.status(), reqwest::StatusCode::NOT_FOUND); } + #[tokio::test] + async fn initialize_accepts_json_only_response_negotiation() { + let registry = test_registry().await; + publish_agent_with_bindings(®istry, "sales-json-accept", vec![]).await; + let api_key = + create_platform_api_key(®istry, "mcp-json-accept", &[PlatformApiKeyScope::Read]) + .await; + let base_url = spawn_mcp_server(build_test_app( + registry, + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let response = client + .post(agent_mcp_url(&base_url, "sales-json-accept")) + .header(header::ACCEPT, "application/json") + .header(header::AUTHORIZATION, format!("Bearer {api_key}")) + .json(&json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": "2025-11-25" + } + })) + .send() + .await + .unwrap(); + + assert_eq!(response.status(), reqwest::StatusCode::OK); + assert_eq!( + response + .headers() + .get(header::CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .unwrap(), + "application/json" + ); + } + #[tokio::test] async fn rejects_get_with_protocol_version_mismatch() { let registry = test_registry().await; @@ -989,6 +1030,194 @@ mod tests { ); } + #[tokio::test] + async fn rejects_cross_agent_session_poll() { + let registry = test_registry().await; + let server_addr = grpc_test_support::spawn_unary_echo_server().await; + let operation_a = test_grpc_session_operation(&server_addr, "echo_stream_session_a"); + let operation_b = test_grpc_session_operation(&server_addr, "echo_stream_session_b"); + + for operation in [&operation_a, &operation_b] { + registry + .create_operation(&test_workspace_id(), operation, Some("alice")) + .await + .unwrap(); + registry + .publish_operation(PublishRequest { + workspace_id: &test_workspace_id(), + operation_id: &operation.id, + version: 1, + published_at: "2026-03-26T10:00:00Z", + published_by: Some("alice"), + }) + .await + .unwrap(); + } + + publish_agent_for_operation(®istry, &operation_a, "sales-session-a").await; + publish_agent_for_operation(®istry, &operation_b, "sales-session-b").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-cross-session", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_test_app( + registry, + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let agent_a_url = agent_mcp_url(&base_url, "sales-session-a"); + let agent_b_url = agent_mcp_url(&base_url, "sales-session-b"); + let session_a = initialize_session(&client, &agent_a_url, &api_key).await; + let session_b = initialize_session(&client, &agent_b_url, &api_key).await; + + let start_response = post_jsonrpc( + &client, + &agent_a_url, + &api_key, + Some(&session_a), + json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "echo_stream_session_a_start", + "arguments": { "message": "hello" } + } + }), + ) + .await; + let foreign_session_id = start_response["result"]["structuredContent"]["session_id"] + .as_str() + .unwrap() + .to_owned(); + + let poll_response = post_jsonrpc( + &client, + &agent_b_url, + &api_key, + Some(&session_b), + json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "echo_stream_session_b_poll", + "arguments": { "session_id": foreign_session_id } + } + }), + ) + .await; + + assert_eq!( + poll_response["result"]["structuredContent"]["error"]["code"], + json!("stream_session_not_found") + ); + assert!( + poll_response["result"]["structuredContent"]["error"]["message"] + .as_str() + .is_some_and(|message| message.starts_with("stream session ")) + ); + } + + #[tokio::test] + async fn rejects_cross_agent_async_job_access() { + let registry = test_registry().await; + let upstream_base_url = spawn_upstream_server().await; + let operation_a = test_rest_async_job_operation(&upstream_base_url, "crm_async_a"); + let operation_b = test_rest_async_job_operation(&upstream_base_url, "crm_async_b"); + + for operation in [&operation_a, &operation_b] { + registry + .create_operation(&test_workspace_id(), operation, Some("alice")) + .await + .unwrap(); + registry + .publish_operation(PublishRequest { + workspace_id: &test_workspace_id(), + operation_id: &operation.id, + version: 1, + published_at: "2026-03-26T10:00:00Z", + published_by: Some("alice"), + }) + .await + .unwrap(); + } + + publish_agent_for_operation(®istry, &operation_a, "sales-async-a").await; + publish_agent_for_operation(®istry, &operation_b, "sales-async-b").await; + let api_key = create_platform_api_key( + ®istry, + "mcp-cross-async", + &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write], + ) + .await; + + let base_url = spawn_mcp_server(build_test_app( + registry, + Duration::from_millis(0), + Some("https://crank.example.com".to_owned()), + )) + .await; + let client = reqwest::Client::new(); + let agent_a_url = agent_mcp_url(&base_url, "sales-async-a"); + let agent_b_url = agent_mcp_url(&base_url, "sales-async-b"); + let session_a = initialize_session(&client, &agent_a_url, &api_key).await; + let session_b = initialize_session(&client, &agent_b_url, &api_key).await; + + let start_response = post_jsonrpc( + &client, + &agent_a_url, + &api_key, + Some(&session_a), + json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "tools/call", + "params": { + "name": "crm_async_a_start", + "arguments": { "email": "user@example.com" } + } + }), + ) + .await; + let foreign_job_id = start_response["result"]["structuredContent"]["job_id"] + .as_str() + .unwrap() + .to_owned(); + + let status_response = post_jsonrpc( + &client, + &agent_b_url, + &api_key, + Some(&session_b), + json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": { + "name": "crm_async_b_status", + "arguments": { "job_id": foreign_job_id } + } + }), + ) + .await; + + assert_eq!( + status_response["result"]["structuredContent"]["error"]["code"], + json!("async_job_not_found") + ); + assert!( + status_response["result"]["structuredContent"]["error"]["message"] + .as_str() + .is_some_and(|message| message.starts_with("async job ")) + ); + } + #[tokio::test] async fn exposes_and_runs_async_job_tool_family_via_mcp() { let registry = test_registry().await;