web-terminal/ios/Packages/WireProtocol/Sources/WireProtocol/MessageCodec.swift

import Foundation

/// Pure static codec for the WS text-frame protocol (frozen contract, plan §3.1).
/// NEVER throws in either direction:
/// - `encode` is total — every `ClientMessage` has a JSON representation.
/// - `decodeServer` mirrors the server's "invalid frames are silently discarded"
///   resilience (src/protocol.ts:45-86, src/server.ts:698-701): malformed input → `nil`.
public enum MessageCodec {
    // MARK: - encode (client → server)

    /// Encode a `ClientMessage` as a JSON text frame accepted by the server's
    /// `parseClientMessage` (src/protocol.ts). String escaping matches
    /// `JSON.stringify` byte-for-byte (short escapes for \b \t \n \f \r, and
    /// `\u00XX` lowercase-hex for other C0 control bytes) so a Swift client is
    /// indistinguishable from the web client on the wire.
    ///
    /// Key shapes (each pinned by CodecRoundtripTests):
    /// - `attach` ALWAYS carries an explicit `sessionId` key — JSON `null` when
    ///   nil (the server requires the key's presence, src/protocol.ts:132-134);
    ///   UUIDs are lowercased (crypto.randomUUID style; SESSION_ID_RE is /i).
    /// - `approve` puts `mode` as a TOP-LEVEL key: the server's WS wiring
    ///   re-parses the RAW frame for `obj['mode']` (src/server.ts:91-102).
    /// - `resize` emits bare integers; `input.data` is escaped verbatim.
    public static func encode(_ message: ClientMessage) -> String {
        switch message {
        case let .attach(sessionId, cwd):
            return encodeAttach(sessionId: sessionId, cwd: cwd)
        case let .input(data):
            return "{\"type\":\"input\",\"data\":\(jsonStringLiteral(data))}"
        case let .resize(cols, rows):
            return "{\"type\":\"resize\",\"cols\":\(cols),\"rows\":\(rows)}"
        case let .approve(mode):
            guard let mode else { return "{\"type\":\"approve\"}" }
            return "{\"type\":\"approve\",\"mode\":\"\(mode.rawValue)\"}"
        case .reject:
            return "{\"type\":\"reject\"}"
        }
    }

    private static func encodeAttach(sessionId: UUID?, cwd: String?) -> String {
        let sessionIdJSON = sessionId.map { "\"\($0.uuidString.lowercased())\"" } ?? "null"
        let cwdPart = cwd.map { ",\"cwd\":\(jsonStringLiteral($0))" } ?? ""
        return "{\"type\":\"attach\",\"sessionId\":\(sessionIdJSON)\(cwdPart)}"
    }

    /// JSON string literal with `JSON.stringify`-identical escaping.
    private static func jsonStringLiteral(_ value: String) -> String {
        var out = "\""
        for scalar in value.unicodeScalars {
            switch scalar {
            case "\"": out += "\\\""
            case "\\": out += "\\\\"
            case "\u{08}": out += "\\b"
            case "\t": out += "\\t"
            case "\n": out += "\\n"
            case "\u{0C}": out += "\\f"
            case "\r": out += "\\r"
            default:
                if scalar.value < 0x20 {
                    out += String(format: "\\u%04x", Int(scalar.value))
                } else {
                    out.unicodeScalars.append(scalar)
                }
            }
        }
        return out + "\""
    }

    // MARK: - decodeServer (server → client, untrusted)

    /// Decode a server JSON text frame. Whitelist semantics, never throws:
    /// bad JSON, non-object frames, unknown `type`, or missing/wrong-typed
    /// REQUIRED fields → `nil` (drop the frame). Wrong-typed OPTIONAL fields
    /// are tolerated as absent (see per-case doc on `ServerMessage`).
    public static func decodeServer(_ text: String) -> ServerMessage? {
        guard let data = text.data(using: .utf8),
              let frame = try? JSONDecoder().decode(ServerFrame.self, from: data)
        else { return nil }
        return frame.message
    }
}

// MARK: - Internal decoding scaffolding

/// Decodable shim so one JSONDecoder pass yields `ServerMessage?` without ever
/// surfacing a throw for merely-malformed content (only structurally non-JSON
/// input makes JSONDecoder itself throw, which `decodeServer` catches).
private struct ServerFrame: Decodable {
    let message: ServerMessage?

    private enum Keys: String, CodingKey {
        case type, sessionId, data, code, reason, status, detail, pending, gate, telemetry
    }

    init(from decoder: any Decoder) throws {
        guard let container = try? decoder.container(keyedBy: Keys.self),
              let type = try? container.decode(String.self, forKey: .type)
        else {
            message = nil
            return
        }
        message = Self.decodeBody(type: type, from: container)
    }

    private static func decodeBody(
        type: String, from container: KeyedDecodingContainer<Keys>
    ) -> ServerMessage? {
        switch type {
        case "attached":
            return decodeAttached(container)
        case "output":
            guard let data = try? container.decode(String.self, forKey: .data) else { return nil }
            return .output(data: data)
        case "exit":
            return decodeExit(container)
        case "status":
            return decodeStatus(container)
        case "telemetry":
            guard let telemetry = try? container.decode(StatusTelemetry.self, forKey: .telemetry)
            else { return nil }
            return .telemetry(telemetry)
        default:
            return nil
        }
    }

    /// `attached.sessionId` must pass the server's own SESSION_ID_RE (M7) —
    /// `UUID(uuidString:)` alone would accept non-v4 UUIDs.
    private static func decodeAttached(_ container: KeyedDecodingContainer<Keys>) -> ServerMessage? {
        guard let raw = try? container.decode(String.self, forKey: .sessionId),
              Validation.isValidSessionId(raw),
              let sessionId = UUID(uuidString: raw)
        else { return nil }
        return .attached(sessionId: sessionId)
    }

    private static func decodeExit(_ container: KeyedDecodingContainer<Keys>) -> ServerMessage? {
        guard let code = try? container.decode(Int.self, forKey: .code) else { return nil }
        let reason = try? container.decode(String.self, forKey: .reason)
        return .exit(code: code, reason: reason)
    }

    private static func decodeStatus(_ container: KeyedDecodingContainer<Keys>) -> ServerMessage? {
        guard let rawStatus = try? container.decode(String.self, forKey: .status),
              let status = ClaudeStatus(rawValue: rawStatus)
        else { return nil }
        let detail = try? container.decode(String.self, forKey: .detail)
        let pending = (try? container.decode(Bool.self, forKey: .pending)) ?? false
        let gate = (try? container.decode(String.self, forKey: .gate))
            .flatMap(GateKind.init(rawValue:))
        return .status(status, detail: detail, pending: pending, gate: gate)
    }
}