@librechat/agents 3.3.3 → 3.3.4

package/src/graphs/__tests__/Graph.contextOverflow.test.ts

@@ -1,15 +1,19 @@
+import { z } from 'zod';
+import { tool } from '@langchain/core/tools';
 import { MemorySaver } from '@langchain/langgraph';
-import { describe, expect, it } from '@jest/globals';
 import { Runnable } from '@langchain/core/runnables';
+import { describe, expect, it, jest } from '@jest/globals';
 import {
   AIMessageChunk,
   HumanMessage,
   AIMessage,
+  ToolMessage,
 } from '@langchain/core/messages';
 import type { BaseMessage } from '@langchain/core/messages';
 import type * as t from '@/types';
 import { OVERFLOW_SIGNATURES } from '@/utils/__tests__/fixtures/contextOverflowSignatures';
-import { GraphEvents, Providers } from '@/common';
+import { ContentTypes, GraphEvents, Providers } from '@/common';
+import * as init from '@/llm/init';
 import { Run } from '@/run';
 
 /**
@@ -64,6 +68,36 @@ class OverflowThenSucceedModel extends Runnable<BaseMessage[], AIMessageChunk> {
   }
 }
 
+class SizeBoundModel extends Runnable<BaseMessage[], AIMessageChunk> {
+  lc_namespace = ['tests'];
+  readonly toolContentChars: number[] = [];
+
+  constructor(
+    private readonly maxToolContentChars: number,
+    private readonly error: Record<string, unknown>
+  ) {
+    super();
+  }
+
+  async invoke(messages: BaseMessage[]): Promise<AIMessageChunk> {
+    let toolContentChars = 0;
+    for (const message of messages) {
+      if (message.getType() !== 'tool') {
+        continue;
+      }
+      toolContentChars +=
+        typeof message.content === 'string'
+          ? message.content.length
+          : JSON.stringify(message.content).length;
+    }
+    this.toolContentChars.push(toolContentChars);
+    if (toolContentChars > this.maxToolContentChars) {
+      throw throwable(this.error);
+    }
+    return new AIMessageChunk({ content: 'recovered' });
+  }
+}
+
 function buildConversation(turns: number): BaseMessage[] {
   const messages: BaseMessage[] = [];
   for (let i = 0; i < turns; i++) {
@@ -78,17 +112,31 @@ async function createRun(options: {
   runId: string;
   maxContextTokens: number;
   checkpointer?: boolean;
+  provider?: Providers;
+  tokenCounter?: t.TokenCounter;
+  indexTokenCountMap?: Record<string, number>;
+  tools?: t.GraphTools;
+  maxToolResultChars?: number;
+  model?: string;
+  promptCache?: boolean;
+  toolOutputReferences?: t.ToolOutputReferencesConfig;
+  fallbacks?: t.FallbackConfig[];
 }): Promise<Run<t.IState>> {
   return Run.create<t.IState>({
     runId: options.runId,
     graphConfig: {
       type: 'standard',
       llmConfig: {
-        provider: Providers.ANTHROPIC,
+        provider: options.provider ?? Providers.ANTHROPIC,
+        ...(options.model != null ? { model: options.model } : {}),
+        ...(options.promptCache === true ? { promptCache: true } : {}),
+        ...(options.fallbacks != null ? { fallbacks: options.fallbacks } : {}),
         disableStreaming: true,
         streamUsage: false,
       },
       maxContextTokens: options.maxContextTokens,
+      maxToolResultChars: options.maxToolResultChars,
+      tools: options.tools,
       compileOptions:
         options.checkpointer === true
           ? { checkpointer: new MemorySaver() }
@@ -96,7 +144,9 @@ async function createRun(options: {
     },
     returnContent: true,
     skipCleanup: true,
-    tokenCounter,
+    tokenCounter: options.tokenCounter ?? tokenCounter,
+    indexTokenCountMap: options.indexTokenCountMap,
+    toolOutputReferences: options.toolOutputReferences,
   });
 }
 
@@ -107,6 +157,1080 @@ const streamConfig = {
 };
 
 describe('context overflow recovery', () => {
+  it('projects structured OpenAI tool content before the final payload check', async () => {
+    const toolCallId = 'tc-openai-structured';
+    const toolMessage = new ToolMessage({
+      content: [
+        { type: ContentTypes.TEXT, text: 'rendered chart' },
+        {
+          type: 'image_url',
+          image_url: {
+            url: `data:image/png;base64,${'A'.repeat(2_000)}`,
+          },
+        },
+      ],
+      tool_call_id: toolCallId,
+      name: 'render_chart',
+    });
+    const messages: BaseMessage[] = [
+      new HumanMessage('render the chart'),
+      new AIMessage({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'render_chart',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      toolMessage,
+    ];
+    const measuredToolContents: string[] = [];
+    const projectionCounter: t.TokenCounter = (message) => {
+      if (message.getType() === 'tool' && typeof message.content === 'string') {
+        measuredToolContents.push(message.content);
+      }
+      return typeof message.content === 'string'
+        ? Math.max(1, Math.ceil(message.content.length / 4))
+        : 1;
+    };
+    const run = await createRun({
+      runId: 'openai-structured-final-projection',
+      maxContextTokens: 10_000,
+      maxToolResultChars: 200,
+      provider: Providers.OPENAI,
+      tokenCounter: projectionCounter,
+      indexTokenCountMap: {
+        0: projectionCounter(messages[0]),
+        1: projectionCounter(messages[1]),
+        2: projectionCounter(messages[2]),
+      },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'render_chart',
+          description: 'Renders a chart',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(signatureFor('gpt-4o-mini'), 0);
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    const projectedTool = model.calls[0].find(
+      (message) => message.getType() === 'tool'
+    ) as ToolMessage | undefined;
+    expect(typeof projectedTool?.content).toBe('string');
+    expect((projectedTool?.content as string).length).toBeLessThanOrEqual(200);
+    expect(measuredToolContents).toContain(projectedTool?.content);
+    expect(Array.isArray(toolMessage.content)).toBe(true);
+  });
+
+  it('measures the bounded generic-provider tool payload that is invoked', async () => {
+    const toolCallId = 'tc-google-structured';
+    const structuredContent = Array.from({ length: 25_000 }, () => ({
+      type: ContentTypes.TEXT,
+      text: 'x',
+    }));
+    const toolMessage = new ToolMessage({
+      content: structuredContent,
+      tool_call_id: toolCallId,
+      name: 'dense_result',
+    });
+    const messages: BaseMessage[] = [
+      new HumanMessage('return the dense result'),
+      new AIMessage({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'dense_result',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      toolMessage,
+    ];
+    const measuredToolContents: string[] = [];
+    const projectionCounter: t.TokenCounter = (message) => {
+      if (message.getType() === 'tool' && typeof message.content === 'string') {
+        measuredToolContents.push(message.content);
+      }
+      return typeof message.content === 'string'
+        ? Math.max(1, Math.ceil(message.content.length / 4))
+        : 1;
+    };
+    const run = await createRun({
+      runId: 'google-structured-final-projection',
+      maxContextTokens: 10_000,
+      maxToolResultChars: 2_000,
+      provider: Providers.GOOGLE,
+      tokenCounter: projectionCounter,
+      indexTokenCountMap: {
+        0: projectionCounter(messages[0]),
+        1: projectionCounter(messages[1]),
+        2: projectionCounter(messages[2]),
+      },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'dense_result',
+          description: 'Returns a dense structured result',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(signatureFor('gpt-4o-mini'), 0);
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    const projectedTool = model.calls[0].find(
+      (message) => message.getType() === 'tool'
+    ) as ToolMessage | undefined;
+    expect(typeof projectedTool?.content).toBe('string');
+    expect((projectedTool?.content as string).length).toBeLessThanOrEqual(
+      2_000
+    );
+    expect(measuredToolContents).toContain(projectedTool?.content);
+    expect(toolMessage.content).toBe(structuredContent);
+  });
+
+  it('guards a fallback-specific projection before invoking the fallback', async () => {
+    const toolCallId = 'tc-fallback-structured';
+    const structuredContent = Array.from({ length: 25_000 }, () => ({
+      type: ContentTypes.TEXT,
+      text: '',
+    }));
+    const messages: BaseMessage[] = [
+      new HumanMessage('return the dense fallback result'),
+      new AIMessage({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'dense_result',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: structuredContent,
+        tool_call_id: toolCallId,
+        name: 'dense_result',
+      }),
+    ];
+    const projectionCounter: t.TokenCounter = (message) =>
+      typeof message.content === 'string'
+        ? Math.max(1, message.content.length)
+        : 1;
+    const run = await createRun({
+      runId: 'fallback-structured-final-projection',
+      maxContextTokens: 1_000_000,
+      provider: Providers.ANTHROPIC,
+      tokenCounter: projectionCounter,
+      indexTokenCountMap: {
+        0: projectionCounter(messages[0]),
+        1: projectionCounter(messages[1]),
+        2: projectionCounter(messages[2]),
+      },
+      fallbacks: [
+        {
+          provider: Providers.GOOGLE,
+          maxContextTokens: 5_000,
+        },
+      ],
+      tools: [
+        tool(async () => 'unused', {
+          name: 'dense_result',
+          description: 'Returns a dense structured result',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const primary = new OverflowThenSucceedModel(
+      { message: '503 primary unavailable' },
+      1
+    );
+    let fallbackInvocations = 0;
+    const fallback = {
+      invoke: async (): Promise<AIMessageChunk> => {
+        fallbackInvocations++;
+        return new AIMessageChunk({ content: 'fallback should not run' });
+      },
+    } as unknown as ReturnType<typeof init.initializeModel>;
+    const initializeSpy = jest
+      .spyOn(init, 'initializeModel')
+      .mockReturnValue(fallback);
+    run.Graph.overrideModel = primary;
+
+    try {
+      const content = await run.processStream({ messages }, streamConfig);
+
+      expect(content).toEqual([{ type: 'text', text: 'recovered' }]);
+      expect(primary.calls).toHaveLength(2);
+      expect(fallbackInvocations).toBe(0);
+      expect(
+        run.Graph.agentContexts.get('default')?.overflowRecoveryAttempts
+      ).toBe(1);
+    } finally {
+      initializeSpy.mockRestore();
+    }
+  });
+
+  it('projects unsafe tool-call args before measuring or invoking the provider', async () => {
+    let toJSONCalls = 0;
+    const toolCallId = 'tc-unsafe-input';
+    const unsafeArgs = {
+      query: 'safe',
+      toJSON() {
+        toJSONCalls++;
+        return { query: 'x'.repeat(100_000) };
+      },
+    };
+    const messages: BaseMessage[] = [
+      new HumanMessage('run the lookup'),
+      new AIMessage({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'lookup_records',
+            args: unsafeArgs,
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: 'done',
+        tool_call_id: toolCallId,
+        name: 'lookup_records',
+      }),
+      new AIMessage('The lookup completed.'),
+      new HumanMessage('continue'),
+    ];
+    const run = await createRun({
+      runId: 'unsafe-tool-input-final-projection',
+      maxContextTokens: 10_000,
+      provider: Providers.OPENAI,
+      tokenCounter: () => 1,
+      indexTokenCountMap: { 0: 1, 1: 1, 2: 1, 3: 1, 4: 1 },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'lookup_records',
+          description: 'Looks up records',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(signatureFor('gpt-4o-mini'), 0);
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    const projectedCall = (
+      model.calls[0].find((message) => message.getType() === 'ai') as AIMessage
+    ).tool_calls?.[0];
+    expect(projectedCall?.args).toEqual({ query: 'safe' });
+    expect(toJSONCalls).toBe(0);
+    expect(messages[1]).toBeInstanceOf(AIMessage);
+    expect((messages[1] as AIMessage).tool_calls?.[0].args).toBe(unsafeArgs);
+  });
+
+  it('compacts cached structured tool output before the first provider call', async () => {
+    const toolCallId = 'tc-structured';
+    const messages: BaseMessage[] = [
+      new HumanMessage('query the table'),
+      new AIMessage({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'run_select_query',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: [
+          {
+            type: ContentTypes.TEXT,
+            text: JSON.stringify(
+              Array.from({ length: 240 }, (_, index) => ({
+                id: index,
+                value: `${'x'.repeat(100)}-${index}`,
+              }))
+            ),
+          },
+        ],
+        tool_call_id: toolCallId,
+        name: 'run_select_query',
+      }),
+      new AIMessage('The query returned 240 rows.'),
+      new HumanMessage('compact context'),
+    ];
+    const structuredTokenCounter: t.TokenCounter = (message) => {
+      const content =
+        typeof message.content === 'string'
+          ? message.content
+          : JSON.stringify(message.content);
+      return Math.ceil(content.length / 4);
+    };
+    const indexTokenCountMap: Record<string, number> = {};
+    for (let i = 0; i < messages.length; i++) {
+      indexTokenCountMap[i] = i === 2 ? 0 : structuredTokenCounter(messages[i]);
+    }
+    const run = await createRun({
+      runId: 'structured-output-preflight',
+      maxContextTokens: 5_000,
+      provider: Providers.BEDROCK,
+      model: 'anthropic.claude-sonnet-4-5',
+      tokenCounter: structuredTokenCounter,
+      indexTokenCountMap,
+      tools: [
+        tool(async () => 'unused', {
+          name: 'run_select_query',
+          description: 'Queries ClickHouse',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new SizeBoundModel(
+      1_500,
+      signatureFor('us.anthropic.claude-sonnet-4-5-20250929-v1:0')
+    );
+    run.Graph.overrideModel = model;
+
+    const content = await run.processStream({ messages }, streamConfig);
+
+    expect(content).toEqual([{ type: 'text', text: 'recovered' }]);
+    expect(model.toolContentChars).toHaveLength(1);
+    expect(model.toolContentChars[0]).toBeGreaterThan(0);
+    expect(model.toolContentChars[0]).toBeLessThanOrEqual(1_500);
+    expect(
+      run.Graph.agentContexts.get('default')?.overflowRecoveryAttempts
+    ).toBe(0);
+  });
+
+  it('compacts an unconsumed structured tool result before its first provider call', async () => {
+    const toolCallId = 'tc-unconsumed-structured';
+    const messages: BaseMessage[] = [
+      new HumanMessage('query the table'),
+      new AIMessage({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'run_select_query',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: [
+          {
+            type: ContentTypes.TEXT,
+            text: JSON.stringify(
+              Array.from({ length: 240 }, (_, index) => ({
+                id: index,
+                value: `${'x'.repeat(100)}-${index}`,
+              }))
+            ),
+          },
+        ],
+        tool_call_id: toolCallId,
+        name: 'run_select_query',
+      }),
+    ];
+    const structuredTokenCounter: t.TokenCounter = (message) => {
+      const content =
+        typeof message.content === 'string'
+          ? message.content
+          : JSON.stringify(message.content);
+      return Math.ceil(content.length / 4);
+    };
+    const run = await createRun({
+      runId: 'unconsumed-structured-output-preflight',
+      maxContextTokens: 5_000,
+      maxToolResultChars: 1_500,
+      provider: Providers.BEDROCK,
+      tokenCounter: structuredTokenCounter,
+      indexTokenCountMap: {
+        0: structuredTokenCounter(messages[0]),
+        1: structuredTokenCounter(messages[1]),
+        2: 0,
+      },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'run_select_query',
+          description: 'Queries ClickHouse',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new SizeBoundModel(
+      1_500,
+      signatureFor('us.anthropic.claude-sonnet-4-5-20250929-v1:0')
+    );
+    run.Graph.overrideModel = model;
+
+    const content = await run.processStream({ messages }, streamConfig);
+
+    expect(content).toEqual([{ type: 'text', text: 'recovered' }]);
+    expect(model.toolContentChars).toHaveLength(1);
+    expect(model.toolContentChars[0]).toBeGreaterThan(0);
+    expect(model.toolContentChars[0]).toBeLessThanOrEqual(1_500);
+    expect(
+      run.Graph.agentContexts.get('default')?.overflowRecoveryAttempts
+    ).toBe(0);
+  });
+
+  it('trusts the pruner baseline for an unchanged payload at high calibration', async () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('earlier question'),
+      new AIMessage('provider-counted answer'),
+      new HumanMessage('latest question'),
+    ];
+    const localCounter: t.TokenCounter = (message) =>
+      message.getType() === 'ai' ? 120 : 5;
+    const providerGroundedTokenMap = { 0: 5, 1: 10, 2: 5 };
+    const run = await createRun({
+      runId: 'provider-grounded-final-projection',
+      maxContextTokens: 200,
+      tokenCounter: localCounter,
+      indexTokenCountMap: providerGroundedTokenMap,
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const agentContext = run.Graph.agentContexts.get('default');
+    if (agentContext == null) {
+      throw new Error('Expected default agent context');
+    }
+    agentContext.calibrationRatio = 5;
+    const model = new OverflowThenSucceedModel(
+      signatureFor('claude-haiku-4-5-20251001'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    const content = await run.processStream({ messages }, streamConfig);
+
+    expect(content).toEqual([{ type: 'text', text: 'recovered' }]);
+    expect(model.calls).toHaveLength(1);
+    expect(model.calls[0]).toHaveLength(messages.length);
+    expect(agentContext.overflowRecoveryAttempts).toBe(0);
+  });
+
+  it('does not let an unrelated provider-format shrink hide artifact growth', async () => {
+    const toolCallId = 'tc-artifact-independent-growth';
+    const artifactSentinel = 'ARTIFACT_INDEPENDENT_GROWTH_SENTINEL';
+    const toolMessage = new ToolMessage({
+      content: 'rendered',
+      tool_call_id: toolCallId,
+      name: 'render_report',
+      artifact: {
+        content: [
+          {
+            type: ContentTypes.TEXT,
+            text: artifactSentinel,
+          },
+        ],
+      },
+    });
+    const messages: BaseMessage[] = [
+      new HumanMessage('render the report'),
+      new AIMessageChunk({
+        content: 'calling render_report',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'render_report',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      toolMessage,
+    ];
+    const transformCounter: t.TokenCounter = (message) => {
+      if (
+        message instanceof AIMessageChunk &&
+        typeof message.content === 'string'
+      ) {
+        return 100;
+      }
+      if (
+        message instanceof ToolMessage &&
+        JSON.stringify(message.content).includes(artifactSentinel)
+      ) {
+        return 50;
+      }
+      return 1;
+    };
+    const run = await createRun({
+      runId: 'artifact-independent-transform-growth',
+      maxContextTokens: 50,
+      provider: Providers.BEDROCK,
+      model: 'anthropic.claude-sonnet-4-5',
+      tokenCounter: transformCounter,
+      indexTokenCountMap: { 0: 1, 1: 1, 2: 1 },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'render_report',
+          description: 'Renders a report',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(
+      signatureFor('us.anthropic.claude-sonnet-4-5-20250929-v1:0'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    expect(
+      JSON.stringify(model.calls[0].map((message) => message.content))
+    ).not.toContain(artifactSentinel);
+    expect(toolMessage.artifact.content[0].text).toBe(artifactSentinel);
+  });
+
+  it('preserves provider attribution for an un-IDd AI clone during orphan sanitization', async () => {
+    const droppedCallId = 'tc-dropped-ai';
+    const missingCallId = 'tc-missing-result';
+    const unrelatedResultId = 'tc-unrelated-result';
+    const messages: BaseMessage[] = [
+      new HumanMessage('earlier question'),
+      new AIMessage({
+        content: [
+          {
+            type: 'tool_use',
+            id: droppedCallId,
+            name: 'declared_tool',
+            input: {},
+          },
+        ],
+        tool_calls: [
+          {
+            id: droppedCallId,
+            name: 'declared_tool',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new AIMessage({
+        content: 'provider-counted answer',
+        tool_calls: [
+          {
+            id: missingCallId,
+            name: 'declared_tool',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: 'orphaned result',
+        tool_call_id: unrelatedResultId,
+        name: 'declared_tool',
+      }),
+      new HumanMessage('latest question'),
+    ];
+    const skewedCounter: t.TokenCounter = (message) => {
+      if (message.getType() !== 'ai') {
+        return 5;
+      }
+      return Array.isArray(message.content) ? 1 : 120;
+    };
+    const run = await createRun({
+      runId: 'orphan-sanitize-provider-origin',
+      maxContextTokens: 1_000,
+      provider: Providers.ANTHROPIC,
+      promptCache: true,
+      tokenCounter: skewedCounter,
+      indexTokenCountMap: { 0: 5, 1: 80, 2: 10, 3: 5, 4: 5 },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'declared_tool',
+          description: 'Declared only to prevent tool-less folding',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const agentContext = run.Graph.agentContexts.get('default');
+    if (agentContext == null) {
+      throw new Error('Expected default agent context');
+    }
+    agentContext.calibrationRatio = 5;
+    const model = new OverflowThenSucceedModel(
+      signatureFor('claude-haiku-4-5-20251001'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    const content = await run.processStream({ messages }, streamConfig);
+
+    expect(content).toEqual([{ type: 'text', text: 'recovered' }]);
+    expect(model.calls).toHaveLength(1);
+    expect(model.calls[0]).toHaveLength(3);
+    expect(model.calls[0].some((message) => message.getType() === 'tool')).toBe(
+      false
+    );
+    const sanitizedAI = model.calls[0].find(
+      (message) => message.getType() === 'ai'
+    ) as AIMessage | undefined;
+    expect(sanitizedAI?.content).toBe('provider-counted answer');
+    expect(sanitizedAI?.tool_calls ?? []).toHaveLength(0);
+    expect(agentContext.overflowRecoveryAttempts).toBe(0);
+  });
+
+  it('reserves the reply primer before accepting fast-path context', async () => {
+    const messages: BaseMessage[] = [
+      new HumanMessage('earlier question'),
+      new AIMessage('earlier answer'),
+      new HumanMessage('latest question'),
+    ];
+    const run = await createRun({
+      runId: 'fast-path-reply-primer',
+      maxContextTokens: 100,
+      tokenCounter: () => 31,
+      indexTokenCountMap: { 0: 31, 1: 31, 2: 31 },
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(
+      signatureFor('claude-haiku-4-5-20251001'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    expect(model.calls[0].length).toBeLessThan(messages.length);
+    expect(3 + model.calls[0].length * 31).toBeLessThanOrEqual(95);
+  });
+
+  it('includes artifact expansion when it fits the post-prune budget', async () => {
+    const toolCallId = 'tc-artifact-fits';
+    const artifactSentinel = 'ARTIFACT_FITS_SENTINEL';
+    const toolMessage = new ToolMessage({
+      content: 'rendered',
+      tool_call_id: toolCallId,
+      name: 'render_report',
+      artifact: {
+        content: [
+          {
+            type: ContentTypes.TEXT,
+            text: `${artifactSentinel}:complete`,
+          },
+        ],
+      },
+    });
+    const messages: BaseMessage[] = [
+      new HumanMessage('render the report'),
+      new AIMessageChunk({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'render_report',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      toolMessage,
+    ];
+    const artifactTokenCounter: t.TokenCounter = (message) => {
+      const content =
+        typeof message.content === 'string'
+          ? message.content
+          : JSON.stringify(message.content);
+      return Math.ceil(content.length / 4);
+    };
+    const run = await createRun({
+      runId: 'artifact-budget-control',
+      maxContextTokens: 10_000,
+      maxToolResultChars: 2_000,
+      provider: Providers.BEDROCK,
+      model: 'anthropic.claude-sonnet-4-5',
+      tokenCounter: artifactTokenCounter,
+      indexTokenCountMap: {
+        0: artifactTokenCounter(messages[0]),
+        1: artifactTokenCounter(messages[1]),
+        2: artifactTokenCounter(messages[2]),
+      },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'render_report',
+          description: 'Renders a report',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(
+      signatureFor('us.anthropic.claude-sonnet-4-5-20250929-v1:0'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    expect(
+      JSON.stringify(model.calls[0].map((message) => message.content))
+    ).toContain(artifactSentinel);
+    expect(toolMessage.content).toBe('rendered');
+    expect(toolMessage.artifact.content[0].text).toContain(artifactSentinel);
+  });
+
+  it('rechecks artifact expansion after provider message transforms', async () => {
+    const toolCallId = 'tc-artifact-final-transform';
+    const artifactSentinel = 'ARTIFACT_FINAL_TRANSFORM_SENTINEL';
+    const toolMessage = new ToolMessage({
+      content: 'rendered',
+      tool_call_id: toolCallId,
+      name: 'render_report',
+      artifact: {
+        content: [
+          {
+            type: ContentTypes.TEXT,
+            text: `${artifactSentinel}:complete`,
+          },
+        ],
+      },
+    });
+    const messages: BaseMessage[] = [
+      new HumanMessage('render the report'),
+      new AIMessageChunk({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'render_report',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      toolMessage,
+    ];
+    const transformSensitiveCounter: t.TokenCounter = (message) => {
+      const content =
+        typeof message.content === 'string'
+          ? message.content
+          : JSON.stringify(message.content);
+      if (
+        message instanceof HumanMessage &&
+        content.includes(artifactSentinel)
+      ) {
+        return 10_000;
+      }
+      return Math.max(1, Math.ceil(content.length / 4));
+    };
+    const run = await createRun({
+      runId: 'artifact-final-transform-guard',
+      maxContextTokens: 5_000,
+      maxToolResultChars: 2_000,
+      provider: Providers.BEDROCK,
+      model: 'anthropic.claude-sonnet-4-5',
+      tokenCounter: transformSensitiveCounter,
+      indexTokenCountMap: {
+        0: transformSensitiveCounter(messages[0]),
+        1: transformSensitiveCounter(messages[1]),
+        2: transformSensitiveCounter(messages[2]),
+      },
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(
+      signatureFor('us.anthropic.claude-sonnet-4-5-20250929-v1:0'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    const providerContent = JSON.stringify(
+      model.calls[0].map((message) => message.content)
+    );
+    expect(providerContent).toContain('[Previous tool interaction]');
+    expect(providerContent).not.toContain(artifactSentinel);
+    expect(toolMessage.artifact.content[0].text).toContain(artifactSentinel);
+  });
+
+  it('compacts expanded synthetic context without an artifact', async () => {
+    const toolCallId = 'tc-final-transform-without-artifact';
+    const messages: BaseMessage[] = [
+      new HumanMessage('query the table'),
+      new AIMessageChunk({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'run_select_query',
+            args: { query: `SELECT '${'x'.repeat(5_000)}'` },
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: 'query complete',
+        tool_call_id: toolCallId,
+        name: 'run_select_query',
+      }),
+    ];
+    const transformSensitiveCounter: t.TokenCounter = (message) => {
+      const content =
+        typeof message.content === 'string'
+          ? message.content
+          : JSON.stringify(message.content);
+      if (
+        message instanceof HumanMessage &&
+        content.includes('[Previous tool interaction]')
+      ) {
+        return content.length * 10;
+      }
+      return 1;
+    };
+    const run = await createRun({
+      runId: 'final-transform-without-artifact',
+      maxContextTokens: 500,
+      provider: Providers.BEDROCK,
+      model: 'anthropic.claude-sonnet-4-5',
+      tokenCounter: transformSensitiveCounter,
+      indexTokenCountMap: { 0: 1, 1: 1, 2: 1 },
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(
+      signatureFor('us.anthropic.claude-sonnet-4-5-20250929-v1:0'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    const humanMessages = model.calls[0].filter(
+      (message) => message instanceof HumanMessage
+    );
+    expect(humanMessages).toHaveLength(2);
+    expect(
+      JSON.stringify(humanMessages[humanMessages.length - 1].content).length
+    ).toBeLessThan(100);
+    expect(
+      JSON.stringify(humanMessages[humanMessages.length - 1].content)
+    ).not.toContain('x'.repeat(1_000));
+  });
+
+  it('counts unresolved-reference annotations before invoking the provider', async () => {
+    const toolCallId = 'tc-unresolved-projection';
+    const unresolvedRefs = Array.from(
+      { length: 1_200 },
+      (_, index) => `missing_tool_${index}_turn_${index}`
+    );
+    const messages: BaseMessage[] = [
+      new HumanMessage(`old question ${'q'.repeat(5_500)}`),
+      new AIMessage({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'lookup_records',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      new ToolMessage({
+        content: `old result ${'r'.repeat(3_000)}`,
+        tool_call_id: toolCallId,
+        name: 'lookup_records',
+        additional_kwargs: { _unresolvedRefs: unresolvedRefs },
+      }),
+      new AIMessage(`old answer ${'a'.repeat(4_500)}`),
+      new HumanMessage(`latest question ${'n'.repeat(3_500)}`),
+    ];
+    const projectionCounter: t.TokenCounter = (message) => {
+      const content =
+        typeof message.content === 'string'
+          ? message.content
+          : JSON.stringify(message.content);
+      return Math.max(1, content.length);
+    };
+    const indexTokenCountMap: Record<string, number> = {};
+    for (let i = 0; i < messages.length; i++) {
+      indexTokenCountMap[i] = projectionCounter(messages[i]);
+    }
+    const run = await createRun({
+      runId: 'unresolved-reference-final-projection',
+      maxContextTokens: 20_000,
+      provider: Providers.ANTHROPIC,
+      tokenCounter: projectionCounter,
+      indexTokenCountMap,
+      toolOutputReferences: { enabled: true },
+      tools: [
+        tool(async () => 'unused', {
+          name: 'lookup_records',
+          description: 'Looks up records',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(
+      signatureFor('claude-haiku-4-5-20251001'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    const content = await run.processStream({ messages }, streamConfig);
+
+    expect(content).toEqual([{ type: 'text', text: 'recovered' }]);
+    expect(model.calls).toHaveLength(1);
+    expect(
+      run.Graph.agentContexts.get('default')?.overflowRecoveryAttempts
+    ).toBeGreaterThan(0);
+    const providerPayload = JSON.stringify(
+      model.calls[0].map((message) => message.content)
+    );
+    expect(providerPayload).not.toContain(unresolvedRefs[0]);
+    expect(providerPayload).not.toContain(unresolvedRefs.at(-1));
+    const sentMessageTokens =
+      3 +
+      model.calls[0].reduce(
+        (total, message) => total + projectionCounter(message),
+        0
+      );
+    expect(sentMessageTokens).toBeLessThan(
+      run.Graph.agentContexts.get('default')?.maxContextTokens ?? 0
+    );
+  });
+
+  it('omits artifact expansion that would exceed the post-prune budget', async () => {
+    const toolCallId = 'tc-artifact';
+    const artifactSentinel = 'ARTIFACT_SENTINEL';
+    const toolMessage = new ToolMessage({
+      content: 'result'.repeat(100),
+      tool_call_id: toolCallId,
+      name: 'render_report',
+      artifact: {
+        content: [
+          {
+            type: ContentTypes.TEXT,
+            text: `${artifactSentinel}:${'a'.repeat(5_000)}`,
+          },
+        ],
+      },
+    });
+    const messages: BaseMessage[] = [
+      new HumanMessage('h'.repeat(2_400)),
+      new AIMessageChunk({
+        content: '',
+        tool_calls: [
+          {
+            id: toolCallId,
+            name: 'render_report',
+            args: {},
+            type: 'tool_call',
+          },
+        ],
+      }),
+      toolMessage,
+    ];
+    const artifactTokenCounter: t.TokenCounter = (message) => {
+      const content =
+        typeof message.content === 'string'
+          ? message.content
+          : JSON.stringify(message.content);
+      return Math.ceil(content.length / 4);
+    };
+    const indexTokenCountMap: Record<string, number> = {};
+    for (let i = 0; i < messages.length; i++) {
+      indexTokenCountMap[i] = artifactTokenCounter(messages[i]);
+    }
+    const run = await createRun({
+      runId: 'artifact-budget-guard',
+      maxContextTokens: 1_000,
+      maxToolResultChars: 2_000,
+      provider: Providers.BEDROCK,
+      model: 'anthropic.claude-sonnet-4-5',
+      tokenCounter: artifactTokenCounter,
+      indexTokenCountMap,
+      tools: [
+        tool(async () => 'unused', {
+          name: 'render_report',
+          description: 'Renders a report',
+          schema: z.object({}),
+        }),
+      ],
+    });
+    if (!run.Graph) {
+      throw new Error('Expected graph to be initialized');
+    }
+    const model = new OverflowThenSucceedModel(
+      signatureFor('us.anthropic.claude-sonnet-4-5-20250929-v1:0'),
+      0
+    );
+    run.Graph.overrideModel = model;
+
+    await run.processStream({ messages }, streamConfig);
+
+    expect(model.calls).toHaveLength(1);
+    expect(
+      JSON.stringify(model.calls[0].map((message) => message.content))
+    ).not.toContain(artifactSentinel);
+    expect(toolMessage.content).toBe('result'.repeat(100));
+    expect(toolMessage.artifact.content[0].text).toContain(artifactSentinel);
+  });
+
   it('preserves masked tool originals while checkpointed messages survive', async () => {
     const run = await createRun({
       runId: 'overflow-originals-checkpoint',