Conversation Pruner — Swift source

// Conversation Pruner — compute a deterministic pruning plan for a
// token-budgeted chat history.
//
// Language: Swift (5.9+, 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.

/// Chat roles recognized by the pruner.
enum ChatRole: String {
    case system, user, assistant, tool
}

/// One per-message decision.
enum PruneAction: String {
    case keep, summarize, drop
}

/// One chat message with its caller-supplied token count.
struct ConversationMessage {
    var role: ChatRole
    var content: String
    var tokens: Int
    /// Pinned messages are never dropped or summarized.
    var pinned: Bool = false
}

/// One per-message decision record.
struct PruneDecision {
    var index: Int
    var role: ChatRole
    var action: PruneAction
    var tokens: Int
}

/// The full plan: decisions, tallies, projection, warnings.
struct PrunePlan {
    var decisions: [PruneDecision]
    var keptTokens: Int
    var summarizedTokens: Int
    var droppedTokens: Int
    /// Tokens the summary placeholder itself will cost in the prompt.
    var summaryCostTokens: Int
    var projectedTokens: Int
    var fitsBudget: Bool
    var warnings: [String]
}

/// Summary compression model: fixed framing tokens.
let SUMMARY_FIXED_TOKENS = 60

/// Summary compression model: share of the folded content.
let SUMMARY_RATIO = 0.1

/// Group an integer with thousands separators (toLocaleString stand-in).
private func fmt(_ v: Int) -> String {
    let s = String(v)
    var out = ""
    for (offset, ch) in s.enumerated().reversed() {
        out.append(ch)
        let remaining = s.count - offset - 1
        if remaining > 0 && remaining % 3 == 0 { out.append(",") }
    }
    return String(out.reversed())
}

/// Errors mirroring the TS suite's RangeError contract.
enum PruneError: Error, CustomStringConvertible {
    case negativeBudget
    case negativeTokens

    var description: String {
        switch self {
        case .negativeBudget: return "budgetTokens must be >= 0"
        case .negativeTokens: return "message tokens must be >= 0"
        }
    }
}

/// Compute the pruning plan for `messages` under `budgetTokens`.
func planPrune(_ messages: [ConversationMessage], budgetTokens: Int) throws -> PrunePlan {
    var warnings: [String] = []
    if budgetTokens < 0 { throw PruneError.negativeBudget }
    if messages.contains(where: { $0.tokens < 0 }) { throw PruneError.negativeTokens }

    let n = messages.count
    let lastUser = messages.lastIndex(where: { $0.role == .user }) ?? -1

    // 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 = Set<Int>()
    for (i, m) in messages.enumerated() {
        if m.role == .system || m.pinned { protectedIdx.insert(i) }
    }
    if n > 0 { protectedIdx.insert(0) }
    if let firstTurn = messages.firstIndex(where: { $0.role != .system }) {
        protectedIdx.insert(firstTurn)
    }
    for i in max(lastUser == -1 ? n - 1 : lastUser, 0)..<n { protectedIdx.insert(i) }

    let protectedTokens = protectedIdx.reduce(0) { $0 + messages[$1].tokens }
    if protectedTokens > budgetTokens {
        warnings.append(
            "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 = Array(repeating: PruneAction.drop, count: n)
    for i in protectedIdx { actions[i] = .keep }
    var used = protectedTokens
    for i in stride(from: n - 1, through: 0, by: -1) {
        if actions[i] != .drop { continue }
        if used + messages[i].tokens <= budgetTokens {
            actions[i] = .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.
    let summarizeIdx = (0..<n).filter { actions[$0] == .drop && !protectedIdx.contains($0) }
    let summarizeTokens = summarizeIdx.reduce(0) { $0 + messages[$1].tokens }
    let attempted = summarizeIdx.isEmpty
        ? 0
        : SUMMARY_FIXED_TOKENS + Int((Double(summarizeTokens) * SUMMARY_RATIO).rounded(.up))

    // The summary only costs anything when it is actually applied.
    var summaryCost = 0
    if attempted > 0 && used + attempted <= budgetTokens {
        for i in summarizeIdx { actions[i] = .summarize }
        summaryCost = attempted
        used += summaryCost
    } else if attempted > 0 {
        warnings.append(
            "Even the compressed summary (\(fmt(attempted)) tokens) does not fit the " +
            "remaining budget — the oldest turns are dropped instead.")
    }

    let decisions = messages.enumerated().map { i, m in
        PruneDecision(index: i, role: m.role, action: actions[i], tokens: m.tokens)
    }

    var kept = 0, dropped = 0, folded = 0
    for d in decisions {
        switch d.action {
        case .keep: kept += d.tokens
        case .drop: dropped += d.tokens
        case .summarize: folded += d.tokens
        }
    }

    return PrunePlan(
        decisions: decisions,
        keptTokens: kept,
        summarizedTokens: folded,
        droppedTokens: dropped,
        summaryCostTokens: summaryCost,
        projectedTokens: kept + summaryCost,
        fitsBudget: kept + summaryCost <= budgetTokens,
        warnings: warnings)
}

/// Human-readable one-line summary of a plan.
func describePrune(_ plan: PrunePlan) -> String {
    if !plan.fitsBudget {
        return "Does not fit: \(fmt(plan.projectedTokens)) tokens projected against the budget."
    }
    var parts = ["\(fmt(plan.keptTokens)) kept"]
    if plan.summarizedTokens > 0 {
        parts.append(
            "\(fmt(plan.summarizedTokens)) folded into a \(fmt(plan.summaryCostTokens))-token summary")
    }
    if plan.droppedTokens > 0 { parts.append("\(fmt(plan.droppedTokens)) dropped") }
    return parts.joined(separator: " · ") + " — fits the budget."
}