Conversation Pruner — Java source

// Conversation Pruner — compute a deterministic pruning plan for a
// token-budgeted chat history.
//
// Language: Java (17+, zero dependencies)
// Port of src/lib/conversationPruner.ts (the canonical TypeScript
//           implementation) for the CosmoDev polyglot showcase
//           (slug: conversation-pruner).
//
// Given per-message token counts and a context budget, decide which messages
// to keep verbatim, which to fold into one running summary, and which to drop
// outright — protecting system messages, pinned turns, the first turn, and
// the current (last user) request.

import java.util.ArrayList;
import java.util.List;
import java.util.Locale;

public final class ConversationPruner {

    /** Summary compression model: fixed framing tokens. */
    public static final int SUMMARY_FIXED_TOKENS = 60;

    /** Summary compression model: share of the folded content. */
    public static final double SUMMARY_RATIO = 0.1;

    public enum ChatRole { SYSTEM, USER, ASSISTANT, TOOL }

    public enum PruneAction { KEEP, SUMMARIZE, DROP }

    /** One chat message with its caller-supplied token count. */
    public record ConversationMessage(ChatRole role, String content, int tokens, boolean pinned) {
        public ConversationMessage(ChatRole role, String content, int tokens) {
            this(role, content, tokens, false);
        }
    }

    /** One per-message decision. */
    public record PruneDecision(int index, ChatRole role, PruneAction action, int tokens) {}

    /** The full plan: decisions, tallies, projection, warnings. */
    public record PrunePlan(
            List<PruneDecision> decisions,
            int keptTokens,
            int summarizedTokens,
            int droppedTokens,
            int summaryCostTokens,
            int projectedTokens,
            boolean fitsBudget,
            List<String> warnings) {}

    /**
     * Compute the pruning plan for {@code messages} under {@code budgetTokens}.
     * Throws {@link IllegalArgumentException} on a negative budget or any
     * negative per-message token count.
     */
    public static PrunePlan planPrune(List<ConversationMessage> messages, int budgetTokens) {
        List<String> warnings = new ArrayList<>();
        if (budgetTokens < 0)
            throw new IllegalArgumentException("budgetTokens must be >= 0");
        for (ConversationMessage m : messages)
            if (m.tokens() < 0)
                throw new IllegalArgumentException("message tokens must be >= 0");

        int n = messages.size();
        int lastUser = -1;
        for (int i = n - 1; i >= 0; i--)
            if (messages.get(i).role() == ChatRole.USER) { lastUser = i; break; }

        // Untouchable: every system message, pinned messages, the first turn
        // (the opening user request), and the current request (the last user
        // message and everything after it).
        java.util.Set<Integer> protectedIdx = new java.util.HashSet<>();
        for (int i = 0; i < n; i++)
            if (messages.get(i).role() == ChatRole.SYSTEM || messages.get(i).pinned())
                protectedIdx.add(i);
        if (n > 0) protectedIdx.add(0);
        int firstTurn = -1;
        for (int i = 0; i < n; i++)
            if (messages.get(i).role() != ChatRole.SYSTEM) { firstTurn = i; break; }
        if (firstTurn != -1) protectedIdx.add(firstTurn);
        for (int i = Math.max(lastUser == -1 ? n - 1 : lastUser, 0); i < n; i++)
            protectedIdx.add(i);

        int protectedTokens = 0;
        for (int i : protectedIdx) protectedTokens += messages.get(i).tokens();
        if (protectedTokens > budgetTokens)
            warnings.add(String.format(Locale.US,
                    "Protected messages alone are %,d tokens against a %,d budget — " +
                    "raise the budget (or reserve less for the reply) before pruning anything else.",
                    protectedTokens, budgetTokens));

        // Fill the remaining budget newest-to-oldest through the middle.
        PruneAction[] actions = new PruneAction[n];
        java.util.Arrays.fill(actions, PruneAction.DROP);
        for (int i : protectedIdx) actions[i] = PruneAction.KEEP;
        int used = protectedTokens;
        for (int i = n - 1; i >= 0; i--) {
            if (actions[i] != PruneAction.DROP) continue;
            if (used + messages.get(i).tokens() <= budgetTokens) {
                actions[i] = PruneAction.KEEP;
                used += messages.get(i).tokens();
            } else break; // oldest-unfilled remain drop/summarize candidates
        }

        // Everything still 'drop' in the middle folds into ONE running summary
        // when the compressed form fits where the raw turns did not.
        List<Integer> summarizeIdx = new ArrayList<>();
        for (int i = 0; i < n; i++)
            if (actions[i] == PruneAction.DROP && !protectedIdx.contains(i))
                summarizeIdx.add(i);
        int summarizeTokens = 0;
        for (int i : summarizeIdx) summarizeTokens += messages.get(i).tokens();
        int attempted = summarizeIdx.isEmpty() ? 0
                : SUMMARY_FIXED_TOKENS + (int) Math.ceil(summarizeTokens * SUMMARY_RATIO);

        // The summary only costs anything when it is actually applied.
        int summaryCost = 0;
        if (attempted > 0 && used + attempted <= budgetTokens) {
            for (int i : summarizeIdx) actions[i] = PruneAction.SUMMARIZE;
            summaryCost = attempted;
            used += summaryCost;
        } else if (attempted > 0)
            warnings.add(String.format(Locale.US,
                    "Even the compressed summary (%,d tokens) does not fit the remaining " +
                    "budget — the oldest turns are dropped instead.", attempted));

        List<PruneDecision> decisions = new ArrayList<>(n);
        for (int i = 0; i < n; i++)
            decisions.add(new PruneDecision(i, messages.get(i).role(), actions[i],
                    messages.get(i).tokens()));

        int kept = 0, dropped = 0, folded = 0;
        for (PruneDecision d : decisions) {
            if (d.action() == PruneAction.KEEP) kept += d.tokens();
            else if (d.action() == PruneAction.DROP) dropped += d.tokens();
            else folded += d.tokens();
        }

        return new PrunePlan(decisions, kept, folded, dropped, summaryCost,
                kept + summaryCost, kept + summaryCost <= budgetTokens, warnings);
    }

    /** Human-readable one-line summary of a plan. */
    public static String describePrune(PrunePlan plan) {
        if (!plan.fitsBudget())
            return String.format(Locale.US, "Does not fit: %,d tokens projected against the budget.",
                    plan.projectedTokens());
        List<String> parts = new ArrayList<>();
        parts.add(String.format(Locale.US, "%,d kept", plan.keptTokens()));
        if (plan.summarizedTokens() > 0)
            parts.add(String.format(Locale.US, "%,d folded into a %,d-token summary",
                    plan.summarizedTokens(), plan.summaryCostTokens()));
        if (plan.droppedTokens() > 0)
            parts.add(String.format(Locale.US, "%,d dropped", plan.droppedTokens()));
        return String.join(" · ", parts) + " — fits the budget.";
    }

    private ConversationPruner() {}
}