fix: include http_request_id in request-wise priming event IDs (#799)
* fix: include http_request_id in request-wise priming event IDs * refactor: use Option::into_iter and usize::from for priming * fix: retain event cache for completed request-wise channels * fix: track completed_at for cache eviction and resume * fix: log resume failures at warn level * test: add completed_cache_ttl eviction test * fix: return empty stream on failed resume * test: add resume after completion test
This commit is contained in:
Dale Seo
GitHub
parent
6603c1ff15
commit
3e56d52764
1 changed files with 467 additions and 89 deletions
|
|
@ -2,7 +2,8 @@
|
|||
use std::time::Duration;
|
||||
|
||||
use rmcp::transport::streamable_http_server::{
|
||||
StreamableHttpServerConfig, StreamableHttpService, session::local::LocalSessionManager,
|
||||
StreamableHttpServerConfig, StreamableHttpService,
|
||||
session::{SessionId, local::LocalSessionManager},
|
||||
};
|
||||
use tokio_util::sync::CancellationToken;
|
||||
|
||||
|
|
@ -54,7 +55,7 @@ async fn test_priming_on_stream_start() -> anyhow::Result<()> {
|
|||
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
|
||||
assert!(events.len() >= 2);
|
||||
|
||||
// Verify priming event (first event)
|
||||
// Verify priming event (first event) — initialize uses "0" (no http_request_id)
|
||||
let priming_event = events[0];
|
||||
assert!(priming_event.contains("id: 0"));
|
||||
assert!(priming_event.contains("retry: 3000"));
|
||||
|
|
@ -71,6 +72,341 @@ async fn test_priming_on_stream_start() -> anyhow::Result<()> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_request_wise_priming_includes_http_request_id() -> anyhow::Result<()> {
|
||||
let ct = CancellationToken::new();
|
||||
|
||||
let service: StreamableHttpService<Calculator, LocalSessionManager> =
|
||||
StreamableHttpService::new(
|
||||
|| Ok(Calculator::new()),
|
||||
Default::default(),
|
||||
StreamableHttpServerConfig::default()
|
||||
.with_sse_keep_alive(None)
|
||||
.with_cancellation_token(ct.child_token()),
|
||||
);
|
||||
|
||||
let router = axum::Router::new().nest_service("/mcp", service);
|
||||
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
|
||||
let addr = tcp_listener.local_addr()?;
|
||||
|
||||
let handle = tokio::spawn({
|
||||
let ct = ct.clone();
|
||||
async move {
|
||||
let _ = axum::serve(tcp_listener, router)
|
||||
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
// Initialize the session
|
||||
let response = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-11-25","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
|
||||
.send()
|
||||
.await?;
|
||||
assert_eq!(response.status(), 200);
|
||||
let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
|
||||
|
||||
// Send notifications/initialized
|
||||
let status = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
|
||||
.send()
|
||||
.await?
|
||||
.status();
|
||||
assert_eq!(status, 202);
|
||||
|
||||
// First tool call — should get http_request_id 0
|
||||
let body = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
|
||||
.send()
|
||||
.await?
|
||||
.text()
|
||||
.await?;
|
||||
|
||||
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
|
||||
assert!(
|
||||
events.len() >= 2,
|
||||
"expected priming + response, got: {body}"
|
||||
);
|
||||
|
||||
// Priming event should encode the http_request_id (0)
|
||||
let priming = events[0];
|
||||
assert!(
|
||||
priming.contains("id: 0/0"),
|
||||
"first request priming should be 0/0, got: {priming}"
|
||||
);
|
||||
assert!(priming.contains("retry: 3000"));
|
||||
|
||||
// Response event should use index 1 (since priming occupies index 0)
|
||||
let response_event = events[1];
|
||||
assert!(
|
||||
response_event.contains("id: 1/0"),
|
||||
"first response event id should be 1/0, got: {response_event}"
|
||||
);
|
||||
assert!(response_event.contains(r#""id":2"#));
|
||||
|
||||
// Second tool call — should get http_request_id 1
|
||||
let body = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"sum","arguments":{"a":3,"b":4}}}"#)
|
||||
.send()
|
||||
.await?
|
||||
.text()
|
||||
.await?;
|
||||
|
||||
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
|
||||
assert!(
|
||||
events.len() >= 2,
|
||||
"expected priming + response, got: {body}"
|
||||
);
|
||||
|
||||
let priming = events[0];
|
||||
assert!(
|
||||
priming.contains("id: 0/1"),
|
||||
"second request priming should be 0/1, got: {priming}"
|
||||
);
|
||||
|
||||
let response_event = events[1];
|
||||
assert!(
|
||||
response_event.contains("id: 1/1"),
|
||||
"second response event id should be 1/1, got: {response_event}"
|
||||
);
|
||||
assert!(response_event.contains(r#""id":3"#));
|
||||
|
||||
ct.cancel();
|
||||
handle.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_resume_after_request_wise_channel_completed() -> anyhow::Result<()> {
|
||||
let ct = CancellationToken::new();
|
||||
|
||||
let service: StreamableHttpService<Calculator, LocalSessionManager> =
|
||||
StreamableHttpService::new(
|
||||
|| Ok(Calculator::new()),
|
||||
Default::default(),
|
||||
StreamableHttpServerConfig::default()
|
||||
.with_sse_keep_alive(None)
|
||||
.with_cancellation_token(ct.child_token()),
|
||||
);
|
||||
|
||||
let router = axum::Router::new().nest_service("/mcp", service);
|
||||
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
|
||||
let addr = tcp_listener.local_addr()?;
|
||||
|
||||
let handle = tokio::spawn({
|
||||
let ct = ct.clone();
|
||||
async move {
|
||||
let _ = axum::serve(tcp_listener, router)
|
||||
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
// Initialize session
|
||||
let response = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
|
||||
.send()
|
||||
.await?;
|
||||
assert_eq!(response.status(), 200);
|
||||
let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
|
||||
|
||||
// Complete handshake
|
||||
let status = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
|
||||
.send()
|
||||
.await?
|
||||
.status();
|
||||
assert_eq!(status, 202);
|
||||
|
||||
// Call a tool and consume the full response (channel completes)
|
||||
let body = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
|
||||
.send()
|
||||
.await?
|
||||
.text()
|
||||
.await?;
|
||||
|
||||
let events: Vec<&str> = body.split("\n\n").filter(|e| !e.is_empty()).collect();
|
||||
assert!(
|
||||
events.len() >= 2,
|
||||
"expected priming + response, got: {body}"
|
||||
);
|
||||
assert!(events[0].contains("id: 0/0"));
|
||||
assert!(events[1].contains(r#""id":2"#));
|
||||
|
||||
// Resume with Last-Event-ID after the channel has completed.
|
||||
// The server returns 200 — either with replayed cached events
|
||||
// (if the channel is still retained) or an empty stream (if the
|
||||
// session worker hasn't processed the completion yet).
|
||||
let resume = client
|
||||
.get(format!("http://{addr}/mcp"))
|
||||
.header("Accept", "text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.header("last-event-id", "0/0")
|
||||
.send()
|
||||
.await?;
|
||||
assert_eq!(resume.status(), 200);
|
||||
|
||||
let resume_body = resume.text().await?;
|
||||
// The stream should complete (not hang), regardless of whether
|
||||
// it contains replayed events or is empty.
|
||||
assert!(
|
||||
!resume_body.contains("standalone"),
|
||||
"should not receive events from a different stream"
|
||||
);
|
||||
|
||||
ct.cancel();
|
||||
handle.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_completed_cache_ttl_eviction() -> anyhow::Result<()> {
|
||||
use std::sync::Arc;
|
||||
|
||||
let ct = CancellationToken::new();
|
||||
let mut session_manager = LocalSessionManager::default();
|
||||
session_manager.session_config.completed_cache_ttl = Duration::from_millis(200);
|
||||
let session_manager = Arc::new(session_manager);
|
||||
|
||||
let service = StreamableHttpService::new(
|
||||
|| Ok(Calculator::new()),
|
||||
session_manager.clone(),
|
||||
StreamableHttpServerConfig::default()
|
||||
.with_sse_keep_alive(None)
|
||||
.with_cancellation_token(ct.child_token()),
|
||||
);
|
||||
|
||||
let router = axum::Router::new().nest_service("/mcp", service);
|
||||
let tcp_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await?;
|
||||
let addr = tcp_listener.local_addr()?;
|
||||
|
||||
let handle = tokio::spawn({
|
||||
let ct = ct.clone();
|
||||
async move {
|
||||
let _ = axum::serve(tcp_listener, router)
|
||||
.with_graceful_shutdown(async move { ct.cancelled_owned().await })
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
|
||||
// Initialize session
|
||||
let response = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.body(r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{"protocolVersion":"2025-06-18","capabilities":{},"clientInfo":{"name":"test","version":"1.0"}}}"#)
|
||||
.send()
|
||||
.await?;
|
||||
assert_eq!(response.status(), 200);
|
||||
let session_id: SessionId = response.headers()["mcp-session-id"].to_str()?.into();
|
||||
|
||||
// Complete handshake
|
||||
client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
// Call a tool and consume the response (channel completes)
|
||||
let body = client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"name":"sum","arguments":{"a":1,"b":2}}}"#)
|
||||
.send()
|
||||
.await?
|
||||
.text()
|
||||
.await?;
|
||||
assert!(body.contains(r#""id":2"#));
|
||||
|
||||
// Wait for TTL to expire (200ms) plus margin
|
||||
tokio::time::sleep(Duration::from_millis(400)).await;
|
||||
|
||||
// Send a notification to trigger an event loop iteration (runs eviction)
|
||||
client
|
||||
.post(format!("http://{addr}/mcp"))
|
||||
.header("Content-Type", "application/json")
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.body(r#"{"jsonrpc":"2.0","method":"notifications/initialized"}"#)
|
||||
.send()
|
||||
.await?;
|
||||
|
||||
// Small delay to ensure the eviction ran
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
// Resume after TTL — channel should be evicted. The server returns
|
||||
// 200 with an empty stream (no events from a different stream).
|
||||
let resume = client
|
||||
.get(format!("http://{addr}/mcp"))
|
||||
.header("Accept", "text/event-stream")
|
||||
.header("mcp-session-id", session_id.to_string())
|
||||
.header("Mcp-Protocol-Version", "2025-06-18")
|
||||
.header("last-event-id", "0/0")
|
||||
.send()
|
||||
.await?;
|
||||
assert_eq!(resume.status(), 200);
|
||||
|
||||
let body = resume.text().await?;
|
||||
assert!(
|
||||
!body.contains(r#""id":2"#),
|
||||
"should NOT contain the old tool response after eviction, got: {body}"
|
||||
);
|
||||
|
||||
ct.cancel();
|
||||
handle.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_priming_on_stream_close() -> anyhow::Result<()> {
|
||||
use std::sync::Arc;
|
||||
|
|
|
|||
Loading…
Reference in a new issue