@librechat/agents 3.3.11 → 3.3.12

package/src/llm/invoke.ts

@@ -12,6 +12,7 @@ import type { ChatGeneration } from '@langchain/core/outputs';
 import type { ToolCall } from '@langchain/core/messages/tool';
 import type { BaseMessage } from '@langchain/core/messages';
 import type { ToolOutputReferenceRegistry } from '@/tools/toolOutputReferences';
+import type { StreamLimitState } from '@/llm/streamLimits';
 import type { ContextOverflowContext } from '@/utils/errors';
 import type * as t from '@/types';
 import {
@@ -37,6 +38,16 @@ import {
 } from '@/messages/cache';
 import { ChatModelStreamHandler, dispatchesChatModelStream } from '@/stream';
 import { Constants, ContentTypes, GraphEvents, Providers } from '@/common';
+import {
+  enforceStreamLimitsForWireChunk,
+  registerActiveStreamLimitGeneration,
+  releaseStreamLimitGeneration,
+  resolveGenerationKey,
+  streamLimitAccountingEnabled,
+  StreamLimitExceededError,
+  STREAM_LIMIT_REDISPATCH_KEY,
+  STREAM_LIMIT_ATTEMPT_KEY,
+} from '@/llm/streamLimits';
 import { annotateMessagesForLLM } from '@/tools/toolOutputReferences';
 import { assertNotTruncatedToolCall } from '@/llm/truncation';
 import { manualToolStreamProviders } from '@/llm/providers';
@@ -83,8 +94,18 @@ export type InvokeContext = NonNullable<
 /**
  * Per-chunk callback for custom stream processing.
  * When provided, replaces the default `ChatModelStreamHandler`.
+ *
+ * `metadata` is the attempt's callback metadata (carrying the provider and
+ * stream-limit attempt stamps), so consumers that count against the stream
+ * limits key each model attempt separately.
  */
-export type OnChunk = (chunk: AIMessageChunk) => void | Promise<void>;
+export type OnChunk = (
+  chunk: AIMessageChunk,
+  metadata?: Record<string, unknown>
+) => void | Promise<void>;
+
+/** Unique per-model-attempt sequence; see the stamp in `attemptInvoke`. */
+let streamLimitAttemptSeq = 0;
 
 export function usesNativeOpenAIResponses(
   model: t.ChatModel,
@@ -659,21 +680,81 @@ function appendStreamChunk({
  * Pass an `onChunk` callback to override this with custom chunk processing
  * (e.g. summarization delta events).
  */
+interface AttemptInvokeParams {
+  model: t.ChatModel;
+  messages: BaseMessage[];
+  provider: Providers;
+  context?: InvokeContext;
+  onChunk?: OnChunk;
+  /** Accounting owner for callers that deliberately pass no `context`
+   * (summarization) — used ONLY for the attempt's accounting lease, never
+   * for charge claims. */
+  streamLimitState?: StreamLimitState;
+}
+
+/**
+ * One model attempt. Stamps the attempt identity into callback metadata
+ * (see the generation-key notes in `streamLimits.ts`), leases the attempt's
+ * accounting for its LIFETIME, and releases both from `finally`: retention
+ * must follow the attempt — a cancellation-ignoring straggler keeps its
+ * original budget no matter how many runs start and reset while it drains.
+ */
 export async function attemptInvoke(
+  params: AttemptInvokeParams,
+  config?: RunnableConfig
+): Promise<Partial<t.BaseGraphState>> {
+  const stampedConfig: RunnableConfig = {
+    ...config,
+    metadata: {
+      ...(config?.metadata ?? {}),
+      [Constants.INVOKED_PROVIDER]: params.provider,
+      /**
+       * One `attemptInvoke` call is one model attempt; primary, fallback,
+       * and retry attempts within a node otherwise share the same langgraph
+       * metadata, so without a unique attempt stamp a fallback re-streaming
+       * a tool call from scratch would be charged the failed primary's
+       * partial bytes (or a same-named sibling fallback's) and could
+       * falsely trip the stream limits. The stamp rides the same metadata
+       * rebuild that already attributes the serving provider.
+       */
+      [STREAM_LIMIT_ATTEMPT_KEY]: ++streamLimitAttemptSeq,
+    },
+  };
+  const rawLeaseTarget = params.context ?? params.streamLimitState;
+  /** No lease when no guard can fire: the lease only protects accounting
+   * entries, and fully disabled guards must allocate no bookkeeping at
+   * all — per-attempt included. */
+  const leaseTarget =
+    rawLeaseTarget != null && streamLimitAccountingEnabled(rawLeaseTarget)
+      ? rawLeaseTarget
+      : undefined;
+  const generationKey =
+    leaseTarget != null
+      ? resolveGenerationKey(
+        stampedConfig.metadata as Record<string, unknown>
+      )
+      : undefined;
+  if (leaseTarget != null && generationKey != null) {
+    registerActiveStreamLimitGeneration(leaseTarget, generationKey);
+  }
+  try {
+    return await attemptInvokeBody(params, stampedConfig);
+  } finally {
+    if (leaseTarget != null && generationKey != null) {
+      releaseStreamLimitGeneration(leaseTarget, generationKey);
+    }
+  }
+}
+
+async function attemptInvokeBody(
   {
     model,
     messages,
     provider,
     context,
     onChunk,
-  }: {
-    model: t.ChatModel;
-    messages: BaseMessage[];
-    provider: Providers;
-    context?: InvokeContext;
-    onChunk?: OnChunk;
-  },
-  config?: RunnableConfig
+  }: AttemptInvokeParams,
+  config: RunnableConfig
 ): Promise<Partial<t.BaseGraphState>> {
   /**
    * Pull the run-scoped tool output registry off the graph (when one
@@ -689,7 +770,7 @@ export async function attemptInvoke(
     callOptions: config,
   });
   const registry = context?.getOrCreateToolOutputRegistry();
-  const runId = config?.configurable?.run_id as string | undefined;
+  const runId = config.configurable?.run_id as string | undefined;
   const annotated = annotateMessagesForLLM(invocationMessages, registry, runId);
   /**
    * Keyed on the provider ACTUALLY serving this call, not the agent's primary.
@@ -748,14 +829,6 @@ export async function attemptInvoke(
    * wrong for fallback-served calls — or `ls_provider` — which derived
    * providers inherit from their base class.
    */
-  config = {
-    ...config,
-    metadata: {
-      ...(config?.metadata ?? {}),
-      [Constants.INVOKED_PROVIDER]: provider,
-    },
-  };
-
   if (model.stream) {
     /**
      * Observed, not dictated. `handleChatModelStart` fires with the chat
@@ -787,10 +860,40 @@ export async function attemptInvoke(
     let preempted = false;
     const registeredStreamHandler =
       getRegisteredDefaultChatStreamHandler(context);
+    /** A sibling's trip aborts the composed signal, but an adapter that
+     * ignores cancellation keeps yielding — and text-only chunks with the
+     * event cap off never throw in enforcement, so nothing else would stop
+     * the drain. Checked on every yielded chunk in all three loops;
+     * throwing closes the iterator and tears down the provider stream. */
+    const throwIfBreakerTripped = (): void => {
+      const signal = config.signal;
+      if (
+        signal?.aborted === true &&
+        signal.reason instanceof StreamLimitExceededError
+      ) {
+        throw signal.reason;
+      }
+    };
 
     if (onChunk) {
+      const attemptMetadata = config.metadata as
+        | Record<string, unknown>
+        | undefined;
       for await (const chunk of stream) {
-        await onChunk(chunk);
+        throwIfBreakerTripped();
+        /** An onChunk consumer replaces the stream handler entirely, so
+         * stream limits are enforced here for every such caller — public
+         * package consumers get no other accounting. The internal
+         * summarization onChunk charges producer-side itself and passes no
+         * context, precisely so this claim and its own never stack. */
+        if (context != null) {
+          enforceStreamLimitsForWireChunk({
+            graph: context,
+            metadata: attemptMetadata,
+            chunk,
+          });
+        }
+        await onChunk(chunk, attemptMetadata);
         finalChunk = appendStreamChunk({
           current: finalChunk,
           next: chunk,
@@ -801,6 +904,7 @@ export async function attemptInvoke(
       const metadata = config.metadata as Record<string, unknown> | undefined;
       const streamHandler = new ChatModelStreamHandler();
       for await (const chunk of stream) {
+        throwIfBreakerTripped();
         const handlingChunk = getStreamHandlingChunk({
           current: finalChunk,
           next: chunk,
@@ -813,6 +917,23 @@ export async function attemptInvoke(
             metadata,
             context
           );
+        } else if (context != null) {
+          /**
+           * A replay-skipped chunk yields no handling chunk, and in this
+           * local branch no `streamEvents` consumer judges the wire event
+           * either — yet a cumulative OpenRouter replay can still carry
+           * `tool_call_chunks` or complete `tool_calls` that are appended
+           * below. Charge the full limits (event budget and argument bytes)
+           * directly so neither cap can be bypassed. Consumer side: the
+           * local handler.handle call above claims as consumer, and one
+           * reused chunk object can alternate between these two arms.
+           */
+          enforceStreamLimitsForWireChunk({
+            graph: context,
+            metadata,
+            chunk,
+            side: 'consumer',
+          });
         }
         finalChunk = appendStreamChunk({
           current: finalChunk,
@@ -842,17 +963,40 @@ export async function attemptInvoke(
       }
     } else {
       const metadata = config.metadata as Record<string, unknown> | undefined;
+      /**
+       * The original wire chunk still reaches the registered handler through
+       * `streamEvents` (where the late-reasoning skip discards it AFTER the
+       * event guard counts it), so this inline re-dispatch of the transformed
+       * chunk is marked to not consume a second event-budget slot. Allocated
+       * once per attempt, only when a transformation occurs.
+       */
+      let redispatchMetadata: Record<string, unknown> | undefined;
       for await (const chunk of stream) {
+        throwIfBreakerTripped();
+        /**
+         * Charged synchronously, ahead of the decoupled `streamEvents`
+         * reader that will echo this same chunk to the registered handler:
+         * a lagging reader would otherwise let an oversized complete call
+         * return to LangGraph and reach ToolNode before the queued handler
+         * throws. The chunk is marked so the echo skips accounting.
+         */
+        if (context != null) {
+          enforceStreamLimitsForWireChunk({ graph: context, metadata, chunk });
+        }
         const handlingChunk = getStreamHandlingChunk({
           current: finalChunk,
           next: chunk,
           provider,
         });
         if (handlingChunk != null && handlingChunk !== chunk) {
+          redispatchMetadata ??= {
+            ...(metadata ?? {}),
+            [STREAM_LIMIT_REDISPATCH_KEY]: true,
+          };
           await registeredStreamHandler.handle(
             GraphEvents.CHAT_MODEL_STREAM,
             { chunk: handlingChunk },
-            metadata,
+            redispatchMetadata,
             context
           );
         }
@@ -994,6 +1138,7 @@ export async function tryFallbackProviders({
   primaryError,
   context,
   onChunk,
+  streamLimitState,
   overflowContext,
   prepareProviderMessages,
 }: {
@@ -1004,6 +1149,9 @@ export async function tryFallbackProviders({
   primaryError: unknown;
   context?: InvokeContext;
   onChunk?: OnChunk;
+  /** Accounting-lease owner forwarded to each fallback attempt (see
+   * `AttemptInvokeParams.streamLimitState`). */
+  streamLimitState?: StreamLimitState;
   /**
    * Prompt-size corroboration for signatures that are not self-describing.
    * Vertex AI's overflow is a bare `400` with no reason, so without this a
@@ -1074,6 +1222,16 @@ export async function tryFallbackProviders({
           maxContextTokens: fb.maxContextTokens,
           config: fbConfig,
         })) ?? messages;
+      /** A sibling can trip the breaker while the preparation above is
+       * awaited — and the catch below only sees attempts that THROW, so a
+       * provider that ignores an aborted signal and succeeds would resolve
+       * a run that must reject. Check before every fallback invocation. */
+      if (
+        config?.signal?.aborted === true &&
+        config.signal.reason instanceof StreamLimitExceededError
+      ) {
+        throw config.signal.reason;
+      }
       const result = await attemptInvoke(
         {
           model: fbModel as t.ChatModel,
@@ -1081,11 +1239,30 @@ export async function tryFallbackProviders({
           provider: fb.provider,
           context,
           onChunk,
+          streamLimitState,
         },
         fbConfig
       );
       return result;
     } catch (e) {
+      /**
+       * A tripped stream circuit breaker is a deliberate abort, not a
+       * provider failure. Continuing would try the remaining fallbacks and a
+       * succeeding one would resolve a run that must reject.
+       */
+      if (e instanceof StreamLimitExceededError) {
+        throw e;
+      }
+      /** A parallel sibling's trip aborts this branch's composed signal, and
+       * a provider can surface that as a generic abort error; advancing to
+       * the next fallback would start new provider work after the safety
+       * abort. Rethrow the breaker's own reason instead. */
+      if (
+        config?.signal?.aborted === true &&
+        config.signal.reason instanceof StreamLimitExceededError
+      ) {
+        throw config.signal.reason;
+      }
       lastError = e;
       const fallbackOverflowContext: ContextOverflowContext = {
         provider: fb.provider,

package/src/llm/openai/cacheWriteTokens.test.ts

@@ -0,0 +1,112 @@
+import { AIMessage } from '@langchain/core/messages';
+import type { OpenAIClient } from '@langchain/openai';
+
+import {
+  getCacheWriteTokens,
+  attachCacheWriteUsage,
+  attachCacheWriteMetadata,
+} from './index';
+
+/**
+ * Regression coverage for a crash reported against OpenAI-compatible
+ * third-party servers (e.g. mlx_vlm.server) whose `/v1/responses` usage
+ * payload omits `input_tokens_details` entirely — a shape the OpenAI API
+ * itself always populates, but which `ResponsesUsageWithCacheWrite`
+ * declares optional. Reading `.cache_write_tokens` off that missing object
+ * without a second `?.` threw "Cannot read properties of undefined
+ * (reading 'cache_write_tokens')" on every completion from such a server.
+ */
+describe('cache write token extraction (Responses API)', () => {
+  describe('getCacheWriteTokens', () => {
+    it('returns undefined without throwing when usage has no input_tokens_details', () => {
+      const message = new AIMessage({
+        content: 'hi',
+        response_metadata: {
+          usage: {
+            input_tokens: 10,
+            output_tokens: 2,
+            total_tokens: 12,
+            // input_tokens_details intentionally omitted, mirroring a
+            // minimal OpenAI-compatible server's usage payload.
+          },
+        },
+      });
+
+      expect(() => getCacheWriteTokens(message)).not.toThrow();
+      expect(getCacheWriteTokens(message)).toBeUndefined();
+    });
+
+    it('returns undefined without throwing when there is no usage at all', () => {
+      const message = new AIMessage({
+        content: 'hi',
+        response_metadata: {},
+      });
+
+      expect(() => getCacheWriteTokens(message)).not.toThrow();
+      expect(getCacheWriteTokens(message)).toBeUndefined();
+    });
+
+    it('still reports cache_write_tokens when the field is present', () => {
+      const message = new AIMessage({
+        content: 'hi',
+        response_metadata: {
+          usage: {
+            input_tokens: 10,
+            output_tokens: 2,
+            total_tokens: 12,
+            input_tokens_details: { cache_write_tokens: 5 },
+          },
+        },
+      });
+
+      expect(getCacheWriteTokens(message)).toBe(5);
+    });
+
+    it('falls back to the serialized metadata key when usage is absent', () => {
+      const message = new AIMessage({
+        content: 'hi',
+        response_metadata: {
+          metadata: { __librechat_cache_write_tokens: '7' },
+        },
+      });
+
+      expect(getCacheWriteTokens(message)).toBe(7);
+    });
+  });
+
+  describe('attachCacheWriteUsage', () => {
+    it('leaves usage_metadata untouched without throwing when input_tokens_details is missing', () => {
+      const message = new AIMessage({
+        content: 'hi',
+        response_metadata: {
+          usage: { input_tokens: 10, output_tokens: 2, total_tokens: 12 },
+        },
+        usage_metadata: {
+          input_tokens: 10,
+          output_tokens: 2,
+          total_tokens: 12,
+        },
+      });
+
+      expect(() => attachCacheWriteUsage(message)).not.toThrow();
+      expect(message.usage_metadata?.input_token_details).toBeUndefined();
+    });
+  });
+
+  describe('attachCacheWriteMetadata', () => {
+    it('returns the response unmodified without throwing when input_tokens_details is missing', () => {
+      const response = {
+        usage: { input_tokens: 10, output_tokens: 2, total_tokens: 12 },
+      } as OpenAIClient.Responses.Response;
+
+      expect(() => attachCacheWriteMetadata(response)).not.toThrow();
+      expect(attachCacheWriteMetadata(response)).toBe(response);
+    });
+
+    it('returns the response unmodified without throwing when usage is missing entirely', () => {
+      const response = {} as OpenAIClient.Responses.Response;
+
+      expect(() => attachCacheWriteMetadata(response)).not.toThrow();
+    });
+  });
+});

package/src/llm/openai/index.ts

@@ -245,7 +245,15 @@ type CacheableResponsePart = (
 ) & {
   prompt_cache_breakpoint?: { mode: 'explicit' };
 };
-type ResponsesUsageWithCacheWrite = OpenAIClient.Responses.ResponseUsage & {
+// `Omit` before re-adding `input_tokens_details` as optional matters: the SDK's own
+// `ResponseUsage` declares it required (true for OpenAI itself), so a plain intersection
+// would keep it required in the merged type despite the `?:` here — masking, at the type
+// level, that OpenAI-*compatible* servers (e.g. mlx_vlm.server) may omit it entirely.
+// Mirrors `CompletionUsageWithCacheWrite`'s handling of the analogous Completions API field.
+type ResponsesUsageWithCacheWrite = Omit<
+  OpenAIClient.Responses.ResponseUsage,
+  'input_tokens_details'
+> & {
   input_tokens_details?: OpenAIClient.Responses.ResponseUsage['input_tokens_details'] & {
     cache_write_tokens?: number;
   };
@@ -509,13 +517,13 @@ export function shouldIncludeEncryptedReasoning(
   );
 }
 
-function getCacheWriteTokens(message: BaseMessage): number | undefined {
+export function getCacheWriteTokens(message: BaseMessage): number | undefined {
   const responseMetadata = message.response_metadata as {
     usage?: ResponsesUsageWithCacheWrite;
     metadata?: Record<string, string>;
   };
   const reported =
-    responseMetadata.usage?.input_tokens_details.cache_write_tokens;
+    responseMetadata.usage?.input_tokens_details?.cache_write_tokens;
   if (reported != null) {
     return reported;
   }
@@ -527,7 +535,7 @@ function getCacheWriteTokens(message: BaseMessage): number | undefined {
   return Number.isFinite(parsed) ? parsed : undefined;
 }
 
-function attachCacheWriteUsage(message: BaseMessage): void {
+export function attachCacheWriteUsage(message: BaseMessage): void {
   const cacheWriteTokens = getCacheWriteTokens(message);
   if (
     cacheWriteTokens == null ||
@@ -554,11 +562,11 @@ function attachCacheWriteUsage(message: BaseMessage): void {
   };
 }
 
-function attachCacheWriteMetadata(
+export function attachCacheWriteMetadata(
   response: OpenAIClient.Responses.Response
 ): OpenAIClient.Responses.Response {
   const usage = response.usage as ResponsesUsageWithCacheWrite | undefined;
-  const cacheWriteTokens = usage?.input_tokens_details.cache_write_tokens;
+  const cacheWriteTokens = usage?.input_tokens_details?.cache_write_tokens;
   if (cacheWriteTokens == null) {
     return response;
   }