Skip to content

Commit 3f25d44

Browse files
bdshadowclaude
bdshadow
and
claude
authored
fix: sync MCP sessions to Redis for multi-replica deployments (#3491)
WebMvcStreamableServerTransportProvider stores sessions in an in-memory ConcurrentHashMap, causing 404 "Session not found" errors when requests scatter across replicas. Add a servlet filter that persists session metadata to Redis and reconstructs sessions on replicas that haven't seen the client. Tracked upstream: modelcontextprotocol/java-sdk#201 <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * MCP sessions are now optionally persisted to Redis and synchronized across instances, enabling automatic session recovery and continuity for up to 48 hours (no-op when Redis is not configured). * **Tests** * Added integration tests validating end-to-end Redis-backed session persistence and recovery across replicas. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 09b56c1 commit 3f25d44

5 files changed

Lines changed: 380 additions & 1 deletion

File tree

‎backend/app/build.gradle‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ dependencies {
139139
implementation "org.springframework.boot:spring-boot-properties-migrator"
140140

141141
implementation libs.springDocOpenApiCommon
142-
testApi libs.redissonSpringBootStarter
142+
implementation libs.redissonSpringBootStarter
143143

144144
/**
145145
* OPENTELEMETRY

‎backend/app/src/main/kotlin/io/tolgee/mcp/McpConfig.kt‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,13 @@
11
package io.tolgee.mcp
22

3+
import com.fasterxml.jackson.databind.ObjectMapper
34
import io.modelcontextprotocol.server.McpServer
45
import io.modelcontextprotocol.server.McpSyncServer
56
import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider
67
import io.modelcontextprotocol.spec.McpSchema
78
import io.tolgee.util.VersionProvider
9+
import org.redisson.api.RedissonClient
10+
import org.springframework.boot.web.servlet.FilterRegistrationBean
811
import org.springframework.context.annotation.Bean
912
import org.springframework.context.annotation.Configuration
1013
import org.springframework.web.servlet.function.RouterFunction
@@ -50,4 +53,16 @@ class McpConfig {
5053
): RouterFunction<ServerResponse> {
5154
return transportProvider.routerFunction
5255
}
56+
57+
@Bean
58+
fun mcpSessionRedisFilter(
59+
transportProvider: WebMvcStreamableServerTransportProvider,
60+
redissonClient: RedissonClient?,
61+
objectMapper: ObjectMapper,
62+
): FilterRegistrationBean<McpSessionRedisFilter> {
63+
val filter = McpSessionRedisFilter(transportProvider, redissonClient, objectMapper)
64+
val registration = FilterRegistrationBean(filter)
65+
registration.addUrlPatterns("/mcp/*")
66+
return registration
67+
}
5368
}

‎backend/app/src/main/kotlin/io/tolgee/mcp/McpSessionData.kt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
package io.tolgee.mcp
2+
3+
data class McpSessionData(
4+
val clientCapabilitiesJson: String?,
5+
val clientInfoJson: String?,
6+
)

‎backend/app/src/main/kotlin/io/tolgee/mcp/McpSessionRedisFilter.kt‎

Lines changed: 250 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,250 @@
1+
package io.tolgee.mcp
2+
3+
import com.fasterxml.jackson.databind.ObjectMapper
4+
import io.modelcontextprotocol.server.McpNotificationHandler
5+
import io.modelcontextprotocol.server.McpRequestHandler
6+
import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider
7+
import io.modelcontextprotocol.spec.DefaultMcpStreamableServerSessionFactory
8+
import io.modelcontextprotocol.spec.McpSchema
9+
import io.modelcontextprotocol.spec.McpStreamableServerSession
10+
import jakarta.servlet.FilterChain
11+
import jakarta.servlet.http.HttpServletRequest
12+
import jakarta.servlet.http.HttpServletResponse
13+
import jakarta.servlet.http.HttpServletResponseWrapper
14+
import org.redisson.api.RedissonClient
15+
import org.slf4j.LoggerFactory
16+
import org.springframework.web.filter.OncePerRequestFilter
17+
import java.time.Duration
18+
import java.util.concurrent.ConcurrentHashMap
19+
20+
/**
21+
* Servlet filter that syncs MCP sessions to Redis for multi-replica deployments.
22+
*
23+
* The MCP Java SDK's [WebMvcStreamableServerTransportProvider] stores sessions in an in-memory
24+
* `ConcurrentHashMap`. When running multiple replicas behind a load balancer, requests with an
25+
* `Mcp-Session-Id` header may land on a replica that doesn't have the session, resulting in
26+
* 404 "Session not found" errors.
27+
*
28+
* This filter works around the issue by:
29+
* - **Before handler**: If the request carries an `Mcp-Session-Id` not present in the local map,
30+
* reconstructs the session from Redis and injects it into the transport provider's sessions map.
31+
* - **After handler**: If the response sets a new `Mcp-Session-Id` (initialize), persists the
32+
* session metadata to Redis.
33+
*
34+
* Uses reflection because [WebMvcStreamableServerTransportProvider] has a private constructor and
35+
* all private fields/methods — it cannot be extended or configured with custom session storage.
36+
*
37+
* When Redis is unavailable (e.g. single-instance deployment), the filter is a no-op.
38+
*
39+
* This is a temporary workaround until the SDK provides a proper session storage API.
40+
* Tracked upstream: https://github.com/modelcontextprotocol/java-sdk/issues/201
41+
* SDK maintainer confirmed persistent session storage is planned:
42+
* https://github.com/modelcontextprotocol/java-sdk/issues/201#issuecomment-3915069460
43+
*/
44+
class McpSessionRedisFilter(
45+
private val transportProvider: WebMvcStreamableServerTransportProvider,
46+
private val redissonClient: RedissonClient?,
47+
private val objectMapper: ObjectMapper,
48+
) : OncePerRequestFilter() {
49+
private val log = LoggerFactory.getLogger(McpSessionRedisFilter::class.java)
50+
51+
private val sessionsMap: ConcurrentHashMap<String, McpStreamableServerSession> by lazy {
52+
loadSessionsMap()
53+
}
54+
55+
private val factoryFields: FactoryFields by lazy {
56+
loadFactoryFields()
57+
}
58+
59+
override fun doFilterInternal(
60+
request: HttpServletRequest,
61+
response: HttpServletResponse,
62+
filterChain: FilterChain,
63+
) {
64+
if (redissonClient == null) {
65+
filterChain.doFilter(request, response)
66+
return
67+
}
68+
69+
val sessionId = request.getHeader(MCP_SESSION_ID_HEADER)
70+
71+
// Pre-handle: recover session from Redis if not in local map
72+
if (sessionId != null && !sessionsMap.containsKey(sessionId)) {
73+
recoverSessionFromRedis(sessionId)
74+
}
75+
76+
// Wrap response to capture new session ID from initialize responses
77+
val responseWrapper = SessionIdCapturingResponseWrapper(response)
78+
filterChain.doFilter(request, responseWrapper)
79+
80+
// Post-handle: persist new session to Redis
81+
val newSessionId = responseWrapper.capturedSessionId
82+
if (newSessionId != null) {
83+
persistSessionToRedis(newSessionId)
84+
}
85+
}
86+
87+
private fun recoverSessionFromRedis(sessionId: String) {
88+
try {
89+
val bucket = redissonClient!!.getBucket<String>("$REDIS_KEY_PREFIX$sessionId")
90+
val json = bucket.get() ?: return
91+
92+
val sessionData = objectMapper.readValue(json, McpSessionData::class.java)
93+
94+
val clientCapabilities =
95+
sessionData.clientCapabilitiesJson?.let {
96+
objectMapper.readValue(it, McpSchema.ClientCapabilities::class.java)
97+
}
98+
val clientInfo =
99+
sessionData.clientInfoJson?.let {
100+
objectMapper.readValue(it, McpSchema.Implementation::class.java)
101+
}
102+
103+
@Suppress("UNCHECKED_CAST")
104+
val session =
105+
McpStreamableServerSession(
106+
sessionId,
107+
clientCapabilities,
108+
clientInfo,
109+
factoryFields.requestTimeout,
110+
factoryFields.requestHandlers as Map<String, McpRequestHandler<*>>,
111+
factoryFields.notificationHandlers as Map<String, McpNotificationHandler>,
112+
)
113+
114+
val existing = sessionsMap.putIfAbsent(sessionId, session)
115+
if (existing != null) {
116+
log.debug("MCP session {} was already recovered by another thread", sessionId)
117+
} else {
118+
log.debug("Recovered MCP session {} from Redis", sessionId)
119+
}
120+
} catch (e: Exception) {
121+
log.warn("Failed to recover MCP session {} from Redis", sessionId, e)
122+
}
123+
}
124+
125+
private fun persistSessionToRedis(sessionId: String) {
126+
try {
127+
val session = sessionsMap[sessionId] ?: return
128+
129+
val clientCapabilities = getPrivateField<Any?>(session, "clientCapabilities")
130+
val clientInfo = getPrivateField<Any?>(session, "clientInfo")
131+
132+
val capabilitiesValue = unwrapField(clientCapabilities, "clientCapabilities")
133+
val infoValue = unwrapField(clientInfo, "clientInfo")
134+
135+
val sessionData =
136+
McpSessionData(
137+
clientCapabilitiesJson = capabilitiesValue?.let { objectMapper.writeValueAsString(it) },
138+
clientInfoJson = infoValue?.let { objectMapper.writeValueAsString(it) },
139+
)
140+
141+
val bucket = redissonClient!!.getBucket<String>("$REDIS_KEY_PREFIX$sessionId")
142+
bucket.set(objectMapper.writeValueAsString(sessionData), SESSION_TTL)
143+
log.debug("Persisted MCP session {} to Redis", sessionId)
144+
} catch (e: Exception) {
145+
log.warn("Failed to persist MCP session {} to Redis", sessionId, e)
146+
}
147+
}
148+
149+
private fun unwrapField(
150+
fieldValue: Any?,
151+
fieldName: String,
152+
): Any? {
153+
if (fieldValue == null) return null
154+
if (fieldValue is java.util.concurrent.atomic.AtomicReference<*>) {
155+
return fieldValue.get()
156+
}
157+
log.warn(
158+
"Expected AtomicReference for field '{}' but got {}; using value directly",
159+
fieldName,
160+
fieldValue.javaClass.name,
161+
)
162+
return fieldValue
163+
}
164+
165+
@Suppress("UNCHECKED_CAST")
166+
private fun loadSessionsMap(): ConcurrentHashMap<String, McpStreamableServerSession> {
167+
val field = WebMvcStreamableServerTransportProvider::class.java.getDeclaredField("sessions")
168+
field.isAccessible = true
169+
return field.get(transportProvider) as ConcurrentHashMap<String, McpStreamableServerSession>
170+
}
171+
172+
private fun loadFactoryFields(): FactoryFields {
173+
val factoryField = WebMvcStreamableServerTransportProvider::class.java.getDeclaredField("sessionFactory")
174+
factoryField.isAccessible = true
175+
val factory = factoryField.get(transportProvider)
176+
177+
if (factory !is DefaultMcpStreamableServerSessionFactory) {
178+
throw IllegalStateException(
179+
"Expected DefaultMcpStreamableServerSessionFactory but got ${factory?.javaClass?.name}",
180+
)
181+
}
182+
183+
val timeoutField = DefaultMcpStreamableServerSessionFactory::class.java.getDeclaredField("requestTimeout")
184+
timeoutField.isAccessible = true
185+
val requestTimeout = timeoutField.get(factory) as Duration
186+
187+
val handlersField = DefaultMcpStreamableServerSessionFactory::class.java.getDeclaredField("requestHandlers")
188+
handlersField.isAccessible = true
189+
190+
@Suppress("UNCHECKED_CAST")
191+
val requestHandlers = handlersField.get(factory) as Map<String, Any>
192+
193+
val notifField = DefaultMcpStreamableServerSessionFactory::class.java.getDeclaredField("notificationHandlers")
194+
notifField.isAccessible = true
195+
196+
@Suppress("UNCHECKED_CAST")
197+
val notificationHandlers = notifField.get(factory) as Map<String, Any>
198+
199+
return FactoryFields(requestTimeout, requestHandlers, notificationHandlers)
200+
}
201+
202+
@Suppress("UNCHECKED_CAST")
203+
private fun <T> getPrivateField(
204+
obj: Any,
205+
fieldName: String,
206+
): T {
207+
val field = obj.javaClass.getDeclaredField(fieldName)
208+
field.isAccessible = true
209+
return field.get(obj) as T
210+
}
211+
212+
private class SessionIdCapturingResponseWrapper(
213+
response: HttpServletResponse,
214+
) : HttpServletResponseWrapper(response) {
215+
var capturedSessionId: String? = null
216+
private set
217+
218+
override fun setHeader(
219+
name: String,
220+
value: String?,
221+
) {
222+
if (name == MCP_SESSION_ID_HEADER && value != null) {
223+
capturedSessionId = value
224+
}
225+
super.setHeader(name, value)
226+
}
227+
228+
override fun addHeader(
229+
name: String,
230+
value: String?,
231+
) {
232+
if (name == MCP_SESSION_ID_HEADER && value != null) {
233+
capturedSessionId = value
234+
}
235+
super.addHeader(name, value)
236+
}
237+
}
238+
239+
private data class FactoryFields(
240+
val requestTimeout: Duration,
241+
val requestHandlers: Map<String, Any>,
242+
val notificationHandlers: Map<String, Any>,
243+
)
244+
245+
companion object {
246+
private const val MCP_SESSION_ID_HEADER = "Mcp-Session-Id"
247+
private const val REDIS_KEY_PREFIX = "mcp_session:"
248+
private val SESSION_TTL = Duration.ofHours(48)
249+
}
250+
}

0 commit comments

Comments
 (0)