oai-swift/oAI/Services/ExternalMCPManager.swift

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

import Foundation

// MARK: - ExternalMCPManager

@Observable
@MainActor
final class ExternalMCPManager {
    nonisolated static let shared = ExternalMCPManager()

    private(set) var clientStates: [UUID: MCPClientState] = [:]
    private(set) var cachedToolSchemas: [Tool] = []

    // Keep server config alongside client so we can access slug without await
    private var clients: [UUID: ExternalMCPClient] = [:]
    private var serverConfigs: [UUID: ExternalMCPServer] = [:]
    private var restartTasks: [UUID: Task<Void, Never>] = [:]
    private var restartAttempts: [UUID: Int] = [:]

    private nonisolated init() {}

    // MARK: - Lifecycle

    func startAll() {
        for server in SettingsService.shared.externalMCPServers where server.isEnabled {
            startClient(for: server)
        }
    }

    func stopAll() {
        for client in clients.values { client.stop() }
        clients.removeAll()
        serverConfigs.removeAll()
        clientStates.removeAll()
        cachedToolSchemas.removeAll()
        for task in restartTasks.values { task.cancel() }
        restartTasks.removeAll()
        restartAttempts.removeAll()
    }

    func reconfigure(servers: [ExternalMCPServer]) {
        let activeIds = Set(servers.filter { $0.isEnabled }.map { $0.id })
        for id in clients.keys where !activeIds.contains(id) {
            clients[id]?.stop()
            clients.removeValue(forKey: id)
            serverConfigs.removeValue(forKey: id)
            clientStates.removeValue(forKey: id)
            restartTasks[id]?.cancel()
            restartTasks.removeValue(forKey: id)
            restartAttempts.removeValue(forKey: id)
            removeCachedSchemas(for: id)
        }
        for server in servers where server.isEnabled && clients[server.id] == nil {
            startClient(for: server)
        }
    }

    private func startClient(for server: ExternalMCPServer) {
        // Stop any existing client for this ID before creating a new one
        clients[server.id]?.stop()
        let client = ExternalMCPClient(server: server, stateDelegate: self)
        clients[server.id] = client
        serverConfigs[server.id] = server
        clientStates[server.id] = .connecting
        Task {
            do {
                try await client.start()
            } catch {
                // start() already notified stateDelegate with .crashed (uniformly, for both
                // transports and every failure kind) before throwing — restart-with-backoff is
                // already scheduled via clientDidChangeState. Nothing further to do but log.
                Log.extMcp.warning("'\(server.name)' start failed: \(error.localizedDescription)")
            }
        }
    }

    /// Manually retries one server's connection, resetting its restart-attempt counter — for use
    /// after the user fixes an external condition (e.g. installs a missing command) rather than
    /// waiting for `.error("Maximum restart attempts reached")`/`.error("Command not found: …")`
    /// to somehow resolve on their own; neither state is retried automatically.
    func retryClient(id: UUID) {
        guard let server = serverConfigs[id] else { return }
        restartAttempts.removeValue(forKey: id)
        restartTasks[id]?.cancel()
        restartTasks.removeValue(forKey: id)
        startClient(for: server)
    }

    private func scheduleRestart(for server: ExternalMCPServer, attempt: Int) {
        let delays: [Double] = [5, 15, 30]
        let delay = delays[min(attempt - 1, delays.count - 1)]
        Log.extMcp.warning("MCP server '\(server.name)' crashed — restarting in \(Int(delay))s (attempt \(attempt)/3)")

        restartTasks[server.id]?.cancel()
        let id = server.id
        restartTasks[id] = Task { [weak self, id] in
            try? await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
            guard !Task.isCancelled, let self,
                  self.clients[id] != nil,
                  SettingsService.shared.externalMCPServers.contains(where: { $0.id == id && $0.isEnabled })
            else { return }
            // startClient is the single place that creates and launches clients.
            // It handles processLaunchFailed by calling clientDidChangeState(.crashed),
            // and all other failures let the termination handler drive the .crashed callback.
            self.startClient(for: server)
        }
    }

    // MARK: - Tool Schema Integration (synchronous)

    func getToolSchemas() -> [Tool] { cachedToolSchemas }

    func isExternalTool(_ name: String) -> Bool {
        cachedToolSchemas.contains { $0.function.name == name }
    }

    // MARK: - Tool Execution

    func executeTool(name: String, argumentsJSON: String) async -> [String: Any] {
        for (id, client) in clients {
            guard clientStates[id] == .ready,
                  let server = serverConfigs[id] else { continue }
            let prefix = "\(server.slug)_"
            if name.hasPrefix(prefix) {
                let originalName = String(name.dropFirst(prefix.count))
                return await client.callTool(originalName: originalName, argumentsJSON: argumentsJSON)
            }
        }
        return ["error": "No external MCP server found for tool: \(name)"]
    }

    // MARK: - Schema Cache

    private func rebuildCache(for server: ExternalMCPServer, tools: [MCPToolDefinition]) {
        removeCachedSchemas(for: server.id, slug: server.slug)
        let prefixed = tools.compactMap { Self.convertToolDefinition($0, server: server) }
        cachedToolSchemas.append(contentsOf: prefixed)
        Log.extMcp.info("[\(server.name)] cached \(prefixed.count) tools: \(prefixed.map { $0.function.name }.joined(separator: ", "))")
    }

    private func removeCachedSchemas(for id: UUID) {
        guard let server = serverConfigs[id] else { return }
        removeCachedSchemas(for: id, slug: server.slug)
    }

    private func removeCachedSchemas(for id: UUID, slug: String) {
        cachedToolSchemas.removeAll { $0.function.name.hasPrefix("\(slug)_") }
    }

    nonisolated static func convertToolDefinition(_ def: MCPToolDefinition, server: ExternalMCPServer) -> Tool? {
        Tool(
            type: "function",
            function: Tool.Function(
                name: "\(server.slug)_\(def.name)",
                description: "[\(server.name)] \(def.description ?? "")",
                parameters: convertInputSchema(def.inputSchema)
            )
        )
    }

    /// A schema property with no `"type"` at all is valid JSON Schema (e.g. an `enum`-only or
    /// composed property) — not every MCP server's tool schemas set it, so this must not assume
    /// it's present. (Found via a real crash: Obsidian's Local REST API plugin sends at least one
    /// tool parameter with no `type`, which a `prop.type!` force-unwrap here used to crash on.)
    nonisolated static func convertInputSchema(_ schema: MCPInputSchema) -> Tool.Function.Parameters {
        var properties: [String: Tool.Function.Parameters.Property] = [:]
        for (key, prop) in schema.properties ?? [:] {
            let effectiveType = prop.type ?? "string"
            let normalized: String
            switch effectiveType {
            case "integer":                                        normalized = "number"
            case "string", "number", "boolean", "array", "object": normalized = effectiveType
            default:                                               normalized = "string"
            }
            var items: Tool.Function.Parameters.Property.Items? = nil
            if normalized == "array", let t = prop.items?.type { items = .init(type: t) }
            properties[key] = Tool.Function.Parameters.Property(
                type: normalized,
                description: prop.description ?? "",
                enum: prop.enum,
                items: items
            )
        }
        return Tool.Function.Parameters(type: "object", properties: properties, required: schema.required)
    }
}

// MARK: - ExternalMCPStateDelegate

extension ExternalMCPManager: ExternalMCPStateDelegate {
    func clientDidBecomeReady(id: UUID, tools: [MCPToolDefinition], server: ExternalMCPServer) {
        clientStates[id] = .ready
        restartAttempts.removeValue(forKey: id)
        rebuildCache(for: server, tools: tools)
    }

    func clientDidChangeState(id: UUID, state: MCPClientState) {
        clientStates[id] = state
        if case .crashed = state,
           let server = serverConfigs[id],
           SettingsService.shared.externalMCPServers.contains(where: { $0.id == id && $0.isEnabled }) {
            removeCachedSchemas(for: id, slug: server.slug)
            let attempt = (restartAttempts[id] ?? 0) + 1
            guard attempt <= 3 else {
                Log.extMcp.error("MCP server '\(server.name)' gave up after 3 restart attempts")
                clientStates[id] = .error("Maximum restart attempts reached")
                restartAttempts.removeValue(forKey: id)
                return
            }
            restartAttempts[id] = attempt
            scheduleRestart(for: server, attempt: attempt)
        }
    }
}