Context Window Planner — Ruby source

# Context Window Planner — plan labeled prompt sections against a model's
# context window.
#
# Language: Ruby (Ruby 3.2, standard library only — `json` is a default gem)
# 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 raises (public API returns plain values).
#   - Functionally equivalent to the TS reference: same inputs -> same outputs.
#   - Self-contained: stdlib only (no gems beyond the default set).
#
# 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 Ruby'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). Ruby's
#     `String#length` counts code points, so line arithmetic goes through
#     `utf16_len` to count the same unit.
#   - JSON: `JSON.parse`'s default strictness matches `JSON.parse` in JS —
#     it rejects NaN/Infinity literals (allow_nan defaults to false), rejects
#     trailing garbage, and accepts any top-level scalar — so the whole-text
#     gate needs no hand-rolled validator here, unlike the Rust/C ports.
#   - Rounding: `Float#round` rounds halfway away from zero — equal to JS
#     `Math.round` over the non-negative inputs used here; `js_round` states
#     the contract with `floor(x + 0.5)`.

require 'json'

module ContextWindowPlanner
  module_function

  # One labeled block of the prompt (system / docs / history / ...).
  # Mirrors the TS PlanSection interface. Use `sec` for the TS object
  # literal `{ label, text }`.
  PlanSection = Struct.new(:label, :text)

  def sec(label, text)
    PlanSection.new(label, text)
  end

  # The subset of the TS AiModel record the planner reads. Production code
  # passes the full snapshot entry; only these fields influence the plan.
  Model = Struct.new(:id, :context_window, :max_output)

  # Sample table for standalone use (mirrors the shared test fixtures).
  # Production code passes the model snapshot instead.
  SAMPLE_MODELS = [
    Model.new('alpha-mini',   200_000, 10_000),
    Model.new('beta-pro',   1_000_000, 10_000),
    Model.new('gamma-open',   100_000, 10_000)
  ].freeze

  # Result of `plan_window`. Field-for-field twin of the TS WindowPlan
  # interface (same keys, same meanings).
  WindowPlan = Struct.new(
    :id,                # the model id planned against
    :input_tokens,      # sum of per-section token estimates
    :context_window,    # the model's context window
    :free,              # window - input; negative on overflow
    :fits,              # raw fit: free >= 0
    :output_reserve_ok, # room for the output reserve: free >= reserve
    :max_output,        # the model's output cap (informational)
    keyword_init: true
  )

  # Average characters per token, by content type. Mirrors CHARS_PER_TOKEN
  # in src/lib/tokenEstimator.ts.
  CHARS_PER_TOKEN = {
    prose: 4,
    code: 3.5,
    json: 3,
    cjk: 1.5
  }.freeze

  # CJK ideographs (U+4E00..U+9FFF), kana (U+3040..U+30FF), Hangul syllables
  # (U+AC00..U+D7AF). Mirrors CJK_RE = /[一-鿿぀-ヿ가-힯]/ in the TS lib.
  CJK_RE = /[一-鿿぀-ヿ가-힯]/

  # Code-flavored symbols, counted over the raw line. Mirrors CODE_SYMBOL_RE.
  CODE_SYMBOL_RE = /[{}();=<>\[\]#]/

  # 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.
  def utf16_len(s)
    s.each_char.sum { |ch| ch.ord > 0xFFFF ? 2 : 1 }
  end

  # JS Math.round: halfway cases round up (floor(x + 0.5)).
  def js_round(x)
    (x + 0.5).floor
  end

  # Classifies a single line by its shape. Order: json, cjk, code, prose.
  # Inlined from detectLineType() in src/lib/tokenEstimator.ts.
  def detect_line_type(line)
    trimmed = line.strip
    # JSON-ish: opens like a JSON fragment AND carries a separator.
    first = trimmed[0]
    jsonish = first == '{' || first == '}' || first == '[' || first == '"'
    return :json if jsonish && (line.include?(':') || line.include?(','))
    # CJK ideographs / kana / Hangul pack roughly one token per 1.5 chars.
    return :cjk if CJK_RE.match?(line)
    # Code: symbol-dense, or a statement terminator / block opener at EOL.
    length = utf16_len(line)
    density = length.zero? ? 0.0 : line.scan(CODE_SYMBOL_RE).size.to_f / length
    return :code if density > 0.08 || trimmed.end_with?(';', '{', '}')

    :prose
  end

  # Whole-text JSON gate: a document that parses as JSON is json all the way
  # down. Mirrors isValidJson() (JSON.parse in a begin/rescue);
  # empty/whitespace text is not.
  def valid_json?(text)
    return false if text.strip.empty?

    begin
      JSON.parse(text)
    rescue JSON::ParserError
      false
    else
      true
    end
  end

  # Token count of `text` under auto content detection — exactly the slice of
  # estimateTokens() the planner consumes (`.tokens`): per non-empty line,
  # max(1, round(utf16_len / chars_per_token)). Framing tokens are the
  # caller's job.
  def estimate_tokens(text)
    # AUTO + whole-text JSON: json's 3 chars/token rate applies to every
    # line, not just the reported content type.
    whole_text_json = valid_json?(text)
    tokens = 0
    # Split on LF or CRLF (text.split(/\r?\n/)); a lone CR is NOT a break.
    text.split(/\r?\n/, -1).each do |line|
      next if line.strip.empty?

      type = whole_text_json ? :json : detect_line_type(line)
      tokens += [1, js_round(utf16_len(line) / CHARS_PER_TOKEN[type])].max
    end
    tokens
  end

  # Sum of per-section token estimates (framing tokens are the caller's job).
  # Mirrors inputTokenTotal() in the TS lib.
  def input_token_total(sections)
    sections.sum { |s| estimate_tokens(s.text) }
  end

  # 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.
  def plan_window(sections, model_id, output_reserve: 0, models: [])
    input_tokens = input_token_total(sections)
    # Fit check inlined from fitsWindow() in src/lib/ai/models.ts.
    m = models.find { |model| model.id == model_id }
    return nil if m.nil?

    free = m.context_window - input_tokens
    WindowPlan.new(
      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
    )
  end

  # Plan against several models; unknown ids are dropped from the result.
  # Mirrors planAll() in the TS lib.
  def plan_all(sections, model_ids, output_reserve: 0, models: [])
    model_ids.filter_map { |id| plan_window(sections, id, output_reserve: output_reserve, models: models) }
  end
end

# ---------- tests (showcase-only; the canonical suite lives in src/lib) ----------
if $PROGRAM_NAME == __FILE__
  # A 1600-char single line of 'a' is pure prose: 1600 / 4 = 400 tokens.
  line_a = 'a' * 1600
  two_sections = [
    ContextWindowPlanner.sec('sys', line_a),
    ContextWindowPlanner.sec('docs', line_a)
  ]
  cp = ContextWindowPlanner

  # input totals
  raise 'input totals' unless cp.input_token_total(two_sections) == 800
  raise unless cp.input_token_total([]) == 0
  raise unless cp.input_token_total([cp.sec('sys', '')]) == 0

  # plans two 400-token sections against beta-pro
  plan = cp.plan_window(two_sections, 'beta-pro', models: cp::SAMPLE_MODELS)
  raise unless plan.id == 'beta-pro'
  raise unless plan.input_tokens == 800
  raise unless plan.context_window == 1_000_000
  raise unless plan.free == 999_200
  raise unless plan.fits
  raise unless plan.output_reserve_ok
  raise unless plan.max_output == 10_000

  # reserve larger than free leaves raw fit true
  plan = cp.plan_window(two_sections, 'beta-pro', output_reserve: 1_000_000, models: cp::SAMPLE_MODELS)
  raise unless plan.fits
  raise if plan.output_reserve_ok

  # reserve exactly equal to free is ok
  plan = cp.plan_window(two_sections, 'beta-pro', output_reserve: 999_200, models: cp::SAMPLE_MODELS)
  raise unless plan.output_reserve_ok

  # smaller window leaves 199,200 free
  plan = cp.plan_window(two_sections, 'alpha-mini', models: cp::SAMPLE_MODELS)
  raise unless plan.context_window == 200_000
  raise unless plan.free == 199_200
  raise unless plan.fits

  # unknown model id returns nil
  raise if cp.plan_window(two_sections, 'ghost', models: cp::SAMPLE_MODELS)

  # no sections: full window free
  plan = cp.plan_window([], 'beta-pro', models: cp::SAMPLE_MODELS)
  raise unless plan.input_tokens.zero?
  raise unless plan.free == 1_000_000
  raise unless plan.fits

  # overflow: fits false, reserve false
  plan = cp.plan_window([cp.sec('big', 'z' * 4_400_000)], 'beta-pro', models: cp::SAMPLE_MODELS)
  raise unless plan.input_tokens == 1_100_000
  raise unless plan.free == -100_000
  raise if plan.fits
  raise if plan.output_reserve_ok

  # plan_all drops unknown ids and keeps order
  plans = cp.plan_all(two_sections, %w[beta-pro alpha-mini ghost], models: cp::SAMPLE_MODELS)
  raise unless plans.length == 2
  raise unless plans[0].id == 'beta-pro'
  raise unless plans[1].id == 'alpha-mini'
  raise unless plans[1].free == 199_200
  raise unless cp.plan_all(two_sections, [], models: cp::SAMPLE_MODELS).empty?

  puts 'context-window-planner (Ruby): all tests passed'
end