core/src/js/tokens.js

// ==========================================
// Chat Switchboard – Token Estimation
// ==========================================
// Context tracking, token estimation, and context warning.
//
// Exports: window.Tokens,window.updateInputTokens,window.updateContextWarning


// ── Token Estimation + Context Tracking ─────

const Tokens = {
    // Rough heuristic: ~4 chars per token for English (GPT/Claude average).
    // Not exact, but good enough for a UI indicator.
    estimate(text) {
        if (!text) return 0;
        return Math.ceil(text.length / 4);
    },

    // Estimate tokens for the full conversation context sent to the model
    estimateConversation(messages, systemPrompt) {
        let total = 0;
        // System prompt
        if (systemPrompt) total += this.estimate(systemPrompt) + 4; // +4 for role/delimiters
        // Messages
        for (const m of messages) {
            total += this.estimate(m.content) + 4; // +4 per message overhead (role, delimiters)
        }
        return total;
    },

    // Estimate tokens for a set of files (staged or sent).
    // Images: ~765 tokens (base tile estimate, reasonable cross-provider).
    // Documents with extracted text: text length / 4.
    // Other/unknown: raw size / 4 (rough fallback).
    estimateFiles(files) {
        if (!files?.length) return 0;
        let total = 0;
        for (const att of files) {
            if (att.content_type?.startsWith('image/')) {
                total += 765; // base tile estimate
            } else if (att.extracted_text) {
                total += this.estimate(att.extracted_text);
            } else if (att.size_bytes) {
                total += Math.ceil(att.size_bytes / 4);
            }
        }
        return total;
    },

    // Get context budget for current model
    getContextBudget() {
        const caps = UI.getSelectedModelCaps();
        return {
            maxContext: caps.max_context || 0,
            maxOutput: caps.max_output_tokens || 0,
        };
    },

    // Format token count for display
    format(n) {
        if (n >= 100000) return (n / 1000).toFixed(0) + 'K';
        if (n >= 10000) return (n / 1000).toFixed(1) + 'K';
        if (n >= 1000) return (n / 1000).toFixed(1) + 'K';
        return String(n);
    },

    _warningDismissed: false,
};

// Update the token counter below the input
function updateInputTokens() {
    const el = document.getElementById('inputTokenCount');
    if (!el) return;

    const inputText = (typeof ChatInput !== 'undefined') ? ChatInput.getValue() : '';
    const inputTokens = Tokens.estimate(inputText);

    if (!inputText.trim()) {
        el.textContent = '';
        el.className = 'input-token-count';
        return;
    }

    const budget = Tokens.getContextBudget();
    if (budget.maxContext > 0) {
        // Show relative to available context
        const chat = App.chats.find(c => c.id === App.activeId);
        const convTokens = Tokens.estimateConversation(chat?.messages || [], App.settings.systemPrompt);
        // Include staged file estimates
        const fileTokens = Tokens.estimateFiles(
            (App.stagedFiles || []).filter(a => a.status !== 'error').map(a => ({
                content_type: a.contentType,
                size_bytes: a.sizeBytes,
            }))
        );
        const totalWithInput = convTokens + inputTokens + fileTokens;
        const pct = totalWithInput / budget.maxContext;
        const attLabel = fileTokens > 0 ? ` +${Tokens.format(fileTokens)} files` : '';
        el.textContent = `~${Tokens.format(inputTokens)} tokens${attLabel} · ${Tokens.format(totalWithInput)} / ${Tokens.format(budget.maxContext)} context`;
        el.className = 'input-token-count' + (pct > 0.9 ? ' danger' : pct > 0.75 ? ' warning' : '');
    } else {
        el.textContent = `~${Tokens.format(inputTokens)} tokens`;
        el.className = 'input-token-count';
    }
}

// Check conversation length and show/hide warning
function updateContextWarning() {
    const warning = document.getElementById('contextWarning');
    const text = document.getElementById('contextWarningText');
    const summarizeBtn = document.getElementById('summarizeBtn');
    if (!warning || !text) return;

    const budget = Tokens.getContextBudget();
    if (budget.maxContext <= 0) {
        warning.style.display = 'none';
        return;
    }

    const chat = App.chats.find(c => c.id === App.activeId);
    if (!chat || !chat.messages?.length) {
        warning.style.display = 'none';
        return;
    }

    const convTokens = Tokens.estimateConversation(chat.messages, App.settings.systemPrompt);
    const pct = convTokens / budget.maxContext;

    // Show summarize button at ≥75% context usage (if enough messages)
    const showSummarize = pct >= 0.75 && chat.messages.length >= 4;
    if (summarizeBtn) {
        summarizeBtn.style.display = showSummarize ? 'inline-block' : 'none';
    }

    // Show auto-compact toggle next to summarize button
    const autoToggle = document.getElementById('autoCompactToggle');
    const autoCheck = document.getElementById('autoCompactCheck');
    if (autoToggle && autoCheck) {
        autoToggle.style.display = showSummarize ? 'inline-flex' : 'none';
        // Reflect current channel setting
        const chatSettings = chat?.settings || {};
        autoCheck.checked = chatSettings.auto_compaction !== false; // default: on
    }

    if (pct >= 0.9 && !Tokens._warningDismissed) {
        warning.style.display = 'flex';
        warning.className = 'context-warning danger';
        text.textContent = `This conversation is using ~${Math.round(pct * 100)}% of the model's context window (~${Tokens.format(convTokens)} / ${Tokens.format(budget.maxContext)} tokens). Responses may lose earlier context.`;
    } else if (pct >= 0.75 && !Tokens._warningDismissed) {
        warning.style.display = 'flex';
        warning.className = 'context-warning';
        text.textContent = `Conversation is getting long (~${Math.round(pct * 100)}% of context window). The model may start losing track of earlier messages.`;
    } else {
        warning.style.display = 'none';
    }
}

function dismissContextWarning() {
    Tokens._warningDismissed = true;
    const el = document.getElementById('contextWarning');
    if (el) el.style.display = 'none';
}

// ── Exports ─────────────────────────────────
sb.ns('Tokens', Tokens);
sb.register('updateInputTokens', updateInputTokens);
sb.register('updateContextWarning', updateContextWarning);