77import java .io .IOException ;
88import java .time .Duration ;
99import java .util .List ;
10- import java .util .UUID ;
1110import java .util .concurrent .ConcurrentHashMap ;
1211import java .util .concurrent .locks .ReentrantLock ;
1312
@@ -256,39 +255,31 @@ private ServerResponse handleSseConnection(ServerRequest request) {
256255 return ServerResponse .status (HttpStatus .SERVICE_UNAVAILABLE ).body ("Server is shutting down" );
257256 }
258257
259- String sessionId = UUID .randomUUID ().toString ();
260- logger .debug ("Creating new SSE connection for session: {}" , sessionId );
261-
262258 // Send initial endpoint event
263- try {
264- return ServerResponse . sse ( sseBuilder -> {
265- sseBuilder . onComplete (() -> {
266- logger . debug ( "SSE connection completed for session: {}" , sessionId );
267- sessions . remove ( sessionId );
268- });
269- sseBuilder . onTimeout (() -> {
270- logger . debug ( "SSE connection timed out for session: {}" , sessionId );
271- sessions . remove ( sessionId );
272- });
273-
274- WebMvcMcpSessionTransport sessionTransport = new WebMvcMcpSessionTransport (sessionId , sseBuilder );
275- McpServerSession session = sessionFactory . create ( sessionTransport );
276- this .sessions .put (sessionId , session );
259+ return ServerResponse . sse ( sseBuilder -> {
260+ WebMvcMcpSessionTransport sessionTransport = new WebMvcMcpSessionTransport ( sseBuilder );
261+ McpServerSession session = sessionFactory . create ( sessionTransport );
262+ String sessionId = session . getId ( );
263+ logger . debug ( "Creating new SSE connection for session: {}" , sessionId );
264+ sseBuilder . onComplete (() -> {
265+ logger . debug ( "SSE connection completed for session: {}" , sessionId );
266+ sessions . remove ( sessionId );
267+ } );
268+ sseBuilder . onTimeout (() -> {
269+ logger . debug ( "SSE connection timed out for session: {}" , sessionId );
270+ sessions . remove (sessionId );
271+ } );
272+ this .sessions .put (sessionId , session );
277273
278- try {
279- sseBuilder .id (sessionId ).event (ENDPOINT_EVENT_TYPE ).data (buildEndpointUrl (sessionId ));
280- }
281- catch (Exception e ) {
282- logger .error ("Failed to send initial endpoint event: {}" , e .getMessage ());
283- sseBuilder .error (e );
284- }
285- }, Duration .ZERO );
286- }
287- catch (Exception e ) {
288- logger .error ("Failed to send initial endpoint event to session {}: {}" , sessionId , e .getMessage ());
289- sessions .remove (sessionId );
290- return ServerResponse .status (HttpStatus .INTERNAL_SERVER_ERROR ).build ();
291- }
274+ try {
275+ sseBuilder .event (ENDPOINT_EVENT_TYPE ).data (buildEndpointUrl (sessionId ));
276+ }
277+ catch (Exception e ) {
278+ logger .error ("Failed to send initial endpoint event: {}" , e .getMessage ());
279+ this .sessions .remove (sessionId );
280+ sseBuilder .error (e );
281+ }
282+ }, Duration .ZERO );
292283 }
293284
294285 /**
@@ -363,8 +354,6 @@ private ServerResponse handleMessage(ServerRequest request) {
363354 */
364355 private class WebMvcMcpSessionTransport implements McpServerTransport {
365356
366- private final String sessionId ;
367-
368357 private final SseBuilder sseBuilder ;
369358
370359 /**
@@ -374,14 +363,11 @@ private class WebMvcMcpSessionTransport implements McpServerTransport {
374363 private final ReentrantLock sseBuilderLock = new ReentrantLock ();
375364
376365 /**
377- * Creates a new session transport with the specified ID and SSE builder.
378- * @param sessionId The unique identifier for this session
366+ * Creates a new session transport with the specified SSE builder.
379367 * @param sseBuilder The SSE builder for sending server events to the client
380368 */
381- WebMvcMcpSessionTransport (String sessionId , SseBuilder sseBuilder ) {
382- this .sessionId = sessionId ;
369+ WebMvcMcpSessionTransport (SseBuilder sseBuilder ) {
383370 this .sseBuilder = sseBuilder ;
384- logger .debug ("Session transport {} initialized with SSE builder" , sessionId );
385371 }
386372
387373 /**
@@ -395,11 +381,10 @@ public Mono<Void> sendMessage(McpSchema.JSONRPCMessage message) {
395381 sseBuilderLock .lock ();
396382 try {
397383 String jsonText = jsonMapper .writeValueAsString (message );
398- sseBuilder .id (sessionId ).event (MESSAGE_EVENT_TYPE ).data (jsonText );
399- logger .debug ("Message sent to session {}" , sessionId );
384+ sseBuilder .event (MESSAGE_EVENT_TYPE ).data (jsonText );
400385 }
401386 catch (Exception e ) {
402- logger .error ("Failed to send message to session {} : {}" , sessionId , e .getMessage ());
387+ logger .error ("Failed to send message: {}" , e .getMessage ());
403388 sseBuilder .error (e );
404389 }
405390 finally {
@@ -427,14 +412,12 @@ public <T> T unmarshalFrom(Object data, TypeRef<T> typeRef) {
427412 @ Override
428413 public Mono <Void > closeGracefully () {
429414 return Mono .fromRunnable (() -> {
430- logger .debug ("Closing session transport: {}" , sessionId );
431415 sseBuilderLock .lock ();
432416 try {
433417 sseBuilder .complete ();
434- logger .debug ("Successfully completed SSE builder for session {}" , sessionId );
435418 }
436419 catch (Exception e ) {
437- logger .warn ("Failed to complete SSE builder for session {} : {}" , sessionId , e .getMessage ());
420+ logger .warn ("Failed to complete SSE builder: {}" , e .getMessage ());
438421 }
439422 finally {
440423 sseBuilderLock .unlock ();
@@ -450,10 +433,9 @@ public void close() {
450433 sseBuilderLock .lock ();
451434 try {
452435 sseBuilder .complete ();
453- logger .debug ("Successfully completed SSE builder for session {}" , sessionId );
454436 }
455437 catch (Exception e ) {
456- logger .warn ("Failed to complete SSE builder for session {} : {}" , sessionId , e .getMessage ());
438+ logger .warn ("Failed to complete SSE builder: {}" , e .getMessage ());
457439 }
458440 finally {
459441 sseBuilderLock .unlock ();
0 commit comments