@librechat/agents 3.8.0 → 3.8.2

package/src/utils/tokens.ts

@@ -16,6 +16,7 @@ export type UnsafeTokenMeasurementReason =
   | 'content_proxy'
   | 'metadata_proxy'
   | 'metadata_accessor'
+  | 'metadata_limit'
   | 'invalid_count';
 
 export class UnsafeTokenMeasurementError extends Error {
@@ -121,6 +122,7 @@ const MAX_STRUCTURED_TOKENIZATION_CHARS = 200_000;
 /** One UTF-16 character can encode to several tokenizer pieces. */
 const MAX_TOKENS_PER_OMITTED_STRUCTURED_CHAR = 4;
 const MAX_STRUCTURED_SAFETY_INSPECTION_WORK = 10_000;
+const MAX_RAW_TOOL_CALLS = MAX_STRUCTURED_SAFETY_INSPECTION_WORK;
 /** Bounds BPE work per plain-text segment without omitting any text. */
 const MAX_TEXT_TOKENIZATION_CHARS = 8_192;
 /** Covers tokenizer segmentation changes on both sides of a chunk boundary. */
@@ -946,6 +948,63 @@ function getBoundedStructuredTokenCount(
   );
 }
 
+export function readTokenMetadataProperty(
+  value: object,
+  key: PropertyKey,
+  path: string
+): unknown {
+  const seen = new Set<object>();
+  let current: object | null = value;
+  for (let depth = 0; current != null && depth < 100; depth++) {
+    if (isProxy(current)) {
+      throw new UnsafeTokenMeasurementError({
+        reason: 'metadata_proxy',
+        path,
+      });
+    }
+    if (seen.has(current)) {
+      throw new UnsafeTokenMeasurementError({
+        reason: 'metadata_accessor',
+        path,
+      });
+    }
+    seen.add(current);
+    let descriptor: PropertyDescriptor | undefined;
+    try {
+      descriptor = Object.getOwnPropertyDescriptor(current, key);
+    } catch {
+      throw new UnsafeTokenMeasurementError({
+        reason: 'metadata_accessor',
+        path,
+      });
+    }
+    if (descriptor != null) {
+      if (!('value' in descriptor)) {
+        throw new UnsafeTokenMeasurementError({
+          reason: 'metadata_accessor',
+          path,
+        });
+      }
+      return descriptor.value;
+    }
+    try {
+      current = Object.getPrototypeOf(current) as object | null;
+    } catch {
+      throw new UnsafeTokenMeasurementError({
+        reason: 'metadata_accessor',
+        path,
+      });
+    }
+  }
+  if (current != null) {
+    throw new UnsafeTokenMeasurementError({
+      reason: 'metadata_accessor',
+      path,
+    });
+  }
+  return undefined;
+}
+
 export function getTokenCountForMessage(
   message: BaseMessage,
   getTokenCount: (text: string) => number,
@@ -1187,6 +1246,130 @@ export function getTokenCountForMessage(
             : getBoundedStructuredTokenCount(args, countText);
       }
     }
+    const rawCalls =
+      toolCalls.length > 0 || additionalKwargs == null
+        ? undefined
+        : readTokenMetadataProperty(
+          additionalKwargs,
+          'tool_calls',
+          'additional_kwargs.tool_calls'
+        );
+    if (rawCalls != null) {
+      if (isProxy(rawCalls)) {
+        throw new UnsafeTokenMeasurementError({
+          reason: 'metadata_proxy',
+          path: 'additional_kwargs.tool_calls',
+        });
+      }
+      if (!Array.isArray(rawCalls)) {
+        throw new UnsafeTokenMeasurementError({
+          reason: 'metadata_accessor',
+          path: 'additional_kwargs.tool_calls',
+        });
+      }
+      if (rawCalls.length > MAX_RAW_TOOL_CALLS) {
+        throw new UnsafeTokenMeasurementError({
+          reason: 'metadata_limit',
+          path: 'additional_kwargs.tool_calls',
+        });
+      }
+      for (let i = 0; i < rawCalls.length; i++) {
+        const rawCall = readTokenMetadataProperty(
+          rawCalls,
+          String(i),
+          `additional_kwargs.tool_calls[${i}]`
+        );
+        if (
+          rawCall == null ||
+          typeof rawCall !== 'object' ||
+          isProxy(rawCall)
+        ) {
+          throw new UnsafeTokenMeasurementError({
+            reason: 'metadata_proxy',
+            path: `additional_kwargs.tool_calls[${i}]`,
+          });
+        }
+        const callPath = `additional_kwargs.tool_calls[${i}]`;
+        if (hasUnsafeStructuredSerialization(rawCall)) {
+          throw new UnsafeTokenMeasurementError({
+            reason: 'metadata_accessor',
+            path: callPath,
+          });
+        }
+        const callId = readTokenMetadataProperty(
+          rawCall,
+          'id',
+          `${callPath}.id`
+        );
+        const callType = readTokenMetadataProperty(
+          rawCall,
+          'type',
+          `${callPath}.type`
+        );
+        const rawFunction = readTokenMetadataProperty(
+          rawCall,
+          'function',
+          `${callPath}.function`
+        );
+        const rawCustom = readTokenMetadataProperty(
+          rawCall,
+          'custom',
+          `${callPath}.custom`
+        );
+        let projectedFunction = rawFunction;
+        if (rawFunction != null && typeof rawFunction === 'object') {
+          if (hasUnsafeStructuredSerialization(rawFunction)) {
+            throw new UnsafeTokenMeasurementError({
+              reason: 'metadata_accessor',
+              path: `${callPath}.function`,
+            });
+          }
+          const name = readTokenMetadataProperty(
+            rawFunction,
+            'name',
+            `${callPath}.function.name`
+          );
+          const args = readTokenMetadataProperty(
+            rawFunction,
+            'arguments',
+            `${callPath}.function.arguments`
+          );
+          projectedFunction = { name, arguments: args };
+        }
+        let projectedCustom = rawCustom;
+        if (rawCustom != null && typeof rawCustom === 'object') {
+          if (hasUnsafeStructuredSerialization(rawCustom)) {
+            throw new UnsafeTokenMeasurementError({
+              reason: 'metadata_accessor',
+              path: `${callPath}.custom`,
+            });
+          }
+          const name = readTokenMetadataProperty(
+            rawCustom,
+            'name',
+            `${callPath}.custom.name`
+          );
+          const input = readTokenMetadataProperty(
+            rawCustom,
+            'input',
+            `${callPath}.custom.input`
+          );
+          projectedCustom = { name, input };
+        }
+        if (typeof callId === 'string' && representedToolCallIds.has(callId)) {
+          continue;
+        }
+        numTokens += getBoundedStructuredTokenCount(
+          {
+            id: callId,
+            type: callType,
+            custom: projectedCustom,
+            function: projectedFunction,
+          },
+          countText
+        );
+      }
+    }
     let legacyFunctionCall: PropertyDescriptor | undefined;
     try {
       legacyFunctionCall =