Skip to content

Commit 5a2bd2f

Browse files
bdshadowclaude
bdshadow
and
claude
committed
fix: sync MCP sessions to Redis for multi-replica deployments
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 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 007f2f6 commit 5a2bd2f

5 files changed

Lines changed: 355 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: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
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+
sessionsMap[sessionId] = session
115+
log.debug("Recovered MCP session {} from Redis", sessionId)
116+
} catch (e: Exception) {
117+
log.warn("Failed to recover MCP session {} from Redis", sessionId, e)
118+
}
119+
}
120+
121+
private fun persistSessionToRedis(sessionId: String) {
122+
try {
123+
val session = sessionsMap[sessionId] ?: return
124+
125+
val clientCapabilities = getPrivateField<Any?>(session, "clientCapabilities")
126+
val clientInfo = getPrivateField<Any?>(session, "clientInfo")
127+
128+
// These are AtomicReference fields, so we need to get the value
129+
val capabilitiesValue = (clientCapabilities as? java.util.concurrent.atomic.AtomicReference<*>)?.get()
130+
val infoValue = (clientInfo as? java.util.concurrent.atomic.AtomicReference<*>)?.get()
131+
132+
val sessionData =
133+
McpSessionData(
134+
clientCapabilitiesJson = capabilitiesValue?.let { objectMapper.writeValueAsString(it) },
135+
clientInfoJson = infoValue?.let { objectMapper.writeValueAsString(it) },
136+
)
137+
138+
val bucket = redissonClient!!.getBucket<String>("$REDIS_KEY_PREFIX$sessionId")
139+
bucket.set(objectMapper.writeValueAsString(sessionData), SESSION_TTL)
140+
log.debug("Persisted MCP session {} to Redis", sessionId)
141+
} catch (e: Exception) {
142+
log.warn("Failed to persist MCP session {} to Redis", sessionId, e)
143+
}
144+
}
145+
146+
@Suppress("UNCHECKED_CAST")
147+
private fun loadSessionsMap(): ConcurrentHashMap<String, McpStreamableServerSession> {
148+
val field = WebMvcStreamableServerTransportProvider::class.java.getDeclaredField("sessions")
149+
field.isAccessible = true
150+
return field.get(transportProvider) as ConcurrentHashMap<String, McpStreamableServerSession>
151+
}
152+
153+
private fun loadFactoryFields(): FactoryFields {
154+
val factoryField = WebMvcStreamableServerTransportProvider::class.java.getDeclaredField("sessionFactory")
155+
factoryField.isAccessible = true
156+
val factory = factoryField.get(transportProvider)
157+
158+
if (factory !is DefaultMcpStreamableServerSessionFactory) {
159+
throw IllegalStateException(
160+
"Expected DefaultMcpStreamableServerSessionFactory but got ${factory?.javaClass?.name}",
161+
)
162+
}
163+
164+
val timeoutField = DefaultMcpStreamableServerSessionFactory::class.java.getDeclaredField("requestTimeout")
165+
timeoutField.isAccessible = true
166+
val requestTimeout = timeoutField.get(factory) as Duration
167+
168+
val handlersField = DefaultMcpStreamableServerSessionFactory::class.java.getDeclaredField("requestHandlers")
169+
handlersField.isAccessible = true
170+
171+
@Suppress("UNCHECKED_CAST")
172+
val requestHandlers = handlersField.get(factory) as Map<String, Any>
173+
174+
val notifField = DefaultMcpStreamableServerSessionFactory::class.java.getDeclaredField("notificationHandlers")
175+
notifField.isAccessible = true
176+
177+
@Suppress("UNCHECKED_CAST")
178+
val notificationHandlers = notifField.get(factory) as Map<String, Any>
179+
180+
return FactoryFields(requestTimeout, requestHandlers, notificationHandlers)
181+
}
182+
183+
@Suppress("UNCHECKED_CAST")
184+
private fun <T> getPrivateField(
185+
obj: Any,
186+
fieldName: String,
187+
): T {
188+
val field = obj.javaClass.getDeclaredField(fieldName)
189+
field.isAccessible = true
190+
return field.get(obj) as T
191+
}
192+
193+
private class SessionIdCapturingResponseWrapper(
194+
response: HttpServletResponse,
195+
) : HttpServletResponseWrapper(response) {
196+
var capturedSessionId: String? = null
197+
private set
198+
199+
override fun setHeader(
200+
name: String,
201+
value: String?,
202+
) {
203+
if (name == MCP_SESSION_ID_HEADER && value != null) {
204+
capturedSessionId = value
205+
}
206+
super.setHeader(name, value)
207+
}
208+
209+
override fun addHeader(
210+
name: String,
211+
value: String?,
212+
) {
213+
if (name == MCP_SESSION_ID_HEADER && value != null) {
214+
capturedSessionId = value
215+
}
216+
super.addHeader(name, value)
217+
}
218+
}
219+
220+
private data class FactoryFields(
221+
val requestTimeout: Duration,
222+
val requestHandlers: Map<String, Any>,
223+
val notificationHandlers: Map<String, Any>,
224+
)
225+
226+
companion object {
227+
private const val MCP_SESSION_ID_HEADER = "Mcp-Session-Id"
228+
private const val REDIS_KEY_PREFIX = "mcp_session:"
229+
private val SESSION_TTL = Duration.ofHours(48)
230+
}
231+
}

‎backend/app/src/test/kotlin/io/tolgee/mcp/McpRedisSessionRecoveryTest.kt‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
package io.tolgee.mcp
2+
3+
import io.modelcontextprotocol.server.transport.WebMvcStreamableServerTransportProvider
4+
import io.modelcontextprotocol.spec.McpStreamableServerSession
5+
import io.tolgee.AbstractMcpTest
6+
import io.tolgee.fixtures.RedisRunner
7+
import io.tolgee.testing.ContextRecreatingTest
8+
import io.tolgee.testing.assertions.Assertions.assertThat
9+
import org.junit.jupiter.api.AfterAll
10+
import org.junit.jupiter.api.BeforeEach
11+
import org.junit.jupiter.api.Test
12+
import org.redisson.api.RedissonClient
13+
import org.springframework.beans.factory.annotation.Autowired
14+
import org.springframework.boot.test.context.SpringBootTest
15+
import org.springframework.boot.test.util.TestPropertyValues
16+
import org.springframework.context.ApplicationContextInitializer
17+
import org.springframework.context.ConfigurableApplicationContext
18+
import org.springframework.test.annotation.DirtiesContext
19+
import org.springframework.test.context.ContextConfiguration
20+
import java.util.concurrent.ConcurrentHashMap
21+
22+
@SpringBootTest(
23+
properties = [
24+
"tolgee.cache.use-redis=true",
25+
"tolgee.cache.enabled=true",
26+
],
27+
webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT,
28+
)
29+
@ContextConfiguration(initializers = [McpRedisSessionRecoveryTest.Companion.Initializer::class])
30+
@DirtiesContext(classMode = DirtiesContext.ClassMode.AFTER_CLASS)
31+
@ContextRecreatingTest
32+
class McpRedisSessionRecoveryTest : AbstractMcpTest() {
33+
companion object {
34+
val redisRunner = RedisRunner()
35+
36+
@AfterAll
37+
@JvmStatic
38+
fun stopRedis() {
39+
redisRunner.stop()
40+
}
41+
42+
class Initializer : ApplicationContextInitializer<ConfigurableApplicationContext> {
43+
override fun initialize(configurableApplicationContext: ConfigurableApplicationContext) {
44+
redisRunner.run()
45+
TestPropertyValues
46+
.of("spring.data.redis.port=${RedisRunner.port}")
47+
.applyTo(configurableApplicationContext)
48+
}
49+
}
50+
}
51+
52+
@Autowired
53+
lateinit var transportProvider: WebMvcStreamableServerTransportProvider
54+
55+
@Autowired
56+
lateinit var redissonClient: RedissonClient
57+
58+
lateinit var data: McpPatTestData
59+
60+
@BeforeEach
61+
fun setup() {
62+
data = createTestDataWithPat()
63+
}
64+
65+
@Test
66+
fun `session is recovered from Redis after local eviction`() {
67+
// 1. Initialize MCP client — this creates a session and the filter persists it to Redis
68+
val client = createMcpClientWithPat(data.pat.token!!)
69+
70+
// Verify the tool call works before eviction
71+
val resultBefore = callTool(client, "list_projects")
72+
assertThat(resultBefore.isError).isFalse()
73+
74+
// 2. Get the sessions map via reflection and find the session ID
75+
val sessionsMap = getSessionsMap()
76+
assertThat(sessionsMap).isNotEmpty
77+
78+
val sessionId = sessionsMap.keys().toList().first()
79+
80+
// Verify session data is in Redis
81+
val bucket = redissonClient.getBucket<String>("mcp_session:$sessionId")
82+
assertThat(bucket.get()).isNotNull
83+
84+
// 3. Remove session from local map to simulate request landing on a different replica
85+
sessionsMap.remove(sessionId)
86+
assertThat(sessionsMap.containsKey(sessionId)).isFalse()
87+
88+
// 4. Call a tool — the filter should recover the session from Redis
89+
val resultAfter = callTool(client, "list_projects")
90+
assertThat(resultAfter.isError).isFalse()
91+
92+
// 5. Verify session is back in local map
93+
assertThat(sessionsMap.containsKey(sessionId)).isTrue()
94+
}
95+
96+
@Suppress("UNCHECKED_CAST")
97+
private fun getSessionsMap(): ConcurrentHashMap<String, McpStreamableServerSession> {
98+
val field = WebMvcStreamableServerTransportProvider::class.java.getDeclaredField("sessions")
99+
field.isAccessible = true
100+
return field.get(transportProvider) as ConcurrentHashMap<String, McpStreamableServerSession>
101+
}
102+
}

0 commit comments

Comments
 (0)