fix: Disable statefulness to attempt to fix initialize race condition (#351)

* fix: Disable statefulness to attempt to fix initialize race condition

* Change stateful_mode to being configuratable on the transport

* Use ..Default syntax

* changeset

* Add new flag to docs config reference page

Author: swcollard

.changesets/config_fix_disable_stateful_sessions.md

@@ -0,0 +1,3 @@
+### fix: Disable statefulness to fix initialize race condition - @swcollard PR #351
+
+We've been seeing errors with state and session handling in the MCP Server. Whether that is requests being sent before the initialized notification is processed. Or running a fleet of MCP Server pods behind a round robin load balancer. A new configuration option under the streamable_http transport `stateful_mode`, allows disabling session handling which appears to fix the race condition issue.

crates/apollo-mcp-server/src/server.rs

@@ -80,6 +80,9 @@ pub enum Transport {
         /// The port to bind to
         #[serde(default = "Transport::default_port")]
         port: u16,
+
+        #[serde(default = "Transport::default_stateful_mode")]
+        stateful_mode: bool,
     },
 }
 
@@ -91,6 +94,10 @@ impl Transport {
     fn default_port() -> u16 {
         5000
     }
+
+    fn default_stateful_mode() -> bool {
+        true
+    }
 }
 
 #[bon]

crates/apollo-mcp-server/src/server/states/starting.rs

@@ -2,8 +2,8 @@ use std::{net::SocketAddr, sync::Arc};
 
 use apollo_compiler::{Name, Schema, ast::OperationType, validation::Valid};
 use axum::{Router, extract::Query, http::StatusCode, response::Json, routing::get};
-use rmcp::transport::StreamableHttpService;
 use rmcp::transport::streamable_http_server::session::local::LocalSessionManager;
+use rmcp::transport::{StreamableHttpServerConfig, StreamableHttpService};
 use rmcp::{
     ServiceExt as _,
     transport::{SseServer, sse_server::SseServerConfig, stdio},
@@ -126,6 +126,7 @@ impl Starting {
                     auth: _,
                     address: _,
                     port: _,
+                    stateful_mode: _,
                 },
                 true,
             ) => Some(HealthCheck::new(self.config.health_check.clone())),
@@ -168,14 +169,18 @@ impl Starting {
                 auth,
                 address,
                 port,
+                stateful_mode,
             } => {
                 info!(port = ?port, address = ?address, "Starting MCP server in Streamable HTTP mode");
                 let running = running.clone();
                 let listen_address = SocketAddr::new(address, port);
                 let service = StreamableHttpService::new(
                     move || Ok(running.clone()),
                     LocalSessionManager::default().into(),
-                    Default::default(),
+                    StreamableHttpServerConfig {
+                        stateful_mode,
+                        ..Default::default()
+                    },
                 );
                 let mut router =
                     with_auth!(axum::Router::new().nest_service("/mcp", service), auth);

docs/source/config-file.mdx

@@ -147,11 +147,13 @@ The available fields depend on the value of the nested `type` key:
 
 ##### Streamable HTTP
 
-| Option    | Value                 | Value Type | Description                                                               |
-| :-------- | :-------------------- | :--------- | :------------------------------------------------------------------------ |
-| `type`    | `"streamable_http"`   |            | Host the MCP server on the configuration, using streamable HTTP messages. |
-| `address` | `127.0.0.1` (default) | `IpAddr`   | The IP address to bind to                                                 |
-| `port`    | `5000` (default)      | `u16`      | The port to bind to                                                       |
+| Option          | Value                 | Value Type | Description                                                               |
+| :-------------- | :-------------------- | :--------- | :------------------------------------------------------------------------ |
+| `type`          | `"streamable_http"`   |            | Host the MCP server on the configuration, using streamable HTTP messages. |
+| `address`       | `127.0.0.1` (default) | `IpAddr`   | The IP address to bind to                                                 |
+| `port`          | `5000` (default)      | `u16`      | The port to bind to                                                       |
+| `stateful_mode` | `true` (default)      | `bool`     | Flag to enable or disable stateful mode and session management.           |
+
 
 ##### SSE (Deprecated, use StreamableHTTP)