oai-swift/oAI/Services/MCPTransport.swift

// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Copyright (C) 2026 Rune Olsen

import Foundation

// MARK: - Transport abstraction

/// One MCP server connection's wire protocol. `ExternalMCPClient` builds JSON-RPC envelopes and
/// decodes typed results; the transport only owns *delivery* — how the bytes get to the server
/// and back. A long-lived subprocess pipe for stdio, discrete HTTP requests for Streamable HTTP.
@MainActor
protocol MCPTransport: AnyObject {
    /// Prepares the transport for use — launches the subprocess for stdio; a no-op for HTTP,
    /// since there's no persistent connection to establish ahead of the first request.
    func prepare() async throws
    /// Sends a JSON-RPC *request* (`message` includes `"id"`) and returns its `result` payload.
    /// Throws `MCPClientError.invalidResponse` if the server returned a JSON-RPC error, or
    /// `.timeout` if `timeoutSeconds` elapses first.
    func sendRequest(_ message: [String: Any], id: Int, timeoutSeconds: Double) async throws -> Data
    /// Sends a JSON-RPC *notification* (no `"id"`, no response expected).
    func sendNotification(_ message: [String: Any]) async throws
    /// Tears down the transport — terminates the subprocess / drops HTTP session state.
    func stop()
}

enum MCPTransportSupport {
    /// Extracts the `result` payload from a decoded top-level JSON-RPC response object, shared
    /// by every transport so error/result semantics stay identical regardless of wire protocol.
    static func extractResult(from json: [String: Any]) throws -> Data {
        if let err = json["error"] as? [String: Any] {
            throw MCPClientError.invalidResponse(err["message"] as? String ?? "Unknown error")
        }
        guard let result = json["result"],
              let resultData = try? JSONSerialization.data(withJSONObject: result)
        else {
            throw MCPClientError.invalidResponse("Missing result field")
        }
        return resultData
    }
}

// MARK: - Stdio Transport

/// Launches the server as a subprocess and speaks newline-delimited JSON-RPC over its
/// stdin/stdout, exactly as the original (pre-transport-split) `ExternalMCPClient` did.
@MainActor
final class StdioMCPTransport: MCPTransport {
    private let server: ExternalMCPServer
    /// Set by the owner after construction (not an init parameter) — the owner (`ExternalMCPClient`)
    /// needs to capture `self` weakly to wire this up, which Swift only allows once `self` is
    /// fully initialized, i.e. after its own `init` has finished assigning all stored properties.
    var onTerminated: (() -> Void)?

    private var process: Process?
    private var stdinHandle: FileHandle?
    private var readTask: Task<Void, Never>?
    private var stderrTask: Task<Void, Never>?
    private var pendingCalls: [Int: CheckedContinuation<Data, Error>] = [:]
    private var lineBuffer = Data()

    init(server: ExternalMCPServer) {
        self.server = server
    }

    func prepare() async throws {
        let proc = Process()
        if server.command.hasPrefix("/") {
            proc.executableURL = URL(fileURLWithPath: server.command)
            proc.arguments = server.args
        } else {
            proc.executableURL = URL(fileURLWithPath: "/usr/bin/env")
            proc.arguments = [server.command] + server.args
        }
        // Inherited environment first (so PATH etc. still resolves — e.g. `npx` needs PATH to
        // find node), then layer the user-configured vars over it so they can override.
        var environment = ProcessInfo.processInfo.environment
        for (key, value) in server.env { environment[key] = value }
        proc.environment = environment

        let stdinPipe = Pipe()
        let stdoutPipe = Pipe()
        let stderrPipe = Pipe()
        proc.standardInput = stdinPipe
        proc.standardOutput = stdoutPipe
        proc.standardError = stderrPipe

        proc.terminationHandler = { [weak self] _ in
            Task { @MainActor [weak self] in self?.handleProcessTerminated() }
        }

        do {
            try proc.run()
        } catch {
            throw MCPClientError.processLaunchFailed(error.localizedDescription)
        }

        process = proc
        stdinHandle = stdinPipe.fileHandleForWriting
        startReadLoop(pipe: stdoutPipe)
        startStderrLoop(pipe: stderrPipe)
    }

    func sendRequest(_ message: [String: Any], id: Int, timeoutSeconds: Double) async throws -> Data {
        try writeJSON(message)

        // Schedule timeout: cancels the specific pending call by ID
        let timeoutId = id
        Task { [weak self, timeoutId] in
            try? await Task.sleep(nanoseconds: UInt64(timeoutSeconds * 1_000_000_000))
            self?.cancelPendingCall(id: timeoutId, with: MCPClientError.timeout)
        }

        return try await withCheckedThrowingContinuation { cont in
            pendingCalls[id] = cont
        }
    }

    func sendNotification(_ message: [String: Any]) async throws {
        try writeJSON(message)
    }

    func stop() {
        readTask?.cancel()
        stderrTask?.cancel()
        process?.terminate()
        process = nil
        stdinHandle = nil
        lineBuffer = Data()
        for (_, cont) in pendingCalls { cont.resume(throwing: MCPClientError.notConnected) }
        pendingCalls.removeAll()
    }

    // MARK: I/O Loops (detached from MainActor)

    private func startReadLoop(pipe: Pipe) {
        readTask = Task.detached { [weak self] in
            let handle = pipe.fileHandleForReading
            while true {
                let data = handle.availableData
                if data.isEmpty { break }
                await self?.receiveData(data)
            }
        }
    }

    private func startStderrLoop(pipe: Pipe) {
        let name = server.name
        stderrTask = Task.detached {
            let handle = pipe.fileHandleForReading
            var buf = Data()
            while true {
                let data = handle.availableData
                if data.isEmpty { break }
                buf.append(data)
                while let idx = buf.firstIndex(of: UInt8(ascii: "\n")) {
                    let line = String(data: buf[buf.startIndex..<idx], encoding: .utf8) ?? ""
                    buf = Data(buf[buf.index(after: idx)...])
                    if !line.trimmingCharacters(in: .whitespaces).isEmpty {
                        Log.extMcp.warning("[\(name)] \(line)")
                    }
                }
            }
        }
    }

    // MARK: Data Processing (MainActor)

    private func receiveData(_ data: Data) {
        lineBuffer.append(data)
        while let idx = lineBuffer.firstIndex(of: UInt8(ascii: "\n")) {
            let lineData = Data(lineBuffer[lineBuffer.startIndex..<idx])
            lineBuffer = Data(lineBuffer[lineBuffer.index(after: idx)...])
            processLine(lineData)
        }
    }

    private func processLine(_ data: Data) {
        guard !data.isEmpty,
              let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any],
              let id = json["id"] as? Int,
              let cont = pendingCalls.removeValue(forKey: id) else { return }

        do {
            let resultData = try MCPTransportSupport.extractResult(from: json)
            cont.resume(returning: resultData)
        } catch {
            cont.resume(throwing: error)
        }
    }

    private func cancelPendingCall(id: Int, with error: Error) {
        pendingCalls.removeValue(forKey: id)?.resume(throwing: error)
    }

    private func writeJSON(_ message: [String: Any]) throws {
        guard let handle = stdinHandle, process?.isRunning == true else {
            throw MCPClientError.writeFailed
        }
        guard let data = try? JSONSerialization.data(withJSONObject: message),
              let line = String(data: data, encoding: .utf8) else {
            throw MCPClientError.writeFailed
        }
        do {
            try handle.write(contentsOf: Data((line + "\n").utf8))
        } catch {
            throw MCPClientError.writeFailed
        }
    }

    private func handleProcessTerminated() {
        for (_, cont) in pendingCalls { cont.resume(throwing: MCPClientError.notConnected) }
        pendingCalls.removeAll()
        onTerminated?()
    }
}

// MARK: - HTTP (Streamable HTTP) Transport

/// Speaks MCP's Streamable HTTP transport (spec revision 2025-06-18): every JSON-RPC message is
/// its own HTTP POST to the server's single MCP endpoint. The server may answer with a plain
/// `application/json` body, or open a `text/event-stream` (SSE) response — this transport
/// supports both, but (since Confab only needs request/response tool calls, not server-initiated
/// push) does not open a standalone listening GET stream for unsolicited server messages.
@MainActor
final class HTTPMCPTransport: MCPTransport {
    private let server: ExternalMCPServer
    private let urlSession: URLSession
    private var sessionId: String?

    init(server: ExternalMCPServer, urlSession: URLSession = .shared) {
        self.server = server
        self.urlSession = urlSession
    }

    func prepare() async throws {
        guard URL(string: server.url) != nil else {
            throw MCPClientError.invalidConfiguration("Invalid server URL: \(server.url)")
        }
        // No persistent connection to establish ahead of time — the first request (`initialize`)
        // both opens the session and confirms the server is reachable.
    }

    func sendRequest(_ message: [String: Any], id: Int, timeoutSeconds: Double) async throws -> Data {
        let request = try buildRequest(for: message, timeoutSeconds: timeoutSeconds)
        let (data, response) = try await urlSession.data(for: request)
        guard let http = response as? HTTPURLResponse else {
            throw MCPClientError.invalidResponse("Non-HTTP response")
        }
        if let newSessionId = http.value(forHTTPHeaderField: "Mcp-Session-Id") {
            sessionId = newSessionId
        }
        guard (200...299).contains(http.statusCode) else {
            throw MCPClientError.invalidResponse("HTTP \(http.statusCode)")
        }
        let json = try Self.parseResponseBody(
            data, contentType: http.value(forHTTPHeaderField: "Content-Type"), expectedId: id
        )
        return try MCPTransportSupport.extractResult(from: json)
    }

    func sendNotification(_ message: [String: Any]) async throws {
        // Notifications carry no "id", so per spec the server responds 202 Accepted, no body.
        let request = try buildRequest(for: message, timeoutSeconds: 15)
        let (_, response) = try await urlSession.data(for: request)
        guard let http = response as? HTTPURLResponse, (200...299).contains(http.statusCode) else {
            throw MCPClientError.invalidResponse("Notification not accepted")
        }
    }

    func stop() {
        sessionId = nil
    }

    private func buildRequest(for message: [String: Any], timeoutSeconds: Double) throws -> URLRequest {
        guard let url = URL(string: server.url) else {
            throw MCPClientError.invalidConfiguration("Invalid server URL: \(server.url)")
        }
        var request = URLRequest(url: url, timeoutInterval: timeoutSeconds)
        request.httpMethod = "POST"
        request.setValue("application/json", forHTTPHeaderField: "Content-Type")
        request.setValue("application/json, text/event-stream", forHTTPHeaderField: "Accept")
        request.setValue("2025-06-18", forHTTPHeaderField: "MCP-Protocol-Version")
        if !server.bearerToken.isEmpty {
            request.setValue("Bearer \(server.bearerToken)", forHTTPHeaderField: "Authorization")
        }
        for (key, value) in server.headers {
            request.setValue(value, forHTTPHeaderField: key)
        }
        if let sessionId {
            request.setValue(sessionId, forHTTPHeaderField: "Mcp-Session-Id")
        }
        request.httpBody = try JSONSerialization.data(withJSONObject: message)
        return request
    }

    /// Parses either a direct `application/json` body, or an SSE (`text/event-stream`) body —
    /// scanning its `data:` lines for the JSON-RPC message whose `id` matches `expectedId` (the
    /// server may send unrelated requests/notifications on the same stream first, per spec).
    nonisolated static func parseResponseBody(_ data: Data, contentType: String?, expectedId: Int) throws -> [String: Any] {
        if contentType?.contains("text/event-stream") == true {
            guard let text = String(data: data, encoding: .utf8) else {
                throw MCPClientError.invalidResponse("Non-UTF8 SSE body")
            }
            for line in text.components(separatedBy: "\n") {
                let trimmed = line.trimmingCharacters(in: .whitespaces)
                guard trimmed.hasPrefix("data:") else { continue }
                let payload = trimmed.dropFirst("data:".count).trimmingCharacters(in: .whitespaces)
                guard let payloadData = payload.data(using: .utf8),
                      let json = try? JSONSerialization.jsonObject(with: payloadData) as? [String: Any],
                      let id = json["id"] as? Int, id == expectedId
                else { continue }
                return json
            }
            throw MCPClientError.invalidResponse("No matching response in SSE stream")
        }
        guard let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            throw MCPClientError.invalidResponse("Malformed JSON response body")
        }
        return json
    }
}