Refine Rust architecture boundaries
This commit is contained in:
@@ -10,12 +10,9 @@ use axum::{
|
||||
extract::{Path, State},
|
||||
http::{
|
||||
HeaderMap, HeaderValue, StatusCode,
|
||||
header::{self, ACCEPT, AUTHORIZATION, HeaderName, RETRY_AFTER},
|
||||
},
|
||||
response::{
|
||||
IntoResponse, Response,
|
||||
sse::{Event, KeepAlive, Sse},
|
||||
header::{AUTHORIZATION, RETRY_AFTER},
|
||||
},
|
||||
response::{IntoResponse, Response, sse::Event},
|
||||
routing::get,
|
||||
};
|
||||
use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
|
||||
@@ -25,8 +22,8 @@ use crank_core::{
|
||||
};
|
||||
use crank_registry::{CreateInvocationLogRequest, PostgresRegistry, PublishedAgentTool};
|
||||
use crank_runtime::{
|
||||
RateLimitRejection, RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutor,
|
||||
RuntimeOperation, RuntimeRequestContext, SecretCrypto,
|
||||
RateLimitRejection, RequestRateLimiter, ResolvedAuth, RuntimeError, RuntimeExecutionRequest,
|
||||
RuntimeExecutor, RuntimeOperation, RuntimeRequestContext, SecretCrypto,
|
||||
};
|
||||
use futures_util::stream;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@@ -39,9 +36,8 @@ use crate::{
|
||||
auth::{SharedMachineCredentialVerifier, VerifiedMachineCredential},
|
||||
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,
|
||||
DEFAULT_PROTOCOL_VERSION, is_notification, is_request, is_response, jsonrpc_error,
|
||||
jsonrpc_result, method_name, negotiated_protocol_version, params, request_id,
|
||||
},
|
||||
manifest::tool_definitions,
|
||||
session::{SessionState, SharedSessionStore},
|
||||
@@ -49,20 +45,16 @@ use crate::{
|
||||
ToolErrorContract, generic_tool_error_contract, runtime_error_code,
|
||||
tool_error_contract_from_runtime, tool_error_text, tool_error_value,
|
||||
},
|
||||
transport::{
|
||||
AllowedOrigins, HEADER_MCP_SESSION_ID, ResponseMode, json_response,
|
||||
negotiate_post_response_mode, protocol_version_from_headers, resolve_request_id,
|
||||
session_id_from_headers, sse_response, transport_response, validate_get_accept_header,
|
||||
validate_origin, validate_session_protocol_version, with_request_id_header,
|
||||
},
|
||||
};
|
||||
|
||||
const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
|
||||
const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";
|
||||
const HEADER_X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
|
||||
const MAX_REQUEST_ID_LEN: usize = 128;
|
||||
const TRANSPORT_SESSION_TTL_MS: u64 = 86_400_000;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
enum ResponseMode {
|
||||
Json,
|
||||
Sse,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AppState {
|
||||
registry: PostgresRegistry,
|
||||
@@ -75,11 +67,6 @@ pub struct AppState {
|
||||
allowed_origins: AllowedOrigins,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct AllowedOrigins {
|
||||
public_origin: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Deserialize)]
|
||||
struct InitializeParams {
|
||||
#[serde(rename = "protocolVersion")]
|
||||
@@ -623,11 +610,10 @@ async fn handle_base_tool_call(
|
||||
Ok(resolved_auth) => {
|
||||
state
|
||||
.runtime
|
||||
.execute_with_auth_and_context(
|
||||
&operation,
|
||||
&arguments,
|
||||
resolved_auth.as_ref(),
|
||||
Some(&runtime_request_context),
|
||||
.execute_request(
|
||||
RuntimeExecutionRequest::new(&operation, &arguments)
|
||||
.with_optional_auth(resolved_auth.as_ref())
|
||||
.with_context(&runtime_request_context),
|
||||
)
|
||||
.await
|
||||
}
|
||||
@@ -1022,141 +1008,6 @@ fn serialize_machine_access_mode(mode: crank_core::MachineAccessMode) -> &'stati
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
@@ -1308,109 +1159,6 @@ fn hash_access_secret(secret: &str) -> String {
|
||||
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 with_request_id_header(mut response: Response, request_id: &str) -> Response {
|
||||
if let Ok(value) = HeaderValue::from_str(request_id) {
|
||||
response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
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 resolve_request_id(headers: &HeaderMap) -> String {
|
||||
headers
|
||||
.get(&HEADER_X_REQUEST_ID)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| is_valid_request_id(value))
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| uuid::Uuid::now_v7().to_string())
|
||||
}
|
||||
|
||||
fn is_valid_request_id(value: &str) -> bool {
|
||||
!value.is_empty()
|
||||
&& value.len() <= MAX_REQUEST_ID_LEN
|
||||
&& value
|
||||
.bytes()
|
||||
.all(|byte| matches!(byte, 0x21..=0x7e) && byte != b',' && byte != b';')
|
||||
}
|
||||
|
||||
fn resolve_generated_tool(
|
||||
tools: &[PublishedAgentTool],
|
||||
tool_name: &str,
|
||||
@@ -1432,42 +1180,6 @@ fn runtime_operation(tool: &PublishedAgentTool) -> RuntimeOperation {
|
||||
operation
|
||||
}
|
||||
|
||||
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
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use axum::body::to_bytes;
|
||||
|
||||
@@ -5,5 +5,6 @@ pub mod jsonrpc;
|
||||
pub mod manifest;
|
||||
pub mod session;
|
||||
pub mod tool_error;
|
||||
mod transport;
|
||||
|
||||
pub use app::build_app;
|
||||
|
||||
@@ -0,0 +1,311 @@
|
||||
use std::{convert::Infallible, time::Duration};
|
||||
|
||||
use axum::{
|
||||
Json,
|
||||
http::{
|
||||
HeaderMap, HeaderValue, StatusCode,
|
||||
header::{self, ACCEPT, HeaderName},
|
||||
},
|
||||
response::{
|
||||
IntoResponse, Response,
|
||||
sse::{Event, KeepAlive, Sse},
|
||||
},
|
||||
};
|
||||
use futures_util::stream;
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::jsonrpc::{
|
||||
CURRENT_PROTOCOL_VERSION, DEFAULT_PROTOCOL_VERSION, negotiated_protocol_version,
|
||||
};
|
||||
|
||||
pub(super) const HEADER_MCP_SESSION_ID: &str = "MCP-Session-Id";
|
||||
pub(super) const HEADER_MCP_PROTOCOL_VERSION: &str = "MCP-Protocol-Version";
|
||||
pub(super) const HEADER_X_REQUEST_ID: HeaderName = HeaderName::from_static("x-request-id");
|
||||
const MAX_REQUEST_ID_LEN: usize = 128;
|
||||
|
||||
#[derive(Clone, Copy)]
|
||||
pub(super) enum ResponseMode {
|
||||
Json,
|
||||
Sse,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub(super) struct AllowedOrigins {
|
||||
public_origin: Option<String>,
|
||||
}
|
||||
|
||||
impl AllowedOrigins {
|
||||
pub(super) fn new(public_base_url: Option<String>) -> Self {
|
||||
Self {
|
||||
public_origin: public_base_url.and_then(|value| extract_origin(&value)),
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) 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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) 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)
|
||||
}
|
||||
|
||||
pub(super) 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)
|
||||
}
|
||||
|
||||
pub(super) 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)
|
||||
}
|
||||
|
||||
pub(super) 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)
|
||||
}
|
||||
|
||||
pub(super) 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(())
|
||||
}
|
||||
|
||||
pub(super) 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()))
|
||||
}
|
||||
|
||||
pub(super) 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
|
||||
}
|
||||
|
||||
pub(super) 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)
|
||||
}
|
||||
|
||||
pub(super) fn with_request_id_header(mut response: Response, request_id: &str) -> Response {
|
||||
if let Ok(value) = HeaderValue::from_str(request_id) {
|
||||
response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
pub(super) 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
|
||||
}
|
||||
|
||||
pub(super) fn resolve_request_id(headers: &HeaderMap) -> String {
|
||||
headers
|
||||
.get(&HEADER_X_REQUEST_ID)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| is_valid_request_id(value))
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| uuid::Uuid::now_v7().to_string())
|
||||
}
|
||||
|
||||
pub(super) fn is_valid_request_id(value: &str) -> bool {
|
||||
!value.is_empty()
|
||||
&& value.len() <= MAX_REQUEST_ID_LEN
|
||||
&& value
|
||||
.bytes()
|
||||
.all(|byte| matches!(byte, 0x21..=0x7e) && byte != b',' && byte != b';')
|
||||
}
|
||||
|
||||
pub(super) 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
|
||||
}
|
||||
Reference in New Issue
Block a user