WIP: Get rid of initialization notification handler

Signed-off-by: Dariusz Jędrzejczyk <[email protected]>

Author: chemicL

mcp/src/main/java/io/modelcontextprotocol/server/McpAsyncServer.java

@@ -145,9 +145,8 @@ public class McpAsyncServer {
 		Map<String, McpRequestHandler<?>> requestHandlers = prepareRequestHandlers();
 		Map<String, McpNotificationHandler> notificationHandlers = prepareNotificationHandlers(features);
 
-		mcpTransportProvider.setSessionFactory(
-				transport -> new McpServerSession(UUID.randomUUID().toString(), requestTimeout, transport,
-						this::asyncInitializeRequestHandler, Mono::empty, requestHandlers, notificationHandlers));
+		mcpTransportProvider.setSessionFactory(transport -> new McpServerSession(UUID.randomUUID().toString(),
+				requestTimeout, transport, this::asyncInitializeRequestHandler, requestHandlers, notificationHandlers));
 	}
 
 	McpAsyncServer(McpStreamableServerTransportProvider mcpTransportProvider, ObjectMapper objectMapper,

mcp/src/main/java/io/modelcontextprotocol/spec/DefaultMcpStreamableServerSessionFactory.java

@@ -33,8 +33,7 @@ public McpStreamableServerSession.McpStreamableServerSessionInit startSession(
 			McpSchema.InitializeRequest initializeRequest) {
 		return new McpStreamableServerSession.McpStreamableServerSessionInit(
 				new McpStreamableServerSession(UUID.randomUUID().toString(), initializeRequest.capabilities(),
-						initializeRequest.clientInfo(), requestTimeout, Mono::empty, requestHandlers,
-						notificationHandlers),
+						initializeRequest.clientInfo(), requestTimeout, requestHandlers, notificationHandlers),
 				this.initRequestHandler.handle(initializeRequest));
 	}
 

mcp/src/main/java/io/modelcontextprotocol/spec/McpServerSession.java

@@ -38,8 +38,6 @@ public class McpServerSession implements McpLoggableSession {
 
 	private final McpInitRequestHandler initRequestHandler;
 
-	private final InitNotificationHandler initNotificationHandler;
-
 	private final Map<String, McpRequestHandler<?>> requestHandlers;
 
 	private final Map<String, McpNotificationHandler> notificationHandlers;
@@ -62,6 +60,27 @@ public class McpServerSession implements McpLoggableSession {
 
 	private volatile McpSchema.LoggingLevel minLoggingLevel = McpSchema.LoggingLevel.INFO;
 
+	/**
+	 * Creates a new server session with the given parameters and the transport to use.
+	 * @param id session id
+	 * @param transport the transport to use
+	 * @param initHandler called when a
+	 * {@link io.modelcontextprotocol.spec.McpSchema.InitializeRequest} is received by the
+	 * server
+	 * @param requestHandlers map of request handlers to use
+	 * @param notificationHandlers map of notification handlers to use
+	 */
+	public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
+			McpInitRequestHandler initHandler, Map<String, McpRequestHandler<?>> requestHandlers,
+			Map<String, McpNotificationHandler> notificationHandlers) {
+		this.id = id;
+		this.requestTimeout = requestTimeout;
+		this.transport = transport;
+		this.initRequestHandler = initHandler;
+		this.requestHandlers = requestHandlers;
+		this.notificationHandlers = notificationHandlers;
+	}
+
 	/**
 	 * Creates a new server session with the given parameters and the transport to use.
 	 * @param id session id
@@ -74,7 +93,10 @@ public class McpServerSession implements McpLoggableSession {
 	 * received.
 	 * @param requestHandlers map of request handlers to use
 	 * @param notificationHandlers map of notification handlers to use
+	 * @deprecated Use
+	 * {@link #McpServerSession(String, Duration, McpServerTransport, McpInitRequestHandler, Map, Map)}
 	 */
+	@Deprecated
 	public McpServerSession(String id, Duration requestTimeout, McpServerTransport transport,
 			McpInitRequestHandler initHandler, InitNotificationHandler initNotificationHandler,
 			Map<String, McpRequestHandler<?>> requestHandlers,
@@ -83,7 +105,6 @@ public McpServerSession(String id, Duration requestTimeout, McpServerTransport t
 		this.requestTimeout = requestTimeout;
 		this.transport = transport;
 		this.initRequestHandler = initHandler;
-		this.initNotificationHandler = initNotificationHandler;
 		this.requestHandlers = requestHandlers;
 		this.notificationHandlers = notificationHandlers;
 	}
@@ -264,7 +285,6 @@ private Mono<Void> handleIncomingNotification(McpSchema.JSONRPCNotification noti
 				// legacy SSE transport.
 				exchangeSink.tryEmitValue(new McpAsyncServerExchange(this.id, this, clientCapabilities.get(),
 						clientInfo.get(), McpTransportContext.EMPTY));
-				return this.initNotificationHandler.handle();
 			}
 
 			var handler = notificationHandlers.get(notification.method());

mcp/src/main/java/io/modelcontextprotocol/spec/McpStreamableServerSession.java

@@ -32,8 +32,6 @@ public class McpStreamableServerSession implements McpLoggableSession {
 
 	private final AtomicLong requestCounter = new AtomicLong(0);
 
-	private final InitNotificationHandler initNotificationHandler;
-
 	private final Map<String, McpRequestHandler<?>> requestHandlers;
 
 	private final Map<String, McpNotificationHandler> notificationHandlers;
@@ -58,15 +56,14 @@ public class McpStreamableServerSession implements McpLoggableSession {
 
 	public McpStreamableServerSession(String id, McpSchema.ClientCapabilities clientCapabilities,
 			McpSchema.Implementation clientInfo, Duration requestTimeout,
-			InitNotificationHandler initNotificationHandler, Map<String, McpRequestHandler<?>> requestHandlers,
+			Map<String, McpRequestHandler<?>> requestHandlers,
 			Map<String, McpNotificationHandler> notificationHandlers) {
 		this.id = id;
 		this.missingMcpTransportSession = new MissingMcpTransportSession(id);
 		this.listeningStreamRef = new AtomicReference<>(this.missingMcpTransportSession);
 		this.clientCapabilities.lazySet(clientCapabilities);
 		this.clientInfo.lazySet(clientInfo);
 		this.requestTimeout = requestTimeout;
-		this.initNotificationHandler = initNotificationHandler;
 		this.requestHandlers = requestHandlers;
 		this.notificationHandlers = notificationHandlers;
 	}
@@ -231,19 +228,6 @@ public interface InitRequestHandler {
 
 	}
 
-	/**
-	 * Notification handler for the initialization notification from the client.
-	 */
-	public interface InitNotificationHandler {
-
-		/**
-		 * Specifies an action to take upon successful initialization.
-		 * @return a Mono that will complete when the initialization is acted upon.
-		 */
-		Mono<Void> handle();
-
-	}
-
 	public interface Factory {
 
 		McpStreamableServerSessionInit startSession(McpSchema.InitializeRequest initializeRequest);

mcp/src/test/java/io/modelcontextprotocol/client/StdioMcpAsyncClientTests.java

@@ -37,7 +37,7 @@ protected McpClientTransport createMcpTransport() {
 	}
 
 	protected Duration getInitializationTimeout() {
-		return Duration.ofSeconds(10);
+		return Duration.ofSeconds(20);
 	}
 
 }