@womp/spark-agent 5.0.2 → 5.1.0

package/dist/context-budget.js

@@ -0,0 +1,308 @@
+// Keep one model request inside the model's context window.
+//
+// A Spark turn sends the whole conversation on every step: for a thread that
+// Swan and Crow both write to, that is every earlier turn, every tool result
+// (each up to MODEL_RESULT_BYTE_LIMIT, ./bounded-output.ts) and every image. Nothing
+// bounded the total, so a long enough thread passed the provider's limit and every
+// later turn failed on its first request ("prompt is too long: 1006655 tokens >
+// 1000000 maximum") -- permanently, since each failed attempt only adds rows.
+//
+// This runs on the messages of every step, after the host's own preparation and
+// before the request is built, in both executors: the shared loop (./agent.ts,
+// which Crow runs) and the browser's own loop in Swan.
+//
+// WHAT, in this order, until the estimate is back under the target:
+//
+//  1. Images before the current turn, oldest first -- the most tokens for the
+//     least text, and the ones a later step least needs to see again.
+//  2. Large tool results before the current turn, oldest first. The call and a
+//     result part stay, with a note in place of the output, so no provider is
+//     handed a call without its result.
+//  3. Whole turns, oldest first, cut at user messages so every remaining call is
+//     still paired with its result. A note in their place tells the model that
+//     earlier messages were left out.
+//
+// The current turn -- the latest user message and everything after it -- is never
+// changed. When it alone is over the window the request goes out as it is and the
+// provider's own error says so; nothing here can make that turn smaller.
+//
+// ESTIMATES, NOT COUNTS. Counting needs the provider's tokenizer or a request.
+// The estimate is deliberately high (CHARS_PER_TOKEN is under the ~3.5-4 of prose)
+// and trimming starts at TRIM_TRIGGER_RATIO of the budget, so dense JSON that
+// tokenizes worse than prose still lands under the real limit. Images count by
+// number, never by their bytes.
+//
+// STABLE ACROSS STEPS. Once over, the conversation is trimmed to
+// TRIM_TARGET_RATIO rather than to the limit, and always oldest first, so the next
+// steps of the run trim the same prefix again instead of a little more each time --
+// which keeps a provider's prompt cache warm.
+import { SPARK_MODELS } from "./models.js";
+/** Under prose's ~3.5-4, so the estimate errs high. */
+export const CHARS_PER_TOKEN = 3;
+/** About what one image at the providers' 1568 px ceiling costs. */
+export const IMAGE_TOKENS = 1_600;
+export const MESSAGE_OVERHEAD_TOKENS = 4;
+/** Per tool: its name, description and input schema around the counted text. */
+export const TOOL_DEFINITION_OVERHEAD_TOKENS = 300;
+/** Start trimming once the estimate passes this share of the budget. */
+export const TRIM_TRIGGER_RATIO = 0.9;
+/** Then trim down to this share, so the next steps do not trim again. */
+export const TRIM_TARGET_RATIO = 0.75;
+/** A tool result at or below this is left whole: not worth the detail it loses. */
+export const SMALL_TOOL_OUTPUT_TOKENS = 500;
+export const OMITTED_HISTORY_TEXT = "[Earlier messages in this conversation were left out of this request to fit the model's context window.]";
+export const OMITTED_HISTORY_IMAGE_TEXT = "[An earlier image was left out of this request to fit the model's context window.]";
+export const OMITTED_TOOL_OUTPUT_TEXT = "[This earlier tool result was left out of this request to fit the model's context window.]";
+/**
+ * The model's context window, or undefined when it is not known.
+ *
+ * Listed models first, then model families whose window is certain. Anything
+ * else -- a third-party model routed through OpenRouter -- is undefined, and the
+ * conversation is then left exactly as it was: overestimating a window can only
+ * repeat today's failure, but underestimating one would silently hide history
+ * from a model that could read it.
+ */
+export function contextWindowFor(modelId) {
+    const listed = SPARK_MODELS.find(({ id }) => id === modelId)?.contextWindow;
+    if (listed !== undefined)
+        return listed;
+    if (/^claude-haiku-/.test(modelId))
+        return 200_000;
+    if (/^claude-/.test(modelId))
+        return 1_000_000;
+    if (/^gpt-6-/.test(modelId))
+        return 1_050_000;
+    if (/^gemini-/.test(modelId))
+        return 1_048_576;
+    return undefined;
+}
+export function fitMessagesToContext(messages, options) {
+    if (options.contextWindow === undefined)
+        return { messages };
+    const budgetTokens = Math.max(0, options.contextWindow -
+        (options.maxOutputTokens ?? 0) -
+        textTokens(options.system ?? "") -
+        toolsTokens(options.tools));
+    const sizes = messages.map(messageTokens);
+    const originalTokens = sum(sizes);
+    const report = {
+        budgetTokens,
+        originalTokens,
+        estimatedTokens: originalTokens,
+        omittedImages: 0,
+        omittedToolOutputs: 0,
+        omittedMessages: 0,
+        fits: originalTokens <= budgetTokens,
+    };
+    if (originalTokens <= budgetTokens * TRIM_TRIGGER_RATIO)
+        return { messages, report };
+    // The current turn starts at the latest user message and is never changed.
+    const current = lastIndex(messages, ({ role }) => role === "user");
+    if (current <= 0)
+        return { messages, report };
+    const target = budgetTokens * TRIM_TARGET_RATIO;
+    const out = messages.slice();
+    let total = originalTokens;
+    /**
+     * Rewrites the parts of each earlier message, oldest first, until the estimate
+     * is under the target. `rewrite` answers the same part to keep it, or its
+     * replacement and roughly what that saves -- used only to stop partway through a
+     * message, which is then measured again.
+     */
+    const rewriteEarlier = (rewrite) => {
+        for (let i = 0; i < current && total > target; i += 1) {
+            const message = out[i];
+            if (message === undefined || !Array.isArray(message.content))
+                continue;
+            let projected = total;
+            let changed = false;
+            const content = message.content.map((part) => {
+                if (projected <= target)
+                    return part;
+                const result = rewrite(part, projected);
+                if (result.part === part)
+                    return part;
+                changed = true;
+                projected -= result.saved;
+                return result.part;
+            });
+            if (!changed)
+                continue;
+            const next = { ...message, content };
+            const size = messageTokens(next);
+            total += size - (sizes[i] ?? 0);
+            sizes[i] = size;
+            out[i] = next;
+        }
+    };
+    // 1. Earlier images.
+    rewriteEarlier((part, projected) => {
+        if (isImage(part)) {
+            report.omittedImages += 1;
+            return { part: omittedImage(), saved: IMAGE_TOKENS };
+        }
+        if (!isToolResult(part) || part.output.type !== "content" || !Array.isArray(part.output.value)) {
+            return { part, saved: 0 };
+        }
+        let saved = 0;
+        const value = part.output.value.map((item) => {
+            if (projected - saved <= target || !isImage(item))
+                return item;
+            report.omittedImages += 1;
+            saved += IMAGE_TOKENS;
+            return omittedImage();
+        });
+        return saved === 0 ? { part, saved } : { part: { ...part, output: { ...part.output, value } }, saved };
+    });
+    // 2. Earlier large tool results.
+    rewriteEarlier((part) => {
+        if (!isToolResult(part))
+            return { part, saved: 0 };
+        const saved = outputTokens(part.output);
+        if (saved <= SMALL_TOOL_OUTPUT_TOKENS)
+            return { part, saved: 0 };
+        report.omittedToolOutputs += 1;
+        return { part: { ...part, output: { type: "text", value: OMITTED_TOOL_OUTPUT_TEXT } }, saved };
+    });
+    // 3. Whole earlier turns, cut at user messages.
+    if (total > target) {
+        const note = { role: "user", content: [{ type: "text", text: OMITTED_HISTORY_TEXT }] };
+        const noteTokens = messageTokens(note);
+        const ends = [];
+        out.forEach(({ role }, i) => {
+            if (i > 0 && i < current && role === "user")
+                ends.push(i);
+        });
+        ends.push(current);
+        let cut = 0;
+        for (const end of ends) {
+            if (total + noteTokens <= target)
+                break;
+            for (let i = cut; i < end; i += 1) {
+                const message = out[i];
+                if (message === undefined || message.role === "system")
+                    continue;
+                total -= sizes[i] ?? 0;
+                report.omittedMessages += 1;
+            }
+            cut = end;
+        }
+        if (cut > 0) {
+            const kept = out.slice(0, cut).filter(({ role }) => role === "system");
+            out.splice(0, cut, ...kept, note);
+            total += noteTokens;
+        }
+    }
+    report.estimatedTokens = total;
+    report.fits = total <= budgetTokens;
+    return { messages: out, report };
+}
+/** The estimate for a list of messages, as `fitMessagesToContext` counts them. */
+export function estimateMessagesTokens(messages) {
+    return sum(messages.map(messageTokens));
+}
+// ── estimation ──────────────────────────────────────────────────────────────
+const estimates = new WeakMap();
+function messageTokens(message) {
+    const known = estimates.get(message);
+    if (known !== undefined)
+        return known;
+    const content = message.content;
+    const size = MESSAGE_OVERHEAD_TOKENS +
+        (typeof content === "string"
+            ? textTokens(content)
+            : Array.isArray(content)
+                ? sum(content.map(partTokens))
+                : 0);
+    estimates.set(message, size);
+    return size;
+}
+function partTokens(part) {
+    if (isImage(part))
+        return IMAGE_TOKENS;
+    switch (part.type) {
+        case "text":
+        case "reasoning":
+            return textTokens(String(part.text ?? ""));
+        case "file":
+            return dataTokens(part.data);
+        case "tool-call":
+            return textTokens(String(part.toolName ?? "")) + textTokens(json(part.input));
+        case "tool-result":
+            return isToolResult(part) ? outputTokens(part.output) : 0;
+        default:
+            return textTokens(json(part));
+    }
+}
+function outputTokens(output) {
+    switch (output.type) {
+        case "text":
+        case "error-text":
+            return textTokens(String(output.value ?? ""));
+        case "content":
+            return Array.isArray(output.value)
+                ? sum(output.value.map((item) => isImage(item) ? IMAGE_TOKENS : item.type === "text" ? textTokens(String(item.text ?? "")) : dataTokens(item.data)))
+                : 0;
+        default:
+            return textTokens(json(output.value));
+    }
+}
+function textTokens(text) {
+    return Math.ceil(text.length / CHARS_PER_TOKEN);
+}
+function dataTokens(data) {
+    if (typeof data === "string")
+        return textTokens(data);
+    if (data instanceof Uint8Array || data instanceof ArrayBuffer)
+        return Math.ceil(data.byteLength / CHARS_PER_TOKEN);
+    return 0;
+}
+function toolsTokens(tools) {
+    if (!tools)
+        return 0;
+    return sum(Object.entries(tools).map(([name, tool]) => {
+        const description = tool?.description;
+        return (TOOL_DEFINITION_OVERHEAD_TOKENS + textTokens(name) + textTokens(typeof description === "string" ? description : ""));
+    }));
+}
+/** JSON text, with binary left out: its size is not text the model reads. */
+function json(value) {
+    try {
+        return (JSON.stringify(value, (_key, v) => (v instanceof Uint8Array || v instanceof ArrayBuffer ? "" : v)) ?? "");
+    }
+    catch {
+        return "";
+    }
+}
+// ── shapes ──────────────────────────────────────────────────────────────────
+function isImage(part) {
+    if (part.type === "image" || part.type === "image-data" || part.type === "image-url")
+        return true;
+    const mediaType = part.mediaType;
+    return ((part.type === "file" || part.type === "media" || part.type === "file-data" || part.type === "file-url") &&
+        typeof mediaType === "string" &&
+        mediaType.startsWith("image/"));
+}
+function isToolResult(part) {
+    return (part.type === "tool-result" &&
+        typeof part.output === "object" &&
+        part.output !== null &&
+        typeof part.output.type === "string");
+}
+function sum(values) {
+    let total = 0;
+    for (const value of values)
+        total += value;
+    return total;
+}
+function lastIndex(items, test) {
+    for (let i = items.length - 1; i >= 0; i -= 1) {
+        const item = items[i];
+        if (item !== undefined && test(item))
+            return i;
+    }
+    return -1;
+}
+function omittedImage() {
+    return { type: "text", text: OMITTED_HISTORY_IMAGE_TEXT };
+}