"use strict";
/**
 * Token tracking plugin for input/output tokens.
 *
 * Accumulates per-turn token usage (input_tokens + cache_read_input_tokens
 * + cache_creation_input_tokens = total input; output_tokens = output).
 * Respects the dedup flag from TurnCounterPlugin. Cache breakdown (r/5m/1h)
 * is tracked by CacheTrackerPlugin for reporting and pricing.
 */
Object.defineProperty(exports, "__esModule", { value: true });
exports.TokenTrackerPlugin = void 0;
const index_1 = require("../../observability/index");
class TokenTrackerPlugin {
    name = 'tokens';
    phaseInput = 0;
    phaseOutput = 0;
    totalInput = 0;
    totalOutput = 0;
    lastUsage = null;
    phaseSnapshots = [];
    currentPhase = 'setup';
    phaseMessagesWithUsage = 0;
    onMessage(message, ctx, store) {
        if (message.type !== 'assistant')
            return;
        const turns = ctx.get('turns');
        if (turns?.isDuplicate)
            return;
        const usage = message.message?.usage;
        if (usage) {
            const input = Number(usage.input_tokens ?? 0) +
                Number(usage.cache_read_input_tokens ?? 0) +
                Number(usage.cache_creation_input_tokens ?? 0);
            const output = Number(usage.output_tokens ?? 0);
            this.phaseInput += input;
            this.phaseOutput += output;
            this.totalInput += input;
            this.totalOutput += output;
            this.lastUsage = usage;
            this.phaseMessagesWithUsage += 1;
        }
        store.set('tokens', this.getData());
    }
    onPhaseTransition(fromPhase, toPhase, _ctx, store) {
        this.phaseSnapshots.push({
            phase: fromPhase,
            inputTokens: this.phaseInput,
            outputTokens: this.phaseOutput,
            messagesWithUsage: this.phaseMessagesWithUsage,
        });
        this.currentPhase = toPhase;
        this.phaseInput = 0;
        this.phaseOutput = 0;
        this.phaseMessagesWithUsage = 0;
        store.set('tokens', this.getData());
    }
    onFinalize(resultMessage, _totalDurationMs, _ctx, store) {
        this.phaseSnapshots.push({
            phase: this.currentPhase,
            inputTokens: this.phaseInput,
            outputTokens: this.phaseOutput,
            messagesWithUsage: this.phaseMessagesWithUsage,
        });
        store.set('tokens', this.getData());
        // Promote token totals to Sentry trace measurements so they show up as
        // first-class metrics in the active root span. No-op when there is no
        // active span or telemetry is disabled — purely additive to benchmark
        // JSON, never replaces it.
        //
        // We prefer the resultMessage's `usage` (cumulative SDK total) when
        // present so input + cache breakdowns survive even if individual
        // assistant deltas were missed. Fall back to the running totals.
        const usage = resultMessage.usage ?? this.lastUsage ?? null;
        if (usage) {
            (0, index_1.setSpanMeasurement)('agent.tokens.input', Number(usage.input_tokens ?? 0), 'token');
            (0, index_1.setSpanMeasurement)('agent.tokens.output', Number(usage.output_tokens ?? 0), 'token');
            (0, index_1.setSpanMeasurement)('agent.tokens.cache_read_input', Number(usage.cache_read_input_tokens ?? 0), 'token');
            (0, index_1.setSpanMeasurement)('agent.tokens.cache_creation_input', Number(usage.cache_creation_input_tokens ?? 0), 'token');
        }
        // Always emit the wizard-side totals (input includes cache reads + creates,
        // matching how the rest of the wizard reports token spend).
        (0, index_1.setSpanMeasurement)('agent.tokens.total_input', this.totalInput, 'token');
        (0, index_1.setSpanMeasurement)('agent.tokens.total_output', this.totalOutput, 'token');
    }
    getData() {
        return {
            phaseInput: this.phaseInput,
            phaseOutput: this.phaseOutput,
            totalInput: this.totalInput,
            totalOutput: this.totalOutput,
            lastUsage: this.lastUsage,
            phaseSnapshots: [...this.phaseSnapshots],
        };
    }
}
exports.TokenTrackerPlugin = TokenTrackerPlugin;
