oai-swift/oAI/Services/ExternalMCPClient.swift

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

import Foundation

// MARK: - ExternalMCPClient

/// Owns one MCP server connection's lifecycle and JSON-RPC message framing. Delivery (stdio
/// subprocess vs Streamable HTTP) is delegated to a `MCPTransport` — this class only builds
/// envelopes, decodes typed results, and tracks connection state.
/// All state is MainActor-isolated (consistent with SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor
/// project setting).
@MainActor
final class ExternalMCPClient {
    let server: ExternalMCPServer
    weak var stateDelegate: (any ExternalMCPStateDelegate)?

    private let transport: any MCPTransport
    private var nextRequestId: Int = 1

    private(set) var state: MCPClientState = .idle
    private(set) var discoveredTools: [MCPToolDefinition] = []

    init(server: ExternalMCPServer, stateDelegate: (any ExternalMCPStateDelegate)?) {
        self.server = server
        self.stateDelegate = stateDelegate
        let stdioTransport: StdioMCPTransport?
        switch server.transportKind {
        case .stdio:
            let t = StdioMCPTransport(server: server)
            stdioTransport = t
            self.transport = t
        case .http:
            stdioTransport = nil
            self.transport = HTTPMCPTransport(server: server)
        }
        // `self` is only safe to capture once every stored property above has a value —
        // wire the crash callback here, after `init` would otherwise be considered complete.
        stdioTransport?.onTerminated = { [weak self] in
            self?.handleTransportTerminatedUnexpectedly()
        }
    }

    /// Called by a stdio transport whose subprocess died on its own — as opposed to a
    /// deliberate `stop()` call, or a failure already handled inline within `start()`.
    /// No HTTP equivalent: a Streamable HTTP connection has no persistent process to crash;
    /// its failures surface per-request instead (handled in `start()`/`callTool()` directly).
    private func handleTransportTerminatedUnexpectedly() {
        guard state != .stopped else { return }
        state = .crashed
        stateDelegate?.clientDidChangeState(id: server.id, state: .crashed)
    }

    // MARK: - Lifecycle

    func start() async throws {
        guard state == .idle || state == .stopped || state == .crashed else { return }
        state = .connecting
        stateDelegate?.clientDidChangeState(id: server.id, state: .connecting)

        do {
            try await transport.prepare()

            let _: MCPInitializeResult = try await timedRequest(seconds: 15, method: "initialize", params: [
                "protocolVersion": "2024-11-05",
                "capabilities": [:] as [String: Any],
                "clientInfo": ["name": "Confab", "version": "1.0"] as [String: Any]
            ])
            try await transport.sendNotification(["jsonrpc": "2.0", "method": "notifications/initialized"])

            let toolsResult: MCPToolsListResult = try await timedRequest(seconds: 15, method: "tools/list", params: nil)
            discoveredTools = toolsResult.tools
        } catch {
            // Uniformly route every start() failure (bad config, launch failure, handshake
            // failure — for either transport) through .crashed, not .error, so
            // ExternalMCPManager's restart-with-backoff drives from exactly one place.
            // (Previously, stdio relied on the subprocess's termination handler firing
            // asynchronously to reach .crashed; that path doesn't exist for HTTP, so failures
            // there would otherwise get stuck at .error with no retry.)
            state = .crashed
            transport.stop()
            stateDelegate?.clientDidChangeState(id: server.id, state: .crashed)
            throw error
        }

        state = .ready
        stateDelegate?.clientDidBecomeReady(id: server.id, tools: discoveredTools, server: server)
    }

    func stop() {
        state = .stopped
        transport.stop()
    }

    // MARK: - Tool Execution

    func callTool(originalName: String, argumentsJSON: String) async -> [String: Any] {
        guard state == .ready else {
            return ["error": "MCP server '\(server.name)' is not connected"]
        }
        guard let argData = argumentsJSON.data(using: .utf8),
              let argsDict = try? JSONSerialization.jsonObject(with: argData) as? [String: Any] else {
            return ["error": "Invalid arguments JSON for tool \(originalName)"]
        }
        do {
            let result: MCPToolCallResult = try await timedRequest(
                seconds: server.timeout,
                method: "tools/call",
                params: ["name": originalName, "arguments": argsDict]
            )
            return convertMCPResult(result)
        } catch MCPClientError.timeout {
            return ["error": "MCP server '\(server.name)' timed out after \(Int(server.timeout))s"]
        } catch {
            return ["error": "MCP call '\(originalName)' failed: \(error.localizedDescription)"]
        }
    }

    // MARK: - JSON-RPC

    private func timedRequest<T: Decodable>(seconds: Double, method: String, params: [String: Any]?) async throws -> T {
        let id = nextRequestId
        nextRequestId += 1
        var message: [String: Any] = ["jsonrpc": "2.0", "method": method, "id": id]
        if let params { message["params"] = params }
        let resultData = try await transport.sendRequest(message, id: id, timeoutSeconds: seconds)
        return try JSONDecoder().decode(T.self, from: resultData)
    }

    // MARK: - Result conversion

    private func convertMCPResult(_ result: MCPToolCallResult) -> [String: Any] {
        let isError = result.isError ?? false
        var parts: [String] = []
        for content in result.content {
            switch content.type {
            case "text":
                if let text = content.text { parts.append(text) }
            case "image":
                if let base64 = content.data, let imageData = Data(base64Encoded: base64) {
                    parts.append("[Image saved to: \(writeTempImage(imageData, mimeType: content.mimeType))]")
                }
            case "resource":
                if let text = content.text { parts.append(text) }
                else if let uri = content.uri { parts.append("[Resource: \(uri)]") }
            default:
                if let text = content.text { parts.append(text) }
            }
        }
        let combined = parts.joined(separator: "\n")
        return isError ? ["error": combined.isEmpty ? "Tool returned an error" : combined] : ["output": combined]
    }

    private func writeTempImage(_ data: Data, mimeType: String?) -> String {
        let ext = mimeType?.contains("png") == true ? "png" : "jpg"
        let path = "/tmp/oai_mcp_\(Int(Date().timeIntervalSince1970 * 1000)).\(ext)"
        try? data.write(to: URL(fileURLWithPath: path))
        return path
    }
}