Context Window Planner — Swift source

// Context Window Planner — plan labeled prompt sections against a model's
// context window.
//
// Language: Swift (Swift 5.9, standard library only)
// Source:   CosmoDev polyglot showcase port of the Context Window Planner
//           tool, ported from src/lib/contextPlanner.ts (the canonical
//           TypeScript implementation).
// Live at:  https://dev.cosmolabs.org/tools/context-window-planner
// License:  display source — part of CosmoDev's polyglot tool pages.
//
// Design goals:
//   - Pure + deterministic; never traps (public API returns plain values).
//   - Functionally equivalent to the TS reference: same inputs -> same outputs.
//   - Self-contained: the Swift stdlib only (Foundation's JSONSerialization
//     is the framework equivalent of serde_json; this port instead ships the
//     same small strict recursive-descent validator the other polyglot
//     siblings use, so the JSON grammar matches JSON.parse exactly rather
//     than approximately).
//
// Port notes: the TS lib delegates to two siblings — `estimateTokens` from
// src/lib/tokenEstimator.ts and `fitsWindow` from src/lib/ai/models.ts (which
// defaults to the bundled pricing snapshot, src/data/ai-models.json). A
// dependency-free port cannot load that file, so the estimator is inlined
// below in the exact form the planner uses it (`estimateTokens(text).tokens`,
// auto content type — the full heuristic lives in the token-estimator port),
// window math is inlined from `fitsWindow` and `models` is an explicit
// parameter, never re-derived.
//
// Faithfulness notes (the places Swift's stdlib silently differs from JS):
//   - Length: TS's `String.length` counts UTF-16 code units (an astral-plane
//     character — emoji, rare CJK ext-B ideographs — counts as 2). Swift's
//     `String.count` counts Characters (grapheme clusters), so line
//     arithmetic goes through `utf16Len` (`s.utf16.count`) to count the same
//     unit.
//   - Rounding: `jsRound` is `floor(x + 0.5)` — JS `Math.round` rounds
//     halfway cases up; `.toNearestOrAwayFromZero` matches only over
//     non-negative inputs, so the helper pins the exact formula.

/// Namespace for the planner's pure logic — no state, only functions and
/// value types.
enum ContextWindowPlanner {

    /// One labeled block of the prompt (system / docs / history / ...).
    /// Mirrors the TS `PlanSection` interface.
    struct PlanSection: Equatable {
        /// Section label, e.g. "system" or "docs".
        let label: String
        /// The section's raw text.
        let text: String
    }

    /// Convenience constructor mirroring the TS object literal `{ label, text }`.
    static func sec(_ label: String, _ text: String) -> PlanSection {
        PlanSection(label: label, text: text)
    }

    /// The subset of the TS `AiModel` record the planner reads. Production
    /// code passes the full snapshot entry; only these fields influence the
    /// plan.
    struct Model: Equatable {
        /// Model id, e.g. "beta-pro".
        let id: String
        /// Total context window in tokens.
        let contextWindow: Int
        /// The model's output cap (informational).
        let maxOutput: Int
    }

    /// Sample table for standalone use (mirrors the shared test fixtures).
    /// Production code passes the model snapshot instead.
    static let sampleModels: [Model] = [
        Model(id: "alpha-mini", contextWindow: 200_000, maxOutput: 10_000),
        Model(id: "beta-pro", contextWindow: 1_000_000, maxOutput: 10_000),
        Model(id: "gamma-open", contextWindow: 100_000, maxOutput: 10_000),
    ]

    /// Result of `planWindow`. Field-for-field twin of the TS `WindowPlan`
    /// interface.
    struct WindowPlan: Equatable {
        /// The model id planned against.
        let id: String
        /// Sum of per-section token estimates.
        let inputTokens: Int
        /// The model's context window.
        let contextWindow: Int
        /// Context tokens left after the request; negative on overflow.
        let free: Int
        /// Raw fit: free >= 0.
        let fits: Bool
        /// Room for the output reserve: free >= reserve.
        let outputReserveOk: Bool
        /// The model's output cap (informational).
        let maxOutput: Int
    }

    // MARK: - estimator internals

    /// Content classification of a single line. The planner only needs each
    /// type's chars-per-token rate (mirrors `CHARS_PER_TOKEN` in
    /// src/lib/tokenEstimator.ts: prose 4, code 3.5, json 3, cjk 1.5).
    private enum ContentType {
        case prose, code, json, cjk

        var charsPerToken: Double {
            switch self {
            case .prose: return 4.0
            case .code: return 3.5
            case .json: return 3.0
            case .cjk: return 1.5
            }
        }
    }

    /// Length of `s` in UTF-16 code units — the unit TS's `String.length`
    /// counts. BMP code points are one unit, astral-plane ones two.
    private static func utf16Len(_ s: Substring) -> Int {
        s.utf16.count
    }

    /// Reports whether `s` contains a CJK ideograph (U+4E00–U+9FFF), kana
    /// (U+3040–U+30FF), or a Hangul syllable (U+AC00–U+D7AF). Mirrors
    /// `CJK_RE` in the TS lib.
    private static func hasCjk(_ s: Substring) -> Bool {
        s.unicodeScalars.contains { c in
            (0x4E00...0x9FFF).contains(c.value) ||
            (0x3040...0x30FF).contains(c.value) ||
            (0xAC00...0xD7AF).contains(c.value)
        }
    }

    /// Reports whether `c` is one of the code-flavored symbols counted by
    /// `CODE_SYMBOL_RE` (`{}();=<>[]#`).
    private static func isCodeSymbol(_ c: Character) -> Bool {
        "{}();=<>[]#".contains(c)
    }

    /// JS `Math.round`: halfway cases round up (`floor(x + 0.5)`).
    private static func jsRound(_ x: Double) -> Int {
        Int((x + 0.5).rounded(.down))
    }

    /// Splits `text` on LF or CRLF, mirroring `text.split(/\r?\n/)`: strip
    /// the optional CR that belongs to the newline, then split on LF. A lone
    /// CR is NOT a line break.
    private static func splitLines(_ text: String) -> [Substring] {
        text.split(separator: "\n", omittingEmptySubsequences: false)
            .map { line in line.hasSuffix("\r") ? line.dropLast() : line }
    }

    /// Classifies a single line by its shape. Order: json, cjk, code, prose.
    /// Inlined from `detectLineType()` in src/lib/tokenEstimator.ts.
    private static func detectLineType(_ line: Substring) -> ContentType {
        let trimmed = line.trimmingWhitespace
        // JSON-ish: opens like a JSON fragment AND carries a separator.
        let first = trimmed.first
        let startsJsonish = first == "{" || first == "}" || first == "[" || first == "\""
        if startsJsonish, line.contains(":") || line.contains(",") {
            return .json
        }
        // CJK ideographs / kana / Hangul pack roughly one token per 1.5 chars.
        if hasCjk(line) {
            return .cjk
        }
        // Code: symbol-dense, or a statement terminator / block opener at EOL.
        let length = utf16Len(line)
        let symbols = line.filter(isCodeSymbol).count
        let density = length > 0 ? Double(symbols) / Double(length) : 0.0
        if density > 0.08 || trimmed.hasSuffix(";") || trimmed.hasSuffix("{") ||
            trimmed.hasSuffix("}") {
            return .code
        }
        return .prose
    }

    /// A strict JSON syntax validator — the exact grammar `JSON.parse`
    /// accepts, walked with a cursor. Same shape as the Rust/C/C++/C#/Java
    /// siblings' validators, so whole-text JSON detection behaves
    /// identically across every port.
    private struct JsonParser {
        let scalars: [Unicode.Scalar]
        var pos = 0

        init(_ text: String) {
            scalars = Array(text.unicodeScalars)
        }

        /// value := ws* (object | array | string | number | 'true' | 'false' | 'null') ws*
        mutating func value() -> Bool {
            skipWs()
            guard pos < scalars.count else { return false }
            switch scalars[pos] {
            case "{": return object()
            case "[": return array()
            case "\"": return string()
            case "-", "0", "1", "2", "3", "4", "5", "6", "7", "8", "9":
                return number()
            case "t": return literal("true")
            case "f": return literal("false")
            case "n": return literal("null")
            default: return false
            }
        }

        mutating func atEnd() -> Bool {
            skipWs()
            return pos == scalars.count // reject trailing garbage
        }

        private mutating func skipWs() {
            while pos < scalars.count,
                  scalars[pos] == " " || scalars[pos] == "\t" ||
                  scalars[pos] == "\n" || scalars[pos] == "\r" {
                pos += 1
            }
        }

        private func peek() -> Unicode.Scalar? {
            pos < scalars.count ? scalars[pos] : nil
        }

        private mutating func eat(_ b: Unicode.Scalar) -> Bool {
            if pos < scalars.count && scalars[pos] == b { pos += 1; return true }
            return false
        }

        private mutating func literal(_ lit: String) -> Bool {
            let l = Array(lit.unicodeScalars)
            guard scalars.count - pos >= l.count else { return false }
            for (i, c) in l.enumerated() where scalars[pos + i] != c { return false }
            pos += l.count
            return true
        }

        /// object := '{' ws* (string ws* ':' value (ws* ',' ...)*)? ws* '}'
        private mutating func object() -> Bool {
            guard eat("{") else { return false }
            skipWs()
            if eat("}") { return true }
            while true {
                guard string() else { return false }
                skipWs()
                guard eat(":") else { return false }
                guard value() else { return false }
                skipWs()
                if eat(",") { skipWs() }
                else { return eat("}") }
            }
        }

        /// array := '[' ws* (value (ws* ',' ws* value)*)? ws* ']'
        private mutating func array() -> Bool {
            guard eat("[") else { return false }
            skipWs()
            if eat("]") { return true }
            while true {
                guard value() else { return false }
                skipWs()
                if eat(",") { skipWs() }
                else { return eat("]") }
            }
        }

        /// string := '"' (escape | any scalar >= 0x20)* '"'
        /// escape := '\' ('"' | '/' | '\' | 'b' | 'f' | 'n' | 'r' | 't' | 'u' hex4)
        private mutating func string() -> Bool {
            guard eat("\"") else { return false }
            while pos < scalars.count {
                let c = scalars[pos]
                if c == "\"" { pos += 1; return true }
                if c == "\\" {
                    pos += 1
                    guard pos < scalars.count else { return false }
                    let esc = scalars[pos]
                    pos += 1
                    switch esc {
                    case "\"", "/", "\\", "b", "f", "n", "r", "t":
                        break
                    case "u":
                        for _ in 0..<4 {
                            guard let h = peek(),
                                  ("0"..."9").contains(h) || ("a"..."f").contains(h) ||
                                  ("A"..."F").contains(h) else { return false }
                            pos += 1
                        }
                    default:
                        return false
                    }
                } else if c.value < 0x20 {
                    return false // raw control characters not allowed in strings
                } else {
                    pos += 1
                }
            }
            return false // unterminated string
        }

        /// number := '-'? int frac? exp? — no leading zeros, like JSON.parse.
        private mutating func number() -> Bool {
            _ = eat("-")
            switch peek() {
            case "0": pos += 1
            case .some(let c) where ("1"..."9").contains(c):
                while let p = peek(), ("0"..."9").contains(p) { pos += 1 }
            default: return false
            }
            if peek() == "." {
                pos += 1
                var digits = 0
                while let p = peek(), ("0"..."9").contains(p) { pos += 1; digits += 1 }
                guard digits > 0 else { return false }
            }
            if peek() == "e" || peek() == "E" {
                pos += 1
                if peek() == "+" || peek() == "-" { pos += 1 }
                var digits = 0
                while let p = peek(), ("0"..."9").contains(p) { pos += 1; digits += 1 }
                guard digits > 0 else { return false }
            }
            return true
        }
    }

    /// Whole-text JSON gate: a document that parses as JSON is json all the
    /// way down. Mirrors `isValidJson()` (`JSON.parse` in a try/catch);
    /// empty/whitespace text is not.
    static func isValidJson(_ text: String) -> Bool {
        if text.trimmingWhitespace.isEmpty { return false }
        var p = JsonParser(text)
        return p.value() && p.atEnd()
    }

    /// Token count of `text` under auto content detection — exactly the
    /// slice of `estimateTokens()` the planner consumes (`.tokens`): per
    /// non-empty line, `max(1, round(utf16Len / charsPerToken))`. Framing
    /// tokens are the caller's job.
    static func estimateTokens(_ text: String) -> Int {
        // AUTO + whole-text JSON: json's 3 chars/token rate applies to every
        // line, not just the reported content type.
        let wholeTextJson = isValidJson(text)
        var tokens = 0
        for line in splitLines(text) {
            if line.trimmingWhitespace.isEmpty { continue }
            let t = wholeTextJson ? ContentType.json : detectLineType(line)
            tokens += max(1, jsRound(Double(utf16Len(line)) / t.charsPerToken))
        }
        return tokens
    }

    // MARK: - planner API

    /// Sum of per-section token estimates (framing tokens are the caller's
    /// job). Mirrors `inputTokenTotal()` in the TS lib.
    static func inputTokenTotal(_ sections: [PlanSection]) -> Int {
        sections.reduce(0) { $0 + estimateTokens($1.text) }
    }

    /// Plan one section set against one model's context window. Returns
    /// `nil` for an unknown model id (window math is `fitsWindow`'s, never
    /// re-derived). Mirrors `planWindow()` in the TS lib.
    static func planWindow(
        _ sections: [PlanSection], _ modelId: String,
        outputReserve: Int = 0, models: [Model] = []
    ) -> WindowPlan? {
        let inputTokens = inputTokenTotal(sections)
        // Fit check inlined from fitsWindow() in src/lib/ai/models.ts.
        guard let m = models.first(where: { $0.id == modelId }) else { return nil }
        let free = m.contextWindow - inputTokens
        return WindowPlan(
            id: modelId,
            inputTokens: inputTokens,
            contextWindow: m.contextWindow,
            free: free,
            fits: free >= 0,
            outputReserveOk: free >= outputReserve,
            maxOutput: m.maxOutput)
    }

    /// Plan against several models; unknown ids are dropped from the result.
    /// Mirrors `planAll()` in the TS lib.
    static func planAll(
        _ sections: [PlanSection], _ modelIds: [String],
        outputReserve: Int = 0, models: [Model] = []
    ) -> [WindowPlan] {
        modelIds.compactMap {
            planWindow(sections, $0, outputReserve: outputReserve, models: models)
        }
    }
}

/// Minimal ASCII-whitespace trim (the stdlib's `trimmingCharacters(in:)`
/// lives in Foundation; this port is stdlib-only). On StringProtocol so both
/// `String` and `Substring` get it.
private extension StringProtocol {
    var trimmingWhitespace: Substring {
        var s = Substring(self)
        while let f = s.first, f == " " || f == "\t" || f == "\n" || f == "\r" || f == "\u{0B}" || f == "\u{0C}" {
            s = s.dropFirst()
        }
        while let l = s.last, l == " " || l == "\t" || l == "\n" || l == "\r" || l == "\u{0B}" || l == "\u{0C}" {
            s = s.dropLast()
        }
        return s
    }
}

// ---------- tests (showcase-only; the canonical suite lives in src/lib) ----------
// Build with -DCWP_TEST and call ContextWindowPlannerTests.run().
#if CWP_TEST
enum ContextWindowPlannerTests {

    /// A 1600-char single line of 'a' is pure prose: 1600 / 4 = 400 tokens.
    private static func twoSections() -> [ContextWindowPlanner.PlanSection] {
        let lineA = String(repeating: "a", count: 1600)
        return [.init(label: "sys", text: lineA), .init(label: "docs", text: lineA)]
    }

    static func run() {
        let two = twoSections()

        // input totals
        assert(ContextWindowPlanner.inputTokenTotal(two) == 800)
        assert(ContextWindowPlanner.inputTokenTotal([]) == 0)
        assert(ContextWindowPlanner.inputTokenTotal([.init(label: "sys", text: "")]) == 0)

        // plans two 400-token sections against beta-pro
        let p = ContextWindowPlanner.planWindow(two, "beta-pro", models: ContextWindowPlanner.sampleModels)!
        assert(p.id == "beta-pro")
        assert(p.inputTokens == 800)
        assert(p.contextWindow == 1_000_000)
        assert(p.free == 999_200)
        assert(p.fits)
        assert(p.outputReserveOk)
        assert(p.maxOutput == 10_000)

        // reserve larger than free leaves raw fit true
        let q = ContextWindowPlanner.planWindow(two, "beta-pro", outputReserve: 1_000_000, models: ContextWindowPlanner.sampleModels)!
        assert(q.fits)
        assert(!q.outputReserveOk)

        // reserve exactly equal to free is ok
        let r = ContextWindowPlanner.planWindow(two, "beta-pro", outputReserve: 999_200, models: ContextWindowPlanner.sampleModels)!
        assert(r.outputReserveOk)

        // smaller window leaves 199,200 free
        let s = ContextWindowPlanner.planWindow(two, "alpha-mini", models: ContextWindowPlanner.sampleModels)!
        assert(s.contextWindow == 200_000)
        assert(s.free == 199_200)
        assert(s.fits)

        // unknown model id returns nil
        assert(ContextWindowPlanner.planWindow(two, "ghost", models: ContextWindowPlanner.sampleModels) == nil)

        // no sections: full window free
        let t = ContextWindowPlanner.planWindow([], "beta-pro", models: ContextWindowPlanner.sampleModels)!
        assert(t.inputTokens == 0)
        assert(t.free == 1_000_000)
        assert(t.fits)

        // overflow: fits false, reserve false
        let big = [ContextWindowPlanner.sec("big", String(repeating: "z", count: 4_400_000))]
        let u = ContextWindowPlanner.planWindow(big, "beta-pro", models: ContextWindowPlanner.sampleModels)!
        assert(u.inputTokens == 1_100_000)
        assert(u.free == -100_000)
        assert(!u.fits)
        assert(!u.outputReserveOk)

        // plan_all drops unknown ids and keeps order
        let plans = ContextWindowPlanner.planAll(two, ["beta-pro", "alpha-mini", "ghost"], models: ContextWindowPlanner.sampleModels)
        assert(plans.count == 2)
        assert(plans[0].id == "beta-pro")
        assert(plans[1].id == "alpha-mini")
        assert(plans[1].free == 199_200)
        assert(ContextWindowPlanner.planAll(two, [], models: ContextWindowPlanner.sampleModels).isEmpty)

        print("context-window-planner (Swift): all tests passed")
    }
}
#endif