@sema-agent/core/dist/core/tool-result-budget.js

Stateless, task-oriented AI agent core

import { isToolResult } from "./message-utils.js";
import { buildPreview, buildToolResultRef, PERSISTED_OUTPUT_PREFIX, toolResultContentSegment, toolResultProvenanceOf, } from "./tool-result-store.js";
import { budgetExemptToolNames } from "./tool-registry.js";
import { wireOf } from "./tool-name-readings.js";
export const AGGREGATE_TOOL_RESULT_BUDGET_CHARS = 200_000;
const HEAD = 1_000;
const TAIL = 1_000;
const PREVIEW_OVERHEAD = 256;
const COMPACT_HEAD = 200;
const COMPACT_TAIL = 200;
function textChars(content) {
    let n = 0;
    for (const b of content)
        if (b.type === "text" && typeof b.text === "string")
            n += b.text.length;
    return n;
}
function textOf(content) {
    return content.filter((b) => b.type === "text").map((b) => b.text).join("\n");
}
function isOffloadedPreview(content) {
    const first = content[0];
    return first?.type === "text" && typeof first.text === "string" && first.text.startsWith(PERSISTED_OUTPUT_PREFIX);
}
function replaceTextBlock(m, text) {
    if (!isToolResult(m))
        return m;
    const images = m.content.filter((b) => b.type !== "text");
    return { ...m, content: [{ type: "text", text }, ...images] };
}
function truncationPreview(full, head = HEAD, tail = TAIL) {
    const total = full.length;
    const tailStart = Math.max(head, total - tail);
    return (`${full.slice(0, head)}\n` +
        `…[${total} chars total; ${tailStart - head} omitted to fit the per-turn result budget — re-run with a narrower query to see the rest.]\n` +
        `${full.slice(tailStart)}`);
}
const CATALOG_BUDGET_EXEMPT_TOOLS = budgetExemptToolNames();
export async function capAggregateToolResults(messages, opts) {
    const budget = opts.budgetChars ?? AGGREGATE_TOOL_RESULT_BUDGET_CHARS;
    const exempt = opts.exemptTools ?? CATALOG_BUDGET_EXEMPT_TOOLS;
    const exemptResult = (m) => m.toolName !== undefined && exempt.has(wireOf("budget.exemptResult", { toolName: opts.toolNameOf !== undefined ? opts.toolNameOf(m.toolName) : m.toolName }));
    let out;
    let i = 0;
    while (i < messages.length) {
        if (!isToolResult(messages[i])) {
            i++;
            continue;
        }
        let j = i;
        while (j < messages.length && isToolResult(messages[j]))
            j++;
        const arr = out ?? messages;
        let total = 0;
        for (let k = i; k < j; k++) {
            const m = arr[k];
            if (exemptResult(m))
                continue;
            total += textChars(m.content);
        }
        if (total > budget) {
            out = out ?? messages.slice();
            const firedFallback = new Map();
            const previewOne = async (k, before, head, tail, fromOriginal) => {
                const m = out[k];
                const full = textOf((fromOriginal ? messages[k] : m).content);
                const ref = buildToolResultRef(opts.sessionId, m.toolCallId ?? `idx${k}`, toolResultContentSegment(full));
                const sizes = head === HEAD && tail === TAIL ? undefined : { head, tail };
                let previewText;
                let storeFallback;
                if (opts.store) {
                    try {
                        await opts.store.put(ref, full, toolResultProvenanceOf(opts.sessionId));
                        previewText = buildPreview(full, ref, sizes);
                        storeFallback = false;
                    }
                    catch {
                        previewText = truncationPreview(full, head, tail);
                        storeFallback = true;
                    }
                }
                else {
                    previewText = truncationPreview(full, head, tail);
                    storeFallback = true;
                }
                if (previewText.length >= before)
                    return 0;
                out[k] = replaceTextBlock(out[k], previewText);
                const prior = firedFallback.get(k);
                if (prior === undefined || (storeFallback && !prior)) {
                    firedFallback.set(k, storeFallback);
                    opts.onCapped?.({ ...(m.toolName !== undefined ? { tool: m.toolName } : {}), sizeChars: before, storeFallback });
                }
                return before - previewText.length;
            };
            const cand = [];
            for (let k = i; k < j; k++) {
                const m = out[k];
                if (exemptResult(m))
                    continue;
                if (isOffloadedPreview(m.content))
                    continue;
                const size = textChars(m.content);
                if (size > HEAD + TAIL + PREVIEW_OVERHEAD)
                    cand.push({ k, size });
            }
            cand.sort((a, b) => b.size - a.size || a.k - b.k);
            for (const { k, size } of cand) {
                if (total <= budget)
                    break;
                total -= await previewOne(k, size, HEAD, TAIL, false);
            }
            if (total > budget) {
                const cand2 = [];
                for (let k = i; k < j; k++) {
                    const m = out[k];
                    if (exemptResult(m))
                        continue;
                    if (isOffloadedPreview(messages[k].content))
                        continue;
                    const size = textChars(m.content);
                    if (size > COMPACT_HEAD + COMPACT_TAIL + PREVIEW_OVERHEAD)
                        cand2.push({ k, size });
                }
                cand2.sort((a, b) => b.size - a.size || a.k - b.k);
                for (const { k, size } of cand2) {
                    if (total <= budget)
                        break;
                    total -= await previewOne(k, size, COMPACT_HEAD, COMPACT_TAIL, true);
                }
            }
            if (total > budget)
                opts.onBudgetUnreachable?.({ batchChars: total, budgetChars: budget, results: j - i });
        }
        i = j;
    }
    return out ?? messages;
}