nucleic/Sources/NucleicCore/JSONRPCConnection.swift

import Foundation

/// A small JSON-RPC 2.0 client layered over a `ProcessHandle`'s stdio (ADAPTERS §0.2), shared by
/// the two duplex-stdio backends:
///
///  - **`codex app-server`** ([`CodexAppServerBackend`](Codex/CodexAppServerBackend.swift)) —
///    `includeVersionHeader: false`: codex omits the `"jsonrpc":"2.0"` header on the wire (its
///    source-confirmed dialect — it neither emits nor requires it), so we match that exactly.
///  - **`grok agent stdio`** ([`GrokACPBackend`](Grok/GrokACPBackend.swift)) —
///    `includeVersionHeader: true`: the Agent Client Protocol is *standards* JSON-RPC 2.0, so the
///    header is present on every message.
///
/// Three message classes flow over one newline-delimited (JSONL) channel:
///
///  - outbound **requests** (`id` + `method`) → correlated to an `await`ed continuation by `id`;
///  - inbound **notifications** (`method`, no `id`) → published on `notifications`;
///  - inbound **server→client requests** (`id` + `method`) → dispatched to the handler registered
///    in `start`, which must reply with a `result` keyed by the same `id`. **Approvals arrive this
///    way** — codex's `…/requestApproval` and ACP's `session/request_permission` alike (ADAPTERS §2.4).
public actor JSONRPCConnection {
    public struct Notification: Sendable, Equatable {
        public let method: String
        public let params: JSONValue
        public init(method: String, params: JSONValue) {
            self.method = method
            self.params = params
        }
    }

    /// A server→client request we must answer (e.g. an approval). `id` is echoed
    /// verbatim into the reply — it may be a number or a string on the wire.
    public struct InboundRequest: Sendable, Equatable {
        public let id: JSONValue
        public let method: String
        public let params: JSONValue
        public init(id: JSONValue, method: String, params: JSONValue) {
            self.id = id
            self.method = method
            self.params = params
        }
    }

    public enum RPCError: Error, Sendable, Equatable {
        /// The peer returned an `error` object for one of our requests.
        case server(code: Int, message: String)
        /// The connection closed before the response arrived.
        case connectionClosed
    }

    private let handle: any ProcessHandle
    /// Emit the `"jsonrpc":"2.0"` header on outbound messages (standards JSON-RPC, e.g. ACP).
    /// Off for codex app-server, whose confirmed dialect omits it.
    private let includeVersionHeader: Bool
    private var nextID = 0
    private var pending: [Int: CheckedContinuation<JSONValue, Error>] = [:]
    private var notificationContinuation: AsyncStream<Notification>.Continuation?
    /// Buffered (unbounded) so notifications that arrive during `initialize`/`session/new`
    /// — before the backend begins draining — are never dropped.
    public nonisolated let notifications: AsyncStream<Notification>
    private var requestHandler: (@Sendable (InboundRequest) async -> Void)?
    private var readTask: Task<Void, Never>?
    private var closed = false

    public init(handle: any ProcessHandle, includeVersionHeader: Bool = false) {
        self.handle = handle
        self.includeVersionHeader = includeVersionHeader
        var captured: AsyncStream<Notification>.Continuation!
        self.notifications = AsyncStream(bufferingPolicy: .unbounded) { captured = $0 }
        self.notificationContinuation = captured
    }

    /// Begin reading the peer's stdout. `onRequest` handles inbound server→client
    /// requests; it runs detached so a suspended approval handler can't block the
    /// read loop (and thus block our own request responses or further notifications).
    public func start(onRequest: @escaping @Sendable (InboundRequest) async -> Void) {
        requestHandler = onRequest
        guard readTask == nil else { return }
        let handle = self.handle
        readTask = Task { [weak self] in
            do {
                for try await line in handle.stdoutLines {
                    await self?.ingest(line)
                }
            } catch {
                await self?.terminate(reason: error)
                return
            }
            await self?.terminate(reason: RPCError.connectionClosed)
        }
    }

    // MARK: - Outbound

    /// Send a request and suspend until the correlated response arrives.
    public func request(_ method: String, params: JSONValue? = nil) async throws -> JSONValue {
        if closed { throw RPCError.connectionClosed }
        nextID += 1
        let id = nextID
        var fields: [String: JSONValue] = ["id": .number(Double(id)), "method": .string(method)]
        if includeVersionHeader { fields["jsonrpc"] = .string("2.0") }
        if let params { fields["params"] = params }
        let payload = JSONValue.object(fields)
        return try await withCheckedThrowingContinuation { continuation in
            pending[id] = continuation
            do {
                try handle.writeLine(payload.encodedData())
            } catch {
                pending.removeValue(forKey: id)
                continuation.resume(throwing: error)
            }
        }
    }

    /// Fire a notification (no response expected).
    public func notify(_ method: String, params: JSONValue? = nil) throws {
        var fields: [String: JSONValue] = ["method": .string(method)]
        if includeVersionHeader { fields["jsonrpc"] = .string("2.0") }
        if let params { fields["params"] = params }
        try handle.writeLine(JSONValue.object(fields).encodedData())
    }

    /// Answer an inbound server→client request with a `result`.
    public func reply(to id: JSONValue, result: JSONValue) throws {
        var fields: [String: JSONValue] = ["id": id, "result": result]
        if includeVersionHeader { fields["jsonrpc"] = .string("2.0") }
        try handle.writeLine(JSONValue.object(fields).encodedData())
    }

    /// Answer an inbound server→client request with an `error`.
    public func replyError(to id: JSONValue, code: Int, message: String) throws {
        let error = JSONValue.object(["code": .number(Double(code)), "message": .string(message)])
        var fields: [String: JSONValue] = ["id": id, "error": error]
        if includeVersionHeader { fields["jsonrpc"] = .string("2.0") }
        try handle.writeLine(JSONValue.object(fields).encodedData())
    }

    // MARK: - Inbound

    private func ingest(_ line: Data) {
        guard !line.isEmpty, let root = try? JSONValue(parsing: line),
            let object = root.objectValue
        else { return }  // non-JSON noise on stdout is ignored, never fatal

        let idValue = object["id"]
        let hasID = (idValue != nil && idValue != .null)

        if let method = object["method"]?.stringValue {
            let params = object["params"] ?? .object([:])
            if hasID, let idValue {
                // Server→client request — dispatch detached so a suspended handler
                // never stalls the read loop.
                let request = InboundRequest(id: idValue, method: method, params: params)
                if let requestHandler {
                    Task { await requestHandler(request) }
                }
            } else {
                notificationContinuation?.yield(Notification(method: method, params: params))
            }
            return
        }

        // No method → a response to one of our requests.
        guard hasID, let id = idValue?.intValue,
            let continuation = pending.removeValue(forKey: id)
        else { return }
        if let error = object["error"] {
            continuation.resume(
                throwing: RPCError.server(
                    code: error["code"]?.intValue ?? -1,
                    message: error["message"]?.stringValue ?? "RPC error"))
        } else {
            continuation.resume(returning: object["result"] ?? .null)
        }
    }

    /// Finish the notification stream *without* failing in-flight requests or stopping the read
    /// loop. The ACP turn terminal is the `session/prompt` *response*, not a notification, so the
    /// backend awaits that response and then calls this — the unbounded buffer guarantees every
    /// already-yielded `session/update` is still delivered to the drain loop before it ends, so no
    /// trailing chunk is dropped. Idempotent.
    public func endNotifications() {
        notificationContinuation?.finish()
        notificationContinuation = nil
    }

    /// Stream closed (peer exited or read error): fail every in-flight request and
    /// end the notification stream so the backend's drain loop terminates.
    private func terminate(reason: Error) {
        guard !closed else { return }
        closed = true
        for (_, continuation) in pending {
            continuation.resume(throwing: reason)
        }
        pending.removeAll()
        notificationContinuation?.finish()
        notificationContinuation = nil
    }

    /// Stop reading and release continuations (called on backend shutdown).
    public func close() {
        readTask?.cancel()
        terminate(reason: RPCError.connectionClosed)
    }
}