estimateContextTokens function

ContextUsageEstimate estimateContextTokens(List<Message> messages)

Estimate context tokens for messages using provider usage when available.

Ported from pi's estimateContextTokens: the last assistant message with valid usage (not errored/aborted, non-zero tokens) anchors the estimate; everything after it is estimated heuristically. Without any usage, the whole list is estimated.

Implementation

ContextUsageEstimate estimateContextTokens(List<Message> messages) {
  ({Usage usage, int index})? usageInfo;
  for (var i = messages.length - 1; i >= 0; i--) {
    final usage = _assistantUsage(messages[i]);
    if (usage != null) {
      usageInfo = (usage: usage, index: i);
      break;
    }
  }

  // F5: the registry dedups repeated images on the wire to short
  // `[Image N]` labels — mirror that here, or the transcript-side
  // estimate diverges from the request the provider actually prices.
  final seenImages = <String>{};
  int charged(Message message) => _estimateTokensDedup(message, seenImages);
  if (usageInfo == null) {
    var estimated = 0;
    for (final message in messages) {
      estimated += charged(message);
    }
    return ContextUsageEstimate(
      tokens: estimated,
      usageTokens: 0,
      trailingTokens: estimated,
      lastUsageIndex: null,
    );
  }

  // The anchor era still seeds the seen-set so trailing repeats stay
  // cheap even when their first occurrence predates the anchor.
  for (var i = 0; i <= usageInfo.index; i++) {
    charged(messages[i]);
  }
  var trailingTokens = 0;
  for (var i = usageInfo.index + 1; i < messages.length; i++) {
    trailingTokens += charged(messages[i]);
  }

  final usageTokens = calculateContextTokens(usageInfo.usage);
  return ContextUsageEstimate(
    tokens: usageTokens + trailingTokens,
    usageTokens: usageTokens,
    trailingTokens: trailingTokens,
    lastUsageIndex: usageInfo.index,
  );
}