Context Window Planner — TypeScript source

// Context-window planning over labeled prompt sections. Token counting is
// delegated to the estimator; window math is delegated to the model accessor
// (src/lib/ai/models.ts) — this module only combines the two into a plan.
import { estimateTokens } from './tokenEstimator';
import { fitsWindow, type AiModel } from './ai/models';

/** One labeled block of the prompt (system / docs / history / …). */
export interface PlanSection {
  label: string;
  text: string;
}

export interface WindowPlan {
  id: string;
  inputTokens: number;
  contextWindow: number;
  /** Context tokens left after the request; negative on overflow. */
  free: number;
  /** Raw fit: free >= 0. */
  fits: boolean;
  /** Room for the output reserve: free >= outputReserve. */
  outputReserveOk: boolean;
  /** Informational: the model's output cap. */
  maxOutput: number;
}

/** Sum of per-section token estimates (framing tokens are the caller's job). */
export function inputTokenTotal(sections: PlanSection[]): number {
  return sections.reduce((n, s) => n + estimateTokens(s.text).tokens, 0);
}

/**
 * Plan one section set against one model's context window. Returns undefined
 * for an unknown model id (window math is fitsWindow's, never re-derived).
 */
export function planWindow(
  sections: PlanSection[],
  modelId: string,
  outputReserve = 0,
  models?: readonly AiModel[],
): WindowPlan | undefined {
  const inputTokens = inputTokenTotal(sections);
  const fit = fitsWindow(modelId, inputTokens, models);
  if (fit === undefined) return undefined;
  return {
    id: modelId,
    inputTokens,
    contextWindow: fit.contextWindow,
    free: fit.free,
    fits: fit.fits,
    outputReserveOk: fit.free >= outputReserve,
    maxOutput: fit.model.maxOutput,
  };
}

/** Plan against several models; unknown ids are dropped from the result. */
export function planAll(
  sections: PlanSection[],
  modelIds: string[],
  outputReserve = 0,
  models?: readonly AiModel[],
): WindowPlan[] {
  return modelIds
    .map((id) => planWindow(sections, id, outputReserve, models))
    .filter((p): p is WindowPlan => p !== undefined);
}