🪟 fix: Summarize the Pruned Turn Window Safely

src/agents/AgentContext.ts

@@ -472,6 +472,11 @@ export class AgentContext {
    * disturbing a `maxContextTokens` that no correction ever touched.
    */
   private _preOverflowMaxContextTokens?: number;
+  private correctedContextBudgetMode?: t.ContextBudgetMode;
+
+  get summarizationBudgetMode(): t.ContextBudgetMode {
+    return this.correctedContextBudgetMode ?? this.summarizationConfig?.contextBudgetMode ?? 'input';
+  }
   /**
    * Prompt size, normalized into the local counter's uncalibrated units, at
    * the last overflow correction. Keeping both measurements in the same units
@@ -1772,6 +1777,7 @@ export class AgentContext {
     }
     if (budgetTokens != null) {
       this.maxContextTokens = budgetTokens;
+      this.correctedContextBudgetMode = 'input';
     }
     this.pruneMessages = undefined;
     this.clearProviderProjection(true);
@@ -1981,6 +1987,7 @@ export class AgentContext {
     }
     this.maxContextTokens = this._preOverflowMaxContextTokens;
     this._preOverflowMaxContextTokens = undefined;
+    this.correctedContextBudgetMode = undefined;
     this._lastOverflowPromptTokens = undefined;
     this._overflowRecoveryAttempts = 0;
   }

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

@@ -1,8 +1,8 @@
 import { AIMessage, HumanMessage } from '@langchain/core/messages';
 import type * as t from '@/types';
+import { messagesStateReducer } from '@/messages/reducer';
 import { AgentContext } from '@/agents/AgentContext';
 import { Providers } from '@/common';
-import { messagesStateReducer } from '@/messages/reducer';
 
 /**
  * The overflow-recovery bookkeeping on AgentContext: the budget correction,
@@ -28,6 +28,28 @@ describe('AgentContext overflow recovery state', () => {
     expect(context.pruneMessages).toBeUndefined();
   });
 
+  it('tracks input-ceiling units after correction and restores configured units on reset', () => {
+    const context = createContext(32_000);
+    context.summarizationConfig = {contextBudgetMode: 'window'};
+    expect(context.summarizationBudgetMode).toBe('window');
+    context.applyContextBudgetCorrection(26_600);
+    expect(context.maxContextTokens).toBe(26_600);
+    expect(context.summarizationBudgetMode).toBe('input');
+    expect(context.summarizationConfig.contextBudgetMode).toBe('window');
+    context.reset();
+    expect(context.maxContextTokens).toBe(32_000);
+    expect(context.summarizationBudgetMode).toBe('window');
+  });
+
+  it('does not change budget units when an overflow supplies no numeric correction', () => {
+    const context = createContext(32_000);
+    context.summarizationConfig = {contextBudgetMode: 'window'};
+    context.applyContextBudgetCorrection(undefined);
+    expect(context.summarizationBudgetMode).toBe('window');
+    context.reset();
+    expect(context.summarizationBudgetMode).toBe('window');
+  });
+
   it('keeps one provider projection per run and rebuilds it after overflow', () => {
     const context = AgentContext.fromConfig(
       {

src/graphs/Graph.ts

@@ -3133,6 +3133,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
     config,
     originalToolContent,
     estimatedPromptTokens,
+    windowMessages,
   }: {
     recovery: OverflowRecoveryPlan;
     agentContext: AgentContext;
@@ -3142,6 +3143,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
     originalToolContent?: Map<number, string>;
     /** Size of the rejected prompt, recorded to detect a correction that changed nothing. */
     estimatedPromptTokens?: number;
+    windowMessages: BaseMessage[];
   }): Partial<t.AgentSubgraphState> {
     const previousBudget = agentContext.maxContextTokens;
     /**
@@ -3189,6 +3191,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         remainingContextTokens: 0,
         agentId: agentId || agentContext.agentId,
         reason: 'overflow',
+        windowMessages,
         allowSummarization,
       },
     };
@@ -3225,12 +3228,14 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
     agentContext,
     agentId,
     messageCount,
+    windowMessages,
     contextUsage,
     config,
   }: {
     agentContext: AgentContext;
     agentId: string;
     messageCount: number;
+    windowMessages: BaseMessage[];
     contextUsage: t.ContextUsageEvent | null;
     config: RunnableConfig;
   }): Promise<Partial<t.AgentSubgraphState> | null> {
@@ -3262,6 +3267,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
           remainingContextTokens: contextUsage?.remainingContextTokens ?? 0,
           agentId: agentId || agentContext.agentId,
           reason: 'manual',
+          windowMessages,
         },
         /** A checkpointed thread may still hold the previous compaction's
          *  summary; only the one this run produces may be reported. */
@@ -3484,6 +3490,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
           agentContext,
           agentId,
           messageCount: messages.length,
+          windowMessages: context,
           contextUsage,
           config,
         });
@@ -3541,6 +3548,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
             return {
               summarizationRequest: {
                 remainingContextTokens: remainingContextTokens ?? 0,
+                windowMessages: context,
                 agentId: agentId || agentContext.agentId,
               },
             };
@@ -3570,6 +3578,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
           agentContext,
           agentId,
           messageCount: messages.length,
+          windowMessages: messagesToUse,
           contextUsage: null,
           config,
         });
@@ -4596,6 +4605,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
             agentId,
             config,
             originalToolContent: prunedOriginalToolContent,
+            windowMessages: messagesToUse,
             estimatedPromptTokens: recoveryPromptEstimate,
           });
         }
@@ -4830,6 +4840,7 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
             agentId,
             config,
             originalToolContent: prunedOriginalToolContent,
+            windowMessages: messagesToUse,
             estimatedPromptTokens: fallbackRecoveryPromptEstimate,
           });
         }

src/graphs/__tests__/Graph.compactionWindow.test.ts

@@ -0,0 +1,49 @@
+import { Runnable } from '@langchain/core/runnables';
+import { AIMessage, AIMessageChunk, HumanMessage } from '@langchain/core/messages';
+import type { BaseMessage } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { StandardGraph } from '@/graphs/Graph';
+import { Providers } from '@/common';
+
+class RecordingModel extends Runnable<BaseMessage[], AIMessageChunk> {
+  lc_namespace = ['tests'];
+  calls: BaseMessage[][] = [];
+  async invoke(messages: BaseMessage[]): Promise<AIMessageChunk> {
+    this.calls.push(messages);
+    return new AIMessageChunk('answer');
+  }
+}
+
+const count: t.TokenCounter = () => 10;
+
+function makeGraph(mode: 'normal' | 'manual' | 'trigger'): StandardGraph {
+  return new StandardGraph({
+    runId: mode, tokenCounter: count,
+    agents: [{
+      agentId: 'agent', provider: Providers.OPENAI,
+      instructions: 'Same system prompt', maxContextTokens: 1_000,
+      summarizeOnly: mode === 'manual', summarizationEnabled: true,
+      summarizationConfig: {trigger: {type: 'messages_to_refine', value: mode === 'trigger' ? 1 : 10_000}},
+      clientOptions: {disableStreaming: true},
+    }],
+  });
+}
+
+it('manual and automatic requests carry the normal turn’s pruned window, not 720 stored messages', async () => {
+  const messages: BaseMessage[] = [];
+  for (let i = 0; i < 360; i++) messages.push(new HumanMessage(`q${i}`), new AIMessage(`a${i}`));
+  const normal = makeGraph('normal');
+  const model = new RecordingModel();
+  normal.overrideModel = model;
+  const config = {metadata: {langgraph_node: 'agent=agent', langgraph_step: 1}, configurable: {thread_id: 'window'}};
+  await normal.createCallModel('agent')({messages}, config);
+  expect(model.calls).toHaveLength(1);
+  const pruned = model.calls[0].filter(message => message.getType() !== 'system');
+  expect(pruned.length).toBeLessThan(720);
+  expect(pruned.length).toBeGreaterThan(0);
+  for (const mode of ['manual', 'trigger'] as const) {
+    const graph = makeGraph(mode);
+    const result = await graph.createCallModel('agent')({messages}, config);
+    expect(result.summarizationRequest?.windowMessages).toEqual(pruned);
+  }
+});