Conversation Pruner — Zig source

// Conversation Pruner — compute a deterministic pruning plan for a
// token-budgeted chat history.
//
// Language: Zig (0.13+, stdlib only)
// 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. All output slices are allocated from the
// caller's allocator; free `warnings` strings and slices when done.

const std = @import("std");

/// Chat roles recognized by the pruner.
pub const Role = enum { system, user, assistant, tool };

/// One per-message decision.
pub const Action = enum { keep, summarize, drop };

/// One chat message with its caller-supplied token count.
pub const Message = struct {
    role: Role,
    content: []const u8,
    tokens: i64,
    /// Pinned messages are never dropped or summarized.
    pinned: bool = false,
};

/// One per-message decision record.
pub const Decision = struct {
    index: usize,
    role: Role,
    action: Action,
    tokens: i64,
};

/// The full plan: decisions, tallies, projection, warnings.
pub const Plan = struct {
    decisions: []Decision,
    kept_tokens: i64,
    summarized_tokens: i64,
    dropped_tokens: i64,
    /// Tokens the summary placeholder itself will cost in the prompt.
    summary_cost_tokens: i64,
    projected_tokens: i64,
    fits_budget: bool,
    /// Each string is allocated from the caller's allocator.
    warnings: [][]const u8,
};

/// Summary compression model: fixed framing tokens.
pub const summary_fixed_tokens: i64 = 60;

/// Summary compression model: share of the folded content.
pub const summary_ratio: f64 = 0.1;

pub const PruneError = error{
    NegativeBudget,
    NegativeTokens,
    OutOfMemory,
};

/// Group a non-negative integer with thousands separators
/// (toLocaleString stand-in). `buf` must hold 16 bytes.
fn commaFormat(buf: []u8, v: i64) []const u8 {
    var tmp: [24]u8 = undefined;
    const digits = std.fmt.bufPrint(&tmp, "{d}", .{v}) catch unreachable;
    var out_len: usize = 0;
    const n = digits.len;
    for (digits, 0..) |ch, i| {
        buf[out_len] = ch;
        out_len += 1;
        const remaining = n - i - 1;
        if (remaining > 0 and remaining % 3 == 0) {
            buf[out_len] = ',';
            out_len += 1;
        }
    }
    return buf[0..out_len];
}

fn warnProtectedTooBig(alloc: std.mem.Allocator, protected: i64, budget: i64) ![]const u8 {
    var b1: [24]u8 = undefined;
    var b2: [24]u8 = undefined;
    return std.fmt.allocPrint(alloc,
        "Protected messages alone are {s} tokens against a {s} budget — raise the budget " ++
            "(or reserve less for the reply) before pruning anything else.",
        .{ commaFormat(&b1, protected), commaFormat(&b2, budget) });
}

fn warnSummaryDoesNotFit(alloc: std.mem.Allocator, attempted: i64) ![]const u8 {
    var b: [24]u8 = undefined;
    return std.fmt.allocPrint(alloc,
        "Even the compressed summary ({s} tokens) does not fit the remaining budget — " ++
            "the oldest turns are dropped instead.",
        .{commaFormat(&b, attempted)});
}

/// Compute the pruning plan for `messages` under `budget_tokens`.
/// Mirrors planPrune() in the TS lib, including the newest-to-oldest fill
/// with its early `break`, and the single-fold summary attempt.
pub fn planPrune(
    alloc: std.mem.Allocator,
    messages: []const Message,
    budget_tokens: i64,
) PruneError!Plan {
    if (budget_tokens < 0) return error.NegativeBudget;
    for (messages) |m| {
        if (m.tokens < 0) return error.NegativeTokens;
    }

    const n = messages.len;
    var warnings: std.ArrayList([]const u8) = .init(alloc);
    errdefer {
        for (warnings.items) |w| alloc.free(w);
        warnings.deinit();
    }

    var last_user: ?usize = null;
    var i: usize = n;
    while (i > 0) {
        i -= 1;
        if (messages[i].role == .user) {
            last_user = 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).
    const protected = try alloc.alloc(bool, n);
    defer alloc.free(protected);
    @memset(protected, false);
    for (messages, 0..) |m, idx| {
        if (m.role == .system or m.pinned) protected[idx] = true;
    }
    if (n > 0) protected[0] = true;
    for (messages, 0..) |m, idx| {
        if (m.role != .system) {
            protected[idx] = true;
            break;
        }
    }
    const tail_from = last_user orelse (if (n > 0) n - 1 else 0);
    var t: usize = tail_from;
    while (t < n) : (t += 1) protected[t] = true;

    var protected_tokens: i64 = 0;
    for (messages, 0..) |m, idx| {
        if (protected[idx]) protected_tokens += m.tokens;
    }
    if (protected_tokens > budget_tokens) {
        try warnings.append(try warnProtectedTooBig(alloc, protected_tokens, budget_tokens));
    }

    // Fill the remaining budget newest-to-oldest through the middle.
    const actions = try alloc.alloc(Action, n);
    defer alloc.free(actions);
    @memset(actions, .drop);
    for (protected, 0..) |p, idx| {
        if (p) actions[idx] = .keep;
    }
    var used = protected_tokens;
    var j: usize = n;
    while (j > 0) {
        j -= 1;
        if (actions[j] != .drop) continue;
        if (used + messages[j].tokens <= budget_tokens) {
            actions[j] = .keep;
            used += messages[j].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 summarize_tokens: i64 = 0;
    var summarize_count: usize = 0;
    for (actions, 0..) |a, idx| {
        if (a == .drop and !protected[idx]) {
            summarize_tokens += messages[idx].tokens;
            summarize_count += 1;
        }
    }
    const attempted: i64 = if (summarize_count > 0)
        summary_fixed_tokens + @as(i64, @intFromFloat(@ceil(@as(f64, @floatFromInt(summarize_tokens)) * summary_ratio)))
    else
        0;

    // The summary only costs anything when it is actually applied.
    var summary_cost: i64 = 0;
    if (attempted > 0 and used + attempted <= budget_tokens) {
        for (actions, 0..) |*a, idx| {
            if (a.* == .drop and !protected[idx]) a.* = .summarize;
        }
        summary_cost = attempted;
        used += summary_cost;
    } else if (attempted > 0) {
        try warnings.append(try warnSummaryDoesNotFit(alloc, attempted));
    }

    const decisions = try alloc.alloc(Decision, n);
    for (messages, 0..) |m, idx| {
        decisions[idx] = .{
            .index = idx,
            .role = m.role,
            .action = actions[idx],
            .tokens = m.tokens,
        };
    }

    var kept: i64 = 0;
    var dropped: i64 = 0;
    var folded: i64 = 0;
    for (decisions) |d| {
        switch (d.action) {
            .keep => kept += d.tokens,
            .drop => dropped += d.tokens,
            .summarize => folded += d.tokens,
        }
    }

    return .{
        .decisions = decisions,
        .kept_tokens = kept,
        .summarized_tokens = folded,
        .dropped_tokens = dropped,
        .summary_cost_tokens = summary_cost,
        .projected_tokens = kept + summary_cost,
        .fits_budget = kept + summary_cost <= budget_tokens,
        .warnings = try warnings.toOwnedSlice(),
    };
}

/// Free a plan's allocated slices (decisions + warnings strings).
pub fn freePlan(alloc: std.mem.Allocator, plan: Plan) void {
    alloc.free(plan.decisions);
    for (plan.warnings) |w| alloc.free(w);
    alloc.free(plan.warnings);
}

/// Human-readable one-line summary of a plan. Caller owns the returned string.
pub fn describePrune(alloc: std.mem.Allocator, plan: Plan) ![]const u8 {
    if (!plan.fits_budget) {
        var b: [24]u8 = undefined;
        return std.fmt.allocPrint(alloc, "Does not fit: {s} tokens projected against the budget.", .{
            commaFormat(&b, plan.projected_tokens),
        });
    }
    var b1: [24]u8 = undefined;
    var b2: [24]u8 = undefined;
    var out: std.ArrayList(u8) = .init(alloc);
    errdefer out.deinit();
    try out.writer().print("{s} kept", .{commaFormat(&b1, plan.kept_tokens)});
    if (plan.summarized_tokens > 0) {
        try out.writer().print(" · {s} folded into a {s}-token summary", .{
            commaFormat(&b1, plan.summarized_tokens),
            commaFormat(&b2, plan.summary_cost_tokens),
        });
    }
    if (plan.dropped_tokens > 0) {
        try out.writer().print(" · {s} dropped", .{commaFormat(&b1, plan.dropped_tokens)});
    }
    try out.appendSlice(" — fits the budget.");
    return out.toOwnedSlice();
}