Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
275 changes: 242 additions & 33 deletions Cargo.lock

Large diffs are not rendered by default.

7 changes: 6 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ os_pipe = "1"

# MCP Protocol
# NOTE: Never use local path dependency - E:\one-mcp\rust-sdk is for source lookup only
rmcp = { version = "0.14.0", features = [
rmcp = { version = "0.15.0", features = [
"client",
"server",
"transport-io",
Expand All @@ -84,4 +84,9 @@ lto = true
codegen-units = 1
strip = true

# Temporary patch: fixes SSE channel replacement bug (notifications lost on reconnect)
# Remove once upstream merges https://github.com/modelcontextprotocol/rust-sdk/pull/660
[patch.crates-io]
rmcp = { git = "https://github.com/ion-ash/rust-sdk.git", branch = "fix/sse-channel-replacement-conflict" }


59 changes: 34 additions & 25 deletions crates/mcpmux-gateway/src/consumers/mcp_notifier.rs
Original file line number Diff line number Diff line change
Expand Up @@ -415,7 +415,7 @@ impl MCPNotifier {
feature_set_id = %feature_set_id,
"[MCPNotifier] 📨 GrantIssued - notifying all clients in space"
);
self.notify_all_list_changed(space_id).await;
self.notify_all_list_changed(space_id, true).await;
}

DomainEvent::GrantRevoked {
Expand All @@ -429,7 +429,7 @@ impl MCPNotifier {
feature_set_id = %feature_set_id,
"[MCPNotifier] 📨 GrantRevoked - notifying all clients in space"
);
self.notify_all_list_changed(space_id).await;
self.notify_all_list_changed(space_id, true).await;
}

DomainEvent::ClientGrantsUpdated {
Expand All @@ -443,7 +443,7 @@ impl MCPNotifier {
feature_sets = feature_set_ids.len(),
"[MCPNotifier] 📨 ClientGrantsUpdated - notifying all clients in space"
);
self.notify_all_list_changed(space_id).await;
self.notify_all_list_changed(space_id, true).await;
}

DomainEvent::FeatureSetMembersChanged {
Expand All @@ -456,7 +456,7 @@ impl MCPNotifier {
feature_set_id = %feature_set_id,
"[MCPNotifier] 📨 FeatureSetMembersChanged - notifying all clients in space"
);
self.notify_all_list_changed(space_id).await;
self.notify_all_list_changed(space_id, true).await;
}

// ============ Backend Server Notifications (Pass-through with Throttling) ============
Expand Down Expand Up @@ -523,7 +523,7 @@ impl MCPNotifier {
status = ?status,
"[MCPNotifier] ServerStatusChanged (Disconnected) - notifying clients to clear features"
);
self.notify_all_list_changed(space_id).await;
self.notify_all_list_changed(space_id, false).await;
} else {
debug!(
server_id = %server_id,
Expand All @@ -549,7 +549,7 @@ impl MCPNotifier {
removed = removed.len(),
"[MCPNotifier] ServerFeaturesRefreshed"
);
self.notify_all_list_changed(space_id).await;
self.notify_all_list_changed(space_id, false).await;
}

// Other events that affect MCP capabilities are handled above
Expand All @@ -571,8 +571,13 @@ impl MCPNotifier {
/// **Important**: This method handles throttling at the batch level and marks
/// all individual notification types as sent, preventing double-notifications
/// when individual DomainEvent::ToolsChanged/etc. events arrive shortly after.
async fn notify_all_list_changed(&self, space_id: Uuid) {
// 1. Content-Based Deduping
///
/// **`force` parameter**: When `true`, skips content-based hash dedup. Used for
/// grant-related events where the total features in the space haven't changed but
/// the *effective* features visible to clients have (due to grant/feature set changes).
/// The hash is computed from all features in the space, so it can't detect grant changes.
async fn notify_all_list_changed(&self, space_id: Uuid, force: bool) {
// 1. Content-Based Deduping (skipped when force=true)
let tools_hash = self
.calculate_feature_hash(space_id, FeatureType::Tool)
.await;
Expand All @@ -583,23 +588,27 @@ impl MCPNotifier {
.calculate_feature_hash(space_id, FeatureType::Resource)
.await;

let any_changed = {
let hashes = self.state_hashes.read();
let t_changed = hashes
.get(&(space_id, NotificationType::Tools))
.is_none_or(|&h| h != tools_hash);
let p_changed = hashes
.get(&(space_id, NotificationType::Prompts))
.is_none_or(|&h| h != prompts_hash);
let r_changed = hashes
.get(&(space_id, NotificationType::Resources))
.is_none_or(|&h| h != resources_hash);
t_changed || p_changed || r_changed
};

if !any_changed {
debug!(space_id = %space_id, "[MCPNotifier] 🛑 Batch content unchanged, skipping");
return;
if !force {
let any_changed = {
let hashes = self.state_hashes.read();
let t_changed = hashes
.get(&(space_id, NotificationType::Tools))
.is_none_or(|&h| h != tools_hash);
let p_changed = hashes
.get(&(space_id, NotificationType::Prompts))
.is_none_or(|&h| h != prompts_hash);
let r_changed = hashes
.get(&(space_id, NotificationType::Resources))
.is_none_or(|&h| h != resources_hash);
t_changed || p_changed || r_changed
};

if !any_changed {
debug!(space_id = %space_id, "[MCPNotifier] 🛑 Batch content unchanged, skipping");
return;
}
} else {
info!(space_id = %space_id, "[MCPNotifier] 🔓 Force-sending (grant/feature set change, bypassing hash dedup)");
}

let now = Instant::now();
Expand Down
35 changes: 25 additions & 10 deletions crates/mcpmux-gateway/src/mcp/handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -102,14 +102,14 @@ impl ServerHandler for McpMuxGatewayHandler {
protocol_version: Default::default(),
capabilities: ServerCapabilities::builder()
.enable_tools_with(ToolsCapability {
list_changed: Some(false), // Stateless mode - no notifications
list_changed: Some(true),
})
.enable_prompts_with(PromptsCapability {
list_changed: Some(false), // Stateless mode - no notifications
list_changed: Some(true),
})
.enable_resources_with(ResourcesCapability {
subscribe: Some(false),
list_changed: Some(false), // Stateless mode - no notifications
list_changed: Some(true),
})
.build(),
server_info: Implementation {
Expand Down Expand Up @@ -150,19 +150,34 @@ impl ServerHandler for McpMuxGatewayHandler {
}

async fn on_initialized(&self, context: NotificationContext<RoleServer>) {
// Silently process - entry already logged in oauth_middleware
let _oauth_ctx = match self.get_oauth_context(&context.extensions) {
let oauth_ctx = match self.get_oauth_context(&context.extensions) {
Ok(ctx) => ctx,
Err(e) => {
warn!("Failed to extract OAuth context: {}", e);
warn!("Failed to extract OAuth context on_initialized: {}", e);
return;
}
};

// In stateless mode:
// - No session tracking
// - No notification registration
// - Each request is independent
// Register peer with MCPNotifier for list_changed notification delivery
let peer = std::sync::Arc::new(context.peer);
self.notification_bridge
.register_peer(oauth_ctx.client_id.clone(), peer);

// Mark the client stream as active immediately - RMCP's session transport
// handles SSE streaming and message caching internally
self.notification_bridge
.mark_client_stream_active(&oauth_ctx.client_id);

// Pre-populate feature hashes to prevent spurious first notifications
self.notification_bridge
.prime_hashes_for_space(oauth_ctx.space_id)
.await;

info!(
client_id = %oauth_ctx.client_id,
space_id = %oauth_ctx.space_id,
"Client initialized - peer registered for notifications"
);
}

async fn list_tools(
Expand Down
5 changes: 5 additions & 0 deletions crates/mcpmux-gateway/src/pool/credential_store.rs
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,7 @@ impl CredentialStore for DatabaseCredentialStore {
Some(StoredCredentials {
client_id: reg.client_id,
token_response: Some(token_response),
granted_scopes: Vec::new(),
})
}
(Some(reg), None) => {
Expand All @@ -217,6 +218,7 @@ impl CredentialStore for DatabaseCredentialStore {
Some(StoredCredentials {
client_id: reg.client_id,
token_response: None,
granted_scopes: Vec::new(),
})
}
(None, Some(access)) => {
Expand All @@ -228,6 +230,7 @@ impl CredentialStore for DatabaseCredentialStore {
Some(StoredCredentials {
client_id: String::new(),
token_response: Some(token_response),
granted_scopes: Vec::new(),
})
}
(None, None) => {
Expand Down Expand Up @@ -584,6 +587,7 @@ mod tests {
let credentials = StoredCredentials {
client_id: "new-client-id".to_string(),
token_response: Some(token_response),
granted_scopes: Vec::new(),
};

store.save(credentials).await.unwrap();
Expand Down Expand Up @@ -641,6 +645,7 @@ mod tests {
let credentials = StoredCredentials {
client_id: "client-id".to_string(),
token_response: Some(token_response),
granted_scopes: Vec::new(),
};

store.save(credentials).await.unwrap();
Expand Down
1 change: 1 addition & 0 deletions crates/mcpmux-gateway/src/pool/instance.rs
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ impl McpClientHandler {
title: Some("McpMux Gateway".to_string()),
icons: None,
website_url: None,
..Default::default()
},
meta: None,
},
Expand Down
1 change: 1 addition & 0 deletions crates/mcpmux-gateway/src/pool/oauth_utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,7 @@ pub fn convert_from_stored_metadata(stored: &StoredOAuthMetadata) -> Authorizati
scopes_supported: stored.scopes_supported.clone(),
response_types_supported: stored.response_types_supported.clone(),
additional_fields: stored.additional_fields.clone(),
..Default::default()
}
}

Expand Down
19 changes: 18 additions & 1 deletion crates/mcpmux-gateway/src/server/logging_middleware.rs
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ fn redact_headers_compact(headers: &axum::http::HeaderMap) -> String {
| "user-agent"
| "mcp-session-id"
| "mcp-protocol-version"
| "last-event-id"
)
})
.map(|(name, value)| {
Expand Down Expand Up @@ -208,7 +209,23 @@ pub async fn http_logging_middleware(request: Request, next: Next) -> Result<Res
let response = next.run(request).await;
let status = response.status().as_u16();

// Extract response body to log it
// Check if this is a streaming response (SSE) — NEVER buffer these.
// SSE responses have Content-Type: text/event-stream and are infinite
// streams. Calling body.collect() would block forever, preventing VS Code
// from receiving any SSE events (notifications, keep-alive pings).
let is_streaming = response
.headers()
.get("content-type")
.and_then(|v| v.to_str().ok())
.is_some_and(|ct| ct.contains("text/event-stream"));

if is_streaming {
// For SSE streams, log entry only and pass through without touching body
RequestSpan::log_exit(&ctx, status, None);
return Ok(response);
}

// For non-streaming responses, capture body for logging
let (parts, body) = response.into_parts();
let body_bytes = match body.collect().await {
Ok(collected) => collected.to_bytes(),
Expand Down
20 changes: 9 additions & 11 deletions crates/mcpmux-gateway/src/server/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -240,29 +240,27 @@ impl GatewayServer {
let handler =
McpMuxGatewayHandler::new(Arc::new(self.services.clone()), notification_bridge.clone());

// Create STATELESS MCP service
// stateful_mode: false means:
// - No Mcp-Session-Id header
// - GET/DELETE return 405 automatically (no notification streams)
// - Each POST is independent
// - Avoids the stream management issues that caused connection loops
// Trade-off: Cannot send list_changed notifications
// Create STATEFUL MCP service (full Streamable HTTP per spec 2025-11-25)
// stateful_mode: true means:
// - Mcp-Session-Id header for session management
// - GET endpoint for SSE streams (server-initiated notifications)
// - DELETE endpoint for session termination
// - list_changed notifications delivered via SSE
let mcp_service = StreamableHttpService::new(
move || {
debug!("[Gateway] Creating handler instance for MCP request");
debug!("[Gateway] Creating handler instance for MCP session");
Ok(handler.clone())
},
LocalSessionManager::default().into(),
StreamableHttpServerConfig {
stateful_mode: false,
stateful_mode: true,
sse_keep_alive: Some(std::time::Duration::from_secs(30)),
sse_retry: None,
sse_retry: Some(std::time::Duration::from_secs(3)),
cancellation_token: CancellationToken::new(),
},
);

// Wrap MCP service with OAuth middleware
// In stateless mode, no session healing needed - rmcp handles 405 for GET/DELETE
let mcp_routes =
Router::new()
.nest_service("/mcp", mcp_service)
Expand Down
1 change: 1 addition & 0 deletions crates/mcpmux-mcp/src/transports.rs
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ impl McpClientHandler {
title: Some("McpMux Gateway".to_string()),
icons: None,
website_url: None,
..Default::default()
},
meta: None,
},
Expand Down
Loading