@sema-agent/core/dist/core/runner/prepare-offload-wrappers.js

Stateless, task-oriented AI agent core

import { DEFAULT_TOOL_RESULT_THRESHOLD_CHARS, InMemoryToolResultStore, RunnerSharedToolResultStore, ScopedToolResultStore, withToolResultOffload, } from "../tool-result-store.js";
import { faceOf } from "../tool-roster.js";
const FIRST_PARTY_OFFLOAD_THRESHOLD_CHARS = 50_000;
export function prepareOffloadWrappers(input) {
    const { spec, deps, resume, offloadScope, sessionId } = input;
    const nestedStats = { tokens: 0, turns: 0, tasks: 0, costMicroUsd: 0, anyUnpriced: false };
    if (resume) {
        nestedStats.tokens = resume.seed.nestedStats.tokens;
        nestedStats.turns = resume.seed.nestedStats.turns;
        nestedStats.tasks = resume.seed.nestedStats.tasks;
        nestedStats.costMicroUsd = resume.seed.nestedStats.costMicroUsd;
        nestedStats.anyUnpriced = resume.seed.nestedStats.anyUnpriced ?? true;
    }
    const reportUsage = (u) => {
        nestedStats.tokens += u.tokens;
        nestedStats.turns += u.turns;
        nestedStats.tasks += u.tasks;
        nestedStats.costMicroUsd += u.costMicroUsd ?? 0;
        if (u.costMicroUsd === undefined)
            nestedStats.anyUnpriced = true;
    };
    const explicitGlobalThreshold = spec.toolResultThresholdChars ?? deps.toolResultThresholdChars;
    const offloadThreshold = explicitGlobalThreshold ?? DEFAULT_TOOL_RESULT_THRESHOLD_CHARS;
    const offloadEnabled = Number.isFinite(offloadThreshold) && offloadThreshold > 0;
    const rawOffloadStore = offloadEnabled ? (deps.toolResultStore ?? new InMemoryToolResultStore()) : undefined;
    const offloadStore = rawOffloadStore instanceof RunnerSharedToolResultStore
        ? new ScopedToolResultStore(rawOffloadStore, offloadScope)
        : rawOffloadStore;
    const offloadReachableToolsRef = {};
    const maybeOffload = (tool, perTool) => {
        if (!offloadStore || perTool?.offload === false)
            return tool;
        return withToolResultOffload(tool, offloadStore, perTool?.offloadThresholdChars ?? offloadThreshold, sessionId, () => offloadReachableToolsRef.current?.());
    };
    const firstPartyOffload = (tool) => {
        const face = faceOf(tool);
        if (face.offload === false)
            return maybeOffload(tool, { offload: false });
        return maybeOffload(tool, explicitGlobalThreshold === undefined ? { offloadThresholdChars: face.offloadThresholdChars ?? FIRST_PARTY_OFFLOAD_THRESHOLD_CHARS } : undefined);
    };
    const remoteToolOffload = (tool) => {
        if (explicitGlobalThreshold !== undefined)
            return maybeOffload(tool);
        return maybeOffload(tool, { offloadThresholdChars: tool.mcpMaxResultSizeChars ?? 50_000 });
    };
    return { nestedStats, reportUsage, offloadStore, offloadReachableToolsRef, maybeOffload, firstPartyOffload, remoteToolOffload };
}