diff --git a/Sources/MCP/Base/Transports/StdioTransport.swift b/Sources/MCP/Base/Transports/StdioTransport.swift index 45522fac..fbada31e 100644 --- a/Sources/MCP/Base/Transports/StdioTransport.swift +++ b/Sources/MCP/Base/Transports/StdioTransport.swift @@ -54,8 +54,19 @@ import struct Foundation.Data public nonisolated let logger: Logger private var isConnected = false + private var readTask: Task? + private var disconnectTask: Task? + private var activeSends = 0 + private var sendsDrained: CheckedContinuation? private let messageStream: AsyncThrowingStream private let messageContinuation: AsyncThrowingStream.Continuation + private struct DescriptorState { + let flags: CInt + let device: UInt64 + let inode: UInt64 + } + private var inputState: DescriptorState? + private var outputState: DescriptorState? /// Creates a new stdio transport with the specified file descriptors /// @@ -89,17 +100,24 @@ import struct Foundation.Data /// /// - Throws: Error if the file descriptors cannot be configured public func connect() async throws { + await disconnectTask?.value guard !isConnected else { return } // Set non-blocking mode - try setNonBlocking(fileDescriptor: input) - try setNonBlocking(fileDescriptor: output) + let originalInputState = try setNonBlocking(fileDescriptor: input) + do { + outputState = try setNonBlocking(fileDescriptor: output) + } catch { + restoreFlags(originalInputState, for: input) + throw error + } + inputState = originalInputState isConnected = true logger.debug("Transport connected successfully") // Start reading loop in background - Task { + readTask = Task { await readLoop() } } @@ -108,19 +126,25 @@ import struct Foundation.Data /// /// - Parameter fileDescriptor: The file descriptor to configure /// - Throws: Error if the operation fails - private func setNonBlocking(fileDescriptor: FileDescriptor) throws { + private func setNonBlocking(fileDescriptor: FileDescriptor) throws -> DescriptorState { #if canImport(Darwin) || canImport(Glibc) || canImport(Musl) // Get current flags let flags = fcntl(fileDescriptor.rawValue, F_GETFL) guard flags >= 0 else { throw MCPError.transportError(Errno(rawValue: CInt(errno))) } + var status = stat() + guard fstat(fileDescriptor.rawValue, &status) == 0 else { + throw MCPError.transportError(Errno(rawValue: CInt(errno))) + } // Set non-blocking flag let result = fcntl(fileDescriptor.rawValue, F_SETFL, flags | O_NONBLOCK) guard result >= 0 else { throw MCPError.transportError(Errno(rawValue: CInt(errno))) } + return DescriptorState( + flags: flags, device: UInt64(status.st_dev), inode: UInt64(status.st_ino)) #else // For platforms where non-blocking operations aren't supported throw MCPError.internalError( @@ -128,6 +152,20 @@ import struct Foundation.Data #endif } + private func restoreFlags(_ state: DescriptorState, for fileDescriptor: FileDescriptor) { + guard state.flags & O_NONBLOCK == 0 else { return } + var status = stat() + guard fstat(fileDescriptor.rawValue, &status) == 0, + UInt64(status.st_dev) == state.device, + UInt64(status.st_ino) == state.inode + else { return } + let currentFlags = fcntl(fileDescriptor.rawValue, F_GETFL) + guard currentFlags >= 0 else { return } + if fcntl(fileDescriptor.rawValue, F_SETFL, currentFlags & ~O_NONBLOCK) < 0 { + logger.warning("Failed to restore file descriptor blocking mode") + } + } + /// Continuous loop that reads and processes incoming messages /// /// This method runs in the background while the transport is connected, @@ -180,9 +218,31 @@ import struct Foundation.Data /// /// This stops the message reading loop and releases associated resources. public func disconnect() async { + if let disconnectTask { + await disconnectTask.value + return + } guard isConnected else { return } isConnected = false - messageContinuation.finish() + let read = readTask + read?.cancel() + let cleanup = Task { + await read?.value + if activeSends > 0 { + await withCheckedContinuation { continuation in + sendsDrained = continuation + } + } + if let inputState { restoreFlags(inputState, for: input) } + if let outputState { restoreFlags(outputState, for: output) } + inputState = nil + outputState = nil + readTask = nil + messageContinuation.finish() + } + disconnectTask = cleanup + await cleanup.value + disconnectTask = nil logger.debug("Transport disconnected") } @@ -198,6 +258,14 @@ import struct Foundation.Data guard isConnected else { throw MCPError.transportError(Errno(rawValue: ENOTCONN)) } + activeSends += 1 + defer { + activeSends -= 1 + if activeSends == 0 { + sendsDrained?.resume() + sendsDrained = nil + } + } // Add newline as delimiter var messageWithNewline = message @@ -205,6 +273,9 @@ import struct Foundation.Data var remaining = messageWithNewline while !remaining.isEmpty { + guard isConnected else { + throw MCPError.transportError(Errno(rawValue: ENOTCONN)) + } do { let written = try remaining.withUnsafeBytes { buffer in try output.write(UnsafeRawBufferPointer(buffer)) diff --git a/Tests/MCPTests/StdioTransportTests.swift b/Tests/MCPTests/StdioTransportTests.swift index e395fbe0..08331411 100644 --- a/Tests/MCPTests/StdioTransportTests.swift +++ b/Tests/MCPTests/StdioTransportTests.swift @@ -1,6 +1,14 @@ import Foundation import Testing +#if canImport(Darwin) + import Darwin.POSIX +#elseif canImport(Glibc) + import Glibc +#elseif canImport(Musl) + import Musl +#endif + @testable import MCP #if canImport(System) @@ -20,6 +28,121 @@ struct StdioTransportTests { await transport.disconnect() } + @Test("Disconnect restores file descriptor blocking mode") + func testDisconnectRestoresBlockingMode() async throws { + let (input, writer) = try FileDescriptor.pipe() + let (reader, output) = try FileDescriptor.pipe() + defer { + try? input.close() + try? writer.close() + try? reader.close() + try? output.close() + } + + let inputFlags = fcntl(input.rawValue, F_GETFL) + let outputFlags = fcntl(output.rawValue, F_GETFL) + #expect(inputFlags & O_NONBLOCK == 0) + #expect(outputFlags & O_NONBLOCK == 0) + + let transport = StdioTransport(input: input, output: output, logger: nil) + try await transport.connect() + #expect(fcntl(input.rawValue, F_GETFL) & O_NONBLOCK != 0) + #expect(fcntl(output.rawValue, F_GETFL) & O_NONBLOCK != 0) + + await transport.disconnect() + #expect(fcntl(input.rawValue, F_GETFL) & O_NONBLOCK == 0) + #expect(fcntl(output.rawValue, F_GETFL) & O_NONBLOCK == 0) + } + + @Test("Disconnect cancels a backpressured send before restoring descriptor flags") + func testDisconnectDuringBackpressuredSend() async throws { + let (input, writer) = try FileDescriptor.pipe() + let (reader, output) = try FileDescriptor.pipe() + defer { + try? input.close() + try? writer.close() + try? reader.close() + try? output.close() + } + + let transport = StdioTransport(input: input, output: output, logger: nil) + try await transport.connect() + let send = Task { try await transport.send(Data(repeating: 65, count: 512 * 1024)) } + try await Task.sleep(for: .milliseconds(50)) + + await transport.disconnect() + do { + try await send.value + Issue.record("Expected the incomplete send to fail on disconnect") + } catch { + // The caller must be released before the descriptor becomes blocking. + } + #expect(fcntl(output.rawValue, F_GETFL) & O_NONBLOCK == 0) + } + + @Test("Disconnect preserves an already nonblocking descriptor") + func testDisconnectPreservesNonblockingMode() async throws { + let (input, writer) = try FileDescriptor.pipe() + let (reader, output) = try FileDescriptor.pipe() + defer { + try? input.close() + try? writer.close() + try? reader.close() + try? output.close() + } + let flags = fcntl(output.rawValue, F_GETFL) + #expect(fcntl(output.rawValue, F_SETFL, flags | O_NONBLOCK) == 0) + + let transport = StdioTransport(input: input, output: output, logger: nil) + try await transport.connect() + await transport.disconnect() + #expect(fcntl(input.rawValue, F_GETFL) & O_NONBLOCK == 0) + #expect(fcntl(output.rawValue, F_GETFL) & O_NONBLOCK != 0) + } + + @Test("Disconnect leaves a reused output descriptor unchanged") + func testDisconnectDoesNotChangeReusedOutputDescriptor() async throws { + let (input, writer) = try FileDescriptor.pipe() + let (reader, output) = try FileDescriptor.pipe() + let (replacementReader, replacementWriter) = try FileDescriptor.pipe() + defer { + try? input.close() + try? writer.close() + try? reader.close() + try? replacementReader.close() + try? replacementWriter.close() + } + + let transport = StdioTransport(input: input, output: output, logger: nil) + try await transport.connect() + let replacementFlags = fcntl(replacementWriter.rawValue, F_GETFL) + #expect(fcntl(replacementWriter.rawValue, F_SETFL, replacementFlags | O_NONBLOCK) == 0) + + try output.close() + #expect(dup2(replacementWriter.rawValue, output.rawValue) == output.rawValue) + defer { _ = close(output.rawValue) } + + await transport.disconnect() + #expect(fcntl(output.rawValue, F_GETFL) & O_NONBLOCK != 0) + } + + @Test("Failed connect restores the input descriptor mode") + func testFailedConnectRestoresInputMode() async throws { + let (input, writer) = try FileDescriptor.pipe() + defer { + try? input.close() + try? writer.close() + } + + let transport = StdioTransport(input: input, output: FileDescriptor(rawValue: -1), logger: nil) + do { + try await transport.connect() + Issue.record("Expected connect to reject the closed output descriptor") + } catch { + #expect(fcntl(input.rawValue, F_GETFL) & O_NONBLOCK == 0) + } + } + @Test("Send Message") func testStdioTransportSendMessage() async throws { let (reader, output) = try FileDescriptor.pipe()