Conversation Pruner — C# source

// Conversation Pruner — compute a deterministic pruning plan for a
// token-budgeted chat history.
//
// Language: C# (.NET 8, 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.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace CosmoDev.ConversationPruner;

/// <summary>Chat roles recognized by the pruner.</summary>
public enum ChatRole { System, User, Assistant, Tool }
public enum PruneAction { Keep, Summarize, Drop }

public sealed record ConversationMessage(
    ChatRole Role,
    string Content,
    int Tokens,
    bool Pinned = false);

public sealed record PruneDecision(
    int Index,
    ChatRole Role,
    PruneAction Action,
    int Tokens);

public sealed record PrunePlan(
    IReadOnlyList<PruneDecision> Decisions,
    int KeptTokens,
    int SummarizedTokens,
    int DroppedTokens,
    int SummaryCostTokens,
    int ProjectedTokens,
    bool FitsBudget,
    IReadOnlyList<string> Warnings);

public static class ConversationPruner
{
    /// <summary>Summary compression model: fixed framing tokens.</summary>
    public const int SUMMARY_FIXED_TOKENS = 60;

    /// <summary>Summary compression model: share of the folded content.</summary>
    public const double SUMMARY_RATIO = 0.1;

    /// <summary>
    /// Compute the pruning plan for <paramref name="messages"/> under
    /// <paramref name="budgetTokens"/>. Throws <see cref="ArgumentOutOfRangeException"/>
    /// on a negative budget or any negative per-message token count.
    /// </summary>
    public static PrunePlan PlanPrune(
        IReadOnlyList<ConversationMessage> messages, int budgetTokens)
    {
        var warnings = new List<string>();
        if (budgetTokens < 0)
            throw new ArgumentOutOfRangeException(nameof(budgetTokens),
                "budgetTokens must be >= 0");
        if (messages.Any(m => m.Tokens < 0))
            throw new ArgumentException("message tokens must be >= 0");

        int n = messages.Count;
        int LastIndexWithRole(ChatRole role)
        {
            for (int i = n - 1; i >= 0; i--) if (messages[i].Role == role) return i;
            return -1;
        }
        int lastUser = LastIndexWithRole(ChatRole.User);

        // 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).
        var protectedIdx = new HashSet<int>();
        for (int i = 0; i < n; i++)
            if (messages[i].Role == ChatRole.System || messages[i].Pinned)
                protectedIdx.Add(i);
        if (n > 0) protectedIdx.Add(0);
        int firstTurn = Enumerable.Range(0, n).FirstOrDefault(i => messages[i].Role != ChatRole.System, -1);
        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 = protectedIdx.Sum(i => messages[i].Tokens);
        if (protectedTokens > budgetTokens)
            warnings.Add(
                $"Protected messages alone are {Fmt(protectedTokens)} tokens against a " +
                $"{Fmt(budgetTokens)} budget — raise the budget (or reserve less for the reply) before pruning anything else.");

        // Fill the remaining budget newest-to-oldest through the middle.
        var actions = new PruneAction[n];
        Array.Fill(actions, PruneAction.Drop);
        foreach (int i in 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[i].Tokens <= budgetTokens)
            {
                actions[i] = PruneAction.Keep;
                used += messages[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.
        var summarizeIdx = Enumerable.Range(0, n)
            .Where(i => actions[i] == PruneAction.Drop && !protectedIdx.Contains(i))
            .ToList();
        int summarizeTokens = summarizeIdx.Sum(i => messages[i].Tokens);
        int attempted =
            summarizeIdx.Count > 0
                ? SUMMARY_FIXED_TOKENS + (int)Math.Ceiling(summarizeTokens * SUMMARY_RATIO)
                : 0;

        // The summary only costs anything when it is actually applied.
        int summaryCost = 0;
        if (attempted > 0 && used + attempted <= budgetTokens)
        {
            foreach (int i in summarizeIdx) actions[i] = PruneAction.Summarize;
            summaryCost = attempted;
            used += summaryCost;
        }
        else if (attempted > 0)
        {
            warnings.Add(
                $"Even the compressed summary ({Fmt(attempted)} tokens) does not fit " +
                "the remaining budget — the oldest turns are dropped instead.");
        }

        var decisions = messages.Select((m, i) =>
            new PruneDecision(i, m.Role, actions[i], m.Tokens)).ToList();

        int kept = 0, dropped = 0, folded = 0;
        foreach (var d in 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);
    }

    /// <summary>Human-readable one-line summary of a plan.</summary>
    public static string DescribePrune(PrunePlan plan)
    {
        if (!plan.FitsBudget)
            return $"Does not fit: {Fmt(plan.ProjectedTokens)} tokens projected against the budget.";
        var parts = new List<string> { $"{Fmt(plan.KeptTokens)} kept" };
        if (plan.SummarizedTokens > 0)
            parts.Add($"{Fmt(plan.SummarizedTokens)} folded into a {Fmt(plan.SummaryCostTokens)}-token summary");
        if (plan.DroppedTokens > 0) parts.Add($"{Fmt(plan.DroppedTokens)} dropped");
        return $"{string.Join(" · ", parts)} — fits the budget.";
    }

    private static string Fmt(int v) => v.ToString("N0");
}