@librechat/agents 3.2.34 → 3.2.36

package/src/agents/AgentContext.ts

@@ -7,7 +7,6 @@ import type {
   BaseMessageFields,
 } from '@langchain/core/messages';
 import type { RunnableConfig, Runnable } from '@langchain/core/runnables';
-import type { createPruneMessages } from '@/messages';
 import type * as t from '@/types';
 import {
   ANTHROPIC_TOOL_TOKEN_MULTIPLIER,
@@ -19,9 +18,16 @@ import {
 import {
   addCacheControl,
   addCacheControlToStablePrefixMessages,
+  cloneMessage,
 } from '@/messages/cache';
 import { createSchemaOnlyTools } from '@/tools/schema';
-import { DEFAULT_RESERVE_RATIO } from '@/messages';
+import { apportionTokenCounts } from '@/utils/tokens';
+import {
+  DEFAULT_RESERVE_RATIO,
+  createPruneMessages,
+  syncBudgetDerivedFields,
+} from '@/messages';
+import { isThinkingEnabled } from '@/llm/request';
 import { toJsonSchema } from '@/utils/schema';
 
 type AgentSystemTextBlock = {
@@ -191,6 +197,11 @@ export class AgentContext {
   dynamicInstructionTokens: number = 0;
   /** Token count for tool schemas only. */
   toolSchemaTokens: number = 0;
+  /** Per-tool schema token counts (post-multiplier), keyed by tool name.
+   *  `undefined` when not calculated (e.g. cached aggregate schema tokens). */
+  toolTokenCounts?: Record<string, number>;
+  /** Names of counted tools that are deferred (`defer_loading`) and discovered. */
+  deferredToolNames: string[] = [];
   /** Running calibration ratio from the pruner — persisted across runs via contextMeta. */
   calibrationRatio: number = 1;
   /** Provider-observed instruction overhead from the pruner's best-variance turn. */
@@ -894,6 +905,8 @@ export class AgentContext {
     this.systemMessageTokens = 0;
     this.dynamicInstructionTokens = 0;
     this.toolSchemaTokens = 0;
+    this.toolTokenCounts = undefined;
+    this.deferredToolNames = [];
     this.cachedSystemRunnable = undefined;
     this.systemRunnableStale = true;
     this.lastToken = undefined;
@@ -1006,6 +1019,10 @@ export class AgentContext {
   ): Promise<void> {
     let toolTokens = 0;
     const countedToolNames = new Set<string>();
+    /** Prototype-free: external tool names like `toString` must not hit
+     *  inherited properties during accumulation */
+    const rawToolTokenCounts: Record<string, number> = Object.create(null);
+    const deferredCountedNames = new Set<string>();
 
     /**
      * Iterate both `tools` (user-provided instance tools) and `graphTools`
@@ -1040,11 +1057,14 @@ export class AgentContext {
             toolName,
             (genericTool.description as string | undefined) ?? ''
           );
-          toolTokens += tokenCounter(
+          const schemaTokens = tokenCounter(
             new SystemMessage(JSON.stringify(jsonSchema))
           );
+          toolTokens += schemaTokens;
           if (toolName) {
             countedToolNames.add(toolName);
+            rawToolTokenCounts[toolName] =
+              (rawToolTokenCounts[toolName] ?? 0) + schemaTokens;
           }
         }
       }
@@ -1062,7 +1082,16 @@ export class AgentContext {
           parameters: def.parameters ?? {},
         },
       };
-      toolTokens += tokenCounter(new SystemMessage(JSON.stringify(schema)));
+      const schemaTokens = tokenCounter(
+        new SystemMessage(JSON.stringify(schema))
+      );
+      toolTokens += schemaTokens;
+      countedToolNames.add(def.name);
+      rawToolTokenCounts[def.name] =
+        (rawToolTokenCounts[def.name] ?? 0) + schemaTokens;
+      if (def.defer_loading === true) {
+        deferredCountedNames.add(def.name);
+      }
     }
 
     const isAnthropic =
@@ -1077,6 +1106,25 @@ export class AgentContext {
       ? ANTHROPIC_TOOL_TOKEN_MULTIPLIER
       : DEFAULT_TOOL_TOKEN_MULTIPLIER;
     this.toolSchemaTokens = Math.ceil(toolTokens * toolTokenMultiplier);
+
+    /** Largest-remainder apportionment keeps the per-tool counts summing
+     *  exactly to the aggregate despite per-entry rounding */
+    const toolTokenCounts = apportionTokenCounts(
+      rawToolTokenCounts,
+      toolTokenMultiplier,
+      this.toolSchemaTokens
+    );
+    const deferredToolNames: string[] = [];
+    for (const name of Object.keys(rawToolTokenCounts)) {
+      if (
+        deferredCountedNames.has(name) ||
+        this.toolRegistry?.get(name)?.defer_loading === true
+      ) {
+        deferredToolNames.push(name);
+      }
+    }
+    this.toolTokenCounts = toolTokenCounts;
+    this.deferredToolNames = deferredToolNames;
   }
 
   /**
@@ -1212,9 +1260,8 @@ export class AgentContext {
    * Returns a structured breakdown of how the context token budget is consumed.
    * Useful for diagnostics when context overflow or pruning issues occur.
    *
-   * Note: `toolCount` reflects discoveries immediately, but `toolSchemaTokens`
-   * is a snapshot taken during `calculateInstructionTokens` and is not
-   * recomputed when `markToolsAsDiscovered` is called mid-run.
+   * Note: `markToolsAsDiscovered` re-triggers `calculateInstructionTokens`,
+   * so `toolSchemaTokens`/`toolTokenCounts` refresh before the next call.
    */
   getTokenBudgetBreakdown(messages?: BaseMessage[]): t.TokenBudgetBreakdown {
     const maxContextTokens = this.maxContextTokens ?? 0;
@@ -1238,7 +1285,14 @@ export class AgentContext {
       }
     }
 
-    const reserveTokens = Math.round(maxContextTokens * DEFAULT_RESERVE_RATIO);
+    /** Mirror the pruner's reserve math so availableForMessages agrees
+     *  with the contextBudget computed during pruning */
+    const reserveRatio =
+      this.summarizationConfig?.reserveRatio ?? DEFAULT_RESERVE_RATIO;
+    const reserveTokens =
+      reserveRatio > 0 && reserveRatio < 1
+        ? Math.round(maxContextTokens * reserveRatio)
+        : 0;
     const availableForMessages = Math.max(
       0,
       maxContextTokens - reserveTokens - this.instructionTokens
@@ -1255,6 +1309,12 @@ export class AgentContext {
       messageCount,
       messageTokens,
       availableForMessages,
+      toolTokenCounts:
+        this.toolTokenCounts != null ? { ...this.toolTokenCounts } : undefined,
+      deferredToolNames:
+        this.deferredToolNames.length > 0
+          ? [...this.deferredToolNames]
+          : undefined,
     };
   }
 
@@ -1275,6 +1335,102 @@ export class AgentContext {
     return lines.join('\n');
   }
 
+  /**
+   * Projects the context-usage snapshot for an arbitrary message set WITHOUT
+   * invoking the model — the pre-send / page-load / window-switch counterpart to
+   * the live `ON_CONTEXT_USAGE` snapshot. Runs the same pruner + budget math the
+   * graph uses (`createPruneMessages` → `getTokenBudgetBreakdown` →
+   * `syncBudgetDerivedFields`) so projected numbers match a real call. Returns
+   * null when the context lacks the tokenizer or window needed to prune. Omits
+   * the live post-format reconciliation (provider-specific, invoke-time) — a
+   * small, acceptable delta for a pre-send estimate.
+   *
+   * Safe to call off the hot path: the supplied `messages` are never mutated
+   * (each is passed as a clone — the pruner both replaces tool-result slots and
+   * unshifts reasoning blocks into AI content arrays in place), and this
+   * context's own state is untouched apart from refreshing stale instruction
+   * counts (idempotent, exactly what a real call does). Token counts are
+   * recounted for the supplied messages (the context's `indexTokenCountMap` is
+   * keyed to the live run's branch and would missum an arbitrary branch) unless
+   * the caller passes a map it guarantees matches. Calibration is NOT re-derived
+   * from this context's live usage (a fresh pruner would compare the prior
+   * call's provider input against the whole projected branch); the learned
+   * `calibrationRatio` is applied as a static seed, and callers may override it
+   * with a persisted ratio via `opts.calibrationRatio`.
+   */
+  projectContextUsage(
+    messages: BaseMessage[],
+    opts?: {
+      runId?: string;
+      agentId?: string;
+      calibrationRatio?: number;
+      indexTokenCountMap?: Record<string, number | undefined>;
+    }
+  ): t.ContextUsageEvent | null {
+    const tokenCounter = this.tokenCounter;
+    if (tokenCounter == null || this.maxContextTokens == null) {
+      return null;
+    }
+    /** Refresh stale system overhead (handoff/summary changes) so instruction
+     *  tokens match the prompt a real call would send. */
+    this.initializeSystemRunnable();
+    /** Clone array-content messages: the pruner unshifts reasoning blocks into
+     *  AI content arrays in place, which would otherwise corrupt the caller's
+     *  history. (Slot replacements land on the mapped array, not the caller's.) */
+    const projected = messages.map((message) =>
+      Array.isArray(message.content)
+        ? cloneMessage(message, [...message.content])
+        : message
+    );
+    let indexTokenCountMap = opts?.indexTokenCountMap;
+    if (indexTokenCountMap == null) {
+      indexTokenCountMap = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[String(i)] = tokenCounter(messages[i]);
+      }
+    }
+    const prune = createPruneMessages({
+      startIndex: 0,
+      provider: this.provider,
+      tokenCounter,
+      maxTokens: this.maxContextTokens,
+      thinkingEnabled: isThinkingEnabled(this.provider, this.clientOptions),
+      indexTokenCountMap,
+      contextPruningConfig: this.contextPruningConfig,
+      summarizationEnabled: this.summarizationEnabled,
+      reserveRatio: this.summarizationConfig?.reserveRatio,
+      calibrationRatio: opts?.calibrationRatio ?? this.calibrationRatio,
+      getInstructionTokens: () => this.instructionTokens,
+    });
+    const {
+      context,
+      prePruneContextTokens,
+      remainingContextTokens,
+      contextBudget,
+      effectiveInstructionTokens,
+      calibrationRatio,
+    } = prune({
+      messages: projected,
+      usageMetadata: undefined,
+      lastCallUsage: undefined,
+      totalTokensFresh: false,
+    });
+    const breakdown = this.getTokenBudgetBreakdown(messages);
+    breakdown.messageCount = context.length;
+    const usage: t.ContextUsageEvent = {
+      runId: opts?.runId,
+      agentId: opts?.agentId,
+      breakdown,
+      contextBudget,
+      effectiveInstructionTokens,
+      prePruneContextTokens,
+      remainingContextTokens,
+      calibrationRatio,
+    };
+    syncBudgetDerivedFields(usage);
+    return usage;
+  }
+
   /**
    * Updates the last-call usage with data from the most recent LLM response.
    * Unlike `currentUsage` which accumulates, this captures only the single call.
@@ -1324,6 +1480,14 @@ export class AgentContext {
     }
     if (hasNewDiscoveries) {
       this.systemRunnableStale = true;
+      /** Refresh schema token accounting so the next call's budget and
+       *  per-tool breakdown include the newly discovered tools; awaited
+       *  via tokenCalculationPromise before the next model call */
+      if (this.tokenCounter) {
+        this.tokenCalculationPromise = this.calculateInstructionTokens(
+          this.tokenCounter
+        );
+      }
     }
     return hasNewDiscoveries;
   }

package/src/agents/__tests__/AgentContext.test.ts

@@ -1414,7 +1414,7 @@ describe('AgentContext', () => {
       expect(ctx.getTokenBudgetBreakdown().toolCount).toBe(2);
     });
 
-    it('toolSchemaTokens snapshot does not auto-update after markToolsAsDiscovered', async () => {
+    it('refreshes toolSchemaTokens and per-tool counts after markToolsAsDiscovered', async () => {
       const toolDefinitions: t.LCTool[] = [
         {
           name: 'deferred',
@@ -1431,9 +1431,13 @@ describe('AgentContext', () => {
 
       await ctx.tokenCalculationPromise;
       expect(ctx.toolSchemaTokens).toBe(0);
+      expect(ctx.toolTokenCounts).toEqual({});
 
       ctx.markToolsAsDiscovered(['deferred']);
-      expect(ctx.toolSchemaTokens).toBe(0);
+      await ctx.tokenCalculationPromise;
+      expect(ctx.toolSchemaTokens).toBeGreaterThan(0);
+      expect(ctx.toolTokenCounts?.deferred).toBeGreaterThan(0);
+      expect(ctx.deferredToolNames).toContain('deferred');
     });
   });
 
@@ -2143,4 +2147,233 @@ describe('AgentContext', () => {
       expect(ctx.lastCallUsage!.inputTokens).toBe(8005);
     });
   });
+
+  describe('projectContextUsage', () => {
+    const countByChars = (msg: { content: unknown }): number => {
+      const content =
+        typeof msg.content === 'string'
+          ? msg.content
+          : JSON.stringify(msg.content);
+      return content.length;
+    };
+
+    const buildBranch = (
+      maxContextTokens: number,
+      perMessageTokens: number,
+      count: number,
+    ): { ctx: AgentContext; messages: AIMessage[] } => {
+      const ctx = createBasicContext({ tokenCounter: countByChars });
+      ctx.maxContextTokens = maxContextTokens;
+      const messages: AIMessage[] = [];
+      for (let i = 0; i < count; i++) {
+        // countByChars counts content length, and projectContextUsage recounts
+        // the supplied messages — so size content to the intended per-msg tokens.
+        const content = 'x'.repeat(perMessageTokens);
+        messages.push(
+          i % 2 === 0
+            ? (new HumanMessage(content) as unknown as AIMessage)
+            : new AIMessage(content),
+        );
+      }
+      return { ctx, messages };
+    };
+
+    it('returns null without a tokenizer or a window', () => {
+      const noCounter = createBasicContext({});
+      noCounter.maxContextTokens = 1000;
+      expect(noCounter.projectContextUsage([new HumanMessage('hi')])).toBeNull();
+
+      const noWindow = createBasicContext({ tokenCounter: countByChars });
+      noWindow.maxContextTokens = undefined;
+      expect(noWindow.projectContextUsage([new HumanMessage('hi')])).toBeNull();
+    });
+
+    it('keeps the whole branch and reports headroom when it fits', () => {
+      const { ctx, messages } = buildBranch(100_000, 1_000, 4);
+      const usage = ctx.projectContextUsage(messages);
+
+      expect(usage).not.toBeNull();
+      expect(usage!.breakdown.messageCount).toBe(4);
+      expect(usage!.breakdown.maxContextTokens).toBe(100_000);
+      expect(usage!.remainingContextTokens).toBeGreaterThan(0);
+      expect(usage!.breakdown.messageTokens).toBeGreaterThan(0);
+
+      const max = usage!.contextBudget ?? usage!.breakdown.maxContextTokens;
+      const used = max - (usage!.remainingContextTokens ?? 0);
+      expect(used).toBeLessThanOrEqual(max);
+    });
+
+    it('prunes older messages when the branch exceeds the window', () => {
+      const { ctx, messages } = buildBranch(3_000, 1_000, 6);
+      const usage = ctx.projectContextUsage(messages);
+
+      expect(usage).not.toBeNull();
+      expect(usage!.breakdown.messageCount).toBeGreaterThan(0);
+      expect(usage!.breakdown.messageCount).toBeLessThan(6);
+      expect(usage!.remainingContextTokens).toBeGreaterThanOrEqual(0);
+
+      const max = usage!.contextBudget ?? usage!.breakdown.maxContextTokens;
+      expect(max - (usage!.remainingContextTokens ?? 0)).toBeLessThanOrEqual(max);
+    });
+
+    it('does not mutate the context (local pruner, no field writes)', () => {
+      const { ctx, messages } = buildBranch(3_000, 1_000, 6);
+      const mapBefore = { ...ctx.indexTokenCountMap };
+
+      expect(ctx.pruneMessages).toBeUndefined();
+      ctx.projectContextUsage(messages);
+
+      expect(ctx.pruneMessages).toBeUndefined();
+      expect(ctx.indexTokenCountMap).toEqual(mapBefore);
+    });
+
+    it('does not mutate the caller messages under context pressure', () => {
+      const ctx = createBasicContext({ tokenCounter: countByChars });
+      ctx.maxContextTokens = 400;
+      const consumed = new ToolMessage({
+        content: 'x'.repeat(20_000),
+        tool_call_id: 't1',
+        name: 'tool',
+      });
+      const messages: AIMessage[] = [
+        new HumanMessage('question') as unknown as AIMessage,
+        new AIMessage({
+          content: '',
+          tool_calls: [{ id: 't1', name: 'tool', args: {} }],
+        }),
+        consumed as unknown as AIMessage,
+        new AIMessage('final answer'),
+      ];
+      const originalRef = messages[2];
+      const originalContent = (messages[2] as unknown as ToolMessage).content;
+
+      ctx.projectContextUsage(messages);
+
+      expect(messages[2]).toBe(originalRef);
+      expect((messages[2] as unknown as ToolMessage).content).toBe(
+        originalContent,
+      );
+    });
+
+    it('recounts the supplied branch, ignoring a stale context token map', () => {
+      const ctx = createBasicContext({ tokenCounter: countByChars });
+      ctx.maxContextTokens = 3_000;
+      // Empty/stale map — if it were reused, every message would count as 0 and
+      // nothing would prune. The fresh recount must drive pruning instead.
+      ctx.indexTokenCountMap = {};
+      const messages: AIMessage[] = [];
+      for (let i = 0; i < 6; i++) {
+        messages.push(new HumanMessage('x'.repeat(1_000)) as unknown as AIMessage);
+      }
+
+      const usage = ctx.projectContextUsage(messages);
+
+      expect(usage).not.toBeNull();
+      expect(usage!.breakdown.messageCount).toBeLessThan(6);
+    });
+
+    it('uses a caller-supplied token map when provided', () => {
+      const { ctx, messages } = buildBranch(3_000, 1, 6);
+      // Each message is ~1 char, so a recount would fit all 6. The supplied map
+      // claims 1000 each, forcing a prune — proving the map is honored.
+      const indexTokenCountMap: Record<string, number> = {};
+      for (let i = 0; i < messages.length; i++) {
+        indexTokenCountMap[String(i)] = 1_000;
+      }
+
+      const usage = ctx.projectContextUsage(messages, { indexTokenCountMap });
+
+      expect(usage!.breakdown.messageCount).toBeLessThan(6);
+    });
+
+    it('ignores this context live usage so projections are not recalibrated', () => {
+      const build = (): { ctx: AgentContext; messages: AIMessage[] } => {
+        const ctx = createBasicContext({ tokenCounter: countByChars });
+        ctx.maxContextTokens = 5_000;
+        const messages: AIMessage[] = [0, 1, 2].map(
+          () => new HumanMessage('x'.repeat(1_000)) as unknown as AIMessage,
+        );
+        return { ctx, messages };
+      };
+
+      const clean = build();
+      const cleanUsage = clean.ctx.projectContextUsage(clean.messages);
+
+      const dirty = build();
+      dirty.ctx.currentUsage = {
+        input_tokens: 4_000,
+        output_tokens: 50,
+        total_tokens: 4_050,
+      };
+      dirty.ctx.updateLastCallUsage({ input_tokens: 4_000, output_tokens: 50 });
+      const dirtyUsage = dirty.ctx.projectContextUsage(dirty.messages);
+
+      expect(dirtyUsage!.remainingContextTokens).toBe(
+        cleanUsage!.remainingContextTokens,
+      );
+      expect(dirtyUsage!.calibrationRatio).toBe(cleanUsage!.calibrationRatio);
+    });
+
+    it('does not mutate AI message content arrays during projection', () => {
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.ANTHROPIC,
+          clientOptions: {
+            model: 'claude-x',
+            thinking: { type: 'enabled', budget_tokens: 1024 },
+          } as never,
+        },
+        tokenCounter: countByChars,
+      });
+      ctx.maxContextTokens = 2_000;
+      const aiContent = [
+        { type: 'thinking', thinking: 'step by step', signature: 'sig' },
+        { type: 'text', text: 'the answer' },
+      ];
+      const ai = new AIMessage({ content: aiContent as never });
+      const messages: AIMessage[] = [
+        new HumanMessage('question') as unknown as AIMessage,
+        ai,
+        new HumanMessage('another') as unknown as AIMessage,
+      ];
+      const contentRef = ai.content;
+      const lenBefore = (ai.content as unknown[]).length;
+
+      ctx.projectContextUsage(messages);
+
+      expect(messages[1].content).toBe(contentRef);
+      expect((messages[1].content as unknown[]).length).toBe(lenBefore);
+    });
+
+    it('honors an explicit calibrationRatio seed', () => {
+      const base = buildBranch(100_000, 1_000, 4);
+      const baseUsage = base.ctx.projectContextUsage(base.messages);
+
+      const scaled = buildBranch(100_000, 1_000, 4);
+      const scaledUsage = scaled.ctx.projectContextUsage(scaled.messages, {
+        calibrationRatio: 3,
+      });
+
+      expect(scaledUsage!.calibrationRatio).toBe(3);
+      expect(scaledUsage!.remainingContextTokens).not.toBe(
+        baseUsage!.remainingContextTokens,
+      );
+    });
+
+    it('refreshes a stale system runnable before projecting', () => {
+      const ctx = createBasicContext({
+        agentConfig: { instructions: 'system prompt' },
+        tokenCounter: countByChars,
+      });
+      ctx.maxContextTokens = 5_000;
+      ctx.initializeSystemRunnable();
+      const systemBefore = ctx.systemMessageTokens;
+
+      // Adds a handoff preamble + marks stale, but defers the token recount.
+      ctx.setHandoffContext('PriorAgent', ['SiblingA', 'SiblingB']);
+      ctx.projectContextUsage([new HumanMessage('hi') as unknown as AIMessage]);
+
+      expect(ctx.systemMessageTokens).toBeGreaterThan(systemBefore);
+    });
+  });
 });

package/src/agents/__tests__/projection.test.ts

@@ -0,0 +1,73 @@
+import { AIMessage, HumanMessage } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { Providers } from '@/common';
+import { projectAgentContextUsage } from '../projection';
+
+const countByChars = (msg: { content: unknown }): number => {
+  const content =
+    typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content);
+  return content.length;
+};
+
+const agent = (maxContextTokens: number): t.AgentInputs => ({
+  agentId: 'test-agent',
+  provider: Providers.OPENAI,
+  instructions: 'system prompt',
+  maxContextTokens,
+});
+
+const branch = (perMessageChars: number, count: number): AIMessage[] => {
+  const messages: AIMessage[] = [];
+  for (let i = 0; i < count; i++) {
+    const content = 'x'.repeat(perMessageChars);
+    messages.push(
+      i % 2 === 0
+        ? (new HumanMessage(content) as unknown as AIMessage)
+        : new AIMessage(content),
+    );
+  }
+  return messages;
+};
+
+describe('projectAgentContextUsage', () => {
+  it('returns a budget snapshot for a branch that fits', async () => {
+    const usage = await projectAgentContextUsage({
+      agent: agent(100_000),
+      messages: branch(1_000, 4),
+      tokenCounter: countByChars,
+    });
+
+    expect(usage).not.toBeNull();
+    expect(usage!.breakdown.maxContextTokens).toBe(100_000);
+    expect(usage!.breakdown.messageCount).toBe(4);
+    expect(usage!.remainingContextTokens).toBeGreaterThan(0);
+    expect(usage!.agentId).toBe('test-agent');
+  });
+
+  it('prunes when the branch exceeds the window', async () => {
+    const usage = await projectAgentContextUsage({
+      agent: agent(3_000),
+      messages: branch(1_000, 6),
+      tokenCounter: countByChars,
+    });
+
+    expect(usage).not.toBeNull();
+    expect(usage!.breakdown.messageCount).toBeGreaterThan(0);
+    expect(usage!.breakdown.messageCount).toBeLessThan(6);
+  });
+
+  it('returns null without a context window', async () => {
+    const noWindow: t.AgentInputs = {
+      agentId: 'test-agent',
+      provider: Providers.OPENAI,
+      instructions: 'sys',
+    };
+    const usage = await projectAgentContextUsage({
+      agent: noWindow,
+      messages: branch(100, 2),
+      tokenCounter: countByChars,
+    });
+
+    expect(usage).toBeNull();
+  });
+});

package/src/agents/projection.ts

@@ -0,0 +1,46 @@
+import type { BaseMessage } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { AgentContext } from './AgentContext';
+
+export interface ProjectAgentContextUsageParams {
+  /** Same `AgentInputs` a run is built from (instructions, tools, model, window). */
+  agent: t.AgentInputs;
+  /** Branch messages to project, in send order (no leading system message). */
+  messages: BaseMessage[];
+  tokenCounter: t.TokenCounter;
+  /** Per-message counts aligned to `messages` (e.g. from `formatAgentMessages`).
+   *  When omitted, counts are recounted via `tokenCounter`. */
+  indexTokenCountMap?: Record<string, number>;
+  /** Provider-calibrated ratio from a prior snapshot, applied as a static seed. */
+  calibrationRatio?: number;
+  runId?: string;
+  agentId?: string;
+}
+
+/**
+ * Projects a pre-send context-usage snapshot for a branch under an agent config
+ * WITHOUT invoking the model — the host-side (page-load / branch-switch /
+ * window-switch) counterpart to the live `ON_CONTEXT_USAGE` event. Builds a
+ * throwaway `AgentContext` from the same `AgentInputs` a run uses, awaits its
+ * instruction/tool token accounting, then runs the shared pruner + budget math
+ * via `AgentContext.projectContextUsage` (which never mutates the supplied
+ * messages). Returns null when the config has no tokenizer or context window.
+ */
+export async function projectAgentContextUsage({
+  agent,
+  messages,
+  tokenCounter,
+  indexTokenCountMap,
+  calibrationRatio,
+  runId,
+  agentId,
+}: ProjectAgentContextUsageParams): Promise<t.ContextUsageEvent | null> {
+  const context = AgentContext.fromConfig(agent, tokenCounter, indexTokenCountMap);
+  await context.tokenCalculationPromise;
+  return context.projectContextUsage(messages, {
+    runId,
+    agentId: agentId ?? agent.agentId,
+    calibrationRatio,
+    indexTokenCountMap,
+  });
+}

package/src/common/enum.ts

@@ -31,6 +31,8 @@ export enum GraphEvents {
   ON_SUBAGENT_UPDATE = 'on_subagent_update',
   /** [Custom] Diagnostic logging event for context management observability */
   ON_AGENT_LOG = 'on_agent_log',
+  /** [Custom] Per-model-call context window usage snapshot (post-prune token budget) */
+  ON_CONTEXT_USAGE = 'on_context_usage',
 
   /* Official Events */
 
@@ -185,6 +187,17 @@ export enum Constants {
   /** Anthropic server tool ID prefix (web_search, code_execution, etc.) */
   ANTHROPIC_SERVER_TOOL_PREFIX = 'srvtoolu_',
   SKILL_TOOL = 'skill',
+  /**
+   * Callback-metadata keys stamped by `attemptInvoke` /
+   * `tryFallbackProviders` carrying the provider (SDK `Providers` enum
+   * value) and configured model that actually served a model invocation.
+   * Unlike `ls_provider` — which derived providers inherit from their base
+   * class (e.g. DeepSeek/OpenRouter report `'openai'`) — these reflect the
+   * SDK's own routing, including fallback-provider calls. Consumed by the
+   * subagent usage-capture handler to tag billing events.
+   */
+  INVOKED_PROVIDER = '__invoked_provider',
+  INVOKED_MODEL = '__invoked_model',
   READ_FILE = 'read_file',
   BASH_TOOL = 'bash_tool',
   BASH_PROGRAMMATIC_TOOL_CALLING = 'run_tools_with_bash',