Update Client.connect(transport:) to automatically send initialize request (#100)
* Update Client to automatically call initialize in connect * Make Client.connect(transport:) return value discardable * Fix example usage in README * Update example code in transports to remove initialize call * Deprecate initialize method * Update tests * Update README * Fix sampling tests * Fix race condition tests
This commit is contained in:
parent
3b505ee575
commit
a59b50bf5a
9 changed files with 302 additions and 106 deletions
14
README.md
14
README.md
|
|
@ -54,10 +54,7 @@ let client = Client(name: "MyApp", version: "1.0.0")
|
|||
|
||||
// Create a transport and connect
|
||||
let transport = StdioTransport()
|
||||
try await client.connect(transport: transport)
|
||||
|
||||
// Initialize the connection
|
||||
let result = try await client.initialize()
|
||||
let result = try await client.connect(transport: transport)
|
||||
|
||||
// Check server capabilities
|
||||
if result.capabilities.tools != nil {
|
||||
|
|
@ -65,6 +62,11 @@ if result.capabilities.tools != nil {
|
|||
}
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> The `Client.connect(transport:)` method returns the initialization result.
|
||||
> This return value is discardable,
|
||||
> so you can ignore it if you don't need to check server capabilities.
|
||||
|
||||
### Transport Options for Clients
|
||||
|
||||
#### Stdio Transport
|
||||
|
|
@ -278,7 +280,7 @@ Handle common client errors:
|
|||
|
||||
```swift
|
||||
do {
|
||||
let result = try await client.initialize()
|
||||
try await client.connect(transport: transport)
|
||||
// Success
|
||||
} catch let error as MCPError {
|
||||
print("MCP Error: \(error.localizedDescription)")
|
||||
|
|
@ -402,7 +404,7 @@ The server component allows your application to host model capabilities and resp
|
|||
```swift
|
||||
import MCP
|
||||
|
||||
// Initialize the server with capabilities
|
||||
// Create a server with given capabilities
|
||||
let server = Server(
|
||||
name: "MyModelServer",
|
||||
version: "1.0.0",
|
||||
|
|
|
|||
|
|
@ -41,9 +41,6 @@ import Logging
|
|||
/// let client = Client(name: "MyApp", version: "1.0.0")
|
||||
/// try await client.connect(transport: transport)
|
||||
///
|
||||
/// // Initialize the connection
|
||||
/// let result = try await client.initialize()
|
||||
///
|
||||
/// // The transport will automatically handle SSE events
|
||||
/// // and deliver them through the client's notification handlers
|
||||
/// ```
|
||||
|
|
|
|||
|
|
@ -51,9 +51,6 @@ import Logging
|
|||
/// // Use the transport with an MCP client
|
||||
/// let client = Client(name: "MyApp", version: "1.0.0")
|
||||
/// try await client.connect(transport: transport)
|
||||
///
|
||||
/// // Initialize the connection
|
||||
/// let result = try await client.initialize()
|
||||
/// ```
|
||||
public actor NetworkTransport: Transport {
|
||||
/// Represents a heartbeat message for connection health monitoring.
|
||||
|
|
|
|||
|
|
@ -44,9 +44,6 @@ import struct Foundation.Data
|
|||
/// // Create a transport and connect
|
||||
/// let transport = StdioTransport()
|
||||
/// try await client.connect(transport: transport)
|
||||
///
|
||||
/// // Initialize the connection
|
||||
/// let result = try await client.initialize()
|
||||
/// ```
|
||||
public actor StdioTransport: Transport {
|
||||
private let input: FileDescriptor
|
||||
|
|
|
|||
|
|
@ -168,7 +168,8 @@ public actor Client {
|
|||
}
|
||||
|
||||
/// Connect to the server using the given transport
|
||||
public func connect(transport: any Transport) async throws {
|
||||
@discardableResult
|
||||
public func connect(transport: any Transport) async throws -> Initialize.Result {
|
||||
self.connection = transport
|
||||
try await self.connection?.connect()
|
||||
|
||||
|
|
@ -217,6 +218,9 @@ public actor Client {
|
|||
} while true
|
||||
await self.logger?.info("Client message handling loop task is terminating.")
|
||||
}
|
||||
|
||||
// Automatically initialize after connecting
|
||||
return try await _initialize()
|
||||
}
|
||||
|
||||
/// Disconnect the client and cancel all pending requests
|
||||
|
|
@ -480,7 +484,23 @@ public actor Client {
|
|||
|
||||
// MARK: - Lifecycle
|
||||
|
||||
/// Initialize the connection with the server.
|
||||
///
|
||||
/// - Important: This method is deprecated. Initialization now happens automatically
|
||||
/// when calling `connect(transport:)`. You should use that method instead.
|
||||
///
|
||||
/// - Returns: The server's initialization response containing capabilities and server info
|
||||
@available(
|
||||
*, deprecated,
|
||||
message:
|
||||
"Initialization now happens automatically during connect. Use connect(transport:) instead."
|
||||
)
|
||||
public func initialize() async throws -> Initialize.Result {
|
||||
return try await _initialize()
|
||||
}
|
||||
|
||||
/// Internal initialization implementation
|
||||
private func _initialize() async throws -> Initialize.Result {
|
||||
let request = Initialize.request(
|
||||
.init(
|
||||
protocolVersion: Version.latest,
|
||||
|
|
|
|||
|
|
@ -11,43 +11,33 @@ struct ClientTests {
|
|||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
#expect(await transport.isConnected == false)
|
||||
try await client.connect(transport: transport)
|
||||
#expect(await transport.isConnected == true)
|
||||
await client.disconnect()
|
||||
#expect(await transport.isConnected == false)
|
||||
}
|
||||
|
||||
@Test(
|
||||
"Initialize request",
|
||||
.timeLimit(.minutes(1))
|
||||
)
|
||||
func testClientInitialize() async throws {
|
||||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
// Small delay to ensure message loop is started
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
|
||||
// Create a task for initialize that we'll cancel
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await client.initialize()
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
// Give it a moment to send the request
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
|
||||
#expect(await transport.sentMessages.count == 1)
|
||||
#expect(await transport.sentMessages.first?.contains(Initialize.name) == true)
|
||||
#expect(await transport.sentMessages.first?.contains(client.name) == true)
|
||||
#expect(await transport.sentMessages.first?.contains(client.version) == true)
|
||||
|
||||
// Cancel the initialize task
|
||||
initTask.cancel()
|
||||
|
||||
// Disconnect client to clean up message loop and give time for continuation cleanup
|
||||
let result = try await client.connect(transport: transport)
|
||||
#expect(await transport.isConnected == true)
|
||||
#expect(result.protocolVersion == Version.latest)
|
||||
await client.disconnect()
|
||||
try await Task.sleep(for: .milliseconds(50))
|
||||
#expect(await transport.isConnected == false)
|
||||
initTask.cancel()
|
||||
}
|
||||
|
||||
@Test(
|
||||
|
|
@ -58,24 +48,50 @@ struct ClientTests {
|
|||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
// Small delay to ensure message loop is started
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
// Queue a response for the initialize request
|
||||
try await Task.sleep(for: .milliseconds(10)) // Wait for request to be sent
|
||||
|
||||
// Create a task for the ping that we'll cancel
|
||||
let pingTask = Task {
|
||||
try await client.ping()
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.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 connect call which will automatically initialize
|
||||
let result = try await client.connect(transport: transport)
|
||||
#expect(result.protocolVersion == Version.latest)
|
||||
#expect(result.serverInfo.name == "TestServer")
|
||||
#expect(result.serverInfo.version == "1.0")
|
||||
|
||||
// Small delay to ensure message loop is started
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
|
||||
// Create a task for the ping
|
||||
let pingTask = Task {
|
||||
try await client.ping()
|
||||
}
|
||||
|
||||
// Give it a moment to send the request
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
|
||||
#expect(await transport.sentMessages.count == 2) // Initialize + Ping
|
||||
#expect(await transport.sentMessages.last?.contains(Ping.name) == true)
|
||||
|
||||
// Cancel the ping task
|
||||
pingTask.cancel()
|
||||
}
|
||||
|
||||
// Give it a moment to send the request
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
|
||||
#expect(await transport.sentMessages.count == 1)
|
||||
#expect(await transport.sentMessages.first?.contains(Ping.name) == true)
|
||||
|
||||
// Cancel the ping task
|
||||
pingTask.cancel()
|
||||
|
||||
// Disconnect client to clean up message loop and give time for continuation cleanup
|
||||
await client.disconnect()
|
||||
try await Task.sleep(for: .milliseconds(50))
|
||||
|
|
@ -104,10 +120,35 @@ struct ClientTests {
|
|||
@Test("Send failure handling")
|
||||
func testClientSendFailure() async throws {
|
||||
let transport = MockTransport()
|
||||
await transport.setFailSend(true)
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
// Connect first without failure
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
initTask.cancel()
|
||||
|
||||
// Now set the transport to fail sends
|
||||
await transport.setFailSend(true)
|
||||
|
||||
do {
|
||||
try await client.ping()
|
||||
|
|
@ -116,11 +157,13 @@ struct ClientTests {
|
|||
if case MCPError.transportError = error {
|
||||
#expect(Bool(true))
|
||||
} else {
|
||||
#expect(Bool(false), "Expected transport error")
|
||||
#expect(Bool(false), "Expected transport error, got \(error)")
|
||||
}
|
||||
} catch {
|
||||
#expect(Bool(false), "Expected MCP.Error")
|
||||
}
|
||||
|
||||
await client.disconnect()
|
||||
}
|
||||
|
||||
@Test("Strict configuration - capabilities check")
|
||||
|
|
@ -129,6 +172,26 @@ struct ClientTests {
|
|||
let config = Client.Configuration.strict
|
||||
let client = Client(name: "TestClient", version: "1.0", configuration: config)
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
|
||||
// Create a task for listPrompts
|
||||
|
|
@ -152,6 +215,7 @@ struct ClientTests {
|
|||
|
||||
// Cancel the task if it's still running
|
||||
promptsTask.cancel()
|
||||
initTask.cancel()
|
||||
|
||||
// Disconnect client
|
||||
await client.disconnect()
|
||||
|
|
@ -164,8 +228,31 @@ struct ClientTests {
|
|||
let config = Client.Configuration.default
|
||||
let client = Client(name: "TestClient", version: "1.0", configuration: config)
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
|
||||
// Make sure init task is complete
|
||||
initTask.cancel()
|
||||
|
||||
// Wait a bit for any setup to complete
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
|
||||
|
|
@ -232,8 +319,29 @@ struct ClientTests {
|
|||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10)) // Allow connection tasks
|
||||
initTask.cancel()
|
||||
|
||||
let request1 = Ping.request()
|
||||
let request2 = Ping.request()
|
||||
|
|
@ -245,11 +353,11 @@ struct ClientTests {
|
|||
resultTask2 = try await batch.addRequest(request2)
|
||||
}
|
||||
|
||||
// Check if one batch message was sent
|
||||
// Check if batch message was sent (after initialize and initialized notification)
|
||||
let sentMessages = await transport.sentMessages
|
||||
#expect(sentMessages.count == 1)
|
||||
#expect(sentMessages.count == 3) // Initialize request + Initialized notification + Batch
|
||||
|
||||
guard let batchData = sentMessages.first?.data(using: .utf8) else {
|
||||
guard let batchData = sentMessages.last?.data(using: .utf8) else {
|
||||
#expect(Bool(false), "Failed to get batch data")
|
||||
return
|
||||
}
|
||||
|
|
@ -291,8 +399,29 @@ struct ClientTests {
|
|||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
initTask.cancel()
|
||||
|
||||
let request1 = Ping.request() // Success
|
||||
let request2 = Ping.request() // Error
|
||||
|
|
@ -304,8 +433,8 @@ struct ClientTests {
|
|||
resultTasks.append(try await batch.addRequest(request2))
|
||||
}
|
||||
|
||||
// Check if one batch message was sent
|
||||
#expect(await transport.sentMessages.count == 1)
|
||||
// Check if batch message was sent (after initialize and initialized notification)
|
||||
#expect(await transport.sentMessages.count == 3) // Initialize request + Initialized notification + Batch
|
||||
|
||||
// Prepare batch response (success for 1, error for 2)
|
||||
let response1 = Response<Ping>(id: request1.id, result: .init())
|
||||
|
|
@ -350,16 +479,37 @@ struct ClientTests {
|
|||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
initTask.cancel()
|
||||
|
||||
// Call withBatch but don't add any requests
|
||||
try await client.withBatch { _ in
|
||||
// No requests added
|
||||
}
|
||||
|
||||
// Check that no messages were sent
|
||||
#expect(await transport.sentMessages.isEmpty)
|
||||
// Check that only initialize message and initialized notification were sent
|
||||
#expect(await transport.sentMessages.count == 2) // Initialize request + Initialized notification
|
||||
|
||||
await client.disconnect()
|
||||
}
|
||||
|
|
@ -369,17 +519,38 @@ struct ClientTests {
|
|||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
initTask.cancel()
|
||||
|
||||
// Create a test notification
|
||||
let notification = InitializedNotification.message()
|
||||
try await client.notify(notification)
|
||||
|
||||
// Verify notification was sent
|
||||
#expect(await transport.sentMessages.count == 1)
|
||||
// Verify notification was sent (in addition to initialize and initialized notification)
|
||||
#expect(await transport.sentMessages.count == 3) // Initialize request + Initialized notification + Custom notification
|
||||
|
||||
if let sentMessage = await transport.sentMessages.first,
|
||||
if let sentMessage = await transport.sentMessages.last,
|
||||
let data = sentMessage.data(using: .utf8)
|
||||
{
|
||||
|
||||
|
|
@ -404,9 +575,6 @@ struct ClientTests {
|
|||
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
|
||||
|
|
@ -431,7 +599,8 @@ struct ClientTests {
|
|||
try await transport.queue(response: response)
|
||||
|
||||
// Now complete the initialize call
|
||||
_ = try await client.initialize()
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
|
||||
// Verify that two messages were sent: initialize request and initialized notification
|
||||
#expect(await transport.sentMessages.count == 2)
|
||||
|
|
@ -485,8 +654,29 @@ struct ClientTests {
|
|||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
initTask.cancel()
|
||||
|
||||
// Set up the transport to fail sends from the start
|
||||
await transport.setFailSend(true)
|
||||
|
|
@ -525,8 +715,29 @@ struct ClientTests {
|
|||
let transport = MockTransport()
|
||||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
// Set up a task to handle the initialize response
|
||||
let initTask = Task {
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
if let lastMessage = await transport.sentMessages.last,
|
||||
let data = lastMessage.data(using: .utf8),
|
||||
let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
|
||||
{
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
try await client.connect(transport: transport)
|
||||
try await Task.sleep(for: .milliseconds(10))
|
||||
initTask.cancel()
|
||||
|
||||
// Create a ping request to get the ID
|
||||
let request = Ping.request()
|
||||
|
|
|
|||
|
|
@ -653,11 +653,8 @@ import Testing
|
|||
}
|
||||
}
|
||||
|
||||
// Execute the complete flow
|
||||
try await client.connect(transport: transport)
|
||||
|
||||
// Step 1: Initialize client
|
||||
let initResult = try await client.initialize()
|
||||
let initResult = try await client.connect(transport: transport)
|
||||
#expect(initResult.protocolVersion == Version.latest)
|
||||
#expect(initResult.capabilities.tools != nil)
|
||||
|
||||
|
|
|
|||
|
|
@ -99,10 +99,9 @@ struct RoundtripTests {
|
|||
let client = Client(name: "TestClient", version: "1.0")
|
||||
|
||||
try await server.start(transport: serverTransport)
|
||||
try await client.connect(transport: clientTransport)
|
||||
|
||||
let initTask = Task {
|
||||
let result = try await client.initialize()
|
||||
let result = try await client.connect(transport: clientTransport)
|
||||
|
||||
#expect(result.serverInfo.name == "TestServer")
|
||||
#expect(result.serverInfo.version == "1.0.0")
|
||||
|
|
|
|||
|
|
@ -587,30 +587,6 @@ struct SamplingIntegrationTests {
|
|||
try await server.start(transport: serverTransport)
|
||||
try await client.connect(transport: clientTransport)
|
||||
|
||||
// Test initialization and capability negotiation
|
||||
let initTask = Task {
|
||||
let result = try await client.initialize()
|
||||
|
||||
#expect(result.serverInfo.name == "SamplingTestServer")
|
||||
#expect(result.serverInfo.version == "1.0.0")
|
||||
#expect(
|
||||
result.capabilities.sampling != nil, "Server should advertise sampling capability")
|
||||
#expect(result.protocolVersion == Version.latest)
|
||||
}
|
||||
|
||||
try await withThrowingTaskGroup(of: Void.self) { group in
|
||||
group.addTask {
|
||||
try await Task.sleep(for: .seconds(1))
|
||||
initTask.cancel()
|
||||
throw CancellationError()
|
||||
}
|
||||
group.addTask {
|
||||
try await initTask.value
|
||||
}
|
||||
try await group.next()
|
||||
group.cancelAll()
|
||||
}
|
||||
|
||||
await server.stop()
|
||||
await client.disconnect()
|
||||
try? clientToServerRead.close()
|
||||
|
|
|
|||
Loading…
Reference in a new issue