llm-relay 0.43.0 → 0.44.0

package/dist/usage-observer.js

@@ -1,9 +1,10 @@
 /**
- * Observe provider-reported completion usage without putting an observer in
- * the response path.  In particular, this module never changes (or buffers)
- * a response chunk: it only keeps a small amount of parsing state alongside
+ * Observe provider-reported usage and relay-estimated model output without
+ * changing the response path. In particular, this module never changes (or
+ * buffers) a response chunk: it only keeps bounded parsing state alongside
  * the stream.
  */
+import { estimateTokensFromCharacters } from "./metadata.js";
 export function createUsageAccumulator() {
     return {
         inputTokens: undefined,
@@ -12,11 +13,15 @@ export function createUsageAccumulator() {
         cacheCreationInputTokens: undefined,
         cacheReadInputTokens: undefined,
         completionTokens: undefined,
+        estimatedOutputTokens: undefined,
     };
 }
 const MAX_SSE_FRAME = 16 * 1024;
 const MAX_BUFFERED_JSON = 1024 * 1024;
 const MAX_DECODE_SLICE = 4096;
+const MAX_PENDING_ARGUMENT_CHARS = 64 * 1024;
+const MAX_PENDING_TOOL_CALLS = 128;
+const MAX_STREAMED_TEXT_FIELDS = 128;
 function validCompletionTokens(value) {
     return (typeof value === "number" &&
         Number.isFinite(value) &&
@@ -36,6 +41,365 @@ function recordField(accumulator, field, value) {
 function asRecord(value) {
     return value && typeof value === "object" ? value : undefined;
 }
+function publishOutputEstimate(accumulator, state) {
+    try {
+        accumulator.estimatedOutputTokens = state.overflow
+            || state.pendingArguments.size > 0
+            || state.characters <= 0
+            ? undefined
+            : estimateTokensFromCharacters(state.characters);
+    }
+    catch {
+        // A malformed/hostile accumulator must not affect pass-through.
+    }
+}
+function isBase64DataUrl(value) {
+    if (value.length < 13 || value.slice(0, 5).toLowerCase() !== "data:")
+        return false;
+    const comma = value.indexOf(",", 5);
+    return comma >= 12 && value.slice(comma - 7, comma).toLowerCase() === ";base64";
+}
+function sanitizedArgumentJson(value) {
+    let removedBinary = false;
+    try {
+        const json = JSON.stringify(value, (key, nestedValue) => {
+            if (typeof nestedValue === "string" && (key === "data" || isBase64DataUrl(nestedValue))) {
+                removedBinary = true;
+                return "";
+            }
+            return nestedValue;
+        });
+        return { characters: json?.length ?? 0, removedBinary };
+    }
+    catch {
+        return undefined;
+    }
+}
+function argumentCharacters(value) {
+    if (typeof value === "string") {
+        if (isBase64DataUrl(value))
+            return 0;
+        if (value.length === 0 || value.trim().length === 0)
+            return value.length;
+        try {
+            const parsed = JSON.parse(value);
+            const sanitized = sanitizedArgumentJson(parsed);
+            if (!sanitized)
+                return undefined;
+            return sanitized.removedBinary ? sanitized.characters : value.length;
+        }
+        catch {
+            // Malformed JSON cannot be semantically inspected for escaped binary
+            // fields. Unknown is safer than a knowingly inflated raw-length total.
+            return undefined;
+        }
+    }
+    return sanitizedArgumentJson(value)?.characters;
+}
+function recordOutputCharacters(accumulator, state, characters) {
+    if (state.overflow || characters <= 0 || !Number.isSafeInteger(characters))
+        return;
+    const next = state.characters + characters;
+    if (!Number.isSafeInteger(next)) {
+        taintOutputEstimate(accumulator, state);
+        return;
+    }
+    state.characters = next;
+    publishOutputEstimate(accumulator, state);
+}
+function taintOutputEstimate(accumulator, state) {
+    state.overflow = true;
+    state.pendingArgumentCharacters = 0;
+    state.pendingArguments.clear();
+    state.streamedTexts.clear();
+    try {
+        accumulator.estimatedOutputTokens = undefined;
+    }
+    catch {
+        // A malformed/hostile accumulator must not affect pass-through.
+    }
+}
+function recordOutputText(accumulator, state, value) {
+    if (typeof value !== "string" || isBase64DataUrl(value))
+        return;
+    recordOutputCharacters(accumulator, state, value.length);
+}
+/**
+ * A streamed logical string may be split before `data:...;base64,` becomes
+ * recognizable. Count partial text immediately, but retain only enough state
+ * to roll it back if later fragments prove the whole string is binary.
+ */
+function appendStreamedOutputText(accumulator, state, key, value) {
+    if (state.overflow || typeof value !== "string" || value.length === 0)
+        return;
+    let text = state.streamedTexts.get(key);
+    if (!text) {
+        if (state.streamedTexts.size >= MAX_STREAMED_TEXT_FIELDS) {
+            taintOutputEstimate(accumulator, state);
+            return;
+        }
+        text = {
+            disposition: "candidate",
+            heldCharacters: 0,
+            prefix: "",
+            headerTail: "",
+        };
+        state.streamedTexts.set(key, text);
+    }
+    if (text.disposition === "count") {
+        recordOutputCharacters(accumulator, state, value.length);
+        return;
+    }
+    if (text.disposition === "skip")
+        return;
+    const heldCharacters = text.heldCharacters + value.length;
+    const totalCharacters = state.characters + value.length;
+    if (!Number.isSafeInteger(heldCharacters) || !Number.isSafeInteger(totalCharacters)) {
+        taintOutputEstimate(accumulator, state);
+        return;
+    }
+    text.heldCharacters = heldCharacters;
+    state.characters = totalCharacters;
+    for (const character of value) {
+        if (text.prefix.length < 5) {
+            text.prefix += character;
+            if (text.prefix.toLowerCase() !== "data:".slice(0, text.prefix.length)) {
+                text.disposition = "count";
+                text.heldCharacters = 0;
+                publishOutputEstimate(accumulator, state);
+                return;
+            }
+            continue;
+        }
+        if (character === ",") {
+            if (text.headerTail.toLowerCase() === ";base64") {
+                text.disposition = "skip";
+                state.characters -= text.heldCharacters;
+                text.heldCharacters = 0;
+                publishOutputEstimate(accumulator, state);
+            }
+            else {
+                text.disposition = "count";
+                text.heldCharacters = 0;
+                publishOutputEstimate(accumulator, state);
+            }
+            return;
+        }
+        text.headerTail = `${text.headerTail}${character}`.slice(-7);
+    }
+    publishOutputEstimate(accumulator, state);
+}
+function flushStreamedOutputText(state, key) {
+    const text = state.streamedTexts.get(key);
+    if (!text)
+        return;
+    state.streamedTexts.delete(key);
+}
+function flushStreamedOutputTextPrefix(state, prefix) {
+    for (const key of [...state.streamedTexts.keys()]) {
+        if (key.startsWith(`${prefix}:`))
+            flushStreamedOutputText(state, key);
+    }
+}
+function flushAllStreamedOutputText(state) {
+    for (const key of [...state.streamedTexts.keys()]) {
+        flushStreamedOutputText(state, key);
+    }
+}
+function recordToolArguments(accumulator, state, value) {
+    const characters = argumentCharacters(value);
+    if (characters === undefined) {
+        taintOutputEstimate(accumulator, state);
+        return;
+    }
+    recordOutputCharacters(accumulator, state, characters);
+}
+function appendToolArguments(accumulator, state, key, value) {
+    if (state.overflow)
+        return;
+    if (typeof value !== "string") {
+        recordToolArguments(accumulator, state, value);
+        return;
+    }
+    if (value.length === 0)
+        return;
+    const prior = state.pendingArguments.get(key);
+    if (prior === undefined && state.pendingArguments.size >= MAX_PENDING_TOOL_CALLS) {
+        taintOutputEstimate(accumulator, state);
+        return;
+    }
+    const nextSize = state.pendingArgumentCharacters + value.length;
+    if (!Number.isSafeInteger(nextSize) || nextSize > MAX_PENDING_ARGUMENT_CHARS) {
+        taintOutputEstimate(accumulator, state);
+        return;
+    }
+    state.pendingArguments.set(key, (prior ?? "") + value);
+    state.pendingArgumentCharacters = nextSize;
+    // Until the complete JSON can be sanitized, a partial tool argument could be
+    // base64. Hide the whole estimate rather than publish a knowingly partial total.
+    publishOutputEstimate(accumulator, state);
+}
+function flushToolArguments(accumulator, state, key) {
+    const value = state.pendingArguments.get(key);
+    if (value === undefined)
+        return;
+    state.pendingArguments.delete(key);
+    state.pendingArgumentCharacters -= value.length;
+    recordToolArguments(accumulator, state, value);
+    publishOutputEstimate(accumulator, state);
+}
+function flushAllToolArguments(accumulator, state) {
+    for (const key of [...state.pendingArguments.keys()]) {
+        flushToolArguments(accumulator, state, key);
+    }
+}
+function flushToolArgumentsPrefix(accumulator, state, prefix) {
+    for (const key of [...state.pendingArguments.keys()]) {
+        if (key.startsWith(`${prefix}:`))
+            flushToolArguments(accumulator, state, key);
+    }
+}
+function streamArgumentKey(prefix, value, fallback) {
+    const index = typeof value === "number" && Number.isSafeInteger(value) && value >= 0
+        ? value
+        : fallback;
+    return `${prefix}:${index}`;
+}
+function isAnthropicToolUseBlock(type) {
+    return type === "tool_use" || type === "server_tool_use" || type === "mcp_tool_use";
+}
+function inspectAnthropicOutput(root, accumulator, state) {
+    const inspectBlock = (value, includeToolArguments) => {
+        const block = asRecord(value);
+        if (!block)
+            return;
+        if (block.type === "text")
+            recordOutputText(accumulator, state, block.text);
+        else if (block.type === "thinking")
+            recordOutputText(accumulator, state, block.thinking);
+        else if (includeToolArguments && isAnthropicToolUseBlock(block.type)) {
+            recordToolArguments(accumulator, state, block.input);
+        }
+    };
+    if (root.type === "message") {
+        if (Array.isArray(root.content)) {
+            root.content.forEach((block) => inspectBlock(block, true));
+        }
+        return;
+    }
+    const key = streamArgumentKey("anthropic", root.index, "unknown");
+    if (root.type === "content_block_start") {
+        // Anthropic streams seed tool-use blocks with protocol-owned `input: {}`;
+        // the authored argument JSON arrives in input_json_delta frames.
+        const block = asRecord(root.content_block);
+        if (block?.type === "text") {
+            appendStreamedOutputText(accumulator, state, `${key}:text`, block.text);
+        }
+        else if (block?.type === "thinking") {
+            appendStreamedOutputText(accumulator, state, `${key}:thinking`, block.thinking);
+        }
+        return;
+    }
+    if (root.type === "content_block_stop") {
+        flushToolArguments(accumulator, state, key);
+        flushStreamedOutputTextPrefix(state, key);
+        return;
+    }
+    if (root.type !== "content_block_delta")
+        return;
+    const delta = asRecord(root.delta);
+    if (!delta)
+        return;
+    if (delta.type === "text_delta") {
+        appendStreamedOutputText(accumulator, state, `${key}:text`, delta.text);
+    }
+    else if (delta.type === "thinking_delta") {
+        appendStreamedOutputText(accumulator, state, `${key}:thinking`, delta.thinking);
+    }
+    else if (delta.type === "input_json_delta") {
+        appendToolArguments(accumulator, state, key, delta.partial_json);
+    }
+}
+function inspectOpenAiMessage(value, accumulator, state, streamKeyPrefix) {
+    const message = asRecord(value);
+    if (!message)
+        return;
+    if (typeof message.content === "string") {
+        if (streamKeyPrefix) {
+            appendStreamedOutputText(accumulator, state, `${streamKeyPrefix}:content`, message.content);
+        }
+        else {
+            recordOutputText(accumulator, state, message.content);
+        }
+    }
+    else if (Array.isArray(message.content)) {
+        for (const [partOffset, partValue] of message.content.entries()) {
+            const part = asRecord(partValue);
+            if (part && (part.type === "text" || part.type === "output_text")) {
+                if (streamKeyPrefix) {
+                    const key = streamArgumentKey(`${streamKeyPrefix}:content`, part.index, partOffset);
+                    appendStreamedOutputText(accumulator, state, key, part.text);
+                }
+                else {
+                    recordOutputText(accumulator, state, part.text);
+                }
+            }
+        }
+    }
+    // Providers use these as alternate spellings. Prefer the explicit Chat field
+    // so a host emitting both aliases cannot double the same reasoning text.
+    const reasoningContent = typeof message.reasoning_content === "string"
+        && message.reasoning_content.length > 0
+        ? message.reasoning_content
+        : message.reasoning;
+    if (streamKeyPrefix) {
+        appendStreamedOutputText(accumulator, state, `${streamKeyPrefix}:reasoning`, reasoningContent);
+        appendStreamedOutputText(accumulator, state, `${streamKeyPrefix}:refusal`, message.refusal);
+    }
+    else {
+        recordOutputText(accumulator, state, reasoningContent);
+        recordOutputText(accumulator, state, message.refusal);
+    }
+    const functionCall = asRecord(message.function_call);
+    const toolCalls = Array.isArray(message.tool_calls) ? message.tool_calls : [];
+    if (toolCalls.length === 0 && functionCall) {
+        if (streamKeyPrefix) {
+            appendToolArguments(accumulator, state, `${streamKeyPrefix}:function`, functionCall.arguments);
+        }
+        else {
+            recordToolArguments(accumulator, state, functionCall.arguments);
+        }
+    }
+    for (const [callOffset, callValue] of toolCalls.entries()) {
+        const call = asRecord(callValue);
+        const fn = call && asRecord(call.function);
+        if (!fn)
+            continue;
+        if (streamKeyPrefix) {
+            const key = streamArgumentKey(streamKeyPrefix, call?.index, callOffset);
+            appendToolArguments(accumulator, state, key, fn.arguments);
+        }
+        else {
+            recordToolArguments(accumulator, state, fn.arguments);
+        }
+    }
+}
+function inspectOpenAiOutput(root, accumulator, state, streamed) {
+    if (!Array.isArray(root.choices))
+        return;
+    for (const [choiceOffset, choiceValue] of root.choices.entries()) {
+        const choice = asRecord(choiceValue);
+        if (!choice)
+            continue;
+        inspectOpenAiMessage(choice.message, accumulator, state);
+        const prefix = streamArgumentKey("openai", choice.index, choiceOffset);
+        inspectOpenAiMessage(choice.delta, accumulator, state, streamed ? prefix : undefined);
+        if (streamed && choice.finish_reason !== undefined && choice.finish_reason !== null) {
+            flushToolArgumentsPrefix(accumulator, state, prefix);
+            flushStreamedOutputTextPrefix(state, prefix);
+        }
+    }
+}
 function usageRecords(value, protocol) {
     const root = asRecord(value);
     if (!root)
@@ -73,7 +437,19 @@ function inspectUsageRecord(usage, protocol, accumulator) {
     if (details)
         recordField(accumulator, "cachedInputTokens", details.cached_tokens);
 }
-function inspectJson(value, protocol, accumulator) {
+function inspectJson(value, protocol, accumulator, estimate, streamed = false) {
+    try {
+        const root = asRecord(value);
+        if (root) {
+            if (protocol === "anthropic-messages")
+                inspectAnthropicOutput(root, accumulator, estimate);
+            else
+                inspectOpenAiOutput(root, accumulator, estimate, streamed);
+        }
+    }
+    catch {
+        taintOutputEstimate(accumulator, estimate);
+    }
     try {
         for (const usage of usageRecords(value, protocol)) {
             inspectUsageRecord(usage, protocol, accumulator);
@@ -86,24 +462,45 @@ function inspectJson(value, protocol, accumulator) {
 function newSseState() {
     return { event: undefined, data: [], size: 0, overflow: false };
 }
-function inspectSseFrame(state, protocol, accumulator) {
+function inspectSseFrame(state, protocol, accumulator, estimate) {
     if (state.overflow || state.data.length === 0)
         return;
+    const data = state.data.join("\n");
+    const trimmed = data.trim();
+    if (trimmed.length === 0 || trimmed === "[DONE]")
+        return;
     try {
-        const value = JSON.parse(state.data.join("\n"));
+        const value = JSON.parse(data);
         if (protocol === "anthropic-messages") {
             const root = asRecord(value);
-            if (root && (state.event === "message_start" || state.event === "message_delta" || root.type === "message_start" || root.type === "message_delta")) {
-                inspectJson(root, protocol, accumulator);
+            if (root) {
+                try {
+                    inspectAnthropicOutput(root, accumulator, estimate);
+                }
+                catch {
+                    taintOutputEstimate(accumulator, estimate);
+                }
+                if (state.event === "message_start" || state.event === "message_delta" || root.type === "message_start" || root.type === "message_delta") {
+                    try {
+                        for (const usage of usageRecords(root, protocol)) {
+                            inspectUsageRecord(usage, protocol, accumulator);
+                        }
+                    }
+                    catch {
+                        // A reported-usage observer failure is isolated from pass-through.
+                    }
+                }
             }
         }
         else {
             // The only non-Anthropic protocol is openai-chat: every frame is inspected.
-            inspectJson(value, protocol, accumulator);
+            inspectJson(value, protocol, accumulator, estimate, true);
         }
     }
     catch {
-        // Malformed/incomplete SSE data is simply not a usage report.
+        // A skipped frame might have carried content. Preserve traffic, but do not
+        // leave a partial estimate looking complete.
+        taintOutputEstimate(accumulator, estimate);
     }
 }
 function resetSse(state) {
@@ -112,9 +509,9 @@ function resetSse(state) {
     state.size = 0;
     state.overflow = false;
 }
-function processSseLine(line, terminatorBytes, state, protocol, accumulator) {
+function processSseLine(line, terminatorBytes, state, protocol, accumulator, estimate) {
     if (line === "") {
-        inspectSseFrame(state, protocol, accumulator);
+        inspectSseFrame(state, protocol, accumulator, estimate);
         resetSse(state);
         return;
     }
@@ -124,6 +521,7 @@ function processSseLine(line, terminatorBytes, state, protocol, accumulator) {
     if (state.size > MAX_SSE_FRAME) {
         state.overflow = true;
         state.data = [];
+        taintOutputEstimate(accumulator, estimate);
         return;
     }
     if (state.overflow)
@@ -135,8 +533,9 @@ function processSseLine(line, terminatorBytes, state, protocol, accumulator) {
         state.data.push(line.slice(5).replace(/^ /, ""));
     }
 }
-function makeSseObserver(protocol, accumulator) {
+function makeSseObserver(protocol, accumulator, estimate) {
     const decoder = new TextDecoder();
+    const utf8Validator = new TextDecoder("utf-8", { fatal: true });
     let text = "";
     let discardUntilBreak = false;
     const state = newSseState();
@@ -187,25 +586,28 @@ function makeSseObserver(protocol, accumulator) {
                     state.overflow = true;
                     text = "";
                     discardUntilBreak = true;
+                    taintOutputEstimate(accumulator, estimate);
                 }
                 break;
             }
             const line = text.slice(0, index);
             text = text.slice(index + terminatorLength);
-            processSseLine(line, terminatorLength, state, protocol, accumulator);
+            processSseLine(line, terminatorLength, state, protocol, accumulator, estimate);
         }
         if (final && text.length > 0) {
             if (state.size + new TextEncoder().encode(text).byteLength <= MAX_SSE_FRAME) {
-                processSseLine(text, 0, state, protocol, accumulator);
+                processSseLine(text, 0, state, protocol, accumulator, estimate);
             }
-            else
+            else {
                 state.overflow = true;
+                taintOutputEstimate(accumulator, estimate);
+            }
             text = "";
         }
         if (final) {
             // A final line without the customary trailing blank line is still a
             // useful terminal event in practice.
-            inspectSseFrame(state, protocol, accumulator);
+            inspectSseFrame(state, protocol, accumulator, estimate);
             resetSse(state);
         }
     };
@@ -217,26 +619,41 @@ function makeSseObserver(protocol, accumulator) {
                 // code point split between these slices.
                 for (let offset = 0; offset < bytes.byteLength; offset += MAX_DECODE_SLICE) {
                     const end = Math.min(offset + MAX_DECODE_SLICE, bytes.byteLength);
-                    text += decoder.decode(bytes.subarray(offset, end), { stream: true });
+                    const slice = bytes.subarray(offset, end);
+                    try {
+                        utf8Validator.decode(slice, { stream: true });
+                    }
+                    catch {
+                        taintOutputEstimate(accumulator, estimate);
+                    }
+                    text += decoder.decode(slice, { stream: true });
                     drain(false);
                 }
             }
             catch {
-                // Ignore parser failures, retaining pass-through behavior.
+                taintOutputEstimate(accumulator, estimate);
             }
         },
         finish() {
             try {
+                try {
+                    utf8Validator.decode();
+                }
+                catch {
+                    taintOutputEstimate(accumulator, estimate);
+                }
                 text += decoder.decode();
                 drain(true);
+                flushAllStreamedOutputText(estimate);
+                flushAllToolArguments(accumulator, estimate);
             }
             catch {
-                // Ignore parser failures.
+                taintOutputEstimate(accumulator, estimate);
             }
         },
     };
 }
-function makeJsonObserver(protocol, accumulator) {
+function makeJsonObserver(protocol, accumulator, estimate) {
     const chunks = [];
     let size = 0;
     let overflow = false;
@@ -244,6 +661,8 @@ function makeJsonObserver(protocol, accumulator) {
         push(bytes) {
             if (overflow)
                 return;
+            if (bytes.byteLength === 0)
+                return;
             if (size + bytes.byteLength > MAX_BUFFERED_JSON) {
                 overflow = true;
                 chunks.length = 0;
@@ -262,10 +681,11 @@ function makeJsonObserver(protocol, accumulator) {
                     bytes.set(chunk, offset);
                     offset += chunk.byteLength;
                 }
-                inspectJson(JSON.parse(new TextDecoder().decode(bytes)), protocol, accumulator);
+                const json = new TextDecoder("utf-8", { fatal: true }).decode(bytes);
+                inspectJson(JSON.parse(json), protocol, accumulator, estimate);
             }
             catch {
-                // Empty, malformed, or non-UTF-8 JSON is not a report.
+                taintOutputEstimate(accumulator, estimate);
             }
         },
     };
@@ -277,9 +697,16 @@ function makeJsonObserver(protocol, accumulator) {
 export function observeUsage(response, protocol, accumulator, options = {}) {
     if (!response.body)
         return response;
+    const estimate = {
+        characters: 0,
+        overflow: false,
+        pendingArgumentCharacters: 0,
+        pendingArguments: new Map(),
+        streamedTexts: new Map(),
+    };
     const observer = options.streamed
-        ? makeSseObserver(protocol, accumulator)
-        : makeJsonObserver(protocol, accumulator);
+        ? makeSseObserver(protocol, accumulator, estimate)
+        : makeJsonObserver(protocol, accumulator, estimate);
     const transform = new TransformStream({
         transform(chunk, controller) {
             // Observer failures are isolated; errors from the upstream body are