modelcontextprotocol/rust-sdk
Watch
1
0
Fork
You've already forked rust-sdk
 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 2

feat(extension): extract http request part into rmcp extension (#163)

Browse source 
* feat(extension): extract http request part into rmcp extension

* perf(extension): insert extension from one entrypoint
This commit is contained in:
4t145 2025-05-09 10:36:52 +08:00 • committed by GitHub
commit 787cc015b7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 67 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

6
crates/rmcp/src/transport/sse_server.rs
View file

@ -3,7 +3,7 @@ use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};
use axum::{
Json, Router,
extract::{Query, State},
http::StatusCode,
http::{StatusCode, request::Parts},
response::{
Response,
sse::{Event, KeepAlive, Sse},
@ -64,7 +64,8 @@ pub struct PostEventQuery {
async fn post_event_handler(
State(app): State<App>,
Query(PostEventQuery { session_id }): Query<PostEventQuery>,
Json(message): Json<ClientJsonRpcMessage>,
parts: Parts,
Json(mut message): Json<ClientJsonRpcMessage>,
) -> Result<StatusCode, StatusCode> {
tracing::debug!(session_id, ?message, "new client message");
let tx = {
@ -73,6 +74,7 @@ async fn post_event_handler(
.ok_or(StatusCode::NOT_FOUND)?
.clone()
};
message.insert_extension(parts);
if tx.send(message).await.is_err() {
tracing::error!("send message error");
return Err(StatusCode::GONE);

12
crates/rmcp/src/transport/streamable_http_server/axum.rs
View file

@ -3,7 +3,7 @@ use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};
use axum::{
Json, Router,
extract::State,
http::{HeaderMap, HeaderValue, StatusCode},
http::{HeaderMap, HeaderValue, StatusCode, request::Parts},
response::{
IntoResponse, Response,
sse::{Event, KeepAlive, Sse},
@ -68,11 +68,11 @@ fn receiver_as_stream(
async fn post_handler(
State(app): State<App>,
header_map: HeaderMap,
Json(message): Json<ClientJsonRpcMessage>,
parts: Parts,
Json(mut message): Json<ClientJsonRpcMessage>,
) -> Result<Response, Response> {
use futures::StreamExt;
if let Some(session_id) = header_map.get(HEADER_SESSION_ID) {
if let Some(session_id) = parts.headers.get(HEADER_SESSION_ID).cloned() {
let session_id = session_id
.to_str()
.map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()).into_response())?;
@ -84,6 +84,8 @@ async fn post_handler(
.ok_or((StatusCode::NOT_FOUND, "session not found").into_response())?;
session.handle().clone()
};
// inject request part
message.insert_extension(parts);
match &message {
ClientJsonRpcMessage::Request(_) | ClientJsonRpcMessage::BatchRequest(_) => {
let receiver = handle.establish_request_wise_channel().await.map_err(|e| {
@ -128,6 +130,8 @@ async fn post_handler(
} else {
// expect initialize message
let session_id = session_id();
// inject request part
message.insert_extension(parts);
let (session, transport) =
super::session::create_session(session_id.clone(), Default::default());
let Ok(_) = app.transport_tx.send(transport) else {