From 2510dda70ebdef49aaa89111228a853b6eeb5094 Mon Sep 17 00:00:00 2001
From: Mattt <mattt@loopwork.com>
Date: Tue, 18 Mar 2025 09:47:54 -0700
Subject: [PATCH] Add server initialization hook (#19)

---
 Sources/MCP/Server/Server.swift  | 34 ++++++++-----
 Tests/MCPTests/ServerTests.swift | 85 ++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+), 11 deletions(-)

diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift
index 6db249c..66dd6ba 100644
--- a/Sources/MCP/Server/Server.swift
+++ b/Sources/MCP/Server/Server.swift
@@ -155,9 +155,15 @@ public actor Server {
     }
 
     /// Start the server
-    public func start(transport: any Transport) async throws {
+    /// - Parameters:
+    ///   - transport: The transport to use for the server
+    ///   - initializeHook: An optional hook that runs when the client sends an initialize request
+    public func start(
+        transport: any Transport,
+        initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)? = nil
+    ) async throws {
         self.connection = transport
-        registerDefaultHandlers()
+        registerDefaultHandlers(initializeHook: initializeHook)
         try await transport.connect()
 
         await logger?.info(
@@ -364,7 +370,9 @@ public actor Server {
         }
     }
 
-    private func registerDefaultHandlers() {
+    private func registerDefaultHandlers(
+        initializeHook: (@Sendable (Client.Info, Client.Capabilities) async throws -> Void)?
+    ) {
         // Initialize
         withMethodHandler(Initialize.self) { [weak self] params in
             guard let self = self else {
@@ -381,26 +389,30 @@ public actor Server {
                     "Unsupported protocol version: \(params.protocolVersion)")
             }
 
+            // Call initialization hook if registered
+            if let hook = initializeHook {
+                try await hook(params.clientInfo, params.capabilities)
+            }
+
+            // Set initial state
             await self.setInitialState(
                 clientInfo: params.clientInfo,
                 clientCapabilities: params.capabilities,
                 protocolVersion: params.protocolVersion
             )
 
-            let result = Initialize.Result(
-                protocolVersion: Version.latest,
-                capabilities: await self.capabilities,
-                serverInfo: self.serverInfo,
-                instructions: nil
-            )
-
             // Send initialized notification after a short delay
             Task {
                 try? await Task.sleep(nanoseconds: 100_000_000)  // 100ms
                 try? await self.notify(InitializedNotification.message())
             }
 
-            return result
+            return Initialize.Result(
+                protocolVersion: Version.latest,
+                capabilities: await self.capabilities,
+                serverInfo: self.serverInfo,
+                instructions: nil
+            )
         }
 
         // Ping
diff --git a/Tests/MCPTests/ServerTests.swift b/Tests/MCPTests/ServerTests.swift
index e733c01..97e0260 100644
--- a/Tests/MCPTests/ServerTests.swift
+++ b/Tests/MCPTests/ServerTests.swift
@@ -48,4 +48,89 @@ struct ServerTests {
         await server.stop()
         await transport.disconnect()
     }
+
+    @Test("Initialize hook - successful")
+    func testInitializeHookSuccess() async throws {
+        let transport = MockTransport()
+
+        actor TestState {
+            var hookCalled = false
+            func setHookCalled() { hookCalled = true }
+            func wasHookCalled() -> Bool { hookCalled }
+        }
+
+        let state = TestState()
+        let server = Server(name: "TestServer", version: "1.0")
+
+        // Start with the hook directly
+        try await server.start(transport: transport) { clientInfo, capabilities in
+            #expect(clientInfo.name == "TestClient")
+            #expect(clientInfo.version == "1.0")
+            await state.setHookCalled()
+        }
+        
+        // Wait for server to initialize
+        try await Task.sleep(nanoseconds: 10_000_000)  // 10ms
+
+        // Queue an initialize request
+        try await transport.queueRequest(
+            Initialize.request(
+                .init(
+                    protocolVersion: Version.latest,
+                    capabilities: .init(),
+                    clientInfo: .init(name: "TestClient", version: "1.0")
+                )
+            ))
+
+        // Wait for message processing and hook execution
+        try await Task.sleep(nanoseconds: 200_000_000)  // 200ms
+
+        #expect(await state.wasHookCalled() == true)
+        #expect(await transport.sentMessages.count >= 1)
+
+        let messages = await transport.sentMessages
+        if let response = messages.first {
+            #expect(response.contains("serverInfo"))
+        }
+
+        await server.stop()
+    }
+
+    @Test("Initialize hook - rejection")
+    func testInitializeHookRejection() async throws {
+        let transport = MockTransport()
+
+        let server = Server(name: "TestServer", version: "1.0")
+
+        try await server.start(transport: transport) { clientInfo, _ in
+            if clientInfo.name == "BlockedClient" {
+                throw Error.invalidRequest("Client not allowed")
+            }
+        }
+        
+        // Wait for server to initialize
+        try await Task.sleep(nanoseconds: 10_000_000)  // 10ms
+
+        // Queue an initialize request from blocked client
+        try await transport.queueRequest(
+            Initialize.request(
+                .init(
+                    protocolVersion: Version.latest,
+                    capabilities: .init(),
+                    clientInfo: .init(name: "BlockedClient", version: "1.0")
+                )
+            ))
+
+        // Wait for message processing
+        try await Task.sleep(nanoseconds: 200_000_000)  // 200ms
+
+        #expect(await transport.sentMessages.count >= 2)
+
+        let messages = await transport.sentMessages
+        if let response = messages.first {
+            #expect(response.contains("error"))
+            #expect(response.contains("Client not allowed"))
+        }
+        await server.stop()
+    }
 }