fix(mcp): always disconnect transport on shutdown

Author: Tom-Ryder

Sources/AgentRunKit/MCP/MCPClient.swift

@@ -163,7 +163,7 @@ public actor MCPClient {
     }
 
     func shutdown() async {
-        guard drainPendingRequests() else { return }
+        _ = drainPendingRequests()
         readerTask?.cancel()
         readerTask = nil
         await transport.disconnect()

Tests/AgentRunKitTests/MCP/DynamicMCPTransport.swift

@@ -6,6 +6,7 @@ actor DynamicMCPTransport: MCPTransport {
     private let stream: AsyncThrowingStream<Data, Error>
     private let continuation: AsyncThrowingStream<Data, Error>.Continuation
     private var connected = false
+    private var effectiveDisconnectCount = 0
 
     init(handler: @escaping @Sendable (Data) async throws -> Data?) {
         self.handler = handler
@@ -19,6 +20,9 @@ actor DynamicMCPTransport: MCPTransport {
     }
 
     func disconnect() async {
+        if connected {
+            effectiveDisconnectCount += 1
+        }
         connected = false
         continuation.finish()
     }
@@ -45,4 +49,8 @@ actor DynamicMCPTransport: MCPTransport {
     func terminateStreamWithError(_ error: any Error) {
         continuation.finish(throwing: error)
     }
+
+    func effectiveDisconnectCallCount() -> Int {
+        effectiveDisconnectCount
+    }
 }

Tests/AgentRunKitTests/MCP/MCPClientTests.swift

@@ -55,12 +55,14 @@ struct MCPClientTests {
 
     @Test
     func initializeVersionMismatch() async throws {
-        let transport = ScriptedMCPTransport(responses: [
-            MCPTestHelpers.encodeResponse(
-                id: 1,
+        let transport = DynamicMCPTransport { data in
+            guard let request = decodeRequest(data) else { return nil }
+            let idValue: Int = if case let .int(val) = request.id { val } else { 0 }
+            return MCPTestHelpers.encodeResponse(
+                id: idValue,
                 result: MCPTestHelpers.initializeResult(protocolVersion: "1999-01-01")
-            ),
-        ])
+            )
+        }
         let client = MCPClient(serverName: "test", transport: transport)
         do {
             try await client.connectAndInitialize()
@@ -74,6 +76,7 @@ struct MCPClientTests {
             #expect(supported == "1999-01-01")
         }
         await client.shutdown()
+        await #expect(transport.effectiveDisconnectCallCount() == 1)
     }
 
     @Test
@@ -247,27 +250,6 @@ struct MCPClientTests {
         await client.shutdown()
     }
 
-    @Test
-    func transportClosedMidCall() async throws {
-        let transport = DynamicMCPTransport(handler: standardHandler { _ in nil })
-        let client = MCPClient(serverName: "test", transport: transport)
-        try await client.connectAndInitialize()
-
-        let callTask = Task {
-            try await client.callTool(name: "test", arguments: Data("{}".utf8))
-        }
-
-        try await Task.sleep(for: .milliseconds(50))
-        await transport.terminateStream()
-
-        do {
-            _ = try await callTask.value
-            Issue.record("Expected transportClosed")
-        } catch let error as MCPError {
-            #expect(error == .transportClosed)
-        }
-    }
-
     @Test
     func requestTimeout() async throws {
         let transport = DynamicMCPTransport(handler: standardHandler { _ in nil })
@@ -469,13 +451,37 @@ struct MCPClientTests {
 }
 
 struct MCPClientEdgeCaseTests {
+    @Test
+    func transportClosedMidCall() async throws {
+        let transport = DynamicMCPTransport(handler: standardHandler { _ in nil })
+        let client = MCPClient(serverName: "test", transport: transport)
+        try await client.connectAndInitialize()
+
+        let callTask = Task {
+            try await client.callTool(name: "test", arguments: Data("{}".utf8))
+        }
+
+        try await Task.sleep(for: .milliseconds(50))
+        await transport.terminateStream()
+
+        do {
+            _ = try await callTask.value
+            Issue.record("Expected transportClosed")
+        } catch let error as MCPError {
+            #expect(error == .transportClosed)
+        }
+        await client.shutdown()
+        await #expect(transport.effectiveDisconnectCallCount() == 1)
+    }
+
     @Test
     func shutdownIdempotent() async throws {
         let transport = DynamicMCPTransport(handler: standardHandler())
         let client = MCPClient(serverName: "test", transport: transport)
         try await client.connectAndInitialize()
         await client.shutdown()
         await client.shutdown()
+        await #expect(transport.effectiveDisconnectCallCount() == 1)
 
         await #expect(throws: MCPError.transportClosed) {
             try await client.callTool(name: "test", arguments: Data("{}".utf8))