Context Window Planner — Go source

// Package contextwindowplanner is the Go twin of CosmoDev's
// src/lib/contextPlanner.ts (dual source: the web lib is TypeScript, the CLI
// lib is Go — kept in lock-step). Pure + deterministic, never panics. The
// table-driven tests in context-window-planner_test.go share vectors with
// src/lib/contextPlanner.test.ts so the two implementations are held to the
// same contract.
//
// Token counting is delegated to cosmodev/tokenestimator and window math to
// cosmodev/aimodels (FitsWindow) — this package only combines the two into a
// plan, exactly like the TS lib.
package contextwindowplanner

import (
	"cosmodev/aimodels"
	"cosmodev/token-estimator"
)

// PlanSection is one labeled block of the prompt (system / docs / history / …).
// Mirrors PlanSection in the TS lib.
type PlanSection struct {
	Label string
	Text  string
}

// WindowPlan mirrors the WindowPlan interface of the TS lib.
type WindowPlan struct {
	ID            string
	InputTokens   int // sum of per-section estimates (framing is the caller's job)
	ContextWindow int
	// Free is context tokens left after the request; negative on overflow.
	Free int
	// Fits is the raw fit: Free >= 0.
	Fits bool
	// OutputReserveOk is the room for the output reserve: Free >= reserve.
	OutputReserveOk bool
	// MaxOutput is informational: the model's output cap.
	MaxOutput int
}

// InputTokenTotal sums the per-section token estimates. It is the twin of
// inputTokenTotal() in src/lib/contextPlanner.ts and must agree with it on
// every shared vector.
func InputTokenTotal(sections []PlanSection) int {
	total := 0
	for _, s := range sections {
		total += tokenestimator.EstimateTokens(s.Text, nil).Tokens
	}
	return total
}

// PlanWindow plans one section set against one model's context window. Nil
// (with nil error) for an unknown model id — window math is FitsWindow's,
// never re-derived here. The error only propagates a snapshot load failure
// from aimodels (models nil → embedded catalog); it mirrors the TS lib's
// undefined return with nil.
func PlanWindow(sections []PlanSection, modelID string, outputReserve int, models []aimodels.Model) (*WindowPlan, error) {
	inputTokens := InputTokenTotal(sections)
	fit, err := aimodels.FitsWindow(modelID, inputTokens, models)
	if err != nil {
		return nil, err
	}
	if fit == nil {
		return nil, nil
	}
	return &WindowPlan{
		ID:              modelID,
		InputTokens:     inputTokens,
		ContextWindow:   fit.ContextWindow,
		Free:            fit.Free,
		Fits:            fit.Fits,
		OutputReserveOk: fit.Free >= outputReserve,
		MaxOutput:       fit.Model.MaxOutput,
	}, nil
}

// PlanAll plans against several models; unknown ids are dropped from the
// result, keeping known plans in order. Twin of planAll() in the TS lib.
func PlanAll(sections []PlanSection, modelIDs []string, outputReserve int, models []aimodels.Model) ([]WindowPlan, error) {
	plans := make([]WindowPlan, 0, len(modelIDs))
	for _, id := range modelIDs {
		p, err := PlanWindow(sections, id, outputReserve, models)
		if err != nil {
			return nil, err
		}
		if p == nil {
			continue
		}
		plans = append(plans, *p)
	}
	return plans, nil
}