@librechat/agents 3.1.73 → 3.1.75-dev.0

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

@@ -1,9 +1,15 @@
 // src/agents/__tests__/AgentContext.test.ts
+import { HumanMessage } from '@langchain/core/messages';
 import { AgentContext } from '../AgentContext';
 import { Providers } from '@/common';
+import { addBedrockCacheControl } from '@/messages/cache';
 import type * as t from '@/types';
 
 describe('AgentContext', () => {
+  type TestSystemContentBlock =
+    | { type: 'text'; text: string; cache_control?: { type: 'ephemeral' } }
+    | { cachePoint: { type: 'default' } };
+
   type ContextOptions = {
     agentConfig?: Partial<t.AgentInputs>;
     tokenCounter?: t.TokenCounter;
@@ -59,14 +65,161 @@ describe('AgentContext', () => {
       expect(ctx.systemRunnable).toBeUndefined();
     });
 
-    it('includes additional_instructions in system message', () => {
+    it('keeps additional_instructions after stable instructions', async () => {
       const ctx = createBasicContext({
         agentConfig: {
           instructions: 'Base instructions',
           additional_instructions: 'Additional instructions',
         },
       });
-      expect(ctx.systemRunnable).toBeDefined();
+
+      const result = await ctx.systemRunnable!.invoke([]);
+      expect(result[0].content).toBe(
+        'Base instructions\n\nAdditional instructions'
+      );
+    });
+
+    it('marks only stable system text for Anthropic prompt caching', async () => {
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.ANTHROPIC,
+          clientOptions: { model: 'claude-3-5-sonnet', promptCache: true },
+          instructions: 'Stable instructions',
+          additional_instructions: 'Dynamic instructions',
+        },
+      });
+
+      const result = await ctx.systemRunnable!.invoke([]);
+      const content = result[0].content as TestSystemContentBlock[];
+      expect(content).toHaveLength(2);
+      expect(content[0]).toMatchObject({
+        type: 'text',
+        text: 'Stable instructions',
+        cache_control: { type: 'ephemeral' },
+      });
+      expect(content[1]).toEqual({
+        type: 'text',
+        text: 'Dynamic instructions',
+      });
+    });
+
+    it('omits Anthropic cache control when only dynamic system text exists', async () => {
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.ANTHROPIC,
+          clientOptions: { model: 'claude-3-5-sonnet', promptCache: true },
+          instructions: undefined,
+          additional_instructions: 'Dynamic only',
+        },
+      });
+
+      const result = await ctx.systemRunnable!.invoke([]);
+      const content = result[0].content as TestSystemContentBlock[];
+      expect(content).toEqual([{ type: 'text', text: 'Dynamic only' }]);
+      expect(content[0]).not.toHaveProperty('cache_control');
+    });
+
+    it('keeps cross-run summaries in the dynamic Anthropic system tail', async () => {
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.ANTHROPIC,
+          clientOptions: { model: 'claude-3-5-sonnet', promptCache: true },
+          instructions: 'Stable instructions',
+        },
+      });
+      ctx.setInitialSummary('Prior summary', 13);
+
+      const result = await ctx.systemRunnable!.invoke([]);
+      const content = result[0].content as TestSystemContentBlock[];
+      expect(content).toHaveLength(2);
+      expect(content[0]).toHaveProperty('cache_control');
+      expect(content[1]).toEqual({
+        type: 'text',
+        text: '## Conversation Summary\n\nPrior summary',
+      });
+    });
+
+    it('places the Bedrock cache point before dynamic system text', async () => {
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.BEDROCK,
+          clientOptions: {
+            model: 'anthropic.claude-3-5-sonnet',
+            promptCache: true,
+          },
+          instructions: 'Stable instructions',
+          additional_instructions: 'Dynamic instructions',
+        },
+      });
+
+      const result = await ctx.systemRunnable!.invoke([]);
+      const content = result[0].content as TestSystemContentBlock[];
+      expect(content).toEqual([
+        { type: 'text', text: 'Stable instructions' },
+        { cachePoint: { type: 'default' } },
+        { type: 'text', text: 'Dynamic instructions' },
+      ]);
+    });
+
+    it('uses plain Bedrock system text when only dynamic system text exists', async () => {
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.BEDROCK,
+          clientOptions: {
+            model: 'anthropic.claude-3-5-sonnet',
+            promptCache: true,
+          },
+          instructions: undefined,
+          additional_instructions: 'Dynamic only',
+        },
+      });
+
+      const result = await ctx.systemRunnable!.invoke([]);
+      expect(result[0].content).toBe('Dynamic only');
+    });
+
+    it('keeps non-cache providers as plain system text with promptCache-like options', async () => {
+      const clientOptions: t.OpenAIClientOptions & { promptCache: true } = {
+        modelName: 'gpt-4o-mini',
+        promptCache: true,
+      };
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.OPENAI,
+          clientOptions,
+          instructions: 'Stable instructions',
+          additional_instructions: 'Dynamic instructions',
+        },
+      });
+
+      const result = await ctx.systemRunnable!.invoke([]);
+      expect(result[0].content).toBe(
+        'Stable instructions\n\nDynamic instructions'
+      );
+    });
+
+    it('preserves the Bedrock system cache point through message cache-control pass', async () => {
+      const ctx = createBasicContext({
+        agentConfig: {
+          provider: Providers.BEDROCK,
+          clientOptions: {
+            model: 'anthropic.claude-3-5-sonnet',
+            promptCache: true,
+          },
+          instructions: 'Stable instructions',
+          additional_instructions: 'Dynamic instructions',
+        },
+      });
+
+      const result = await ctx.systemRunnable!.invoke([
+        new HumanMessage('Hello'),
+      ]);
+      const finalMessages = addBedrockCacheControl(result);
+      expect(finalMessages[0].content).toEqual([
+        { type: 'text', text: 'Stable instructions' },
+        { cachePoint: { type: 'default' } },
+        { type: 'text', text: 'Dynamic instructions' },
+      ]);
     });
   });
 
@@ -404,6 +557,141 @@ describe('AgentContext', () => {
       expect(ctxWithDeferred.toolSchemaTokens).toBe(ctxBase.toolSchemaTokens);
     });
 
+    it('excludes programmatic-only toolDefinitions from toolSchemaTokens', async () => {
+      // getEventDrivenToolsForBinding excludes definitions whose
+      // allowed_callers omit 'direct'. Accounting must mirror that — a
+      // programmatic-only definition is never bound to the model and
+      // shouldn't inflate toolSchemaTokens.
+      const activeDef: t.LCTool = {
+        name: 'active_tool',
+        description: 'Always loaded',
+        parameters: { type: 'object', properties: {} },
+      };
+      const programmaticDef: t.LCTool = {
+        name: 'programmatic_tool',
+        description: 'Only callable via code execution',
+        parameters: { type: 'object', properties: {} },
+        allowed_callers: ['code_execution'],
+      };
+
+      const ctxBase = createBasicContext({
+        agentConfig: { toolDefinitions: [activeDef] },
+        tokenCounter: mockTokenCounter,
+      });
+      const ctxWithProgrammatic = createBasicContext({
+        agentConfig: { toolDefinitions: [activeDef, programmaticDef] },
+        tokenCounter: mockTokenCounter,
+      });
+
+      await ctxBase.tokenCalculationPromise;
+      await ctxWithProgrammatic.tokenCalculationPromise;
+
+      expect(ctxWithProgrammatic.toolSchemaTokens).toBe(
+        ctxBase.toolSchemaTokens
+      );
+    });
+
+    it('excludes deferred-undiscovered instance tools from toolSchemaTokens', async () => {
+      const activeTool = createMockTool('active_tool');
+      const deferredTool = createMockTool('deferred_tool');
+      const programmaticTool = createMockTool('programmatic_tool');
+      const toolRegistry: t.LCToolRegistry = new Map([
+        ['active_tool', { name: 'active_tool' }],
+        ['deferred_tool', { name: 'deferred_tool', defer_loading: true }],
+        [
+          'programmatic_tool',
+          {
+            name: 'programmatic_tool',
+            allowed_callers: ['code_execution'],
+          },
+        ],
+      ]);
+
+      const ctxBase = createBasicContext({
+        agentConfig: { tools: [activeTool], toolRegistry },
+        tokenCounter: mockTokenCounter,
+      });
+      const ctxWithExcluded = createBasicContext({
+        agentConfig: {
+          tools: [activeTool, deferredTool, programmaticTool],
+          toolRegistry,
+        },
+        tokenCounter: mockTokenCounter,
+      });
+
+      await ctxBase.tokenCalculationPromise;
+      await ctxWithExcluded.tokenCalculationPromise;
+
+      expect(ctxWithExcluded.toolSchemaTokens).toBe(ctxBase.toolSchemaTokens);
+    });
+
+    it('includes deferred instance tools once discovered via discoveredTools input', async () => {
+      const tools = [createMockTool('deferred_tool')];
+      const toolRegistry: t.LCToolRegistry = new Map([
+        ['deferred_tool', { name: 'deferred_tool', defer_loading: true }],
+      ]);
+
+      const ctxUndiscovered = createBasicContext({
+        agentConfig: { tools, toolRegistry },
+        tokenCounter: mockTokenCounter,
+      });
+      const ctxDiscovered = createBasicContext({
+        agentConfig: {
+          tools,
+          toolRegistry,
+          discoveredTools: ['deferred_tool'],
+        },
+        tokenCounter: mockTokenCounter,
+      });
+
+      await ctxUndiscovered.tokenCalculationPromise;
+      await ctxDiscovered.tokenCalculationPromise;
+
+      expect(ctxUndiscovered.toolSchemaTokens).toBe(0);
+      expect(ctxDiscovered.toolSchemaTokens).toBeGreaterThan(0);
+    });
+
+    it('does not filter instance tools in event-driven mode (matches getEventDrivenToolsForBinding)', async () => {
+      // In event-driven mode, getEventDrivenToolsForBinding appends
+      // `this.tools` UNFILTERED. Accounting must do the same — otherwise we
+      // under-count and risk exceeding the model's context budget.
+      const activeDef: t.LCTool = {
+        name: 'active_def',
+        description: 'Always loaded',
+        parameters: { type: 'object', properties: {} },
+      };
+      const nativeTool = createMockTool('native_tool');
+      // Registry marks the native tool as deferred-undiscovered. In the
+      // non-event-driven path this would exclude it; in event-driven mode
+      // it is still bound and must still be counted.
+      const toolRegistry: t.LCToolRegistry = new Map([
+        ['native_tool', { name: 'native_tool', defer_loading: true }],
+      ]);
+
+      const ctxWithoutNative = createBasicContext({
+        agentConfig: {
+          toolDefinitions: [activeDef],
+          toolRegistry,
+        },
+        tokenCounter: mockTokenCounter,
+      });
+      const ctxWithNative = createBasicContext({
+        agentConfig: {
+          toolDefinitions: [activeDef],
+          tools: [nativeTool],
+          toolRegistry,
+        },
+        tokenCounter: mockTokenCounter,
+      });
+
+      await ctxWithoutNative.tokenCalculationPromise;
+      await ctxWithNative.tokenCalculationPromise;
+
+      expect(ctxWithNative.toolSchemaTokens).toBeGreaterThan(
+        ctxWithoutNative.toolSchemaTokens
+      );
+    });
+
     it('includes deferred toolDefinitions once discovered via discoveredTools input', async () => {
       const toolDefinitions: t.LCTool[] = [
         {
@@ -448,6 +736,36 @@ describe('AgentContext', () => {
       expect(ctx.getTokenBudgetBreakdown().toolCount).toBe(1);
     });
 
+    it('getTokenBudgetBreakdown toolCount excludes deferred-undiscovered instance tools', () => {
+      // Mirrors the toolDefinitions test for the instance-tools path so
+      // toolCount stays aligned with toolSchemaTokens (and with what
+      // getToolsForBinding actually emits) for non-event-driven runs.
+      const tools = [
+        createMockTool('active_tool'),
+        createMockTool('deferred_tool'),
+        createMockTool('programmatic_tool'),
+      ];
+      const toolRegistry: t.LCToolRegistry = new Map([
+        ['active_tool', { name: 'active_tool' }],
+        ['deferred_tool', { name: 'deferred_tool', defer_loading: true }],
+        [
+          'programmatic_tool',
+          {
+            name: 'programmatic_tool',
+            allowed_callers: ['code_execution'],
+          },
+        ],
+      ]);
+
+      const ctx = createBasicContext({
+        agentConfig: { tools, toolRegistry },
+      });
+
+      expect(ctx.getTokenBudgetBreakdown().toolCount).toBe(1);
+      ctx.markToolsAsDiscovered(['deferred_tool']);
+      expect(ctx.getTokenBudgetBreakdown().toolCount).toBe(2);
+    });
+
     it('getTokenBudgetBreakdown toolCount reflects newly discovered deferred tools', () => {
       const toolDefinitions: t.LCTool[] = [
         {
@@ -464,6 +782,19 @@ describe('AgentContext', () => {
       expect(ctx.getTokenBudgetBreakdown().toolCount).toBe(1);
     });
 
+    it('getTokenBudgetBreakdown toolCount includes graphTools', () => {
+      // graphTools (handoff/subagent) are bound to the model alongside
+      // instance tools. Now that toolCount derives from getToolsForBinding(),
+      // graphTools are reflected in the diagnostic just like they're
+      // counted in toolSchemaTokens. Locks in that alignment.
+      const ctx = createBasicContext({
+        agentConfig: { tools: [createMockTool('direct_tool')] },
+      });
+      ctx.graphTools = [createMockTool('handoff_tool')];
+
+      expect(ctx.getTokenBudgetBreakdown().toolCount).toBe(2);
+    });
+
     it('toolSchemaTokens snapshot does not auto-update after markToolsAsDiscovered', async () => {
       const toolDefinitions: t.LCTool[] = [
         {

package/src/agents/__tests__/promptCacheLiveHelpers.ts

@@ -0,0 +1,165 @@
+import { expect } from '@jest/globals';
+import { HumanMessage } from '@langchain/core/messages';
+import type { UsageMetadata } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { GraphEvents, Providers } from '@/common';
+import { AgentContext } from '../AgentContext';
+import { ModelEndHandler } from '@/events';
+import { Run } from '@/run';
+
+type LivePromptCacheProvider = Providers.ANTHROPIC | Providers.BEDROCK;
+
+type PromptCacheExpectedSystemBlock =
+  | { type: 'text'; text: string; cache_control?: { type: 'ephemeral' } }
+  | { cachePoint: { type: 'default' } };
+
+type LivePromptCacheClientOptions =
+  | t.ClientOptions
+  | t.BedrockAnthropicClientOptions;
+
+export function buildStableInstructions({
+  nonce,
+  providerLabel,
+}: {
+  nonce: string;
+  providerLabel: string;
+}): string {
+  const records = Array.from(
+    { length: 360 },
+    (_, index) =>
+      `Stable ${providerLabel} cache record ${index}: nonce ${nonce}; keep this reference in the cacheable prefix and do not use it as the dynamic marker.`
+  );
+  return [
+    `You are a ${providerLabel} prompt-cache verification assistant.`,
+    'When asked for the dynamic marker, answer with only the marker value from the Dynamic Marker line.',
+    ...records,
+  ].join('\n');
+}
+
+export function buildDynamicInstructions({
+  marker,
+  tailDescription,
+}: {
+  marker: string;
+  tailDescription: string;
+}): string {
+  return [`Dynamic Marker: ${marker}`, tailDescription].join('\n');
+}
+
+export function waitForCachePropagation(): Promise<void> {
+  return new Promise((resolve) => setTimeout(resolve, 2000));
+}
+
+export async function assertSystemPayloadShape({
+  agentId,
+  provider,
+  clientOptions,
+  stableInstructions,
+  dynamicInstructions,
+  expectedContent,
+}: {
+  agentId: string;
+  provider: LivePromptCacheProvider;
+  clientOptions: LivePromptCacheClientOptions;
+  stableInstructions: string;
+  dynamicInstructions: string;
+  expectedContent: PromptCacheExpectedSystemBlock[];
+}): Promise<void> {
+  const ctx = AgentContext.fromConfig({
+    agentId,
+    provider,
+    clientOptions,
+    instructions: stableInstructions,
+    additional_instructions: dynamicInstructions,
+  });
+
+  const messages = await ctx.systemRunnable!.invoke([
+    new HumanMessage('What is the dynamic marker?'),
+  ]);
+
+  expect(messages[0].content).toEqual(expectedContent);
+}
+
+function latestUsage({
+  collectedUsage,
+  label,
+  providerLabel,
+}: {
+  collectedUsage: UsageMetadata[];
+  label: string;
+  providerLabel: string;
+}): UsageMetadata {
+  if (collectedUsage.length === 0) {
+    throw new Error(`Missing ${providerLabel} usage metadata for ${label}`);
+  }
+  return collectedUsage[collectedUsage.length - 1];
+}
+
+function collectText(parts: t.MessageContentComplex[] | undefined): string {
+  return (parts ?? []).reduce((text, part) => {
+    if (part.type === 'text') {
+      return text + part.text;
+    }
+    return text;
+  }, '');
+}
+
+export async function runLiveTurn({
+  provider,
+  providerLabel,
+  clientOptions,
+  runId,
+  threadId,
+  stableInstructions,
+  dynamicInstructions,
+}: {
+  provider: LivePromptCacheProvider;
+  providerLabel: string;
+  clientOptions: LivePromptCacheClientOptions;
+  runId: string;
+  threadId: string;
+  stableInstructions: string;
+  dynamicInstructions: string;
+}): Promise<{
+  text: string;
+  usage: UsageMetadata;
+}> {
+  const collectedUsage: UsageMetadata[] = [];
+  const run = await Run.create<t.IState>({
+    runId,
+    graphConfig: {
+      type: 'standard',
+      llmConfig: {
+        provider,
+        ...clientOptions,
+      } as t.LLMConfig,
+      instructions: stableInstructions,
+      additional_instructions: dynamicInstructions,
+    },
+    returnContent: true,
+    skipCleanup: true,
+    customHandlers: {
+      [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(collectedUsage),
+    },
+  });
+
+  const config = {
+    configurable: { thread_id: threadId },
+    streamMode: 'values',
+    version: 'v2' as const,
+  };
+
+  const contentParts = await run.processStream(
+    {
+      messages: [
+        new HumanMessage('What is the dynamic marker? Reply with only it.'),
+      ],
+    },
+    config
+  );
+
+  return {
+    text: collectText(contentParts),
+    usage: latestUsage({ collectedUsage, label: runId, providerLabel }),
+  };
+}

package/src/graphs/Graph.ts

@@ -399,12 +399,25 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
   ): (string | number | undefined)[] {
     if (!metadata) return [];
 
+    const configurable = this.config?.configurable;
+    const runId =
+      (metadata.run_id as string | undefined) ??
+      (configurable?.run_id as string | undefined) ??
+      this.runId;
+    const threadId =
+      (metadata.thread_id as string | undefined) ??
+      (configurable?.thread_id as string | undefined) ??
+      runId;
+    const checkpointNs =
+      (metadata.checkpoint_ns as string | undefined) ??
+      (metadata.langgraph_checkpoint_ns as string | undefined) ??
+      '';
     const keyList = [
-      metadata.run_id as string,
-      metadata.thread_id as string,
+      runId,
+      threadId,
       metadata.langgraph_node as string,
       metadata.langgraph_step as number,
-      metadata.checkpoint_ns as string,
+      checkpointNs,
     ];
 
     const agentContext = this.getAgentContext(metadata);
@@ -1461,7 +1474,14 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
       }),
     });
     const workflow = new StateGraph(StateAnnotation)
-      .addNode(this.defaultAgentId, agentNode, { ends: [END] })
+      .addNode(
+        this.defaultAgentId,
+        agentNode as Runnable<
+          t.AgentSubgraphState,
+          Partial<t.AgentSubgraphState>
+        >,
+        { ends: [END] }
+      )
       .addEdge(START, this.defaultAgentId)
       // LangGraph compile() types are overly strict for opt-in options
       .compile(this.compileOptions as unknown as never);