Stateless streamable-HTTP logs a placeholder client_info ("rmcp") instead of the real client identity #1172

fn peer_info_for_stateless_request(
    request: &crate::model::JsonRpcRequest<ClientRequest>,
    headers: &HeaderMap,
) -> Option<InitializeRequestParams> {
    let version = if let ClientRequest::InitializeRequest(ref init) = request.request {
        init.params.protocol_version.clone()
    } else {
        headers
            .get(HEADER_MCP_PROTOCOL_VERSION)
            .and_then(|v| v.to_str().ok())
            .and_then(|s| serde_json::from_value(serde_json::Value::String(s.to_owned())).ok())
            .unwrap_or(ProtocolVersion::V_2025_03_26)
    };
    Some(InitializeRequestParams {
        meta: None,
        protocol_version: version,             // reconstructed correctly
        capabilities: ClientCapabilities::default(),   // always blank
        client_info: Implementation::default(),        // always "rmcp"/crate version
    })
}

pub fn serve_directly_with_ct<R, S, T, E, A>(service: S, transport: T, peer_info: Option<R::PeerInfo>, ct: ...) -> RunningService<R, S> {
    let (peer, peer_rx) = Peer::new(..., peer_info);   // <- placeholder from above
    R::configure_direct_peer(&peer, &service.get_info());
    serve_inner(service, transport.into_transport(), peer, peer_rx, ct)
}

let peer_info = peer.peer_info();
if R::IS_CLIENT {
    tracing::info!(?peer_info, "Service initialized as client");
} else {
    tracing::info!(?peer_info, "Service initialized as server");
}

use std::sync::{Arc, Mutex};
use axum::{Router, body::Body};
use http::{Request, StatusCode};
use rmcp::ServerHandler;
use rmcp::model::{Implementation, InitializeRequestParams, ServerInfo};
use rmcp::service::{RequestContext, RoleServer};
use rmcp::transport::streamable_http_server::{
    StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
};
use tower::ServiceExt;

#[derive(Debug, Clone)]
struct Observed {
    from_request_arg: Implementation,
    from_peer_snapshot: Option<Implementation>,
}

#[derive(Clone, Default)]
struct ProbeHandler {
    observed: Arc<Mutex<Option<Observed>>>,
}

impl ServerHandler for ProbeHandler {
    fn get_info(&self) -> ServerInfo {
        ServerInfo::default()
    }

    async fn initialize(
        &self,
        request: InitializeRequestParams,
        context: RequestContext<RoleServer>,
    ) -> Result<ServerInfo, rmcp::ErrorData> {
        *self.observed.lock().unwrap() = Some(Observed {
            from_request_arg: request.client_info.clone(),
            from_peer_snapshot: context.peer.peer_info().map(|p| p.client_info.clone()),
        });
        Ok(self.get_info())
    }
}

async fn send_initialize_and_observe(legacy_session_mode: bool) -> Observed {
    let observed = Arc::new(Mutex::new(None));
    let handler = ProbeHandler { observed: observed.clone() };

    let mut config = StreamableHttpServerConfig::default();
    config.legacy_session_mode = legacy_session_mode;

    let service = StreamableHttpService::new(
        move || Ok(handler.clone()),
        LocalSessionManager::default().into(),
        config,
    );
    let app = Router::new().nest_service("/mcp", service);

    let body = serde_json::json!({
        "jsonrpc": "2.0", "id": 1, "method": "initialize",
        "params": {
            "protocolVersion": "2025-06-18",
            "capabilities": {},
            "clientInfo": { "name": "totally-real-client", "version": "42.0.0" }
        }
    });

    let response = app
        .oneshot(
            Request::builder()
                .method("POST")
                .uri("/mcp")
                .header("host", "127.0.0.1:3000")
                .header("content-type", "application/json")
                .header("accept", "application/json, text/event-stream")
                .body(Body::from(body.to_string()))
                .unwrap(),
        )
        .await
        .unwrap();

    assert_eq!(response.status(), StatusCode::OK);
    let _ = axum::body::to_bytes(response.into_body(), usize::MAX).await.unwrap();
    observed.lock().unwrap().clone().expect("handler did not run")
}

#[tokio::test]
async fn stateless_mode_discards_real_client_info_before_handler_runs() {
    let observed = send_initialize_and_observe(false).await;
    assert_eq!(observed.from_request_arg.name, "totally-real-client"); // handler got the real value
    let peer_client_info = observed.from_peer_snapshot.unwrap();
    assert_eq!(peer_client_info.name, "rmcp"); // but Peer::peer_info() — what gets logged — did not
}

#[tokio::test]
async fn stateful_mode_preserves_real_client_info() {
    let observed = send_initialize_and_observe(true).await;
    let peer_client_info = observed.from_peer_snapshot.unwrap();
    assert_eq!(peer_client_info.name, "totally-real-client"); // control: works fine when stateful
}