Context Window Planner — Zig source

//! Context Window Planner — plan labeled prompt sections against a model's
//! context window.
//!
//! Language: Zig (Zig 0.13, 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 panics (public API returns plain values,
//!     no allocator required).
//!   - Functionally equivalent to the TS reference: same inputs -> same outputs.
//!   - Self-contained: std only (std.json exists but pulls the full reader
//!     machinery; this port 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 Zig's std 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). Zig
//!     slices are UTF-8 bytes, so line arithmetic goes through `utf16Len`,
//!     which decodes code points and counts the same unit.
//!   - Rounding: `jsRound` is `@floor(x + 0.5)` — JS `Math.round` rounds
//!     halfway cases up; `@round` rounds away from zero (identical on the
//!     non-negative numbers used here, but the helper pins the formula).

const std = @import("std");

/// One labeled block of the prompt (system / docs / history / ...).
/// Mirrors the TS `PlanSection` interface.
pub const PlanSection = struct {
    /// Section label, e.g. "system" or "docs".
    label: []const u8,
    /// The section's raw text.
    text: []const u8,
};

/// Convenience constructor mirroring the TS object literal `{ label, text }`.
pub fn sec(label: []const u8, text: []const u8) PlanSection {
    return .{ .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.
pub const Model = struct {
    /// Model id, e.g. "beta-pro".
    id: []const u8,
    /// Total context window in tokens.
    context_window: i64,
    /// The model's output cap (informational).
    max_output: i64,
};

/// Sample table for standalone use (mirrors the shared test fixtures).
/// Production code passes the model snapshot instead.
pub const sample_models = [_]Model{
    .{ .id = "alpha-mini", .context_window = 200_000, .max_output = 10_000 },
    .{ .id = "beta-pro", .context_window = 1_000_000, .max_output = 10_000 },
    .{ .id = "gamma-open", .context_window = 100_000, .max_output = 10_000 },
};

/// Result of `planWindow`. Field-for-field twin of the TS `WindowPlan`
/// interface.
pub const WindowPlan = struct {
    /// The model id planned against.
    id: []const u8,
    /// Sum of per-section token estimates.
    input_tokens: i64,
    /// The model's context window.
    context_window: i64,
    /// Context tokens left after the request; negative on overflow.
    free: i64,
    /// Raw fit: free >= 0.
    fits: bool,
    /// Room for the output reserve: free >= reserve.
    output_reserve_ok: bool,
    /// The model's output cap (informational).
    max_output: i64,
};

/// 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).
const ContentType = enum {
    prose,
    code,
    json,
    cjk,

    fn charsPerToken(self: ContentType) f64 {
        return switch (self) {
            .prose => 4.0,
            .code => 3.5,
            .json => 3.0,
            .cjk => 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. Input must
/// be valid UTF-8 (page text always is).
fn utf16Len(s: []const u8) i64 {
    var units: i64 = 0;
    var it = std.unicode.Utf8View.initUnchecked(s).iterator();
    while (it.nextCodepoint()) |cp| {
        units += if (cp > 0xFFFF) 2 else 1;
    }
    return units;
}

/// 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.
fn hasCjk(s: []const u8) bool {
    var it = std.unicode.Utf8View.initUnchecked(s).iterator();
    while (it.nextCodepoint()) |cp| {
        if ((cp >= 0x4E00 and cp <= 0x9FFF) or
            (cp >= 0x3040 and cp <= 0x30FF) or
            (cp >= 0xAC00 and cp <= 0xD7AF))
        {
            return true;
        }
    }
    return false;
}

/// Reports whether `c` is one of the code-flavored symbols counted by
/// `CODE_SYMBOL_RE` (`{}();=<>[]#`).
fn isCodeSymbol(c: u8) bool {
    return c == '{' or c == '}' or c == '(' or c == ')' or c == ';' or
        c == '=' or c == '<' or c == '>' or c == '[' or c == ']' or c == '#';
}

const ws = " \t\n\r\x0b\x0c";

/// JS `Math.round`: halfway cases round up (`@floor(x + 0.5)`).
fn jsRound(x: f64) i64 {
    return @intFromFloat(@floor(x + 0.5));
}

/// Classifies a single line by its shape. Order: json, cjk, code, prose.
/// Inlined from `detectLineType()` in src/lib/tokenEstimator.ts.
fn detectLineType(line: []const u8) ContentType {
    const trimmed = std.mem.trim(u8, line, ws);
    // JSON-ish: opens like a JSON fragment AND carries a separator.
    const starts_jsonish = trimmed.len > 0 and
        (trimmed[0] == '{' or trimmed[0] == '}' or
        trimmed[0] == '[' or trimmed[0] == '"');
    if (starts_jsonish and
        (std.mem.indexOfScalar(u8, line, ':') != null or
        std.mem.indexOfScalar(u8, line, ',') != null))
    {
        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.
    const length = utf16Len(line);
    var symbols: i64 = 0;
    for (line) |c| {
        if (isCodeSymbol(c)) symbols += 1;
    }
    const density = if (length > 0)
        @as(f64, @floatFromInt(symbols)) / @as(f64, @floatFromInt(length))
    else
        0.0;
    const ends_code = std.mem.endsWith(u8, trimmed, ";") or
        std.mem.endsWith(u8, trimmed, "{") or std.mem.endsWith(u8, trimmed, "}");
    if (density > 0.08 or ends_code) {
        return .code;
    }
    return .prose;
}

/// A strict JSON syntax validator — the exact grammar `JSON.parse` accepts,
/// walked with a byte cursor. (Multibyte UTF-8 inside strings never contains
/// an ASCII byte, so byte-level scanning is safe.)
const JsonParser = struct {
    bytes: []const u8,
    pos: usize,

    fn init(text: []const u8) JsonParser {
        return .{ .bytes = text, .pos = 0 };
    }

    fn skipWs(self: *JsonParser) void {
        while (self.pos < self.bytes.len) {
            const b = self.bytes[self.pos];
            if (b == ' ' or b == '\t' or b == '\n' or b == '\r') {
                self.pos += 1;
            } else {
                break;
            }
        }
    }

    fn peek(self: *const JsonParser) ?u8 {
        return if (self.pos < self.bytes.len) self.bytes[self.pos] else null;
    }

    fn eat(self: *JsonParser, b: u8) bool {
        if (self.peek() == b and self.pos < self.bytes.len) {
            self.pos += 1;
            return true;
        }
        return false;
    }

    fn literal(self: *JsonParser, lit: []const u8) bool {
        if (std.mem.startsWith(u8, self.bytes[self.pos..], lit)) {
            self.pos += lit.len;
            return true;
        }
        return false;
    }

    /// value := ws* (object | array | string | number | 'true' | 'false' | 'null') ws*
    fn value(self: *JsonParser) bool {
        self.skipWs();
        const b = self.peek() orelse return false;
        return switch (b) {
            '{' => self.object(),
            '[' => self.array(),
            '"' => self.string(),
            '-', '0'...'9' => self.number(),
            't' => self.literal("true"),
            'f' => self.literal("false"),
            'n' => self.literal("null"),
            else => false,
        };
    }

    /// object := '{' ws* (string ws* ':' value (ws* ',' ...)*)? ws* '}'
    fn object(self: *JsonParser) bool {
        if (!self.eat('{')) return false;
        self.skipWs();
        if (self.eat('}')) return true;
        while (true) {
            if (!self.string()) return false;
            self.skipWs();
            if (!self.eat(':')) return false;
            if (!self.value()) return false;
            self.skipWs();
            if (self.eat(',')) {
                self.skipWs();
            } else {
                return self.eat('}');
            }
        }
    }

    /// array := '[' ws* (value (ws* ',' ws* value)*)? ws* ']'
    fn array(self: *JsonParser) bool {
        if (!self.eat('[')) return false;
        self.skipWs();
        if (self.eat(']')) return true;
        while (true) {
            if (!self.value()) return false;
            self.skipWs();
            if (self.eat(',')) {
                self.skipWs();
            } else {
                return self.eat(']');
            }
        }
    }

    /// string := '"' (escape | any byte >= 0x20)* '"'
    /// escape := '\' ('"' | '/' | '\' | 'b' | 'f' | 'n' | 'r' | 't' | 'u' hex4)
    fn string(self: *JsonParser) bool {
        if (!self.eat('"')) return false;
        while (self.pos < self.bytes.len) {
            const b = self.bytes[self.pos];
            if (b == '"') {
                self.pos += 1;
                return true;
            }
            if (b == '\\') {
                self.pos += 1;
                if (self.pos >= self.bytes.len) return false;
                const esc = self.bytes[self.pos];
                self.pos += 1;
                switch (esc) {
                    '"', '/', '\\', 'b', 'f', 'n', 'r', 't' => {},
                    'u' => {
                        var i: usize = 0;
                        while (i < 4) : (i += 1) {
                            const h = self.peek() orelse return false;
                            const hex = (h >= '0' and h <= '9') or
                                (h >= 'a' and h <= 'f') or (h >= 'A' and h <= 'F');
                            if (!hex) return false;
                            self.pos += 1;
                        }
                    },
                    else => return false,
                }
            } else if (b < 0x20) {
                // Raw control characters are not allowed inside strings.
                return false;
            } else {
                self.pos += 1;
            }
        }
        return false; // unterminated string
    }

    /// number := '-'? int frac? exp? — no leading zeros, like JSON.parse.
    fn number(self: *JsonParser) bool {
        _ = self.eat('-');
        const first = self.peek() orelse return false;
        switch (first) {
            '0' => self.pos += 1,
            '1'...'9' => {
                while (self.peek()) |p| {
                    if (p < '0' or p > '9') break;
                    self.pos += 1;
                }
            },
            else => return false,
        }
        if (self.peek() == @as(?u8, '.')) {
            self.pos += 1;
            var digits: usize = 0;
            while (self.peek()) |p| {
                if (p < '0' or p > '9') break;
                self.pos += 1;
                digits += 1;
            }
            if (digits == 0) return false;
        }
        if (self.peek() == @as(?u8, 'e') or self.peek() == @as(?u8, 'E')) {
            self.pos += 1;
            if (self.peek() == @as(?u8, '+') or self.peek() == @as(?u8, '-')) {
                self.pos += 1;
            }
            var digits: usize = 0;
            while (self.peek()) |p| {
                if (p < '0' or p > '9') break;
                self.pos += 1;
                digits += 1;
            }
            if (digits == 0) 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.
fn isValidJson(text: []const u8) bool {
    if (std.mem.trim(u8, text, ws).len == 0) return false;
    var p = JsonParser.init(text);
    if (!p.value()) return false;
    p.skipWs();
    return p.pos == p.bytes.len; // reject trailing garbage
}

/// 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.
fn estimateTokens(text: []const u8) i64 {
    // AUTO + whole-text JSON: json's 3 chars/token rate applies to every
    // line, not just the reported content type.
    const whole_text_json = isValidJson(text);
    var tokens: i64 = 0;
    // Split on LF or CRLF (text.split(/\r?\n/)): strip the optional CR that
    // belongs to the newline, then split on LF. A lone CR is NOT a break.
    var rest = text;
    while (rest.len > 0) {
        const nl = std.mem.indexOfScalar(u8, rest, '\n');
        var line = if (nl) |i| rest[0..i] else rest;
        if (line.len > 0 and line[line.len - 1] == '\r') {
            line = line[0 .. line.len - 1];
        }
        if (std.mem.trim(u8, line, ws).len > 0) {
            const t: ContentType = if (whole_text_json) .json else detectLineType(line);
            const est = jsRound(
                @as(f64, @floatFromInt(utf16Len(line))) / t.charsPerToken(),
            );
            tokens += if (est >= 1) est else 1;
        }
        if (nl) |i| {
            rest = rest[i + 1 ..];
        } else {
            break;
        }
    }
    return tokens;
}

/// Sum of per-section token estimates (framing tokens are the caller's job).
/// Mirrors `inputTokenTotal()` in the TS lib.
pub fn inputTokenTotal(sections: []const PlanSection) i64 {
    var total: i64 = 0;
    for (sections) |s| total += estimateTokens(s.text);
    return total;
}

/// Plan one section set against one model's context window. Returns `null`
/// for an unknown model id (window math is `fitsWindow`'s, never
/// re-derived). Mirrors `planWindow()` in the TS lib.
pub fn planWindow(
    sections: []const PlanSection,
    model_id: []const u8,
    output_reserve: i64,
    models: []const Model,
) ?WindowPlan {
    const input_tokens = inputTokenTotal(sections);
    // Fit check inlined from fitsWindow() in src/lib/ai/models.ts.
    var found: ?Model = null;
    for (models) |m| {
        if (std.mem.eql(u8, m.id, model_id)) {
            found = m;
            break;
        }
    }
    const m = found orelse return null;
    const free = m.context_window - input_tokens;
    return .{
        .id = model_id,
        .input_tokens = input_tokens,
        .context_window = m.context_window,
        .free = free,
        .fits = free >= 0,
        .output_reserve_ok = free >= output_reserve,
        .max_output = m.max_output,
    };
}

/// Plan against several models; unknown ids are dropped from the result.
/// Writes at most `out.len` plans and returns the count written.
/// Mirrors `planAll()` in the TS lib.
pub fn planAll(
    sections: []const PlanSection,
    model_ids: []const []const u8,
    output_reserve: i64,
    models: []const Model,
    out: []WindowPlan,
) usize {
    var written: usize = 0;
    for (model_ids) |id| {
        if (written >= out.len) break;
        if (planWindow(sections, id, output_reserve, models)) |plan| {
            out[written] = plan;
            written += 1;
        }
    }
    return written;
}

// ---------- tests (showcase-only; the canonical suite lives in src/lib) ----------

/// A 1600-char single line of 'a' is pure prose: 1600 / 4 = 400 tokens.
fn twoSections(buf: []u8) [2]PlanSection {
    @memset(buf, 'a');
    return .{
        sec("sys", buf[0..1600]),
        sec("docs", buf[0..1600]),
    };
}

test "input totals" {
    var buf: [1600]u8 = undefined;
    const two = twoSections(&buf);
    try std.testing.expectEqual(@as(i64, 800), inputTokenTotal(&two));
    try std.testing.expectEqual(@as(i64, 0), inputTokenTotal(&[_]PlanSection{}));
    try std.testing.expectEqual(@as(i64, 0), inputTokenTotal(&[_]PlanSection{sec("sys", "")}));
}

test "plans two 400-token sections against beta-pro" {
    var buf: [1600]u8 = undefined;
    const two = twoSections(&buf);
    const p = planWindow(&two, "beta-pro", 0, &sample_models).?;
    try std.testing.expectEqualStrings("beta-pro", p.id);
    try std.testing.expectEqual(@as(i64, 800), p.input_tokens);
    try std.testing.expectEqual(@as(i64, 1_000_000), p.context_window);
    try std.testing.expectEqual(@as(i64, 999_200), p.free);
    try std.testing.expect(p.fits);
    try std.testing.expect(p.output_reserve_ok);
    try std.testing.expectEqual(@as(i64, 10_000), p.max_output);
}

test "reserve larger than free leaves raw fit true" {
    var buf: [1600]u8 = undefined;
    const two = twoSections(&buf);
    const p = planWindow(&two, "beta-pro", 1_000_000, &sample_models).?;
    try std.testing.expect(p.fits);
    try std.testing.expect(!p.output_reserve_ok);
}

test "reserve exactly equal to free is ok" {
    var buf: [1600]u8 = undefined;
    const two = twoSections(&buf);
    const p = planWindow(&two, "beta-pro", 999_200, &sample_models).?;
    try std.testing.expect(p.output_reserve_ok);
}

test "smaller window leaves 199,200 free" {
    var buf: [1600]u8 = undefined;
    const two = twoSections(&buf);
    const p = planWindow(&two, "alpha-mini", 0, &sample_models).?;
    try std.testing.expectEqual(@as(i64, 200_000), p.context_window);
    try std.testing.expectEqual(@as(i64, 199_200), p.free);
    try std.testing.expect(p.fits);
}

test "unknown model id returns null" {
    var buf: [1600]u8 = undefined;
    const two = twoSections(&buf);
    try std.testing.expect(planWindow(&two, "ghost", 0, &sample_models) == null);
}

test "no sections full window free" {
    const p = planWindow(&[_]PlanSection{}, "beta-pro", 0, &sample_models).?;
    try std.testing.expectEqual(@as(i64, 0), p.input_tokens);
    try std.testing.expectEqual(@as(i64, 1_000_000), p.free);
    try std.testing.expect(p.fits);
}

test "overflow fits false reserve false" {
    const big = [_]PlanSection{sec("big", "z" ** 4_400_000)};
    const p = planWindow(&big, "beta-pro", 0, &sample_models).?;
    try std.testing.expectEqual(@as(i64, 1_100_000), p.input_tokens);
    try std.testing.expectEqual(@as(i64, -100_000), p.free);
    try std.testing.expect(!p.fits);
    try std.testing.expect(!p.output_reserve_ok);
}

test "planAll drops unknown ids and keeps order" {
    var buf: [1600]u8 = undefined;
    const two = twoSections(&buf);
    const ids = [_][]const u8{ "beta-pro", "alpha-mini", "ghost" };
    var plans: [3]WindowPlan = undefined;
    const n = planAll(&two, &ids, 0, &sample_models, &plans);
    try std.testing.expectEqual(@as(usize, 2), n);
    try std.testing.expectEqualStrings("beta-pro", plans[0].id);
    try std.testing.expectEqualStrings("alpha-mini", plans[1].id);
    try std.testing.expectEqual(@as(i64, 199_200), plans[1].free);
    try std.testing.expectEqual(
        @as(usize, 0),
        planAll(&two, &[_][]const u8{}, 0, &sample_models, &plans),
    );
}