Conversation Pruner — C++ source

// conversation-pruner — C++ port: compute a deterministic pruning plan for a
// token-budgeted chat history. C++17, stdlib only. Port of
// src/lib/conversationPruner.ts — same decisions as this dir's javascript.js.
#include <cmath>
#include <cstdio>
#include <stdexcept>
#include <string>
#include <vector>

// Summary compression model: fixed framing tokens, then 10% of the folded
// content's tokens, rounded up.
inline constexpr long long kSummaryFixedTokens = 60;
inline constexpr double kSummaryRatio = 0.1;

enum class ChatRole { System, User, Assistant, Tool };

// One message plus its prompt-side token count. Pinned messages are never
// dropped or summarized (the TypeScript's optional flag defaults to false).
struct ConversationMessage {
    ChatRole role;
    std::string content;
    long long tokens;
    bool pinned = false;
};

enum class PruneAction { Keep, Summarize, Drop };

// The plan's verdict on one message.
struct PruneDecision {
    std::size_t index;
    ChatRole role;
    PruneAction action;
    long long tokens;
};

// The full pruning plan for one conversation.
struct PrunePlan {
    std::vector<PruneDecision> decisions;
    long long keptTokens = 0;
    long long summarizedTokens = 0;
    long long droppedTokens = 0;
    long long summaryCostTokens = 0; // tokens the summary placeholder costs
    long long projectedTokens = 0;
    bool fitsBudget = false;
    std::vector<std::string> warnings;
};

const char *roleName(ChatRole r) {
    switch (r) {
        case ChatRole::System: return "system";
        case ChatRole::User: return "user";
        case ChatRole::Assistant: return "assistant";
        case ChatRole::Tool: return "tool";
    }
    return "?";
}

const char *actionName(PruneAction a) {
    switch (a) {
        case PruneAction::Keep: return "keep";
        case PruneAction::Summarize: return "summarize";
        case PruneAction::Drop: return "drop";
    }
    return "?";
}

// Group an integer with en-US commas: 1234567 -> "1,234,567", matching
// toLocaleString('en-US').
std::string thousands(long long v) {
    std::string digits = std::to_string(v < 0 ? -v : v);
    const std::size_t len = digits.size();
    std::string out;
    out.reserve(len + len / 3 + 1);
    if (v < 0) out.push_back('-');
    for (std::size_t i = 0; i < len; i++) {
        if (i > 0 && (len - i) % 3 == 0) out.push_back(',');
        out.push_back(digits[i]);
    }
    return out;
}

/**
 * Compute the pruning plan. Throws std::invalid_argument — the TypeScript
 * RangeError — on a negative budget or any negative per-message token count.
 */
PrunePlan planPrune(const std::vector<ConversationMessage> &messages, long long budgetTokens) {
    std::vector<std::string> warnings;
    if (budgetTokens < 0) throw std::invalid_argument("budgetTokens must be >= 0");
    for (const auto &m : messages) {
        if (m.tokens < 0) throw std::invalid_argument("message tokens must be >= 0");
    }

    const std::size_t n = messages.size();
    long long lastUser = -1;
    for (std::size_t k = n; k-- > 0;) {
        if (messages[k].role == ChatRole::User) {
            lastUser = static_cast<long long>(k);
            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).
    std::vector<unsigned char> prot(n, 0);
    for (std::size_t i = 0; i < n; i++) {
        if (messages[i].role == ChatRole::System || messages[i].pinned) prot[i] = 1;
    }
    if (n > 0) prot[0] = 1;
    long long firstTurn = -1;
    for (std::size_t i = 0; i < n; i++) {
        if (messages[i].role != ChatRole::System) {
            firstTurn = static_cast<long long>(i);
            break;
        }
    }
    if (firstTurn != -1) prot[static_cast<std::size_t>(firstTurn)] = 1;
    std::size_t tailStart = static_cast<std::size_t>(
        lastUser == -1 ? static_cast<long long>(n > 0 ? n - 1 : 0) : lastUser);
    for (std::size_t i = tailStart; i < n; i++) prot[i] = 1;

    long long protectedTokens = 0;
    for (std::size_t i = 0; i < n; i++) {
        if (prot[i]) protectedTokens += messages[i].tokens;
    }
    if (protectedTokens > budgetTokens) {
        warnings.push_back(
            "Protected messages alone are " + thousands(protectedTokens) + " tokens against a " +
            thousands(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.
    std::vector<PruneAction> actions(n, PruneAction::Drop);
    for (std::size_t i = 0; i < n; i++) {
        if (prot[i]) actions[i] = PruneAction::Keep;
    }
    long long used = protectedTokens;
    for (std::size_t k = n; k-- > 0;) {
        if (actions[k] != PruneAction::Drop) continue;
        if (used + messages[k].tokens <= budgetTokens) {
            actions[k] = PruneAction::Keep;
            used += messages[k].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.
    std::vector<std::size_t> summarizeIdx;
    for (std::size_t i = 0; i < n; i++) {
        if (actions[i] == PruneAction::Drop && !prot[i]) summarizeIdx.push_back(i);
    }
    long long summarizeTokens = 0;
    for (std::size_t i : summarizeIdx) summarizeTokens += messages[i].tokens;
    long long attemptedSummaryCost =
        summarizeIdx.empty()
            ? 0
            : kSummaryFixedTokens +
                  static_cast<long long>(std::ceil(static_cast<double>(summarizeTokens) * kSummaryRatio));

    // The summary only costs anything when it is actually applied — otherwise
    // those turns drop and cost zero.
    long long summaryCost = 0;
    if (attemptedSummaryCost > 0 && used + attemptedSummaryCost <= budgetTokens) {
        for (std::size_t i : summarizeIdx) actions[i] = PruneAction::Summarize;
        summaryCost = attemptedSummaryCost;
        // Mirrors the TypeScript; `used` is not consulted again afterwards.
        used += summaryCost;
        (void)used;
    } else if (attemptedSummaryCost > 0) {
        warnings.push_back("Even the compressed summary (" + thousands(attemptedSummaryCost) +
                           " tokens) does not fit the remaining budget — the oldest turns are "
                           "dropped instead.");
    }

    std::vector<PruneDecision> decisions;
    decisions.reserve(n);
    for (std::size_t i = 0; i < n; i++) {
        decisions.push_back(
            PruneDecision{i, messages[i].role, actions[i], messages[i].tokens});
    }

    long long keptTokens = 0, droppedTokens = 0, foldedTokens = 0;
    for (const auto &d : decisions) {
        if (d.action == PruneAction::Keep) keptTokens += d.tokens;
        else if (d.action == PruneAction::Drop) droppedTokens += d.tokens;
        else foldedTokens += d.tokens;
    }

    PrunePlan plan;
    plan.decisions = std::move(decisions);
    plan.keptTokens = keptTokens;
    plan.summarizedTokens = foldedTokens;
    plan.droppedTokens = droppedTokens;
    plan.summaryCostTokens = summaryCost;
    plan.projectedTokens = keptTokens + summaryCost;
    plan.fitsBudget = keptTokens + summaryCost <= budgetTokens;
    plan.warnings = std::move(warnings);
    return plan;
}

// Human-readable one-line summary of a plan.
std::string describePrune(const PrunePlan &plan) {
    if (!plan.fitsBudget) {
        return "Does not fit: " + thousands(plan.projectedTokens) +
               " tokens projected against the budget.";
    }
    std::vector<std::string> parts{thousands(plan.keptTokens) + " kept"};
    if (plan.summarizedTokens > 0) {
        parts.push_back(thousands(plan.summarizedTokens) + " folded into a " +
                        thousands(plan.summaryCostTokens) + "-token summary");
    }
    if (plan.droppedTokens > 0) {
        parts.push_back(thousands(plan.droppedTokens) + " dropped");
    }
    std::string out;
    for (std::size_t i = 0; i < parts.size(); i++) {
        if (i > 0) out += " · ";
        out += parts[i];
    }
    return out + " — fits the budget.";
}

int main() {
    std::vector<ConversationMessage> msgs = {
        {ChatRole::System, "You are CosmoDev's assistant.", 80, false},
        {ChatRole::User, "Please review my repo layout.", 250, false},
        {ChatRole::Assistant, "(long review of the tree)", 900, false},
        {ChatRole::Tool, "git status output", 400, false},
        {ChatRole::Assistant, "summary of the findings", 700, false},
        {ChatRole::User, "Now prune this conversation.", 350, false},
    };

    PrunePlan plan = planPrune(msgs, 1000);
    std::printf("budget 1,000:\n");
    for (const auto &d : plan.decisions) {
        std::printf("  %zu %-9s %-9s %5lld\n", d.index, roleName(d.role), actionName(d.action),
                    d.tokens);
    }
    for (const auto &w : plan.warnings) std::printf("  warning: %s\n", w.c_str());
    std::printf("  %s\n", describePrune(plan).c_str());

    // Squeezed budget: protected turns alone overflow, and even the
    // compressed summary does not fit.
    plan = planPrune(msgs, 50);
    std::printf("budget 50:\n");
    for (const auto &w : plan.warnings) std::printf("  warning: %s\n", w.c_str());
    std::printf("  %s\n", describePrune(plan).c_str());
    return 0;
}