oai-swift/oAI/Services/ExternalMCPClient.swift

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright (C) 2026 Rune Olsen

import Foundation

// MARK: - ExternalMCPClient

/// Manages one MCP stdio server process. All state is MainActor-isolated
/// (consistent with SWIFT_DEFAULT_ACTOR_ISOLATION = MainActor project setting).
/// Background I/O runs in Task.detached; state mutations hop back to MainActor.
@MainActor
final class ExternalMCPClient {
    let server: ExternalMCPServer
    weak var stateDelegate: (any ExternalMCPStateDelegate)?

    private var process: Process?
    private var stdinHandle: FileHandle?
    private var readTask: Task<Void, Never>?
    private var stderrTask: Task<Void, Never>?

    private var nextRequestId: Int = 1
    private var pendingCalls: [Int: CheckedContinuation<Data, Error>] = [:]
    private var lineBuffer = Data()

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

    init(server: ExternalMCPServer, stateDelegate: (any ExternalMCPStateDelegate)?) {
        self.server = server
        self.stateDelegate = stateDelegate
    }

    // MARK: - Lifecycle

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

        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
        }
        proc.environment = ProcessInfo.processInfo.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 {
            state = .error(error.localizedDescription)
            stateDelegate?.clientDidChangeState(id: server.id, state: .error(error.localizedDescription))
            throw MCPClientError.processLaunchFailed(error.localizedDescription)
        }

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

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

            let toolsResult: MCPToolsListResult = try await timedRequest(seconds: 15, method: "tools/list", params: nil)
            discoveredTools = toolsResult.tools
        } catch {
            state = .error(error.localizedDescription)
            stateDelegate?.clientDidChangeState(id: server.id, state: .error(error.localizedDescription))
            proc.terminate()
            throw error
        }

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

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

    // 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: - 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 }

        if let err = json["error"] as? [String: Any] {
            cont.resume(throwing: MCPClientError.invalidResponse(err["message"] as? String ?? "Unknown error"))
        } else if let result = json["result"],
                  let resultData = try? JSONSerialization.data(withJSONObject: result) {
            cont.resume(returning: resultData)
        } else {
            cont.resume(throwing: MCPClientError.invalidResponse("Missing result field"))
        }
    }

    // MARK: - JSON-RPC

    /// Send a JSON-RPC request with a per-call timeout. The timeout fires a cancellation
    /// directly into the pending-calls table rather than using a task group (which would
    /// pass the generic T through a @Sendable closure and trigger an isolated-conformance warning).
    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 }
        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(seconds * 1_000_000_000))
            self?.cancelPendingCall(id: timeoutId, with: MCPClientError.timeout)
        }

        // Await response data, then decode on MainActor
        let resultData: Data = try await withCheckedThrowingContinuation { cont in
            pendingCalls[id] = cont
        }
        return try JSONDecoder().decode(T.self, from: resultData)
    }

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

    private func sendNotification(method: String) throws {
        try writeJSON(["jsonrpc": "2.0", "method": method])
    }

    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
        }
    }

    // MARK: - Process termination

    private func handleProcessTerminated() {
        guard state != .stopped else { return }
        state = .crashed
        for (_, cont) in pendingCalls { cont.resume(throwing: MCPClientError.notConnected) }
        pendingCalls.removeAll()
        stateDelegate?.clientDidChangeState(id: server.id, state: .crashed)
    }

    // 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
    }
}