From 25dee3bb4ca7001a636af59fd1dcd328e20520b0 Mon Sep 17 00:00:00 2001 From: Mattt Date: Mon, 21 Apr 2025 14:43:33 -0700 Subject: [PATCH] Send initialized notification from client instead of server (#85) --- Sources/MCP/Client/Client.swift | 12 ++++ Sources/MCP/Server/Server.swift | 8 +-- Tests/MCPTests/ClientTests.swift | 116 +++++++++++++++++++++++++++++++ Tests/MCPTests/ServerTests.swift | 6 +- 4 files changed, 130 insertions(+), 12 deletions(-) diff --git a/Sources/MCP/Client/Client.swift b/Sources/MCP/Client/Client.swift index fa4fbe0..f6ec77a 100644 --- a/Sources/MCP/Client/Client.swift +++ b/Sources/MCP/Client/Client.swift @@ -249,6 +249,16 @@ public actor Client { return self } + /// Send a notification to the server + public func notify(_ notification: Message) async throws { + guard let connection = connection else { + throw MCPError.internalError("Client connection not initialized") + } + + let notificationData = try encoder.encode(notification) + try await connection.send(notificationData) + } + // MARK: - Requests /// Send a request and receive its response @@ -454,6 +464,8 @@ public actor Client { self.serverVersion = result.protocolVersion self.instructions = result.instructions + try await notify(InitializedNotification.message()) + return result } diff --git a/Sources/MCP/Server/Server.swift b/Sources/MCP/Server/Server.swift index 9049d75..9d348d2 100644 --- a/Sources/MCP/Server/Server.swift +++ b/Sources/MCP/Server/Server.swift @@ -496,12 +496,6 @@ public actor Server { protocolVersion: params.protocolVersion ) - // Send initialized notification after a short delay - Task { - try? await Task.sleep(for: .milliseconds(10)) - try? await self.notify(InitializedNotification.message()) - } - return Initialize.Result( protocolVersion: Version.latest, capabilities: await self.capabilities, @@ -571,4 +565,4 @@ extension Server.Batch.Item: Codable { try notification.encode(to: encoder) } } -} \ No newline at end of file +} diff --git a/Tests/MCPTests/ClientTests.swift b/Tests/MCPTests/ClientTests.swift index e7e2c02..cfbedeb 100644 --- a/Tests/MCPTests/ClientTests.swift +++ b/Tests/MCPTests/ClientTests.swift @@ -363,4 +363,120 @@ struct ClientTests { await client.disconnect() } + + @Test("Notify method sends notifications") + func testClientNotify() async throws { + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + try await client.connect(transport: transport) + try await Task.sleep(for: .milliseconds(10)) + + // Create a test notification + let notification = InitializedNotification.message() + try await client.notify(notification) + + // Verify notification was sent + #expect(await transport.sentMessages.count == 1) + + if let sentMessage = await transport.sentMessages.first, + let data = sentMessage.data(using: .utf8) + { + + // Decode as Message + let decoder = JSONDecoder() + do { + let decodedNotification = try decoder.decode( + Message.self, from: data) + #expect(decodedNotification.method == InitializedNotification.name) + } catch { + #expect(Bool(false), "Failed to decode notification: \(error)") + } + } else { + #expect(Bool(false), "No message was sent") + } + + await client.disconnect() + } + + @Test("Initialize sends initialized notification") + func testClientInitializeNotification() async throws { + let transport = MockTransport() + let client = Client(name: "TestClient", version: "1.0") + + try await client.connect(transport: transport) + try await Task.sleep(for: .milliseconds(10)) + + // Create a task for initialize + let initTask = Task { + // Queue a response for the initialize request + try await Task.sleep(for: .milliseconds(10)) // Wait for request to be sent + + if let lastMessage = await transport.sentMessages.last, + let data = lastMessage.data(using: .utf8), + let request = try? JSONDecoder().decode(Request.self, from: data) + { + + // Create a valid initialize response + let response = Initialize.response( + id: request.id, + result: .init( + protocolVersion: Version.latest, + capabilities: .init(), + serverInfo: .init(name: "TestServer", version: "1.0"), + instructions: nil + ) + ) + + try await transport.queue(response: response) + + // Now complete the initialize call + _ = try await client.initialize() + + // Verify that two messages were sent: initialize request and initialized notification + #expect(await transport.sentMessages.count == 2) + + // Check that the second message is the initialized notification + let notifications = await transport.sentMessages + if notifications.count >= 2 { + let notificationJson = notifications[1] + if let notificationData = notificationJson.data(using: .utf8) { + do { + let decoder = JSONDecoder() + let decodedNotification = try decoder.decode( + Message.self, from: notificationData) + #expect(decodedNotification.method == InitializedNotification.name) + } catch { + #expect(Bool(false), "Failed to decode notification: \(error)") + } + } else { + #expect(Bool(false), "Could not convert notification to data") + } + } else { + #expect( + Bool(false), "Expected both initialize request and initialized notification" + ) + } + } + } + + // Wait with timeout + let timeoutTask = Task { + try await Task.sleep(for: .seconds(1)) + initTask.cancel() + } + + // Wait for the task to complete + do { + _ = try await initTask.value + } catch is CancellationError { + #expect(Bool(false), "Test timed out") + } catch { + #expect(Bool(false), "Unexpected error: \(error)") + } + + timeoutTask.cancel() + + await client.disconnect() + } } diff --git a/Tests/MCPTests/ServerTests.swift b/Tests/MCPTests/ServerTests.swift index 91b97bd..1a5d5f6 100644 --- a/Tests/MCPTests/ServerTests.swift +++ b/Tests/MCPTests/ServerTests.swift @@ -41,17 +41,13 @@ struct ServerTests { // Wait for message processing and response try await Task.sleep(nanoseconds: 100_000_000) // 100ms - #expect(await transport.sentMessages.count == 2) + #expect(await transport.sentMessages.count == 1) let messages = await transport.sentMessages if let response = messages.first { #expect(response.contains("serverInfo")) } - if let noticiation = messages.last { - #expect(noticiation.contains("initialized")) - } - // Clean up await server.stop() await transport.disconnect()