swift-apple-intelligence-grpc/Sources/AppleIntelligenceApp/ServerManager.swift

import Foundation
import AppleIntelligenceCore
import GRPCCore
import GRPCNIOTransportHTTP2
import GRPCReflectionService

@MainActor
@Observable
final class ServerManager {
    enum ServerState: Equatable {
        case stopped
        case starting
        case running(host: String, port: Int)
        case error(String)

        var isRunning: Bool {
            if case .running = self { return true }
            return false
        }

        var statusText: String {
            switch self {
            case .stopped:
                return "Server offline"
            case .starting:
                return "Starting..."
            case .running(let host, let port):
                return "Running on \(host):\(port)"
            case .error:
                return "Server offline"
            }
        }
    }

    private(set) var state: ServerState = .stopped
    private(set) var modelStatus: String = "Unknown"

    private var serverTask: Task<Void, Never>?
    private var service: AppleIntelligenceService?
    private let settings: AppSettings

    init(settings: AppSettings) {
        self.settings = settings
    }

    func start() {
        guard !state.isRunning else { return }

        state = .starting

        // Capture settings values for use in Task
        let host = settings.host
        let port = settings.port
        let apiKey = settings.apiKey.isEmpty ? nil : settings.apiKey
        let enableReflection = settings.enableReflection

        serverTask = Task {
            do {
                // Initialize Apple Intelligence service
                let aiService = await AppleIntelligenceService()
                self.service = aiService

                let isAvailable = await aiService.isAvailable
                let status = await aiService.getModelStatus()

                await MainActor.run {
                    self.modelStatus = status
                }

                guard isAvailable else {
                    await MainActor.run {
                        self.state = .error("Apple Intelligence not available")
                    }
                    return
                }

                // Create provider
                let provider = AppleIntelligenceProvider(service: aiService, apiKey: apiKey)

                // Create transport and server
                let transport = HTTP2ServerTransport.Posix(
                    address: .ipv4(host: host, port: port),
                    transportSecurity: .plaintext,
                    config: .defaults
                )

                // Build services list with optional reflection
                var services: [any RegistrableRPCService] = [provider]
                if enableReflection {
                    if let descriptorURL = AppleIntelligenceResources.descriptorSetURL {
                        let reflectionService = try ReflectionService(descriptorSetFileURLs: [descriptorURL])
                        services.append(reflectionService)
                    }
                }

                let server = GRPCServer(transport: transport, services: services)

                await MainActor.run {
                    self.state = .running(host: host, port: port)
                }

                // Run server until cancelled
                try await server.serve()

            } catch is CancellationError {
                // Normal shutdown
            } catch {
                await MainActor.run {
                    self.state = .error(error.localizedDescription)
                }
            }

            await MainActor.run {
                if case .running = self.state {
                    self.state = .stopped
                }
            }
        }
    }

    func stop() {
        serverTask?.cancel()
        serverTask = nil
        state = .stopped
    }

    func restart() {
        guard state.isRunning else { return }

        // Stop the current server
        stop()
        state = .starting

        // Start again after a short delay to allow port release
        DispatchQueue.main.asyncAfter(deadline: .now() + 0.5) { [weak self] in
            self?.start()
        }
    }

    func toggle() {
        if state.isRunning {
            stop()
        } else {
            start()
        }
    }
}