enforceToolResultBudget method

Future<({List<Message> messages, List<ContentReplacementRecord> newlyReplaced})> enforceToolResultBudget(List<Message> messages, ContentReplacementState state, { Set<String> skipToolNames = const {}, })

Enforce the per-message budget on aggregate tool result size.

Implementation

Future<
  ({List<Message> messages, List<ContentReplacementRecord> newlyReplaced})
>
enforceToolResultBudget(
  List<Message> messages,
  ContentReplacementState state, {
  Set<String> skipToolNames = const {},
}) async {
  final candidatesByMessage = _collectCandidatesByMessage(messages);
  final nameByToolUseId = skipToolNames.isNotEmpty
      ? _buildToolNameMap(messages)
      : null;

  bool shouldSkip(String id) {
    if (nameByToolUseId == null) return false;
    return skipToolNames.contains(nameByToolUseId[id] ?? '');
  }

  final limit = getPerMessageBudgetLimit();
  final replacementMap = <String, String>{};
  final toPersist = <_ToolResultCandidate>[];

  for (final candidates in candidatesByMessage) {
    final partition = _partitionByPriorDecision(candidates, state);

    // Re-apply cached replacements.
    for (final c in partition.mustReapply) {
      replacementMap[c.toolUseId] = c.replacement;
    }

    if (partition.fresh.isEmpty) {
      for (final c in candidates) {
        state.seenIds.add(c.toolUseId);
      }
      continue;
    }

    // Tools with opt-out — mark as seen.
    final skipped = partition.fresh.where((c) => shouldSkip(c.toolUseId));
    for (final c in skipped) {
      state.seenIds.add(c.toolUseId);
    }
    final eligible = partition.fresh
        .where((c) => !shouldSkip(c.toolUseId))
        .toList();

    final frozenSize = partition.frozen.fold<int>(0, (s, c) => s + c.size);
    final freshSize = eligible.fold<int>(0, (s, c) => s + c.size);

    final selected = (frozenSize + freshSize) > limit
        ? _selectFreshToReplace(eligible, frozenSize, limit)
        : <_ToolResultCandidate>[];

    final selectedIds = selected.map((c) => c.toolUseId).toSet();
    for (final c in candidates) {
      if (!selectedIds.contains(c.toolUseId)) {
        state.seenIds.add(c.toolUseId);
      }
    }

    if (selected.isEmpty) continue;
    toPersist.addAll(selected);
  }

  if (replacementMap.isEmpty && toPersist.isEmpty) {
    return (messages: messages, newlyReplaced: <ContentReplacementRecord>[]);
  }

  // Persist all selected candidates.
  final newlyReplaced = <ContentReplacementRecord>[];
  for (final candidate in toPersist) {
    state.seenIds.add(candidate.toolUseId);
    final result = await persistToolResult(
      candidate.content,
      candidate.toolUseId,
    );
    if (result is PersistToolResultError) continue;
    final persisted = result as PersistedToolResult;
    final replacementContent = buildLargeToolResultMessage(persisted);
    replacementMap[candidate.toolUseId] = replacementContent;
    state.replacements[candidate.toolUseId] = replacementContent;
    newlyReplaced.add(
      ContentReplacementRecord.toolResult(
        toolUseId: candidate.toolUseId,
        replacement: replacementContent,
      ),
    );
  }

  if (replacementMap.isEmpty) {
    return (messages: messages, newlyReplaced: <ContentReplacementRecord>[]);
  }

  return (
    messages: _replaceToolResultContents(messages, replacementMap),
    newlyReplaced: newlyReplaced,
  );
}