Conversation Pruner — Go source

// Package conversationpruner is the Go twin of CosmoDev's
// src/lib/conversationPruner.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 conversationpruner_test.go share vectors with
// src/lib/conversationPruner.test.ts so the two implementations are held to
// the same contract.
//
// Given a conversation with per-message token counts and a context budget,
// PlanPrune computes a deterministic pruning plan: which messages to keep,
// which to fold into a running summary, and which to drop outright —
// protecting system messages, pinned turns, and the current (last user)
// request.
package conversationpruner

import (
	"fmt"
	"strconv"
	"strings"
)

// ChatRole mirrors the TS ChatRole union ('system' | 'user' | 'assistant' | 'tool').
type ChatRole string

const (
	RoleSystem    ChatRole = "system"
	RoleUser      ChatRole = "user"
	RoleAssistant ChatRole = "assistant"
	RoleTool      ChatRole = "tool"
)

// ConversationMessage mirrors the TS ConversationMessage interface.
type ConversationMessage struct {
	Role ChatRole
	// Content is the message body (never inspected by the planner).
	Content string
	// Tokens is the token count for this message (prompt-side framing
	// included by the caller).
	Tokens int
	// Pinned messages are never dropped or summarized. (TS `pinned?: boolean`
	// → zero value false = unpinned.)
	Pinned bool
}

// Options mirrors the TS PruneOptions interface; the zero value
// (Options{BudgetTokens: 0}) is valid, like every other non-negative budget.
type Options struct {
	// BudgetTokens is the total tokens available for the history (context
	// window minus reserved reply). Negative values are rejected, mirroring
	// the TS RangeError.
	BudgetTokens int
}

// PruneAction mirrors the TS PruneAction union.
type PruneAction string

const (
	ActionKeep      PruneAction = "keep"
	ActionSummarize PruneAction = "summarize"
	ActionDrop      PruneAction = "drop"
)

// PruneDecision mirrors the TS PruneDecision interface.
type PruneDecision struct {
	Index  int
	Role   ChatRole
	Action PruneAction
	Tokens int
}

// PrunePlan mirrors the TS PrunePlan interface.
type PrunePlan struct {
	Decisions        []PruneDecision
	KeptTokens       int
	SummarizedTokens int
	DroppedTokens    int
	// SummaryCostTokens is what the summary placeholder itself will cost in
	// the prompt. An unapplied (rejected) summary costs zero.
	SummaryCostTokens int
	ProjectedTokens   int
	FitsBudget        bool
	Warnings          []string
}

// Summary compression model: fixed framing + 10% of the folded content.
// Mirrors SUMMARY_FIXED_TOKENS / SUMMARY_RATIO in the TS lib.
const (
	SummaryFixedTokens = 60
	SummaryRatio       = 0.1
)

// PlanPrune computes the pruning plan for a conversation. It is the Go twin of
// planPrune() in src/lib/conversationPruner.ts and must agree with it on every
// shared vector. Where the TS lib throws a RangeError (negative budget or
// negative per-message tokens), the Go twin returns an error and a zero-value
// plan; it never panics.
func PlanPrune(messages []ConversationMessage, opts Options) (PrunePlan, error) {
	if opts.BudgetTokens < 0 {
		return PrunePlan{}, fmt.Errorf("budgetTokens must be >= 0")
	}
	for _, msg := range messages {
		if msg.Tokens < 0 {
			return PrunePlan{}, fmt.Errorf("message tokens must be >= 0")
		}
	}

	warnings := []string{}
	n := len(messages)

	lastUser := -1
	for i := n - 1; i >= 0; i-- {
		if messages[i].Role == RoleUser {
			lastUser = i
			break
		}
	}

	// Untouchable: every system message, pinned messages, the first turn (the
	// opening user request that anchors the conversation), and the current
	// request (the last user message and everything after it).
	protected := make([]bool, n)
	for i, msg := range messages {
		if msg.Role == RoleSystem || msg.Pinned {
			protected[i] = true
		}
	}
	if n > 0 {
		protected[0] = true
	}
	for i, msg := range messages {
		if msg.Role != RoleSystem {
			protected[i] = true
			break
		}
	}
	start := n - 1 // no user message at all → only the final index anchors the tail
	if lastUser != -1 {
		start = lastUser
	}
	if start < 0 {
		start = 0
	}
	for i := start; i < n; i++ {
		protected[i] = true
	}

	// (reduce seeded: an empty conversation has no protected indexes at all)
	protectedTokens := 0
	for i, p := range protected {
		if p {
			protectedTokens += messages[i].Tokens
		}
	}
	if protectedTokens > opts.BudgetTokens {
		warnings = append(warnings, fmt.Sprintf(
			"Protected messages alone are %s tokens against a %s budget — raise the budget (or reserve less for the reply) before pruning anything else.",
			comma(protectedTokens), comma(opts.BudgetTokens)))
	}

	// Fill the remaining budget newest-to-oldest through the middle.
	actions := make([]PruneAction, n)
	for i := range actions {
		actions[i] = ActionDrop
	}
	for i, p := range protected {
		if p {
			actions[i] = ActionKeep
		}
	}
	used := protectedTokens
	for i := n - 1; i >= 0; i-- {
		if actions[i] != ActionDrop {
			continue
		}
		if used+messages[i].Tokens <= opts.BudgetTokens {
			actions[i] = ActionKeep
			used += messages[i].Tokens
		} else {
			break // oldest-unfilled remain drop/summarize candidates, newest first stopped
		}
	}

	// Everything still 'drop' in the middle folds into ONE running summary when
	// the compressed form fits where the raw turns did not.
	var summarizeIdx []int
	for i, a := range actions {
		if a == ActionDrop && !protected[i] {
			summarizeIdx = append(summarizeIdx, i)
		}
	}
	summarizeTokens := 0
	for _, i := range summarizeIdx {
		summarizeTokens += messages[i].Tokens
	}
	attemptedSummaryCost := 0
	if len(summarizeIdx) > 0 {
		attemptedSummaryCost = SummaryFixedTokens + ceilSummaryShare(summarizeTokens)
	}

	// The summary only costs anything when it is actually applied — otherwise
	// those turns drop and cost zero.
	summaryCost := 0
	if attemptedSummaryCost > 0 && used+attemptedSummaryCost <= opts.BudgetTokens {
		for _, i := range summarizeIdx {
			actions[i] = ActionSummarize
		}
		summaryCost = attemptedSummaryCost
		used += summaryCost
	} else if attemptedSummaryCost > 0 {
		warnings = append(warnings, fmt.Sprintf(
			"Even the compressed summary (%s tokens) does not fit the remaining budget — the oldest turns are dropped instead.",
			comma(attemptedSummaryCost)))
	}

	decisions := make([]PruneDecision, 0, n)
	for i, msg := range messages {
		decisions = append(decisions, PruneDecision{
			Index:  i,
			Role:   msg.Role,
			Action: actions[i],
			Tokens: msg.Tokens,
		})
	}

	var keptTokens, droppedTokens, foldedTokens int
	for _, d := range decisions {
		switch d.Action {
		case ActionKeep:
			keptTokens += d.Tokens
		case ActionDrop:
			droppedTokens += d.Tokens
		default: // ActionSummarize
			foldedTokens += d.Tokens
		}
	}

	return PrunePlan{
		Decisions:         decisions,
		KeptTokens:        keptTokens,
		SummarizedTokens:  foldedTokens,
		DroppedTokens:     droppedTokens,
		SummaryCostTokens: summaryCost,
		ProjectedTokens:   keptTokens + summaryCost,
		FitsBudget:        keptTokens+summaryCost <= opts.BudgetTokens,
		Warnings:          warnings,
	}, nil
}

// DescribePrune renders a human-readable one-line summary of a plan. It is the
// Go twin of describePrune() in src/lib/conversationPruner.ts and must agree
// with it on every shared vector.
func DescribePrune(plan PrunePlan) string {
	if !plan.FitsBudget {
		return fmt.Sprintf(
			"Does not fit: %s tokens projected against the budget.",
			comma(plan.ProjectedTokens))
	}
	parts := []string{fmt.Sprintf("%s kept", comma(plan.KeptTokens))}
	if plan.SummarizedTokens > 0 {
		parts = append(parts, fmt.Sprintf(
			"%s folded into a %s-token summary",
			comma(plan.SummarizedTokens), comma(plan.SummaryCostTokens)))
	}
	if plan.DroppedTokens > 0 {
		parts = append(parts, fmt.Sprintf("%s dropped", comma(plan.DroppedTokens)))
	}
	return strings.Join(parts, " · ") + " — fits the budget."
}

// ceilSummaryShare returns ceil(tokens * SummaryRatio) — the variable share of
// the summary cost — computed in exact integer arithmetic. SummaryRatio is 0.1,
// and (t+9)/10 == ceil(t/10) == the TS Math.ceil(t * 0.1) for non-negative
// integers, avoiding float rounding surprises at multiples of 10.
func ceilSummaryShare(tokens int) int {
	if tokens <= 0 {
		return 0
	}
	return (tokens + 9) / 10
}

// comma formats n with en-US thousands separators, mirroring the TS lib's
// toLocaleString('en-US') in warnings and descriptions.
func comma(n int) string {
	neg := n < 0
	digits := strconv.Itoa(n)
	if neg {
		digits = digits[1:]
	}
	var groups []string
	for len(digits) > 3 {
		groups = append([]string{digits[len(digits)-3:]}, groups...)
		digits = digits[:len(digits)-3]
	}
	groups = append([]string{digits}, groups...)
	s := strings.Join(groups, ",")
	if neg {
		return "-" + s
	}
	return s
}