oai-swift/oAI/Services/ExternalMCPModels.swift

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

import Foundation

// MARK: - Server Configuration

/// Which wire protocol an `ExternalMCPServer` uses. `.stdio` fields are `command`/`args`/`env`;
/// `.http` fields are `url`/`bearerToken`/`headers`. Kept as one flat struct rather than an enum
/// with associated values — simpler `Codable` and simpler settings-JSON storage, at the cost of
/// each server config carrying some always-unused fields for its transport.
enum MCPTransportKind: String, Codable, Sendable, CaseIterable {
    case stdio
    case http
}

struct ExternalMCPServer: Codable, Identifiable, Sendable {
    var id: UUID
    var name: String
    var transportKind: MCPTransportKind
    var command: String
    var args: [String]
    var env: [String: String]
    var url: String
    var bearerToken: String
    var headers: [String: String]
    var isEnabled: Bool
    var timeout: TimeInterval
    var createdAt: Date

    init(
        id: UUID = UUID(),
        name: String,
        transportKind: MCPTransportKind = .stdio,
        command: String = "",
        args: [String] = [],
        env: [String: String] = [:],
        url: String = "",
        bearerToken: String = "",
        headers: [String: String] = [:],
        isEnabled: Bool = true,
        timeout: TimeInterval = 30,
        createdAt: Date = Date()
    ) {
        self.id = id
        self.name = name
        self.transportKind = transportKind
        self.command = command
        self.args = args
        self.env = env
        self.url = url
        self.bearerToken = bearerToken
        self.headers = headers
        self.isEnabled = isEnabled
        self.timeout = timeout
        self.createdAt = createdAt
    }

    enum CodingKeys: String, CodingKey {
        case id, name, transportKind, command, args, env, url, bearerToken, headers, isEnabled, timeout, createdAt
    }

    /// Custom decoding so servers saved before `transportKind`/`env`/`url`/`bearerToken`/`headers`
    /// existed (plain stdio-only configs) still decode — those keys default rather than fail.
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        id = try c.decode(UUID.self, forKey: .id)
        name = try c.decode(String.self, forKey: .name)
        command = try c.decode(String.self, forKey: .command)
        args = try c.decode([String].self, forKey: .args)
        isEnabled = try c.decode(Bool.self, forKey: .isEnabled)
        timeout = try c.decode(TimeInterval.self, forKey: .timeout)
        createdAt = try c.decode(Date.self, forKey: .createdAt)
        transportKind = try c.decodeIfPresent(MCPTransportKind.self, forKey: .transportKind) ?? .stdio
        env = try c.decodeIfPresent([String: String].self, forKey: .env) ?? [:]
        url = try c.decodeIfPresent(String.self, forKey: .url) ?? ""
        bearerToken = try c.decodeIfPresent(String.self, forKey: .bearerToken) ?? ""
        headers = try c.decodeIfPresent([String: String].self, forKey: .headers) ?? [:]
    }

    var slug: String { Self.makeSlug(from: name) }

    static func makeSlug(from name: String) -> String {
        let s = name
            .lowercased()
            .components(separatedBy: CharacterSet.alphanumerics.inverted)
            .filter { !$0.isEmpty }
            .joined(separator: "_")
        return s.isEmpty ? "ext" : s
    }

    /// Splits a raw arguments string into tokens, respecting single/double-quoted
    /// segments so arguments containing spaces (e.g. `--root "/Users/x/My Documents"`)
    /// survive intact instead of being split on every space.
    static func parseArguments(_ input: String) -> [String] {
        var args: [String] = []
        var current = ""
        var inSingleQuotes = false
        var inDoubleQuotes = false

        for char in input {
            if char == "'" && !inDoubleQuotes {
                inSingleQuotes.toggle()
            } else if char == "\"" && !inSingleQuotes {
                inDoubleQuotes.toggle()
            } else if char.isWhitespace && !inSingleQuotes && !inDoubleQuotes {
                if !current.isEmpty {
                    args.append(current)
                    current = ""
                }
            } else {
                current.append(char)
            }
        }
        if !current.isEmpty { args.append(current) }
        return args
    }

    static let reservedSlugs: Set<String> = [
        "anytype", "paperless", "calendar", "reminders",
        "contacts", "location", "maps", "bash", "web", "read", "write",
        "list", "search", "edit", "delete", "create", "move", "copy", "spawn"
    ]

    var isSlugReserved: Bool { Self.reservedSlugs.contains(slug) }
}

// MARK: - Client State

enum MCPClientState: Equatable {
    case idle
    case connecting
    case ready
    case error(String)
    case crashed
    case stopped
}

// MARK: - State Delegate (all callbacks on MainActor)

@MainActor
protocol ExternalMCPStateDelegate: AnyObject {
    func clientDidBecomeReady(id: UUID, tools: [MCPToolDefinition], server: ExternalMCPServer)
    func clientDidChangeState(id: UUID, state: MCPClientState)
}

// MARK: - Client Errors

enum MCPClientError: LocalizedError {
    case notConnected
    case invalidResponse(String)
    case timeout
    case processLaunchFailed(String)
    case handshakeFailed(String)
    case writeFailed
    case invalidConfiguration(String)

    var errorDescription: String? {
        switch self {
        case .notConnected:               return "MCP server is not connected"
        case .invalidResponse(let s):     return "Invalid MCP response: \(s)"
        case .timeout:                    return "MCP request timed out"
        case .processLaunchFailed(let s): return "Failed to launch MCP server: \(s)"
        case .handshakeFailed(let s):     return "MCP handshake failed: \(s)"
        case .writeFailed:                return "Failed to write to MCP server stdin"
        case .invalidConfiguration(let s): return "Invalid MCP server configuration: \(s)"
        }
    }
}

// MARK: - MCP Protocol Types

struct MCPInitializeResult: Decodable {
    let protocolVersion: String
    let capabilities: MCPCapabilities
    let serverInfo: MCPServerInfo?
}

struct MCPCapabilities: Decodable {
    let tools: MCPToolsCapability?
    struct MCPToolsCapability: Decodable { let listChanged: Bool? }
}

struct MCPServerInfo: Decodable {
    let name: String
    let version: String?
}

struct MCPToolsListResult: Decodable {
    let tools: [MCPToolDefinition]
    let nextCursor: String?
}

struct MCPToolDefinition: Decodable {
    let name: String
    let description: String?
    let inputSchema: MCPInputSchema
}

struct MCPInputSchema: Decodable {
    let type: String
    let properties: [String: MCPPropertySchema]?
    let required: [String]?
}

struct MCPPropertySchema: Decodable {
    let type: String?
    let description: String?
    let `enum`: [String]?
    let items: MCPItemsSchema?
    struct MCPItemsSchema: Decodable { let type: String? }
}

struct MCPToolCallResult: Decodable {
    let content: [MCPContent]
    let isError: Bool?
}

struct MCPContent: Decodable {
    let type: String
    let text: String?
    let data: String?
    let mimeType: String?
    let uri: String?
}